Skip to main content

tokenburn_core/
lib.rs

1//! TokenBurn core — UI-agnostic logic shared by the CLI, TUI and GUI front-ends.
2//!
3//! Reads **local logs only** (no network):
4//!
5//! * [pi](https://github.com/earendil-works/pi-coding-agent) session files,
6//!   `~/.pi/agent/sessions/**/*.jsonl` — one row per assistant turn.
7//! * Zed's agent thread database (`threads.db`, sqlite + zstd JSON blobs) —
8//!   one row per thread (input/output tokens only, no `$` cost).
9//! * Claude Code, Codex CLI, Copilot CLI, Gemini CLI, OpenCode and Amp — see
10//!   [`Tool`] and `features/<tool>/`.
11//!
12//! # Typical flow
13//!
14//! ```no_run
15//! use tokenburn_core::{collect, build_report, Query};
16//!
17//! let query = Query::default();
18//! let now = chrono::Local::now();
19//! let collected = collect(query.tool, query.window.start(now));
20//! let report = build_report(&collected.rows, &query);
21//! println!("{} rows", report.len());
22//! ```
23
24pub mod features;
25pub mod utils;
26
27pub use features::report::{
28    axis_label, axis_tokens, bucket_at, bucket_label_of, build_report, by_project, chart_bucket,
29    drill, nice_max, pie, project_key, short_label, short_name, slice_angles, time_series, Bucket,
30    Drill, Metric, Pie, Query, ReportRow, RowFilter, Selection, Slice, TimeSeries, ToolFilter,
31    Window,
32};
33pub use features::settings::{self, Settings};
34pub use features::theme::{
35    find_theme, theme_by_index, theme_index_by_name, AppTheme, Rgb, THEME_COUNT, THEME_NAMES,
36};
37pub use features::usage::{Row, Summary, Tool};
38
39use chrono::{DateTime, Utc};
40
41/// Result of reading every enabled source.
42#[derive(Debug, Default, Clone)]
43pub struct Collected {
44    /// Every usage row at or after the requested start.
45    pub rows: Vec<Row>,
46    /// Non-fatal problems (missing db, unreadable file…), safe to show to the user.
47    pub warnings: Vec<String>,
48}
49
50/// Collect rows for `tool` from the default on-disk locations, keeping only
51/// rows whose timestamp is `>= start`.
52///
53/// Every tool in the filter is read on its own thread (a Zed database or a
54/// year of Claude transcripts can take a moment). A tool that is not installed
55/// simply contributes nothing; a tool whose logs cannot be read adds a warning.
56pub fn collect(tool: ToolFilter, start: DateTime<Utc>) -> Collected {
57    use features::{amp, claude, codex, copilot, gemini, opencode, pi, zed};
58
59    type Job = fn(DateTime<Utc>) -> anyhow::Result<Vec<Row>>;
60    let jobs: [(Tool, Job); 8] = [
61        (Tool::Pi, pi::collect_pi),
62        (Tool::Zed, zed::collect_zed),
63        (Tool::Claude, claude::collect_claude),
64        (Tool::Codex, codex::collect_codex),
65        (Tool::Copilot, copilot::collect_copilot),
66        (Tool::Gemini, gemini::collect_gemini),
67        (Tool::OpenCode, opencode::collect_opencode),
68        (Tool::Amp, amp::collect_amp),
69    ];
70
71    let results: Vec<(Tool, anyhow::Result<Vec<Row>>)> = std::thread::scope(|scope| {
72        let handles: Vec<_> = jobs
73            .iter()
74            .filter(|(t, _)| tool.includes(*t))
75            .map(|(t, job)| (*t, scope.spawn(move || job(start))))
76            .collect();
77        handles
78            .into_iter()
79            .map(|(t, h)| {
80                let r = h
81                    .join()
82                    .unwrap_or_else(|_| Err(anyhow::anyhow!("collector panicked")));
83                (t, r)
84            })
85            .collect()
86    });
87
88    let mut out = Collected::default();
89    for (t, r) in results {
90        match r {
91            Ok(rows) => out.rows.extend(rows),
92            Err(e) => out.warnings.push(format!("{}: {e:#}", t.as_str())),
93        }
94    }
95    out.rows.retain(|r| r.ts >= start);
96    out.rows.sort_by_key(|r| r.ts);
97    out
98}