haki_telemetry/
lib.rs

1//! haki-telemetry — token tracking, cost accumulation, and budget enforcement.
2
3#[derive(Debug, Clone, PartialEq)]
4pub enum BudgetStatus {
5    /// Under 80% of the configured budget.
6    Ok,
7    /// Between 80% and 100% of the configured budget.
8    Warning { used_usd: f64, budget_usd: f64 },
9    /// Budget exceeded.
10    Exceeded { used_usd: f64, budget_usd: f64 },
11}
12
13impl BudgetStatus {
14    pub fn is_exceeded(&self) -> bool {
15        matches!(self, Self::Exceeded { .. })
16    }
17}
18
19/// Accumulates token usage and USD cost for the current session.
20#[derive(Debug, Clone, Default)]
21pub struct UsageTracker {
22    pub input_tokens: u64,
23    pub output_tokens: u64,
24    pub cache_read_tokens: u64,
25    pub cache_write_tokens: u64,
26    pub total_cost_usd: f64,
27    budget_usd: Option<f64>,
28}
29
30impl UsageTracker {
31    pub fn new(budget_usd: Option<f64>) -> Self {
32        Self { budget_usd, ..Default::default() }
33    }
34
35    /// Record a single LLM call. `cost_usd` should be pre-computed by the
36    /// caller using `ModelRegistry::estimate_cost`.
37    pub fn record(
38        &mut self,
39        input: u64,
40        output: u64,
41        cache_read: u64,
42        cache_write: u64,
43        cost_usd: f64,
44    ) -> BudgetStatus {
45        self.input_tokens += input;
46        self.output_tokens += output;
47        self.cache_read_tokens += cache_read;
48        self.cache_write_tokens += cache_write;
49        self.total_cost_usd += cost_usd;
50        self.budget_status()
51    }
52
53    pub fn budget_status(&self) -> BudgetStatus {
54        match self.budget_usd {
55            None => BudgetStatus::Ok,
56            Some(budget) => {
57                let ratio = self.total_cost_usd / budget;
58                if ratio >= 1.0 {
59                    BudgetStatus::Exceeded { used_usd: self.total_cost_usd, budget_usd: budget }
60                } else if ratio >= 0.8 {
61                    BudgetStatus::Warning { used_usd: self.total_cost_usd, budget_usd: budget }
62                } else {
63                    BudgetStatus::Ok
64                }
65            }
66        }
67    }
68
69    /// Format a compact status bar string for the TUI.
70    /// Example: `claude-sonnet-4-5  ↑12.3K ↓4.1K  $0.053  [ask]`
71    pub fn format_bar(&self, model: &str, security_mode: &str) -> String {
72        fn fmt_k(n: u64) -> String {
73            if n >= 1_000 {
74                format!("{:.1}K", n as f64 / 1_000.0)
75            } else {
76                n.to_string()
77            }
78        }
79        format!(
80            "{}  ↑{} ↓{}  ${:.4}  [{}]",
81            model,
82            fmt_k(self.input_tokens),
83            fmt_k(self.output_tokens),
84            self.total_cost_usd,
85            security_mode,
86        )
87    }
88}
89
90// ─── Tests ────────────────────────────────────────────────────────────────────
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn record_accumulates_tokens() {
98        let mut t = UsageTracker::new(None);
99        t.record(1000, 500, 0, 0, 0.01);
100        t.record(2000, 1000, 0, 0, 0.02);
101        assert_eq!(t.input_tokens, 3000);
102        assert_eq!(t.output_tokens, 1500);
103        assert!((t.total_cost_usd - 0.03).abs() < 1e-9);
104    }
105
106    #[test]
107    fn budget_exceeded_when_over_limit() {
108        let mut t = UsageTracker::new(Some(0.05));
109        let status = t.record(0, 0, 0, 0, 0.06);
110        assert!(status.is_exceeded());
111    }
112
113    #[test]
114    fn budget_warning_at_80_percent() {
115        let mut t = UsageTracker::new(Some(0.10));
116        let status = t.record(0, 0, 0, 0, 0.09);
117        assert!(matches!(status, BudgetStatus::Warning { .. }));
118    }
119
120    #[test]
121    fn no_budget_is_always_ok() {
122        let mut t = UsageTracker::new(None);
123        let status = t.record(0, 0, 0, 0, 9999.0);
124        assert_eq!(status, BudgetStatus::Ok);
125    }
126
127    #[test]
128    fn format_bar_contains_model_and_cost() {
129        let mut t = UsageTracker::new(None);
130        t.record(12_300, 4_100, 0, 0, 0.053);
131        let bar = t.format_bar("claude-sonnet-4-5", "ask");
132        assert!(bar.contains("claude-sonnet-4-5"));
133        assert!(bar.contains("↑12.3K"));
134        assert!(bar.contains("↓4.1K"));
135        assert!(bar.contains("[ask]"));
136    }
137}
haki_telemetry/lib.rs

haki_telemetry/
lib.rs
