byokey_proxy/
usage.rs

1//! In-memory usage statistics for request/token tracking.
2
3use serde::Serialize;
4use std::collections::HashMap;
5use std::sync::Mutex;
6use std::sync::atomic::{AtomicU64, Ordering};
7
8/// Global request/token counters.
9#[derive(Default)]
10pub struct UsageStats {
11    /// Total requests received.
12    pub total_requests: AtomicU64,
13    /// Successful requests (2xx from upstream).
14    pub success_requests: AtomicU64,
15    /// Failed requests (non-2xx or internal error).
16    pub failure_requests: AtomicU64,
17    /// Total input tokens across all requests.
18    pub input_tokens: AtomicU64,
19    /// Total output tokens across all requests.
20    pub output_tokens: AtomicU64,
21    /// Per-model request counts.
22    model_counts: Mutex<HashMap<String, ModelStats>>,
23}
24
25/// Per-model usage counters.
26#[derive(Default, Clone, Serialize)]
27pub struct ModelStats {
28    pub requests: u64,
29    pub success: u64,
30    pub failure: u64,
31    pub input_tokens: u64,
32    pub output_tokens: u64,
33}
34
35/// JSON-serializable snapshot of current usage.
36#[derive(Serialize)]
37pub struct UsageSnapshot {
38    pub total_requests: u64,
39    pub success_requests: u64,
40    pub failure_requests: u64,
41    pub input_tokens: u64,
42    pub output_tokens: u64,
43    pub models: HashMap<String, ModelStats>,
44}
45
46impl UsageStats {
47    /// Creates a new empty stats tracker.
48    #[must_use]
49    pub fn new() -> Self {
50        Self::default()
51    }
52
53    /// Record a successful request with optional token counts.
54    pub fn record_success(&self, model: &str, input_tokens: u64, output_tokens: u64) {
55        self.total_requests.fetch_add(1, Ordering::Relaxed);
56        self.success_requests.fetch_add(1, Ordering::Relaxed);
57        self.input_tokens.fetch_add(input_tokens, Ordering::Relaxed);
58        self.output_tokens
59            .fetch_add(output_tokens, Ordering::Relaxed);
60
61        if let Ok(mut map) = self.model_counts.lock() {
62            let entry = map.entry(model.to_string()).or_default();
63            entry.requests += 1;
64            entry.success += 1;
65            entry.input_tokens += input_tokens;
66            entry.output_tokens += output_tokens;
67        }
68    }
69
70    /// Record a failed request.
71    pub fn record_failure(&self, model: &str) {
72        self.total_requests.fetch_add(1, Ordering::Relaxed);
73        self.failure_requests.fetch_add(1, Ordering::Relaxed);
74
75        if let Ok(mut map) = self.model_counts.lock() {
76            let entry = map.entry(model.to_string()).or_default();
77            entry.requests += 1;
78            entry.failure += 1;
79        }
80    }
81
82    /// Take a JSON-serializable snapshot of current stats.
83    #[must_use]
84    pub fn snapshot(&self) -> UsageSnapshot {
85        let models = self
86            .model_counts
87            .lock()
88            .map(|m| m.clone())
89            .unwrap_or_default();
90        UsageSnapshot {
91            total_requests: self.total_requests.load(Ordering::Relaxed),
92            success_requests: self.success_requests.load(Ordering::Relaxed),
93            failure_requests: self.failure_requests.load(Ordering::Relaxed),
94            input_tokens: self.input_tokens.load(Ordering::Relaxed),
95            output_tokens: self.output_tokens.load(Ordering::Relaxed),
96            models,
97        }
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn test_record_success() {
107        let stats = UsageStats::new();
108        stats.record_success("claude-opus-4-5", 100, 200);
109        stats.record_success("claude-opus-4-5", 50, 100);
110        stats.record_success("gpt-4o", 80, 150);
111
112        let snap = stats.snapshot();
113        assert_eq!(snap.total_requests, 3);
114        assert_eq!(snap.success_requests, 3);
115        assert_eq!(snap.failure_requests, 0);
116        assert_eq!(snap.input_tokens, 230);
117        assert_eq!(snap.output_tokens, 450);
118
119        let claude = &snap.models["claude-opus-4-5"];
120        assert_eq!(claude.requests, 2);
121        assert_eq!(claude.success, 2);
122        assert_eq!(claude.input_tokens, 150);
123        assert_eq!(claude.output_tokens, 300);
124    }
125
126    #[test]
127    fn test_record_failure() {
128        let stats = UsageStats::new();
129        stats.record_failure("gpt-4o");
130        stats.record_success("gpt-4o", 10, 20);
131
132        let snap = stats.snapshot();
133        assert_eq!(snap.total_requests, 2);
134        assert_eq!(snap.success_requests, 1);
135        assert_eq!(snap.failure_requests, 1);
136
137        let model = &snap.models["gpt-4o"];
138        assert_eq!(model.requests, 2);
139        assert_eq!(model.failure, 1);
140        assert_eq!(model.success, 1);
141    }
142
143    #[test]
144    fn test_snapshot_empty() {
145        let stats = UsageStats::new();
146        let snap = stats.snapshot();
147        assert_eq!(snap.total_requests, 0);
148        assert!(snap.models.is_empty());
149    }
150}
byokey_proxy/usage.rs

byokey_proxy/
usage.rs
