haki-telemetry 0.1.0

//! haki-telemetry — token tracking, cost accumulation, and budget enforcement.

#[derive(Debug, Clone, PartialEq)]
pub enum BudgetStatus {
    /// Under 80% of the configured budget.
    Ok,
    /// Between 80% and 100% of the configured budget.
    Warning { used_usd: f64, budget_usd: f64 },
    /// Budget exceeded.
    Exceeded { used_usd: f64, budget_usd: f64 },
}

impl BudgetStatus {
    pub fn is_exceeded(&self) -> bool {
        matches!(self, Self::Exceeded { .. })
    }
}

/// Accumulates token usage and USD cost for the current session.
#[derive(Debug, Clone, Default)]
pub struct UsageTracker {
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cache_read_tokens: u64,
    pub cache_write_tokens: u64,
    pub total_cost_usd: f64,
    budget_usd: Option<f64>,
}

impl UsageTracker {
    pub fn new(budget_usd: Option<f64>) -> Self {
        Self { budget_usd, ..Default::default() }
    }

    /// Record a single LLM call. `cost_usd` should be pre-computed by the
    /// caller using `ModelRegistry::estimate_cost`.
    pub fn record(
        &mut self,
        input: u64,
        output: u64,
        cache_read: u64,
        cache_write: u64,
        cost_usd: f64,
    ) -> BudgetStatus {
        self.input_tokens += input;
        self.output_tokens += output;
        self.cache_read_tokens += cache_read;
        self.cache_write_tokens += cache_write;
        self.total_cost_usd += cost_usd;
        self.budget_status()
    }

    pub fn budget_status(&self) -> BudgetStatus {
        match self.budget_usd {
            None => BudgetStatus::Ok,
            Some(budget) => {
                let ratio = self.total_cost_usd / budget;
                if ratio >= 1.0 {
                    BudgetStatus::Exceeded { used_usd: self.total_cost_usd, budget_usd: budget }
                } else if ratio >= 0.8 {
                    BudgetStatus::Warning { used_usd: self.total_cost_usd, budget_usd: budget }
                } else {
                    BudgetStatus::Ok
                }
            }
        }
    }

    /// Format a compact status bar string for the TUI.
    /// Example: `claude-sonnet-4-5  ↑12.3K ↓4.1K  $0.053  [ask]`
    pub fn format_bar(&self, model: &str, security_mode: &str) -> String {
        fn fmt_k(n: u64) -> String {
            if n >= 1_000 {
                format!("{:.1}K", n as f64 / 1_000.0)
            } else {
                n.to_string()
            }
        }
        format!(
            "{}  ↑{} ↓{}  ${:.4}  [{}]",
            model,
            fmt_k(self.input_tokens),
            fmt_k(self.output_tokens),
            self.total_cost_usd,
            security_mode,
        )
    }
}

// ─── Tests ────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn record_accumulates_tokens() {
        let mut t = UsageTracker::new(None);
        t.record(1000, 500, 0, 0, 0.01);
        t.record(2000, 1000, 0, 0, 0.02);
        assert_eq!(t.input_tokens, 3000);
        assert_eq!(t.output_tokens, 1500);
        assert!((t.total_cost_usd - 0.03).abs() < 1e-9);
    }

    #[test]
    fn budget_exceeded_when_over_limit() {
        let mut t = UsageTracker::new(Some(0.05));
        let status = t.record(0, 0, 0, 0, 0.06);
        assert!(status.is_exceeded());
    }

    #[test]
    fn budget_warning_at_80_percent() {
        let mut t = UsageTracker::new(Some(0.10));
        let status = t.record(0, 0, 0, 0, 0.09);
        assert!(matches!(status, BudgetStatus::Warning { .. }));
    }

    #[test]
    fn no_budget_is_always_ok() {
        let mut t = UsageTracker::new(None);
        let status = t.record(0, 0, 0, 0, 9999.0);
        assert_eq!(status, BudgetStatus::Ok);
    }

    #[test]
    fn format_bar_contains_model_and_cost() {
        let mut t = UsageTracker::new(None);
        t.record(12_300, 4_100, 0, 0, 0.053);
        let bar = t.format_bar("claude-sonnet-4-5", "ask");
        assert!(bar.contains("claude-sonnet-4-5"));
        assert!(bar.contains("↑12.3K"));
        assert!(bar.contains("↓4.1K"));
        assert!(bar.contains("[ask]"));
    }
}