1use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10#[derive(Debug, Clone, Default, Serialize, Deserialize)]
12#[serde(rename_all = "camelCase")]
13pub struct StatsCache {
14 #[serde(default)]
16 pub version: u32,
17
18 #[serde(default)]
20 pub last_computed_date: Option<String>,
21
22 #[serde(default)]
24 pub daily_activity: Vec<DailyActivityEntry>,
25
26 #[serde(default)]
28 pub daily_model_tokens: Vec<DailyModelTokens>,
29
30 #[serde(default)]
32 pub model_usage: HashMap<String, ModelUsage>,
33
34 #[serde(default)]
36 pub total_sessions: u64,
37
38 #[serde(default)]
40 pub total_messages: u64,
41
42 #[serde(default)]
44 pub longest_session: Option<LongestSession>,
45
46 #[serde(default)]
48 pub first_session_date: Option<String>,
49
50 #[serde(default)]
52 pub hour_counts: HashMap<String, u64>,
53
54 #[serde(default)]
56 pub total_speculation_time_saved_ms: u64,
57}
58
59#[derive(Debug, Clone, Default, Serialize, Deserialize)]
61#[serde(rename_all = "camelCase")]
62pub struct DailyActivityEntry {
63 pub date: String,
64 #[serde(default)]
65 pub message_count: u64,
66 #[serde(default)]
67 pub session_count: u64,
68 #[serde(default)]
69 pub tool_call_count: u64,
70}
71
72#[derive(Debug, Clone, Default, Serialize, Deserialize)]
74#[serde(rename_all = "camelCase")]
75pub struct DailyModelTokens {
76 pub date: String,
77 #[serde(default)]
78 pub tokens_by_model: HashMap<String, u64>,
79}
80
81#[derive(Debug, Clone, Default, Serialize, Deserialize)]
83#[serde(rename_all = "camelCase")]
84pub struct ModelUsage {
85 #[serde(default)]
86 pub input_tokens: u64,
87 #[serde(default)]
88 pub output_tokens: u64,
89 #[serde(default)]
90 pub cache_read_input_tokens: u64,
91 #[serde(default)]
92 pub cache_creation_input_tokens: u64,
93 #[serde(default)]
94 pub web_search_requests: u64,
95 #[serde(default)]
96 pub cost_usd: f64,
97 #[serde(default)]
98 pub context_window: u64,
99 #[serde(default)]
100 pub max_output_tokens: u64,
101}
102
103impl ModelUsage {
104 pub fn total_tokens(&self) -> u64 {
105 self.input_tokens + self.output_tokens
106 }
107
108 pub fn total_with_cache(&self) -> u64 {
109 self.input_tokens
110 + self.output_tokens
111 + self.cache_read_input_tokens
112 + self.cache_creation_input_tokens
113 }
114}
115
116#[derive(Debug, Clone, Default, Serialize, Deserialize)]
118#[serde(rename_all = "camelCase")]
119pub struct LongestSession {
120 #[serde(default)]
121 pub session_id: Option<String>,
122 #[serde(default)]
123 pub message_count: u64,
124 #[serde(default)]
125 pub date: Option<String>,
126}
127
128#[derive(Debug, Clone, Default, Serialize, Deserialize)]
130#[serde(rename_all = "camelCase")]
131pub struct DailyActivity {
132 #[serde(default)]
133 pub tokens: u64,
134 #[serde(default)]
135 pub input_tokens: u64,
136 #[serde(default)]
137 pub output_tokens: u64,
138 #[serde(default)]
139 pub messages: u64,
140 #[serde(default)]
141 pub sessions: u64,
142}
143
144impl StatsCache {
145 pub fn total_input_tokens(&self) -> u64 {
147 self.model_usage.values().map(|m| m.input_tokens).sum()
148 }
149
150 pub fn total_output_tokens(&self) -> u64 {
152 self.model_usage.values().map(|m| m.output_tokens).sum()
153 }
154
155 pub fn total_tokens(&self) -> u64 {
157 self.total_input_tokens() + self.total_output_tokens()
158 }
159
160 pub fn total_cache_read_tokens(&self) -> u64 {
162 self.model_usage
163 .values()
164 .map(|m| m.cache_read_input_tokens)
165 .sum()
166 }
167
168 pub fn total_cache_write_tokens(&self) -> u64 {
170 self.model_usage
171 .values()
172 .map(|m| m.cache_creation_input_tokens)
173 .sum()
174 }
175
176 pub fn recalculate_costs(&mut self) {
181 for (model_name, usage) in self.model_usage.iter_mut() {
182 usage.cost_usd = crate::pricing::calculate_cost(
183 model_name,
184 usage.input_tokens,
185 usage.output_tokens,
186 usage.cache_creation_input_tokens,
187 usage.cache_read_input_tokens,
188 );
189 }
190 }
191
192 pub fn session_count(&self) -> u64 {
194 self.total_sessions
195 }
196
197 pub fn message_count(&self) -> u64 {
199 self.total_messages
200 }
201
202 pub fn top_models(&self, n: usize) -> Vec<(&str, &ModelUsage)> {
204 let mut models: Vec<_> = self
205 .model_usage
206 .iter()
207 .filter(|(_, usage)| usage.total_tokens() > 0)
208 .map(|(k, v)| (k.as_str(), v))
209 .collect();
210 models.sort_by(|a, b| b.1.total_tokens().cmp(&a.1.total_tokens()));
211 models.truncate(n);
212 models
213 }
214
215 pub fn recent_daily(&self, n: usize) -> Vec<&DailyActivityEntry> {
217 let len = self.daily_activity.len();
218 if len <= n {
219 self.daily_activity.iter().collect()
220 } else {
221 self.daily_activity[len - n..].iter().collect()
222 }
223 }
224
225 pub fn cache_ratio(&self) -> f64 {
227 let cache_read = self.total_cache_read_tokens();
228 let total_input = self.total_input_tokens() + cache_read;
229 if total_input == 0 {
230 return 0.0;
231 }
232 cache_read as f64 / total_input as f64
233 }
234
235 pub const CONTEXT_WINDOW: u64 = 200_000;
237
238 pub fn calculate_context_saturation(
243 session_metadata: &[&crate::models::SessionMetadata],
244 last_n: usize,
245 ) -> ContextWindowStats {
246 if session_metadata.is_empty() {
247 return ContextWindowStats::default();
248 }
249
250 let mut sorted: Vec<_> = session_metadata
252 .iter()
253 .filter(|s| s.last_timestamp.is_some() && s.total_tokens > 0)
254 .collect();
255 sorted.sort_by(|a, b| b.last_timestamp.cmp(&a.last_timestamp));
256
257 let recent: Vec<_> = sorted.into_iter().take(last_n).collect();
259
260 if recent.is_empty() {
261 return ContextWindowStats::default();
262 }
263
264 let mut total_pct = 0.0;
266 let mut high_load_count = 0;
267 let mut peak_pct = 0.0;
268
269 for session in &recent {
270 let saturation_pct =
271 (session.total_tokens as f64 / Self::CONTEXT_WINDOW as f64) * 100.0;
272 total_pct += saturation_pct;
273
274 if saturation_pct > 85.0 {
275 high_load_count += 1;
276 }
277
278 if saturation_pct > peak_pct {
279 peak_pct = saturation_pct;
280 }
281 }
282
283 ContextWindowStats {
284 avg_saturation_pct: total_pct / recent.len() as f64,
285 high_load_count,
286 peak_saturation_pct: peak_pct,
287 }
288 }
289}
290
291#[derive(Debug, Clone, Default)]
293pub struct ContextWindowStats {
294 pub avg_saturation_pct: f64,
296
297 pub high_load_count: usize,
299
300 pub peak_saturation_pct: f64,
302}
303
304#[cfg(test)]
305mod tests {
306 use super::*;
307
308 #[test]
309 fn test_stats_cache_defaults() {
310 let stats = StatsCache::default();
311 assert_eq!(stats.total_tokens(), 0);
312 assert!(stats.model_usage.is_empty());
313 }
314
315 #[test]
316 fn test_model_usage_total() {
317 let usage = ModelUsage {
318 input_tokens: 1000,
319 output_tokens: 500,
320 ..Default::default()
321 };
322 assert_eq!(usage.total_tokens(), 1500);
323 }
324
325 #[test]
326 fn test_cache_ratio() {
327 let mut stats = StatsCache::default();
328 stats.model_usage.insert(
329 "test".into(),
330 ModelUsage {
331 input_tokens: 800,
332 cache_read_input_tokens: 200,
333 ..Default::default()
334 },
335 );
336 assert!((stats.cache_ratio() - 0.2).abs() < 0.001);
337 }
338
339 #[test]
340 fn test_top_models() {
341 let mut stats = StatsCache::default();
342 stats.model_usage.insert(
343 "opus".to_string(),
344 ModelUsage {
345 input_tokens: 1000,
346 output_tokens: 500,
347 ..Default::default()
348 },
349 );
350 stats.model_usage.insert(
351 "sonnet".to_string(),
352 ModelUsage {
353 input_tokens: 2000,
354 output_tokens: 1000,
355 ..Default::default()
356 },
357 );
358
359 let top = stats.top_models(2);
360 assert_eq!(top[0].0, "sonnet");
361 assert_eq!(top[1].0, "opus");
362 }
363
364 #[test]
365 fn test_parse_real_format() {
366 let json = r#"{
367 "version": 2,
368 "lastComputedDate": "2026-01-31",
369 "dailyActivity": [
370 {"date": "2026-01-30", "messageCount": 100, "sessionCount": 5, "toolCallCount": 20}
371 ],
372 "modelUsage": {
373 "claude-opus-4-5": {
374 "inputTokens": 1000,
375 "outputTokens": 500,
376 "cacheReadInputTokens": 200,
377 "cacheCreationInputTokens": 100
378 }
379 },
380 "totalSessions": 10,
381 "totalMessages": 1000,
382 "hourCounts": {"10": 50, "14": 100}
383 }"#;
384
385 let stats: StatsCache = serde_json::from_str(json).unwrap();
386 assert_eq!(stats.version, 2);
387 assert_eq!(stats.total_sessions, 10);
388 assert_eq!(stats.total_messages, 1000);
389 assert_eq!(stats.daily_activity.len(), 1);
390 assert_eq!(stats.total_input_tokens(), 1000);
391 assert_eq!(stats.total_output_tokens(), 500);
392 }
393
394 #[test]
395 fn test_context_saturation_calculation() {
396 use crate::models::SessionMetadata;
397 use chrono::Utc;
398 use std::path::PathBuf;
399
400 let mut sessions = vec![];
401 let now = Utc::now();
402
403 for (i, tokens) in [50_000u64, 100_000, 150_000, 170_000, 190_000]
405 .iter()
406 .enumerate()
407 {
408 let mut meta = SessionMetadata::from_path(
409 PathBuf::from(format!("/test{}.jsonl", i)),
410 "test".into(),
411 );
412 meta.total_tokens = *tokens;
413 meta.last_timestamp = Some(now - chrono::Duration::seconds((4 - i) as i64 * 60));
414 sessions.push(meta);
415 }
416
417 let refs: Vec<_> = sessions.iter().collect();
418 let stats = StatsCache::calculate_context_saturation(&refs, 30);
419
420 assert!((stats.avg_saturation_pct - 66.0).abs() < 1.0);
422
423 assert_eq!(stats.high_load_count, 1);
425
426 assert!((stats.peak_saturation_pct - 95.0).abs() < 1.0);
428 }
429
430 #[test]
431 fn test_context_saturation_empty_sessions() {
432 let stats = StatsCache::calculate_context_saturation(&[], 30);
433 assert_eq!(stats.avg_saturation_pct, 0.0);
434 assert_eq!(stats.high_load_count, 0);
435 }
436
437 #[test]
438 fn test_context_saturation_fewer_than_requested() {
439 use crate::models::SessionMetadata;
440 use chrono::Utc;
441 use std::path::PathBuf;
442
443 let mut sessions = vec![];
444 let now = Utc::now();
445
446 for (i, tokens) in [60_000u64, 80_000, 120_000].iter().enumerate() {
448 let mut meta = SessionMetadata::from_path(
449 PathBuf::from(format!("/test{}.jsonl", i)),
450 "test".into(),
451 );
452 meta.total_tokens = *tokens;
453 meta.last_timestamp = Some(now - chrono::Duration::seconds((2 - i) as i64 * 60));
454 sessions.push(meta);
455 }
456
457 let refs: Vec<_> = sessions.iter().collect();
458 let stats = StatsCache::calculate_context_saturation(&refs, 30);
459
460 assert!((stats.avg_saturation_pct - 43.33).abs() < 0.1);
463 }
464}