1use std::collections::HashMap;
15use std::sync::{Mutex, OnceLock};
16
17use serde::{Deserialize, Serialize};
18
19use crate::core::gain::model_pricing::ModelPricing;
20
21#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
24pub struct ModelUsage {
25 pub requests: u64,
26 pub input_tokens: u64,
27 pub output_tokens: u64,
28 pub cache_read_tokens: u64,
29 pub cache_write_tokens: u64,
30 pub reasoning_tokens: u64,
31 #[serde(default)]
35 pub counterfactual_requests: u64,
36 #[serde(default)]
39 pub counterfactual_input_tokens: u64,
40 #[serde(default)]
43 pub counterfactual_billed_tokens: u64,
44 #[serde(default)]
49 pub measured: MeasuredSlice,
50}
51
52#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq)]
54pub struct MeasuredSlice {
55 pub requests: u64,
56 pub input_tokens: u64,
57 pub output_tokens: u64,
58 pub cache_read_tokens: u64,
59 pub cache_write_tokens: u64,
60 pub cost_usd: f64,
62}
63
64impl MeasuredSlice {
65 fn merge(&mut self, other: &Self) {
66 self.requests += other.requests;
67 self.input_tokens += other.input_tokens;
68 self.output_tokens += other.output_tokens;
69 self.cache_read_tokens += other.cache_read_tokens;
70 self.cache_write_tokens += other.cache_write_tokens;
71 self.cost_usd += other.cost_usd;
72 }
73}
74
75impl ModelUsage {
76 fn add(&mut self, u: &super::usage::RealUsage) {
77 self.requests += 1;
78 self.input_tokens += u.input_tokens;
79 self.output_tokens += u.output_tokens;
80 self.cache_read_tokens += u.cache_read_tokens;
81 self.cache_write_tokens += u.cache_write_tokens;
82 self.reasoning_tokens += u.reasoning_tokens;
83 if let Some(cost) = u.provider_cost_usd {
84 self.measured.merge(&MeasuredSlice {
85 requests: 1,
86 input_tokens: u.input_tokens,
87 output_tokens: u.output_tokens,
88 cache_read_tokens: u.cache_read_tokens,
89 cache_write_tokens: u.cache_write_tokens,
90 cost_usd: cost,
91 });
92 }
93 if let Some(counted) = u
96 .wire
97 .as_deref()
98 .and_then(|w| w.counterfactual.as_ref())
99 .and_then(super::counterfactual::CounterfactualSlot::get)
100 {
101 self.counterfactual_requests += 1;
102 self.counterfactual_input_tokens += counted;
103 self.counterfactual_billed_tokens +=
104 u.input_tokens + u.cache_read_tokens + u.cache_write_tokens;
105 }
106 }
107
108 fn billable_tokens(&self) -> u64 {
109 self.input_tokens + self.output_tokens + self.cache_read_tokens + self.cache_write_tokens
110 }
111}
112
113#[derive(Debug, Clone, Serialize, PartialEq)]
115pub struct ModelSpend {
116 pub model: String,
117 pub requests: u64,
118 pub input_tokens: u64,
119 pub output_tokens: u64,
120 pub cache_read_tokens: u64,
121 pub cache_write_tokens: u64,
122 pub reasoning_tokens: u64,
123 pub cost_usd: f64,
124 pub measured_cost_usd: f64,
127 pub measured_requests: u64,
129 pub pricing_estimated: bool,
132}
133
134#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
138pub struct CohortUsage {
139 pub requests: u64,
140 pub input_tokens: u64,
141 pub output_tokens: u64,
142 #[serde(default)]
146 pub sum_sq_output: u64,
147}
148
149impl CohortUsage {
150 fn add(&mut self, u: &super::usage::RealUsage) {
151 self.requests += 1;
152 self.input_tokens += u.input_tokens;
153 self.output_tokens += u.output_tokens;
154 self.sum_sq_output += u.output_tokens.saturating_mul(u.output_tokens);
155 }
156
157 #[must_use]
159 pub fn avg_output(&self) -> Option<f64> {
160 if self.requests == 0 {
161 None
162 } else {
163 #[allow(clippy::cast_precision_loss)]
164 Some(self.output_tokens as f64 / self.requests as f64)
165 }
166 }
167
168 #[must_use]
172 pub fn variance_output(&self) -> Option<f64> {
173 if self.requests < 2 {
174 return None;
175 }
176 #[allow(clippy::cast_precision_loss)]
177 let n = self.requests as f64;
178 #[allow(clippy::cast_precision_loss)]
179 let sum = self.output_tokens as f64;
180 #[allow(clippy::cast_precision_loss)]
181 let sum_sq = self.sum_sq_output as f64;
182 let var = (sum_sq - sum * sum / n) / (n - 1.0);
183 Some(var.max(0.0))
184 }
185}
186
187#[derive(Debug, Clone, Default, Serialize, Deserialize)]
189pub struct PersistedUsage {
190 pub ts: u64,
191 pub models: HashMap<String, ModelUsage>,
192 #[serde(default)]
195 pub cohorts: HashMap<String, CohortUsage>,
196}
197
198const MAX_TRACKED_MODELS: usize = 256;
201
202const PROXY_USAGE_FILE: &str = "proxy_usage.json";
203
204fn store() -> &'static Mutex<HashMap<String, ModelUsage>> {
205 static STORE: OnceLock<Mutex<HashMap<String, ModelUsage>>> = OnceLock::new();
206 STORE.get_or_init(|| Mutex::new(HashMap::new()))
207}
208
209fn cohort_store() -> &'static Mutex<HashMap<String, CohortUsage>> {
210 static STORE: OnceLock<Mutex<HashMap<String, CohortUsage>>> = OnceLock::new();
211 STORE.get_or_init(|| Mutex::new(HashMap::new()))
212}
213
214pub fn resume_from_disk() {
218 let Some(persisted) = load_persisted() else {
219 return;
220 };
221 let mut map = store()
222 .lock()
223 .unwrap_or_else(std::sync::PoisonError::into_inner);
224 for (model, usage) in persisted.models {
225 let acc = map.entry(model).or_default();
226 acc.requests += usage.requests;
227 acc.input_tokens += usage.input_tokens;
228 acc.output_tokens += usage.output_tokens;
229 acc.cache_read_tokens += usage.cache_read_tokens;
230 acc.cache_write_tokens += usage.cache_write_tokens;
231 acc.reasoning_tokens += usage.reasoning_tokens;
232 acc.counterfactual_requests += usage.counterfactual_requests;
233 acc.counterfactual_input_tokens += usage.counterfactual_input_tokens;
234 acc.counterfactual_billed_tokens += usage.counterfactual_billed_tokens;
235 acc.measured.merge(&usage.measured);
236 }
237 drop(map);
238 let mut cohorts = cohort_store()
239 .lock()
240 .unwrap_or_else(std::sync::PoisonError::into_inner);
241 for (arm, usage) in persisted.cohorts {
242 let acc = cohorts.entry(arm).or_default();
243 acc.requests += usage.requests;
244 acc.input_tokens += usage.input_tokens;
245 acc.output_tokens += usage.output_tokens;
246 }
247}
248
249pub fn record(u: &super::usage::RealUsage) {
252 super::usage_sink::push(u);
255
256 if let Some(wire) = u.wire.as_deref()
262 && (wire.person.is_some() || wire.project.is_some())
263 {
264 let baseline = crate::core::config::Config::load().proxy.baseline.clone();
265 #[allow(clippy::cast_precision_loss)]
266 let cost_usd = if wire.is_local {
267 let billable =
268 u.input_tokens + u.output_tokens + u.cache_read_tokens + u.cache_write_tokens;
269 baseline.effective_local_shadow_rate() / 1_000_000.0 * billable as f64
270 } else if let Some(measured) = u.provider_cost_usd {
271 measured
272 } else {
273 crate::core::gain::model_pricing::ModelPricing::load()
274 .quote(Some(&u.model))
275 .cost
276 .estimate_usd(
277 u.input_tokens,
278 u.output_tokens,
279 u.cache_write_tokens,
280 u.cache_read_tokens,
281 )
282 };
283 super::policy_gate::record_spend(wire.person.as_deref(), wire.project.as_deref(), cost_usd);
284 }
285
286 if let Some(wire) = u.wire.as_deref()
291 && let Some(routed_from) = wire.routed_from.as_deref()
292 {
293 crate::core::savings_ledger::record_routing_event(routed_from, &u.model, u.input_tokens);
294 }
295 if u.cache_read_tokens > 0 {
296 let cost = crate::core::gain::model_pricing::ModelPricing::load()
297 .quote(Some(&u.model))
298 .cost;
299 #[allow(clippy::cast_precision_loss)]
300 let discount_usd =
301 (cost.input_per_m - cost.cache_read_per_m) / 1_000_000.0 * u.cache_read_tokens as f64;
302 crate::core::savings_ledger::record_caching_event(
303 &u.model,
304 u.cache_read_tokens,
305 discount_usd,
306 );
307 }
308
309 let key = normalize_key(&u.model);
310 {
311 let mut map = store()
312 .lock()
313 .unwrap_or_else(std::sync::PoisonError::into_inner);
314 let key = if !map.contains_key(&key) && map.len() >= MAX_TRACKED_MODELS {
315 "unknown".to_string()
316 } else {
317 key
318 };
319 map.entry(key).or_default().add(u);
320 }
321 if let Some(arm) = u.cohort {
322 let mut cohorts = cohort_store()
323 .lock()
324 .unwrap_or_else(std::sync::PoisonError::into_inner);
325 cohorts.entry(arm.as_str().to_string()).or_default().add(u);
326 }
327 persist();
328}
329
330#[must_use]
332pub fn cohort_snapshot() -> HashMap<String, CohortUsage> {
333 cohort_store()
334 .lock()
335 .unwrap_or_else(std::sync::PoisonError::into_inner)
336 .clone()
337}
338
339#[must_use]
341pub fn persisted_cohorts() -> HashMap<String, CohortUsage> {
342 load_persisted().map(|p| p.cohorts).unwrap_or_default()
343}
344
345fn normalize_key(model: &str) -> String {
346 let m = model.trim();
347 if m.is_empty() {
348 "unknown".to_string()
349 } else {
350 m.to_string()
351 }
352}
353
354pub fn snapshot() -> Vec<ModelSpend> {
356 let map = store()
357 .lock()
358 .unwrap_or_else(std::sync::PoisonError::into_inner);
359 price_models(&map)
360}
361
362pub fn total_cost_usd() -> f64 {
364 snapshot().iter().map(|m| m.cost_usd).sum()
365}
366
367#[derive(Debug, Clone, Serialize, PartialEq)]
369pub struct VerifiedSavings {
370 pub requests: u64,
372 pub counterfactual_input_tokens: u64,
375 pub billed_input_tokens: u64,
378 pub verified_saved_tokens: i64,
381}
382
383pub fn verified_savings() -> Option<VerifiedSavings> {
388 let map = store()
389 .lock()
390 .unwrap_or_else(std::sync::PoisonError::into_inner);
391 verified_of(&map)
392}
393
394#[allow(clippy::cast_possible_wrap)]
396fn verified_of(map: &HashMap<String, ModelUsage>) -> Option<VerifiedSavings> {
397 let (mut requests, mut counterfactual, mut billed) = (0u64, 0u64, 0u64);
398 for usage in map.values() {
399 requests += usage.counterfactual_requests;
400 counterfactual += usage.counterfactual_input_tokens;
401 billed += usage.counterfactual_billed_tokens;
402 }
403 (requests > 0).then(|| VerifiedSavings {
404 requests,
405 counterfactual_input_tokens: counterfactual,
406 billed_input_tokens: billed,
407 verified_saved_tokens: counterfactual as i64 - billed as i64,
408 })
409}
410
411pub fn price_models(map: &HashMap<String, ModelUsage>) -> Vec<ModelSpend> {
414 let pricing = ModelPricing::load();
415 let mut rows: Vec<ModelSpend> = map
416 .iter()
417 .map(|(model, usage)| price_one(&pricing, model, usage))
418 .collect();
419 rows.sort_by(|a, b| {
420 b.cost_usd
421 .partial_cmp(&a.cost_usd)
422 .unwrap_or(std::cmp::Ordering::Equal)
423 });
424 rows
425}
426
427fn price_one(pricing: &ModelPricing, model: &str, usage: &ModelUsage) -> ModelSpend {
428 let quote = pricing.quote(Some(model));
429 let m = &usage.measured;
432 let derived_cost = quote.cost.estimate_usd(
433 usage.input_tokens.saturating_sub(m.input_tokens),
434 usage.output_tokens.saturating_sub(m.output_tokens),
435 usage
436 .cache_write_tokens
437 .saturating_sub(m.cache_write_tokens),
438 usage.cache_read_tokens.saturating_sub(m.cache_read_tokens),
439 );
440 let derived_requests = usage.requests.saturating_sub(m.requests);
441 ModelSpend {
442 model: model.to_string(),
443 requests: usage.requests,
444 input_tokens: usage.input_tokens,
445 output_tokens: usage.output_tokens,
446 cache_read_tokens: usage.cache_read_tokens,
447 cache_write_tokens: usage.cache_write_tokens,
448 reasoning_tokens: usage.reasoning_tokens,
449 cost_usd: m.cost_usd + derived_cost,
450 measured_cost_usd: m.cost_usd,
451 measured_requests: m.requests,
452 pricing_estimated: quote.match_kind.is_estimated() && derived_requests > 0,
454 }
455}
456
457fn usage_path() -> Option<std::path::PathBuf> {
458 crate::core::data_dir::lean_ctx_data_dir()
459 .ok()
460 .map(|d| d.join(PROXY_USAGE_FILE))
461}
462
463fn persist() {
465 let Some(path) = usage_path() else {
466 return;
467 };
468 let models = {
469 let map = store()
470 .lock()
471 .unwrap_or_else(std::sync::PoisonError::into_inner);
472 map.clone()
473 };
474 let cohorts = {
475 let map = cohort_store()
476 .lock()
477 .unwrap_or_else(std::sync::PoisonError::into_inner);
478 map.clone()
479 };
480 let payload = PersistedUsage {
481 ts: std::time::SystemTime::now()
482 .duration_since(std::time::UNIX_EPOCH)
483 .unwrap_or_default()
484 .as_secs(),
485 models,
486 cohorts,
487 };
488 let Ok(json) = serde_json::to_string(&payload) else {
489 return;
490 };
491 let tmp = path.with_extension("json.tmp");
492 if std::fs::write(&tmp, json).is_ok() {
493 let _ = std::fs::rename(&tmp, &path);
494 }
495}
496
497pub fn load_persisted() -> Option<PersistedUsage> {
499 let path = usage_path()?;
500 let data = std::fs::read_to_string(path).ok()?;
501 serde_json::from_str(&data).ok()
502}
503
504pub fn persisted_snapshot() -> Vec<ModelSpend> {
506 load_persisted()
507 .map(|p| price_models(&p.models))
508 .unwrap_or_default()
509}
510
511pub fn persisted_dominant_model() -> Option<String> {
515 let persisted = load_persisted()?;
516 persisted
517 .models
518 .iter()
519 .filter(|(m, _)| m.as_str() != "unknown" && !m.trim().is_empty())
520 .max_by_key(|(_, u)| u.billable_tokens())
521 .filter(|(_, u)| u.billable_tokens() > 0)
522 .map(|(m, _)| m.clone())
523}
524
525#[cfg(test)]
526mod tests {
527 use super::*;
528
529 fn usage(
530 model: &str,
531 input: u64,
532 output: u64,
533 cache_read: u64,
534 ) -> super::super::usage::RealUsage {
535 super::super::usage::RealUsage {
536 model: model.to_string(),
537 input_tokens: input,
538 output_tokens: output,
539 cache_read_tokens: cache_read,
540 ..Default::default()
541 }
542 }
543
544 #[test]
547 fn counterfactual_pair_recorded_only_when_probe_answered() {
548 let slot = super::super::counterfactual::CounterfactualSlot::new();
549 slot.set(5_000);
550 let with_probe = super::super::usage::RealUsage {
551 wire: Some(Box::new(super::super::usage::WireContext {
552 counterfactual: Some(slot),
553 ..Default::default()
554 })),
555 cache_write_tokens: 200,
556 ..usage("claude-sonnet-4.5", 1_000, 0, 300)
557 };
558 let mut acc = ModelUsage::default();
559 acc.add(&with_probe);
560 assert_eq!(acc.counterfactual_requests, 1);
561 assert_eq!(acc.counterfactual_input_tokens, 5_000);
562 assert_eq!(acc.counterfactual_billed_tokens, 1_000 + 300 + 200);
564
565 let empty = super::super::usage::RealUsage {
568 wire: Some(Box::new(super::super::usage::WireContext {
569 counterfactual: Some(super::super::counterfactual::CounterfactualSlot::new()),
570 ..Default::default()
571 })),
572 ..usage("claude-sonnet-4.5", 1_000, 0, 0)
573 };
574 acc.add(&empty);
575 assert_eq!(acc.requests, 2);
576 assert_eq!(acc.counterfactual_requests, 1, "empty slot adds no pair");
577
578 acc.add(&usage("claude-sonnet-4.5", 10, 0, 0));
580 assert_eq!(acc.counterfactual_requests, 1);
581 }
582
583 #[test]
585 fn verified_of_aggregates_and_reports_signed_savings() {
586 let mut map = HashMap::new();
587 assert!(verified_of(&map).is_none(), "no coverage → None, not zeros");
588
589 map.insert(
590 "claude-sonnet-4.5".to_string(),
591 ModelUsage {
592 counterfactual_requests: 2,
593 counterfactual_input_tokens: 10_000,
594 counterfactual_billed_tokens: 6_000,
595 ..Default::default()
596 },
597 );
598 map.insert(
601 "claude-haiku-4.5".to_string(),
602 ModelUsage {
603 counterfactual_requests: 1,
604 counterfactual_input_tokens: 1_000,
605 counterfactual_billed_tokens: 1_400,
606 ..Default::default()
607 },
608 );
609 let v = verified_of(&map).expect("covered rows present");
610 assert_eq!(v.requests, 3);
611 assert_eq!(v.counterfactual_input_tokens, 11_000);
612 assert_eq!(v.billed_input_tokens, 7_400);
613 assert_eq!(v.verified_saved_tokens, 3_600);
614
615 map.get_mut("claude-sonnet-4.5")
616 .unwrap()
617 .counterfactual_input_tokens = 0;
618 map.get_mut("claude-sonnet-4.5")
619 .unwrap()
620 .counterfactual_billed_tokens = 0;
621 let negative = verified_of(&map).unwrap();
622 assert_eq!(
623 negative.verified_saved_tokens, -400,
624 "a net-negative verified saving is reported, never clamped"
625 );
626 }
627
628 #[test]
631 fn counterfactual_fields_roundtrip_and_default_for_legacy_files() {
632 let legacy = r#"{"ts":1,"models":{"m":{"requests":1,"input_tokens":10,"output_tokens":5,"cache_read_tokens":0,"cache_write_tokens":0,"reasoning_tokens":0}}}"#;
633 let p: PersistedUsage = serde_json::from_str(legacy).expect("legacy file loads");
634 assert_eq!(p.models["m"].counterfactual_requests, 0);
635
636 let mut p = PersistedUsage::default();
637 p.models.insert(
638 "m".into(),
639 ModelUsage {
640 counterfactual_requests: 4,
641 counterfactual_input_tokens: 9_999,
642 counterfactual_billed_tokens: 5_555,
643 ..Default::default()
644 },
645 );
646 let back: PersistedUsage =
647 serde_json::from_str(&serde_json::to_string(&p).unwrap()).unwrap();
648 assert_eq!(back.models["m"].counterfactual_input_tokens, 9_999);
649 assert_eq!(back.models["m"].counterfactual_billed_tokens, 5_555);
650 }
651
652 #[test]
653 fn prices_known_model_with_cache_split() {
654 let mut map = HashMap::new();
655 let mut acc = ModelUsage::default();
656 acc.add(&usage("claude-sonnet-4.5", 1_000_000, 1_000_000, 1_000_000));
657 map.insert("claude-sonnet-4.5".to_string(), acc);
658
659 let rows = price_models(&map);
660 assert_eq!(rows.len(), 1);
661 let row = &rows[0];
662 assert!(
664 (row.cost_usd - 18.30).abs() < 1e-6,
665 "cost was {}",
666 row.cost_usd
667 );
668 assert!(!row.pricing_estimated, "exact model match");
669 assert_eq!(row.requests, 1);
670 }
671
672 #[test]
673 fn unknown_model_prices_with_fallback_and_is_estimated() {
674 let mut map = HashMap::new();
675 let mut acc = ModelUsage::default();
676 acc.add(&usage("some-novel-model-xyz", 1_000_000, 0, 0));
677 map.insert("some-novel-model-xyz".to_string(), acc);
678
679 let rows = price_models(&map);
680 assert!(rows[0].pricing_estimated, "fallback pricing is estimated");
681 assert!(rows[0].cost_usd > 0.0);
682 }
683
684 #[test]
689 fn measured_provider_cost_replaces_table_estimate() {
690 const MODEL: &str = "vendor/unlisted-model-20990101";
691 let mut acc = ModelUsage::default();
692 let mut measured = usage(MODEL, 431_600, 22_700, 126_700);
693 measured.provider_cost_usd = Some(0.05);
694 acc.add(&measured);
695
696 let mut map = HashMap::new();
697 map.insert(MODEL.to_string(), acc.clone());
698 let row = &price_models(&map)[0];
699 assert!(
700 (row.cost_usd - 0.05).abs() < 1e-12,
701 "the bill, not the table"
702 );
703 assert!((row.measured_cost_usd - 0.05).abs() < 1e-12);
704 assert_eq!(row.measured_requests, 1);
705 assert!(
706 !row.pricing_estimated,
707 "fully measured spend is not an estimate"
708 );
709
710 acc.add(&usage(MODEL, 10_000, 0, 0));
714 let mut map = HashMap::new();
715 map.insert(MODEL.to_string(), acc);
716 let row = &price_models(&map)[0];
717 assert!(row.cost_usd > 0.05, "unmeasured remainder adds table cost");
718 assert!(row.cost_usd < 0.2, "measured slice must not be re-priced");
719 assert!(
720 row.pricing_estimated,
721 "unmeasured remainder is heuristically priced"
722 );
723 }
724
725 #[test]
726 fn measured_slice_roundtrips_and_defaults_for_legacy_files() {
727 let legacy = r#"{"ts":1,"models":{"m":{"requests":1,"input_tokens":10,"output_tokens":5,"cache_read_tokens":0,"cache_write_tokens":0,"reasoning_tokens":0}}}"#;
728 let p: PersistedUsage = serde_json::from_str(legacy).expect("legacy file loads");
729 assert_eq!(p.models["m"].measured, MeasuredSlice::default());
730
731 let mut p = PersistedUsage::default();
732 let mut acc = ModelUsage::default();
733 let mut m = usage("m", 100, 10, 0);
734 m.provider_cost_usd = Some(0.5);
735 acc.add(&m);
736 p.models.insert("m".into(), acc);
737 let back: PersistedUsage =
738 serde_json::from_str(&serde_json::to_string(&p).unwrap()).unwrap();
739 assert_eq!(back.models["m"].measured.requests, 1);
740 assert!((back.models["m"].measured.cost_usd - 0.5).abs() < 1e-12);
741 }
742
743 #[test]
744 fn dominant_model_picks_highest_token_real_model() {
745 let mut models = HashMap::new();
746 models.insert("claude-haiku-4.5".to_string(), {
747 let mut u = ModelUsage::default();
748 u.add(&usage("claude-haiku-4.5", 100, 100, 0));
749 u
750 });
751 models.insert("claude-opus-4.5".to_string(), {
752 let mut u = ModelUsage::default();
753 u.add(&usage("claude-opus-4.5", 10_000, 10_000, 0));
754 u
755 });
756 models.insert("unknown".to_string(), {
757 let mut u = ModelUsage::default();
758 u.add(&usage("unknown", 999_999, 0, 0));
759 u
760 });
761 let dominant = models
762 .iter()
763 .filter(|(m, _)| m.as_str() != "unknown")
764 .max_by_key(|(_, u)| u.billable_tokens())
765 .map(|(m, _)| m.clone());
766 assert_eq!(dominant.as_deref(), Some("claude-opus-4.5"));
767 }
768
769 #[test]
770 fn empty_model_buckets_as_unknown() {
771 assert_eq!(normalize_key(" "), "unknown");
772 assert_eq!(normalize_key(""), "unknown");
773 assert_eq!(normalize_key("gpt-5.4"), "gpt-5.4");
774 }
775
776 #[test]
777 fn cohort_avg_output_is_mean_per_turn() {
778 let mut c = CohortUsage::default();
779 assert_eq!(c.avg_output(), None, "no observations → None");
780 c.add(&usage("m", 10, 100, 0));
781 c.add(&usage("m", 10, 50, 0));
782 assert_eq!(c.requests, 2);
783 assert_eq!(c.output_tokens, 150);
784 assert!((c.avg_output().unwrap() - 75.0).abs() < f64::EPSILON);
785 }
786
787 #[test]
788 fn persisted_usage_without_cohorts_field_loads() {
789 let json = r#"{"ts":1,"models":{"gpt-5.4":{"requests":1,"input_tokens":10,"output_tokens":5,"cache_read_tokens":0,"cache_write_tokens":0,"reasoning_tokens":0}}}"#;
792 let p: PersistedUsage = serde_json::from_str(json).expect("loads legacy file");
793 assert_eq!(p.models.len(), 1);
794 assert!(p.cohorts.is_empty());
795 }
796
797 #[test]
798 fn persisted_usage_roundtrips_cohorts() {
799 let mut p = PersistedUsage::default();
800 p.cohorts.insert(
801 "control".into(),
802 CohortUsage {
803 requests: 3,
804 input_tokens: 30,
805 output_tokens: 300,
806 sum_sq_output: 30_000,
807 },
808 );
809 let json = serde_json::to_string(&p).unwrap();
810 let back: PersistedUsage = serde_json::from_str(&json).unwrap();
811 assert_eq!(back.cohorts.get("control").unwrap().output_tokens, 300);
812 }
813}