use std::collections::{HashMap, VecDeque}; use std::sync::atomic::{AtomicU64, Ordering::Relaxed}; use std::sync::{Arc, Mutex, OnceLock}; use crate::providers::Usage; const WINDOW: usize = 100; const DEGRADE_THRESHOLD: usize = 3; const DEGRADE_WINDOW: usize = 10; #[derive(Default)] struct ProviderStat { model: String, tokens_in: u64, tokens_out: u64, cost: f64, calls: u64, last_latency_ms: u64, latencies: VecDeque, recent_results: VecDeque, } pub struct Metrics { tokens_in: AtomicU64, tokens_out: AtomicU64, turns: AtomicU64, tool_calls: AtomicU64, per_tool: Mutex>, turn_latencies: Mutex>, providers: Mutex>, } #[derive(serde::Serialize)] pub struct MetricsSnapshot { pub tokens_in: u64, pub tokens_out: u64, pub cost: f64, pub turns: u64, pub tool_calls: u64, pub turn_latency_p95_ms: u64, pub per_tool: HashMap, pub providers: Vec, } #[derive(serde::Serialize)] pub struct ProviderSnapshot { pub name: String, pub model: String, pub status: String, pub latency_ms: u64, pub latencies: Vec, pub tokens_in: u64, pub tokens_out: u64, pub cost: f64, } impl Default for Metrics { fn default() -> Self { Self::new() } } impl Metrics { pub fn new() -> Self { Self { tokens_in: AtomicU64::new(0), tokens_out: AtomicU64::new(0), turns: AtomicU64::new(0), tool_calls: AtomicU64::new(0), per_tool: Mutex::new(HashMap::new()), turn_latencies: Mutex::new(VecDeque::new()), providers: Mutex::new(HashMap::new()), } } pub fn record_turn(&self, usage: Option<&Usage>, latency_ms: u64) { self.turns.fetch_add(1, Relaxed); if let Some(u) = usage { self.tokens_in .fetch_add(u64::from(u.prompt_tokens), Relaxed); self.tokens_out .fetch_add(u64::from(u.completion_tokens), Relaxed); } let mut q = self .turn_latencies .lock() .unwrap_or_else(|e| e.into_inner()); q.push_back(latency_ms); while q.len() > WINDOW { q.pop_front(); } } pub fn record_tool_call(&self, name: &str, _success: bool) { self.tool_calls.fetch_add(1, Relaxed); *self .per_tool .lock() .unwrap_or_else(|e| e.into_inner()) .entry(name.to_string()) .or_insert(0) += 1; } pub fn record_provider( &self, name: &str, model: &str, cost: Option, latency_ms: u64, is_error: bool, ) { let mut map = self.providers.lock().unwrap_or_else(|e| e.into_inner()); let stat = map.entry(name.to_string()).or_insert_with(|| ProviderStat { model: model.to_string(), ..Default::default() }); stat.model = model.to_string(); stat.calls += 1; stat.last_latency_ms = latency_ms; stat.latencies.push_back(latency_ms); while stat.latencies.len() > WINDOW { stat.latencies.pop_front(); } stat.recent_results.push_back(is_error); while stat.recent_results.len() > DEGRADE_WINDOW { stat.recent_results.pop_front(); } if let Some(c) = cost { stat.cost += c; } } pub fn record_provider_tokens(&self, name: &str, usage: &Usage) { let mut map = self.providers.lock().unwrap_or_else(|e| e.into_inner()); let stat = map.entry(name.to_string()).or_default(); stat.tokens_in += u64::from(usage.prompt_tokens); stat.tokens_out += u64::from(usage.completion_tokens); } pub fn tool_call_count(&self, name: &str) -> u64 { *self .per_tool .lock() .unwrap_or_else(|e| e.into_inner()) .get(name) .unwrap_or(&0) } pub fn snapshot(&self) -> MetricsSnapshot { let latencies = self .turn_latencies .lock() .unwrap_or_else(|e| e.into_inner()); let p95 = percentile_95(&latencies); let providers = self.providers.lock().unwrap_or_else(|e| e.into_inner()); let mut cost = 0.0; let provider_snaps = providers .iter() .map(|(name, s)| { cost += s.cost; let errors = s.recent_results.iter().filter(|e| **e).count(); ProviderSnapshot { name: name.clone(), model: s.model.clone(), status: if s.recent_results.len() >= DEGRADE_WINDOW && errors >= DEGRADE_THRESHOLD { "degraded".to_string() } else { "ok".to_string() }, latency_ms: s.last_latency_ms, latencies: s.latencies.iter().copied().collect(), tokens_in: s.tokens_in, tokens_out: s.tokens_out, cost: s.cost, } }) .collect(); MetricsSnapshot { tokens_in: self.tokens_in.load(Relaxed), tokens_out: self.tokens_out.load(Relaxed), cost, turns: self.turns.load(Relaxed), tool_calls: self.tool_calls.load(Relaxed), turn_latency_p95_ms: p95, per_tool: self .per_tool .lock() .unwrap_or_else(|e| e.into_inner()) .clone(), providers: provider_snaps, } } } fn percentile_95(values: &VecDeque) -> u64 { if values.is_empty() { return 0; } let mut sorted: Vec = values.iter().copied().collect(); sorted.sort_unstable(); let idx = ((sorted.len() as f64 * 0.95).ceil() as usize) .saturating_sub(1) .min(sorted.len() - 1); sorted[idx] } static GLOBAL: OnceLock> = OnceLock::new(); pub fn global_metrics() -> Arc { GLOBAL.get_or_init(|| Arc::new(Metrics::new())).clone() } #[cfg(test)] mod tests { use super::*; #[test] fn records_turns_tokens_and_p95() { let m = Metrics::new(); for i in 1..=100u64 { m.record_turn( Some(&Usage { prompt_tokens: 10, completion_tokens: 20, total_tokens: 30, ..Default::default() }), i, ); } let s = m.snapshot(); assert_eq!(s.turns, 100); assert_eq!(s.tokens_in, 1000); assert_eq!(s.tokens_out, 2000); assert!(s.turn_latency_p95_ms >= 95); } #[test] fn records_per_tool_counts() { let m = Metrics::new(); m.record_tool_call("bash", true); m.record_tool_call("bash", true); m.record_tool_call("read_file", false); let s = m.snapshot(); assert_eq!(s.tool_calls, 3); assert_eq!(s.per_tool.get("bash"), Some(&2)); assert_eq!(s.per_tool.get("read_file"), Some(&1)); } #[test] fn records_provider_tokens_and_cost() { let m = Metrics::new(); m.record_provider("openai", "gpt-4o", Some(0.05), 120, false); m.record_provider_tokens( "openai", &Usage { prompt_tokens: 100, completion_tokens: 50, total_tokens: 150, ..Default::default() }, ); let s = m.snapshot(); assert_eq!(s.cost, 0.05); let p = s.providers.iter().find(|p| p.name == "openai").unwrap(); assert_eq!(p.tokens_in, 100); assert_eq!(p.tokens_out, 50); assert_eq!(p.latency_ms, 120); assert_eq!(p.status, "ok"); } #[test] fn derives_provider_status() { let m = Metrics::new(); for _ in 0..7 { m.record_provider("openai", "gpt-4o", None, 100, false); } for _ in 0..3 { m.record_provider("openai", "gpt-4o", None, 100, true); } let s = m.snapshot(); let p = s.providers.iter().find(|p| p.name == "openai").unwrap(); assert_eq!(p.status, "degraded"); } #[test] fn tool_call_count_accessor() { let m = Metrics::new(); assert_eq!(m.tool_call_count("bash"), 0); m.record_tool_call("bash", true); assert_eq!(m.tool_call_count("bash"), 1); } #[test] fn global_metrics_returns_same_instance() { let a = global_metrics(); let b = global_metrics(); assert!(Arc::ptr_eq(&a, &b)); } }