tokrs 0.1.0

A CLI tool for counting tokens
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

use super::*;
use std::time::{SystemTime, UNIX_EPOCH};

fn temp_db() -> Connection {
    let path = std::env::temp_dir().join(format!(
        "tokrs-opencode-{}-{}.db",
        std::process::id(),
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    ));
    let conn = Connection::open(&path).unwrap();
    conn.execute_batch(
            "CREATE TABLE session (id TEXT PRIMARY KEY, time_updated INTEGER);
             CREATE TABLE message (id TEXT, session_id TEXT, time_created INTEGER, time_updated INTEGER, data TEXT);",
        )
        .unwrap();
    conn
}

fn insert_message(conn: &Connection, id: &str, session_id: &str, data: &str) {
    conn.execute(
            "INSERT INTO message (id, session_id, time_created, time_updated, data) VALUES (?1, ?2, 1, 1, ?3)",
            [id, session_id, data],
        )
        .unwrap();
}

#[test]
fn test_parse_assistant_messages() {
    let conn = temp_db();
    conn.execute(
        "INSERT INTO session (id, time_updated) VALUES ('s1', 1)",
        [],
    )
    .unwrap();
    insert_message(
        &conn,
        "m1",
        "s1",
        r#"{"role":"assistant","modelID":"deepseek-v4","cost":0.42,"tokens":{"input":10,"output":5,"reasoning":3,"cache":{"read":100,"write":20}},"time":{"created":1788256800000,"completed":1788256801000}}"#,
    );
    insert_message(
        &conn,
        "m2",
        "s1",
        r#"{"role":"assistant","modelID":"deepseek-v4","tokens":{"input":1,"output":1},"time":{"created":1788256800000}}"#,
    );
    insert_message(
        &conn,
        "m3",
        "s1",
        r#"{"role":"user","tokens":{"input":5,"output":5},"time":{"created":1788256800000,"completed":1788256801000}}"#,
    );
    insert_message(
        &conn,
        "m4",
        "s1",
        r#"{"role":"assistant","modelID":"m","tokens":{"input":0,"output":0,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1788256800000,"completed":1788256801000}}"#,
    );

    let db_path = std::path::Path::new(conn.path().unwrap()).to_path_buf();
    drop(conn);
    let entries = collect_from(&db_path).unwrap();
    std::fs::remove_file(&db_path).ok();
    assert_eq!(entries.len(), 1);
    let e = &entries[0];
    assert_eq!(e.model, "deepseek-v4");
    assert_eq!(e.input_tokens, 10);
    assert_eq!(e.output_tokens, 8);
    assert_eq!(e.cache_read_tokens, 100);
    assert_eq!(e.cache_creation_tokens, 20);
    assert_eq!(e.created_at, 1_788_256_800);
    assert_eq!(e.session_id.as_deref(), Some("s1"));
    // 自报聚合成本被捕获
    assert_eq!(e.self_cost_usd, Some(0.42));
}

#[test]
fn test_zero_tokens_with_cost_imported() {
    let conn = temp_db();
    conn.execute(
        "INSERT INTO session (id, time_updated) VALUES ('s2', 1)",
        [],
    )
    .unwrap();
    insert_message(
        &conn,
        "m1",
        "s2",
        r#"{"role":"assistant","modelID":"free","cost":0.9,"tokens":{"input":0,"output":0,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1788256800000,"completed":1788256801000}}"#,
    );
    let db_path = std::path::Path::new(conn.path().unwrap()).to_path_buf();
    drop(conn);
    let entries = collect_from(&db_path).unwrap();
    std::fs::remove_file(&db_path).ok();
    // token 全零但有真实扣费: 保留以记录成本
    assert_eq!(entries.len(), 1);
    assert_eq!(entries[0].total_tokens(), 0);
    assert_eq!(entries[0].self_cost_usd, Some(0.9));
}

#[test]
fn test_missing_db_returns_empty() {
    let path = std::env::temp_dir().join("tokrs-opencode-nonexistent.db");
    assert!(collect_from(&path).unwrap().is_empty());
}