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

pub mod claude;
pub mod codex;
pub mod gemini;
pub mod grok;
pub mod kimi;
pub mod opencode;
pub mod pi;
pub mod prince;

use crate::error::AppError;
use crate::model::{AppKind, UsageEntry};

/// 上游 input 含缓存时归一为 fresh input
///
/// 对齐 cc-switch CACHE_INCLUSIVE_APP_TYPES(codex/gemini/grok): 这三家的
/// input 字段包含 cache read(及 codex 的 cache write), 直接相加会双算.
/// 数据不一致(input < read+write)时保守不扣, 避免把上游异常抹成负数.
pub fn fresh_input(input: u64, cache_read: u64, cache_creation: u64) -> u64 {
    input
        .checked_sub(cache_read.saturating_add(cache_creation))
        .unwrap_or(input)
}

/// 剥离模型名的 provider 前缀, 统一以裸 model id 落表(与 codex normalize 同方向)
///
/// 取最后一个 '/' 之后的段(moonshot-cn/kimi-k3 -> kimi-k3), 使同模型跨渠道合并计价;
/// 斜杠后为空或无斜杠时原样返回, 由调用方处理 unknown 兜底
pub fn strip_provider(model: &str) -> &str {
    match model.rfind('/') {
        Some(pos) if pos + 1 < model.len() => &model[pos + 1..],
        _ => model,
    }
}

pub fn collect(apps: &[AppKind]) -> Result<Vec<UsageEntry>, AppError> {
    let mut entries = Vec::new();
    for &app in apps {
        match app {
            AppKind::Claude => entries.extend(claude::collect()?),
            AppKind::Codex => entries.extend(codex::collect()?),
            AppKind::OpenCode => entries.extend(opencode::collect()?),
            AppKind::Gemini => entries.extend(gemini::collect()?),
            AppKind::Grok => entries.extend(grok::collect()?),
            AppKind::Pi => entries.extend(pi::collect()?),
            AppKind::Kimi => entries.extend(kimi::collect()?),
        }
    }
    Ok(entries)
}

#[cfg(test)]
mod tests {
    use super::{fresh_input, strip_provider};

    #[test]
    fn test_fresh_input() {
        assert_eq!(fresh_input(100, 50, 0), 50);
        assert_eq!(fresh_input(100, 50, 20), 30);
        assert_eq!(fresh_input(50, 0, 0), 50);
        assert_eq!(fresh_input(10, 20, 0), 10); // 上游数据不一致时保守不扣
    }

    #[test]
    fn test_strip_provider() {
        assert_eq!(strip_provider("moonshot-cn/kimi-k3"), "kimi-k3");
        assert_eq!(strip_provider("a/b/kimi-k3"), "kimi-k3"); // 多级取最后一段
        assert_eq!(strip_provider("kimi-k3"), "kimi-k3"); // 无斜杠原样
        assert_eq!(strip_provider("moonshot-cn/"), "moonshot-cn/"); // 斜杠后为空, 原样交给调用方兜底
    }
}