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::currency::{find as currency_by_code, Currency, CURRENCIES};
28pub use features::report::{
29    axis_label, axis_tokens, bucket_at, bucket_label_of, build_report, by_project, chart_bucket,
30    drill, nice_max, pie, project_key, short_label, short_name, slice_angles, time_series, Bucket,
31    Drill, Metric, Pie, Query, ReportRow, RowFilter, Selection, Slice, TimeSeries, ToolFilter,
32    Window,
33};
34pub use features::settings::{self, Settings};
35pub use features::theme::{
36    find_theme, theme_by_index, theme_index_by_name, AppTheme, Rgb, THEME_COUNT, THEME_NAMES,
37};
38pub use features::usage::{Row, Summary, Tool};
39pub use utils::layout::LayoutClass;
40
41use chrono::{DateTime, Utc};
42
43/// Result of reading every enabled source.
44#[derive(Debug, Default, Clone)]
45pub struct Collected {
46    /// Every usage row at or after the requested start.
47    pub rows: Vec<Row>,
48    /// Non-fatal problems (missing db, unreadable file…), safe to show to the user.
49    pub warnings: Vec<String>,
50}
51
52/// Collect rows for `tool` from the default on-disk locations, keeping only
53/// rows whose timestamp is `>= start`.
54///
55/// Every tool in the filter is read on its own thread (a Zed database or a
56/// year of Claude transcripts can take a moment). A tool that is not installed
57/// simply contributes nothing; a tool whose logs cannot be read adds a warning.
58pub fn collect(tool: ToolFilter, start: DateTime<Utc>) -> Collected {
59    use features::{amp, claude, codex, copilot, gemini, opencode, pi, zed};
60
61    type Job = fn(DateTime<Utc>) -> anyhow::Result<Vec<Row>>;
62    let jobs: [(Tool, Job); 8] = [
63        (Tool::Pi, pi::collect_pi),
64        (Tool::Zed, zed::collect_zed),
65        (Tool::Claude, claude::collect_claude),
66        (Tool::Codex, codex::collect_codex),
67        (Tool::Copilot, copilot::collect_copilot),
68        (Tool::Gemini, gemini::collect_gemini),
69        (Tool::OpenCode, opencode::collect_opencode),
70        (Tool::Amp, amp::collect_amp),
71    ];
72
73    let results: Vec<(Tool, anyhow::Result<Vec<Row>>)> = std::thread::scope(|scope| {
74        let handles: Vec<_> = jobs
75            .iter()
76            .filter(|(t, _)| tool.includes(*t))
77            .map(|(t, job)| (*t, scope.spawn(move || job(start))))
78            .collect();
79        handles
80            .into_iter()
81            .map(|(t, h)| {
82                let r = h
83                    .join()
84                    .unwrap_or_else(|_| Err(anyhow::anyhow!("collector panicked")));
85                (t, r)
86            })
87            .collect()
88    });
89
90    let mut out = Collected::default();
91    for (t, r) in results {
92        match r {
93            Ok(rows) => out.rows.extend(rows),
94            Err(e) => out.warnings.push(format!("{}: {e:#}", t.as_str())),
95        }
96    }
97    out.rows.retain(|r| r.ts >= start);
98    out.rows.sort_by_key(|r| r.ts);
99    out
100}