tokrs 0.3.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

//! 命令解析与运行入口
//!
//! 可以指定：
//! - app种类
//! - 按照day, app, model分组
//! - 按照起始日期统计
//! - 按照结束日期统计
//! - 输出json文件
//!
use chrono::NaiveDate;
use clap::{Parser, ValueEnum};
use std::collections::BTreeSet;

use crate::{
    apps,
    error::AppError,
    io::{cli_print, load},
    model::{AppKind, TokenTotals, UsageEntry},
    tokens,
};

/// 分类方法枚举
#[derive(Clone, Copy, Debug, ValueEnum)]
enum GroupBy {
    App,
    Model,
    Day,
}

/// Cli命令结构体
#[derive(Parser)]
#[command(name = "tokrs", about = "Local Token Usage Statistics CLI")]
pub struct Cli {
    #[arg(
        long,
        value_delimiter = ',',
        ignore_case = true,
        help = "Applications of Statistics (Claude, Codex, OpenCode, Gemini, Grok, Pi, Kimi)-all by default"
    )]
    app: Vec<AppKind>,
    #[arg(
        long,
        value_enum,
        default_value_t = GroupBy::App,
        ignore_case = true,
        help = "Grouping"
    )]
    by: GroupBy,
    #[arg(long, short, help = "Start Date YYYY-MM-DD (inclusive)")]
    since: Option<NaiveDate>,
    #[arg(long, short, help = "End Date YYYY-MM-DD (inclusive)")]
    until: Option<NaiveDate>,
    #[arg(long, help = "Output in JSON format")]
    json: bool,
    #[arg(
        long,
        help = "Scan threads per app (default: min(CPU cores, 16, file count); 0 = default; >16 clamped to 16 with a warning; invalid values are rejected by the CLI parser)"
    )]
    threads: Option<usize>,
}

/// 解析 --threads: 0/缺省 → auto; 超上限提示后收敛(意图明确的越界, 资源保护);
/// 负数/非数字由 clap 原生报错终止——显式 CLI 参数是用户意图, fail-fast
/// (与"数据文件失败警告继续"分层: 环境不可控 vs 意图传错, 见 AGENTS.md)
fn resolve_threads(raw: Option<usize>) -> Option<usize> {
    match raw {
        None | Some(0) => None,
        Some(n) if n > load::MAX_THREADS => {
            eprintln!(
                ">_: --threads={n} Exceeded the limit of {}, executed according to the limit",
                load::MAX_THREADS
            );
            Some(load::MAX_THREADS)
        }
        Some(n) => Some(n),
    }
}

/// cli运行函数
pub fn run(cli: Cli) -> Result<(), AppError> {
    let app_kinds: Vec<AppKind> = if cli.app.is_empty() {
        AppKind::ALL.to_vec()
    } else {
        cli.app
            .iter()
            .copied()
            .collect::<BTreeSet<_>>()
            .into_iter()
            .collect()
    };

    let mut entries = apps::collect(&app_kinds, resolve_threads(cli.threads))?;
    // 价目表: 损坏直接报错不回写; 未覆盖模型自动追加 null 模板(用全量模型, 先于日期过滤)
    let pricing_path = apps::prince::pricing_path()?;
    let mut table = apps::prince::load_pricing(&pricing_path)?;
    let added = apps::prince::sync_models(&mut table, &pricing_path, &entries)?;
    if added > 0 {
        eprintln!(
            ">_: added {added} model(s) to {} (prices null until filled in)",
            pricing_path.display()
        );
    }
    apps::prince::resolve(&mut entries, &table);
    let entries = tokens::filter_by_range(entries, cli.since, cli.until);
    let total = tokens::grand_total(&entries);
    let today = tokens::today_total(&entries);
    let rows = build_rows(&entries, cli.by);

    if cli.json {
        cli_print::print_json(&rows, &total, &today);
    } else {
        cli_print::print_report(&rows, &total, &today);
    }
    Ok(())
}

/// 按分组维度聚合出报表行(纯函数, 便于单测覆盖分组分发)
fn build_rows(entries: &[UsageEntry], by: GroupBy) -> Vec<(String, TokenTotals)> {
    match by {
        GroupBy::App => tokens::aggregate_by_app(entries)
            .into_iter()
            .map(|(app, totals)| (app.to_string(), totals))
            .collect(),
        GroupBy::Model => tokens::aggregate_by_model(entries)
            .into_iter()
            .map(|((app, model), totals)| (format!("{app}/{model}"), totals))
            .collect(),
        GroupBy::Day => tokens::aggregate_by_day(entries)
            .into_iter()
            .map(|(day, totals)| (day.format("%Y-%m-%d").to_string(), totals))
            .collect(),
    }
}

#[cfg(test)]
#[path = "tests/commands_test.rs"]
mod tests;