tokrs 0.2.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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259

use super::*;
use crate::model::AppKind;
use chrono::{Local, TimeZone};
use std::time::{SystemTime, UNIX_EPOCH};

fn entry(model: &str, created_at: i64, input: u64, output: u64) -> UsageEntry {
    UsageEntry::new(
        AppKind::Claude,
        model.into(),
        None,
        created_at,
        input,
        output,
        0,
        0,
        None,
    )
}

fn table(json: &str) -> PricingFile {
    serde_json::from_str(json).unwrap()
}

fn local_ts(year: i32, month: u32, day: u32, hour: u32) -> i64 {
    Local
        .with_ymd_and_hms(year, month, day, hour, 0, 0)
        .single()
        .unwrap()
        .timestamp()
}

#[test]
fn test_estimate_basic_and_self_cost_priority() {
    let table = table(
        r#"{"version":1,"models":{"m":[{"input":3.0,"output":15.0,"cache_read":0.3,"cache_write":3.75}]}}"#,
    );
    let m = 1_000_000;
    let mut entries = vec![entry("m", local_ts(2026, 9, 1, 12), m, m)];
    entries[0].cache_read_tokens = m;
    entries[0].cache_creation_tokens = m;
    resolve(&mut entries, &table);
    // 3 + 15 + 0.3 + 3.75 (每百万 tokens)
    assert!((entries[0].cost_usd.unwrap() - 22.05).abs() < 1e-9);
    // 自报成本无条件优先
    let mut reported = entry("m", local_ts(2026, 9, 1, 12), m, m);
    reported.self_cost_usd = Some(0.01);
    resolve(std::slice::from_mut(&mut reported), &table);
    assert_eq!(reported.cost_usd, Some(0.01));
}

#[test]
fn test_null_base_is_unpriced_but_partial_null_zeroes() {
    let table = table(
        r#"{"version":1,"models":{"allnull":[{"input":null,"output":null,"cache_read":null,"cache_write":null}],"partial":[{"input":2.0,"output":null}]}}"#,
    );
    let mut entries = vec![
        entry("allnull", 0, 100, 100),
        entry("partial", 0, 500_000, 999_999_999),
    ];
    resolve(&mut entries, &table);
    assert_eq!(entries[0].cost_usd, None);
    // 未填字段按 0 计: 500k * 2.0 / 1M = 1.0
    assert!((entries[1].cost_usd.unwrap() - 1.0).abs() < 1e-9);
}

#[test]
fn test_longest_prefix_match_with_boundary() {
    let table =
        table(r#"{"version":1,"models":{"gpt-5":[{"input":1.0}],"gpt-5-codex":[{"input":2.0}]}}"#);
    let per_m = |model: &str| {
        let mut e = entry(model, 0, 1_000_000, 0);
        resolve(std::slice::from_mut(&mut e), &table);
        e.cost_usd
    };
    assert_eq!(per_m("gpt-5-codex"), Some(2.0));
    assert_eq!(per_m("gpt-5-codex-2026-01"), Some(2.0));
    assert_eq!(per_m("gpt-5.4"), Some(1.0));
    // "gpt-5" 前缀后紧跟字母数字 => 边界不成立, 不误配
    assert_eq!(per_m("gpt-51x"), None);
}

#[test]
fn test_version_selection_by_local_date() {
    let table = table(
        r#"{"version":1,"models":{"m":[{"since":"2026-01-01","input":1.0},{"since":"2026-03-15","input":2.0}]}}"#,
    );
    let cost_at = |ts: i64| {
        let mut e = entry("m", ts, 1_000_000, 0);
        resolve(std::slice::from_mut(&mut e), &table);
        e.cost_usd
    };
    assert_eq!(cost_at(local_ts(2026, 2, 1, 12)), Some(1.0));
    assert_eq!(cost_at(local_ts(2026, 3, 15, 0)), Some(2.0));
    assert_eq!(cost_at(local_ts(2026, 9, 1, 12)), Some(2.0));
    // 早于所有 since 版本 => 无生效价格
    assert_eq!(cost_at(local_ts(2025, 12, 31, 12)), None);
}

#[test]
fn test_peak_hours_multi_ranges_and_wrapping() {
    // 基础 1.0, 峰时(UTC 8-12, 14-18) 覆盖为 2.0
    let peak_table = table(
        r#"{"version":1,"models":{"m":[{"input":1.0,"peak":{"hours":[[8,12],[14,18]],"utc_offset":0,"input":2.0}}]}}"#,
    );
    let cost_at_utc_hour = |hour: u32| {
        let mut e = entry("m", i64::from(hour) * 3600, 1_000_000, 0);
        resolve(std::slice::from_mut(&mut e), &peak_table);
        e.cost_usd.unwrap()
    };
    assert_eq!(cost_at_utc_hour(9), 2.0);
    assert_eq!(cost_at_utc_hour(12), 1.0); // 区间左闭右开
    assert_eq!(cost_at_utc_hour(15), 2.0);
    assert_eq!(cost_at_utc_hour(20), 1.0);
    // 跨午夜回绕 [22,8) + utc_offset 8 时区换算
    let wrap_table = table(
        r#"{"version":1,"models":{"m":[{"input":1.0,"peak":{"hours":[[22,8]],"utc_offset":8,"input":2.0}}]}}"#,
    );
    let mut e = entry("m", 15 * 3600, 1_000_000, 0); // UTC 15 = 东八区 23 点
    resolve(std::slice::from_mut(&mut e), &wrap_table);
    assert_eq!(e.cost_usd.unwrap(), 2.0);
}

#[test]
fn test_long_context_threshold_and_field_inheritance() {
    let table = table(
        r#"{"version":1,"models":{"m":[{"input":1.0,"output":10.0,"cache_read":0.1,"long_context":{"above":200000,"input":6.0}}]}}"#,
    );
    // 上下文 = fresh input + cache_read + cache_creation = 100k + 150k = 250k >= 200k
    // 覆盖块只改 input 价, cache_read 继承基础价
    let mut big = UsageEntry::new(
        AppKind::Claude,
        "m".into(),
        None,
        0,
        100_000,
        0,
        150_000,
        0,
        None,
    );
    resolve(std::slice::from_mut(&mut big), &table);
    let expect = (100_000.0 * 6.0 + 150_000.0 * 0.1) / 1e6;
    assert!((big.cost_usd.unwrap() - expect).abs() < 1e-9);
    // 50k < 200k: 用基础价 1.0
    let mut small = entry("m", 0, 50_000, 0);
    resolve(std::slice::from_mut(&mut small), &table);
    assert!((small.cost_usd.unwrap() - 0.05).abs() < 1e-9);
}

#[test]
fn test_sync_models_appends_template_without_touching_existing() {
    let dir = std::env::temp_dir().join(format!(
        "tokrs-pricing-{}-{}",
        std::process::id(),
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    ));
    std::fs::create_dir_all(&dir).unwrap();
    let path = dir.join("pricing.json");
    // 用户已填 m1 价格
    std::fs::write(
        &path,
        r#"{"version":1,"models":{"m1":[{"input":5.0,"output":null,"cache_read":null,"cache_write":null}]}}"#,
    )
    .unwrap();
    let mut table = load_pricing(&path).unwrap();
    let entries = vec![
        entry("m1", 0, 1, 1),
        entry("m2", 0, 1, 1),
        entry("m2", 1, 1, 1),
    ];
    assert_eq!(sync_models(&mut table, &path, &entries).unwrap(), 1);
    // 重新加载: 已有条目未被覆盖, 新模型带上可编辑的 null 模板
    let mut reloaded = load_pricing(&path).unwrap();
    assert_eq!(reloaded.models["m1"][0].fields.input, Some(5.0));
    assert!(reloaded.models["m2"][0].fields.is_empty());
    // 再次同步无变化, 不再写文件
    let mtime = std::fs::metadata(&path).unwrap().modified().unwrap();
    assert_eq!(sync_models(&mut reloaded, &path, &entries).unwrap(), 0);
    assert_eq!(std::fs::metadata(&path).unwrap().modified().unwrap(), mtime);
    std::fs::remove_dir_all(&dir).ok();
}

#[test]
fn test_force_overrides_self_cost() {
    let forced_table = table(
        r#"{"version":1,"models":{"m":[{"force":true,"input":1.0,"output":0.0,"cache_read":0.0,"cache_write":0.0}]}}"#,
    );
    let empty_force = table(r#"{"version":1,"models":{"m":[{"force":true}]}}"#);
    let plain = table(r#"{"version":1,"models":{"m":[{"input":1.0}]}}"#);
    // force 且已填价: 无视自报 99, 按表计 1.0
    let mut forced = entry("m", 0, 1_000_000, 0);
    forced.self_cost_usd = Some(99.0);
    resolve(std::slice::from_mut(&mut forced), &forced_table);
    assert_eq!(forced.cost_usd, Some(1.0));
    // force 但基础价未填: 回退自报, 不吞数据
    let mut fallback = entry("m", 0, 1_000_000, 0);
    fallback.self_cost_usd = Some(99.0);
    resolve(std::slice::from_mut(&mut fallback), &empty_force);
    assert_eq!(fallback.cost_usd, Some(99.0));
    // 非 force: 自报照常优先
    let mut normal = entry("m", 0, 1_000_000, 0);
    normal.self_cost_usd = Some(99.0);
    resolve(std::slice::from_mut(&mut normal), &plain);
    assert_eq!(normal.cost_usd, Some(99.0));
}

#[test]
fn test_sync_skips_unknown_and_empty_models() {
    let dir = std::env::temp_dir().join(format!(
        "tokrs-pricing-unknown-{}-{}",
        std::process::id(),
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    ));
    let path = dir.join("pricing.json");
    let mut table = PricingFile::default();
    let entries = vec![
        entry("unknown", 0, 1, 1),
        entry("", 0, 1, 1),
        entry("real-model", 0, 1, 1),
    ];
    assert_eq!(sync_models(&mut table, &path, &entries).unwrap(), 1);
    assert!(table.models.contains_key("real-model"));
    assert!(!table.models.contains_key("unknown"));
    assert!(!table.models.keys().any(|k| k.is_empty()));
    std::fs::remove_dir_all(&dir).ok();
}

#[test]
fn test_load_missing_file_is_empty_but_corrupted_is_err() {
    let dir = std::env::temp_dir().join(format!(
        "tokrs-pricing-missing-{}-{}",
        std::process::id(),
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    ));
    let missing = dir.join("pricing.json");
    assert_eq!(load_pricing(&missing).unwrap(), PricingFile::default());
    std::fs::create_dir_all(&dir).unwrap();
    std::fs::write(&missing, "{ not json").unwrap();
    assert!(matches!(
        load_pricing(&missing),
        Err(AppError::Corrupted { .. })
    ));
    // 版本号不支持同样报错(损坏即终止, 由用户修复)
    std::fs::write(&missing, r#"{"version":99,"models":{}}"#).unwrap();
    assert!(matches!(
        load_pricing(&missing),
        Err(AppError::Corrupted { .. })
    ));
    std::fs::remove_dir_all(&dir).ok();
}