tokrs 0.2.4

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

//! claude日志解析
//!
use serde_json::Value;
use std::{
    collections::{HashMap, hash_map::Entry},
    path::Path,
};

use crate::{
    apps::normalize_model,
    error::AppError,
    io::{load, progress::Progress},
    model::{AppKind, UsageEntry},
};

const MAX_DEPTH: usize = 5;

pub fn collect() -> Result<Vec<UsageEntry>, AppError> {
    let base = load::home_dir()?.join(".claude").join("projects");
    if !base.is_dir() {
        return Ok(Vec::new());
    }
    collect_from(&base)
}

pub fn collect_from(base: &Path) -> Result<Vec<UsageEntry>, AppError> {
    let mut candidates: HashMap<String, Candidate> = HashMap::new();
    let files = load::discover_files(base, "jsonl", MAX_DEPTH);
    let mut progress = Progress::start("claude", load::total_bytes(&files));
    // 流式逐行: 文件可达 GB 级, 整读驻留会耗尽内存;
    // 单文件读取失败警告后跳过, 不中止全局统计
    for file in &files {
        progress.set_file(load::file_name_str(file));
        let mut session_fallback: Option<String> = None;
        if let Err(e) = load::for_each_jsonl_progress(file, &[], &mut progress, |value| {
            if session_fallback.is_none() {
                session_fallback = load::str_get(&value, &["sessionId"]).map(str::to_string);
            }
            parse_assistant_line(&value, session_fallback.as_deref(), &mut candidates);
            true
        }) {
            load::warn_file(&e);
        }
    }
    progress.finish();
    Ok(candidates.into_values().map(|c| c.entry).collect())
}

struct Candidate {
    entry: UsageEntry,
    has_stop_reason: bool,
}

fn parse_assistant_line(
    value: &Value,
    session_fallback: Option<&str>,
    candidates: &mut HashMap<String, Candidate>,
) {
    if load::str_get(value, &["type"]) != Some("assistant") {
        return;
    }
    let Some(message) = value.get("message") else {
        return;
    };
    let Some(msg_id) = load::str_get(message, &["id"]) else {
        return;
    };
    let input = load::u64_get(message, &["usage", "input_tokens"]);
    let output = load::u64_get(message, &["usage", "output_tokens"]);
    let cache_read = load::u64_get(message, &["usage", "cache_read_input_tokens"]);
    let cache_creation = load::u64_get(message, &["usage", "cache_creation_input_tokens"]);
    if input == 0 && output == 0 && cache_read == 0 && cache_creation == 0 {
        return;
    }
    let model =
        load::str_get(message, &["model"]).map_or_else(|| "unknown".to_string(), normalize_model);
    let session_id = load::str_get(value, &["sessionId"])
        .or(session_fallback)
        .map(str::to_string);
    let created_at = value
        .get("timestamp")
        .and_then(load::timestamp_to_epoch)
        .unwrap_or_else(load::now_epoch);
    // claude 无自报成本, 待定价表估价
    let entry = UsageEntry::new(
        AppKind::Claude,
        model,
        session_id,
        created_at,
        input,
        output,
        cache_read,
        cache_creation,
        None,
    );
    let candidate = Candidate {
        has_stop_reason: load::str_get(message, &["stop_reason"]).is_some(),
        entry,
    };
    match candidates.entry(msg_id.to_string()) {
        Entry::Vacant(vacant) => {
            vacant.insert(candidate);
        }
        Entry::Occupied(mut occupied) => {
            let existing = occupied.get();
            let should_replace = (candidate.has_stop_reason && !existing.has_stop_reason)
                || (candidate.has_stop_reason == existing.has_stop_reason
                    && candidate.entry.output_tokens > existing.entry.output_tokens);
            if should_replace {
                occupied.insert(candidate);
            }
        }
    }
}

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