1use std::collections::HashMap;
15use std::sync::{Mutex, OnceLock};
16
17use serde::{Deserialize, Serialize};
18
19use crate::core::gain::model_pricing::ModelPricing;
20use crate::core::ocla::{OclaRegistry, UsageSink};
21
22fn project_ocla_usage(u: &super::usage::RealUsage, sink: &dyn UsageSink) {
28 use futures::FutureExt;
29 let Ok(Some(record)) = u.to_ocla_usage_record() else {
30 return;
31 };
32 let _ = sink.record_usage(record).now_or_never();
33}
34
35#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
38pub struct ModelUsage {
39 pub requests: u64,
40 pub input_tokens: u64,
41 pub output_tokens: u64,
42 pub cache_read_tokens: u64,
43 pub cache_write_tokens: u64,
44 pub reasoning_tokens: u64,
45 #[serde(default)]
49 pub counterfactual_requests: u64,
50 #[serde(default)]
53 pub counterfactual_input_tokens: u64,
54 #[serde(default)]
57 pub counterfactual_billed_tokens: u64,
58 #[serde(default)]
63 pub measured: MeasuredSlice,
64}
65
66#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq)]
68pub struct MeasuredSlice {
69 pub requests: u64,
70 pub input_tokens: u64,
71 pub output_tokens: u64,
72 pub cache_read_tokens: u64,
73 pub cache_write_tokens: u64,
74 pub cost_usd: f64,
76}
77
78impl MeasuredSlice {
79 fn merge(&mut self, other: &Self) {
80 self.requests += other.requests;
81 self.input_tokens += other.input_tokens;
82 self.output_tokens += other.output_tokens;
83 self.cache_read_tokens += other.cache_read_tokens;
84 self.cache_write_tokens += other.cache_write_tokens;
85 self.cost_usd += other.cost_usd;
86 }
87}
88
89impl ModelUsage {
90 fn add(&mut self, u: &super::usage::RealUsage) {
91 self.requests += 1;
92 self.input_tokens += u.input_tokens;
93 self.output_tokens += u.output_tokens;
94 self.cache_read_tokens += u.cache_read_tokens;
95 self.cache_write_tokens += u.cache_write_tokens;
96 self.reasoning_tokens += u.reasoning_tokens;
97 if let Some(cost) = u.provider_cost_usd {
98 self.measured.merge(&MeasuredSlice {
99 requests: 1,
100 input_tokens: u.input_tokens,
101 output_tokens: u.output_tokens,
102 cache_read_tokens: u.cache_read_tokens,
103 cache_write_tokens: u.cache_write_tokens,
104 cost_usd: cost,
105 });
106 }
107 if let Some(counted) = u
110 .wire
111 .as_deref()
112 .and_then(|w| w.counterfactual.as_ref())
113 .and_then(super::counterfactual::CounterfactualSlot::get)
114 {
115 self.counterfactual_requests += 1;
116 self.counterfactual_input_tokens += counted;
117 self.counterfactual_billed_tokens +=
118 u.input_tokens + u.cache_read_tokens + u.cache_write_tokens;
119 }
120 }
121
122 fn billable_tokens(&self) -> u64 {
123 self.input_tokens + self.output_tokens + self.cache_read_tokens + self.cache_write_tokens
124 }
125}
126
127#[derive(Debug, Clone, Serialize, PartialEq)]
129pub struct ModelSpend {
130 pub model: String,
131 pub requests: u64,
132 pub input_tokens: u64,
133 pub output_tokens: u64,
134 pub cache_read_tokens: u64,
135 pub cache_write_tokens: u64,
136 pub reasoning_tokens: u64,
137 pub cost_usd: f64,
138 pub measured_cost_usd: f64,
141 pub measured_requests: u64,
143 pub pricing_estimated: bool,
146}
147
148#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
152pub struct CohortUsage {
153 pub requests: u64,
154 pub input_tokens: u64,
155 pub output_tokens: u64,
156 #[serde(default)]
160 pub sum_sq_output: u64,
161}
162
163impl CohortUsage {
164 fn add(&mut self, u: &super::usage::RealUsage) {
165 self.requests += 1;
166 self.input_tokens += u.input_tokens;
167 self.output_tokens += u.output_tokens;
168 self.sum_sq_output += u.output_tokens.saturating_mul(u.output_tokens);
169 }
170
171 #[must_use]
173 pub fn avg_output(&self) -> Option<f64> {
174 if self.requests == 0 {
175 None
176 } else {
177 #[allow(clippy::cast_precision_loss)]
178 Some(self.output_tokens as f64 / self.requests as f64)
179 }
180 }
181
182 #[must_use]
186 pub fn variance_output(&self) -> Option<f64> {
187 if self.requests < 2 {
188 return None;
189 }
190 #[allow(clippy::cast_precision_loss)]
191 let n = self.requests as f64;
192 #[allow(clippy::cast_precision_loss)]
193 let sum = self.output_tokens as f64;
194 #[allow(clippy::cast_precision_loss)]
195 let sum_sq = self.sum_sq_output as f64;
196 let var = (sum_sq - sum * sum / n) / (n - 1.0);
197 Some(var.max(0.0))
198 }
199}
200
201#[derive(Debug, Clone, Default, Serialize, Deserialize)]
203pub struct PersistedUsage {
204 pub ts: u64,
205 pub models: HashMap<String, ModelUsage>,
206 #[serde(default)]
209 pub cohorts: HashMap<String, CohortUsage>,
210}
211
212const MAX_TRACKED_MODELS: usize = 256;
215
216const PROXY_USAGE_FILE: &str = "proxy_usage.json";
217
218fn store() -> &'static Mutex<HashMap<String, ModelUsage>> {
219 static STORE: OnceLock<Mutex<HashMap<String, ModelUsage>>> = OnceLock::new();
220 STORE.get_or_init(|| Mutex::new(HashMap::new()))
221}
222
223fn cohort_store() -> &'static Mutex<HashMap<String, CohortUsage>> {
224 static STORE: OnceLock<Mutex<HashMap<String, CohortUsage>>> = OnceLock::new();
225 STORE.get_or_init(|| Mutex::new(HashMap::new()))
226}
227
228pub fn resume_from_disk() {
232 let Some(persisted) = load_persisted() else {
233 return;
234 };
235 let mut map = store()
236 .lock()
237 .unwrap_or_else(std::sync::PoisonError::into_inner);
238 for (model, usage) in persisted.models {
239 let acc = map.entry(model).or_default();
240 acc.requests += usage.requests;
241 acc.input_tokens += usage.input_tokens;
242 acc.output_tokens += usage.output_tokens;
243 acc.cache_read_tokens += usage.cache_read_tokens;
244 acc.cache_write_tokens += usage.cache_write_tokens;
245 acc.reasoning_tokens += usage.reasoning_tokens;
246 acc.counterfactual_requests += usage.counterfactual_requests;
247 acc.counterfactual_input_tokens += usage.counterfactual_input_tokens;
248 acc.counterfactual_billed_tokens += usage.counterfactual_billed_tokens;
249 acc.measured.merge(&usage.measured);
250 }
251 drop(map);
252 let mut cohorts = cohort_store()
253 .lock()
254 .unwrap_or_else(std::sync::PoisonError::into_inner);
255 for (arm, usage) in persisted.cohorts {
256 let acc = cohorts.entry(arm).or_default();
257 acc.requests += usage.requests;
258 acc.input_tokens += usage.input_tokens;
259 acc.output_tokens += usage.output_tokens;
260 }
261}
262
263pub fn record(u: &super::usage::RealUsage) {
266 super::usage_sink::push(u);
269
270 {
272 let label = u.wire.as_ref().map_or("Unknown", |w| {
273 match w.provider.as_str() {
276 "Anthropic" => "Anthropic",
277 "OpenAI" | "ChatGPT" => "OpenAI",
278 "Gemini" => "Gemini",
279 "Bedrock" => "Bedrock",
280 "Azure" => "Azure",
281 _ => "Unknown",
282 }
283 });
284 let envelope = super::usage_parity::real_usage_to_envelope(u, label);
285 crate::core::context_kernel::envelope_bridge::record_proxy_envelope(&envelope);
286 }
287
288 project_ocla_usage(u, OclaRegistry::global().usage_sink.as_ref());
292
293 if let Some(wire) = u.wire.as_deref()
299 && (wire.person.is_some() || wire.project.is_some())
300 {
301 let baseline = crate::core::config::Config::load().proxy.baseline.clone();
302 #[allow(clippy::cast_precision_loss)]
303 let cost_usd = if wire.is_local {
304 let billable =
305 u.input_tokens + u.output_tokens + u.cache_read_tokens + u.cache_write_tokens;
306 baseline.effective_local_shadow_rate() / 1_000_000.0 * billable as f64
307 } else if let Some(measured) = u.provider_cost_usd {
308 measured
309 } else {
310 crate::core::gain::model_pricing::ModelPricing::load()
311 .quote(Some(&u.model))
312 .cost
313 .estimate_usd(
314 u.input_tokens,
315 u.output_tokens,
316 u.cache_write_tokens,
317 u.cache_read_tokens,
318 )
319 };
320 super::policy_gate::record_spend(wire.person.as_deref(), wire.project.as_deref(), cost_usd);
321 }
322
323 if let Some(wire) = u.wire.as_deref()
328 && let Some(routed_from) = wire.routed_from.as_deref()
329 {
330 crate::core::savings_ledger::record_routing_event(routed_from, &u.model, u.input_tokens);
331 }
332 if u.cache_read_tokens > 0 {
333 let cost = crate::core::gain::model_pricing::ModelPricing::load()
334 .quote(Some(&u.model))
335 .cost;
336 #[allow(clippy::cast_precision_loss)]
337 let discount_usd =
338 (cost.input_per_m - cost.cache_read_per_m) / 1_000_000.0 * u.cache_read_tokens as f64;
339 crate::core::savings_ledger::record_caching_event(
340 &u.model,
341 u.cache_read_tokens,
342 discount_usd,
343 );
344 }
345
346 let key = normalize_key(&u.model);
347 {
348 let mut map = store()
349 .lock()
350 .unwrap_or_else(std::sync::PoisonError::into_inner);
351 let key = if !map.contains_key(&key) && map.len() >= MAX_TRACKED_MODELS {
352 "unknown".to_string()
353 } else {
354 key
355 };
356 map.entry(key).or_default().add(u);
357 }
358 if let Some(arm) = u.cohort {
359 let mut cohorts = cohort_store()
360 .lock()
361 .unwrap_or_else(std::sync::PoisonError::into_inner);
362 cohorts.entry(arm.as_str().to_string()).or_default().add(u);
363 }
364 persist();
365}
366
367#[must_use]
369pub fn cohort_snapshot() -> HashMap<String, CohortUsage> {
370 cohort_store()
371 .lock()
372 .unwrap_or_else(std::sync::PoisonError::into_inner)
373 .clone()
374}
375
376#[must_use]
378pub fn persisted_cohorts() -> HashMap<String, CohortUsage> {
379 load_persisted().map(|p| p.cohorts).unwrap_or_default()
380}
381
382fn normalize_key(model: &str) -> String {
383 let m = model.trim();
384 if m.is_empty() {
385 "unknown".to_string()
386 } else {
387 m.to_string()
388 }
389}
390
391pub fn snapshot() -> Vec<ModelSpend> {
393 let map = store()
394 .lock()
395 .unwrap_or_else(std::sync::PoisonError::into_inner);
396 price_models(&map)
397}
398
399pub fn total_cost_usd() -> f64 {
401 snapshot().iter().map(|m| m.cost_usd).sum()
402}
403
404#[derive(Debug, Clone, Serialize, PartialEq)]
406pub struct VerifiedSavings {
407 pub requests: u64,
409 pub counterfactual_input_tokens: u64,
412 pub billed_input_tokens: u64,
415 pub verified_saved_tokens: i64,
418}
419
420pub fn verified_savings() -> Option<VerifiedSavings> {
425 let map = store()
426 .lock()
427 .unwrap_or_else(std::sync::PoisonError::into_inner);
428 verified_of(&map)
429}
430
431#[allow(clippy::cast_possible_wrap)]
433fn verified_of(map: &HashMap<String, ModelUsage>) -> Option<VerifiedSavings> {
434 let (mut requests, mut counterfactual, mut billed) = (0u64, 0u64, 0u64);
435 for usage in map.values() {
436 requests += usage.counterfactual_requests;
437 counterfactual += usage.counterfactual_input_tokens;
438 billed += usage.counterfactual_billed_tokens;
439 }
440 (requests > 0).then(|| VerifiedSavings {
441 requests,
442 counterfactual_input_tokens: counterfactual,
443 billed_input_tokens: billed,
444 verified_saved_tokens: counterfactual as i64 - billed as i64,
445 })
446}
447
448pub fn price_models(map: &HashMap<String, ModelUsage>) -> Vec<ModelSpend> {
451 let pricing = ModelPricing::load();
452 let mut rows: Vec<ModelSpend> = map
453 .iter()
454 .map(|(model, usage)| price_one(&pricing, model, usage))
455 .collect();
456 rows.sort_by(|a, b| {
457 b.cost_usd
458 .partial_cmp(&a.cost_usd)
459 .unwrap_or(std::cmp::Ordering::Equal)
460 });
461 rows
462}
463
464fn price_one(pricing: &ModelPricing, model: &str, usage: &ModelUsage) -> ModelSpend {
465 let quote = pricing.quote(Some(model));
466 let m = &usage.measured;
469 let derived_cost = quote.cost.estimate_usd(
470 usage.input_tokens.saturating_sub(m.input_tokens),
471 usage.output_tokens.saturating_sub(m.output_tokens),
472 usage
473 .cache_write_tokens
474 .saturating_sub(m.cache_write_tokens),
475 usage.cache_read_tokens.saturating_sub(m.cache_read_tokens),
476 );
477 let derived_requests = usage.requests.saturating_sub(m.requests);
478 ModelSpend {
479 model: model.to_string(),
480 requests: usage.requests,
481 input_tokens: usage.input_tokens,
482 output_tokens: usage.output_tokens,
483 cache_read_tokens: usage.cache_read_tokens,
484 cache_write_tokens: usage.cache_write_tokens,
485 reasoning_tokens: usage.reasoning_tokens,
486 cost_usd: m.cost_usd + derived_cost,
487 measured_cost_usd: m.cost_usd,
488 measured_requests: m.requests,
489 pricing_estimated: quote.match_kind.is_estimated() && derived_requests > 0,
491 }
492}
493
494fn usage_path() -> Option<std::path::PathBuf> {
495 crate::core::data_dir::lean_ctx_data_dir()
496 .ok()
497 .map(|d| d.join(PROXY_USAGE_FILE))
498}
499
500fn persist() {
502 let Some(path) = usage_path() else {
503 return;
504 };
505 let models = {
506 let map = store()
507 .lock()
508 .unwrap_or_else(std::sync::PoisonError::into_inner);
509 map.clone()
510 };
511 let cohorts = {
512 let map = cohort_store()
513 .lock()
514 .unwrap_or_else(std::sync::PoisonError::into_inner);
515 map.clone()
516 };
517 let payload = PersistedUsage {
518 ts: std::time::SystemTime::now()
519 .duration_since(std::time::UNIX_EPOCH)
520 .unwrap_or_default()
521 .as_secs(),
522 models,
523 cohorts,
524 };
525 let Ok(json) = serde_json::to_string(&payload) else {
526 return;
527 };
528 let tmp = path.with_extension("json.tmp");
529 if std::fs::write(&tmp, json).is_ok() {
530 let _ = std::fs::rename(&tmp, &path);
531 }
532}
533
534pub fn load_persisted() -> Option<PersistedUsage> {
536 let path = usage_path()?;
537 let data = std::fs::read_to_string(path).ok()?;
538 serde_json::from_str(&data).ok()
539}
540
541pub fn persisted_snapshot() -> Vec<ModelSpend> {
543 load_persisted()
544 .map(|p| price_models(&p.models))
545 .unwrap_or_default()
546}
547
548pub fn persisted_dominant_model() -> Option<String> {
552 let persisted = load_persisted()?;
553 persisted
554 .models
555 .iter()
556 .filter(|(m, _)| m.as_str() != "unknown" && !m.trim().is_empty())
557 .max_by_key(|(_, u)| u.billable_tokens())
558 .filter(|(_, u)| u.billable_tokens() > 0)
559 .map(|(m, _)| m.clone())
560}
561
562#[cfg(test)]
563mod tests {
564 use super::*;
565
566 fn usage(
567 model: &str,
568 input: u64,
569 output: u64,
570 cache_read: u64,
571 ) -> super::super::usage::RealUsage {
572 super::super::usage::RealUsage {
573 model: model.to_string(),
574 input_tokens: input,
575 output_tokens: output,
576 cache_read_tokens: cache_read,
577 ..Default::default()
578 }
579 }
580
581 #[test]
584 fn counterfactual_pair_recorded_only_when_probe_answered() {
585 let slot = super::super::counterfactual::CounterfactualSlot::new();
586 slot.set(5_000);
587 let with_probe = super::super::usage::RealUsage {
588 wire: Some(Box::new(super::super::usage::WireContext {
589 counterfactual: Some(slot),
590 ..Default::default()
591 })),
592 cache_write_tokens: 200,
593 ..usage("claude-sonnet-4.5", 1_000, 0, 300)
594 };
595 let mut acc = ModelUsage::default();
596 acc.add(&with_probe);
597 assert_eq!(acc.counterfactual_requests, 1);
598 assert_eq!(acc.counterfactual_input_tokens, 5_000);
599 assert_eq!(acc.counterfactual_billed_tokens, 1_000 + 300 + 200);
601
602 let empty = super::super::usage::RealUsage {
605 wire: Some(Box::new(super::super::usage::WireContext {
606 counterfactual: Some(super::super::counterfactual::CounterfactualSlot::new()),
607 ..Default::default()
608 })),
609 ..usage("claude-sonnet-4.5", 1_000, 0, 0)
610 };
611 acc.add(&empty);
612 assert_eq!(acc.requests, 2);
613 assert_eq!(acc.counterfactual_requests, 1, "empty slot adds no pair");
614
615 acc.add(&usage("claude-sonnet-4.5", 10, 0, 0));
617 assert_eq!(acc.counterfactual_requests, 1);
618 }
619
620 #[test]
622 fn verified_of_aggregates_and_reports_signed_savings() {
623 let mut map = HashMap::new();
624 assert!(verified_of(&map).is_none(), "no coverage → None, not zeros");
625
626 map.insert(
627 "claude-sonnet-4.5".to_string(),
628 ModelUsage {
629 counterfactual_requests: 2,
630 counterfactual_input_tokens: 10_000,
631 counterfactual_billed_tokens: 6_000,
632 ..Default::default()
633 },
634 );
635 map.insert(
638 "claude-haiku-4.5".to_string(),
639 ModelUsage {
640 counterfactual_requests: 1,
641 counterfactual_input_tokens: 1_000,
642 counterfactual_billed_tokens: 1_400,
643 ..Default::default()
644 },
645 );
646 let v = verified_of(&map).expect("covered rows present");
647 assert_eq!(v.requests, 3);
648 assert_eq!(v.counterfactual_input_tokens, 11_000);
649 assert_eq!(v.billed_input_tokens, 7_400);
650 assert_eq!(v.verified_saved_tokens, 3_600);
651
652 map.get_mut("claude-sonnet-4.5")
653 .unwrap()
654 .counterfactual_input_tokens = 0;
655 map.get_mut("claude-sonnet-4.5")
656 .unwrap()
657 .counterfactual_billed_tokens = 0;
658 let negative = verified_of(&map).unwrap();
659 assert_eq!(
660 negative.verified_saved_tokens, -400,
661 "a net-negative verified saving is reported, never clamped"
662 );
663 }
664
665 #[test]
668 fn counterfactual_fields_roundtrip_and_default_for_legacy_files() {
669 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}}}"#;
670 let p: PersistedUsage = serde_json::from_str(legacy).expect("legacy file loads");
671 assert_eq!(p.models["m"].counterfactual_requests, 0);
672
673 let mut p = PersistedUsage::default();
674 p.models.insert(
675 "m".into(),
676 ModelUsage {
677 counterfactual_requests: 4,
678 counterfactual_input_tokens: 9_999,
679 counterfactual_billed_tokens: 5_555,
680 ..Default::default()
681 },
682 );
683 let back: PersistedUsage =
684 serde_json::from_str(&serde_json::to_string(&p).unwrap()).unwrap();
685 assert_eq!(back.models["m"].counterfactual_input_tokens, 9_999);
686 assert_eq!(back.models["m"].counterfactual_billed_tokens, 5_555);
687 }
688
689 #[test]
690 fn prices_known_model_with_cache_split() {
691 let mut map = HashMap::new();
692 let mut acc = ModelUsage::default();
693 acc.add(&usage("claude-sonnet-4.5", 1_000_000, 1_000_000, 1_000_000));
694 map.insert("claude-sonnet-4.5".to_string(), acc);
695
696 let rows = price_models(&map);
697 assert_eq!(rows.len(), 1);
698 let row = &rows[0];
699 assert!(
701 (row.cost_usd - 18.30).abs() < 1e-6,
702 "cost was {}",
703 row.cost_usd
704 );
705 assert!(!row.pricing_estimated, "exact model match");
706 assert_eq!(row.requests, 1);
707 }
708
709 #[test]
710 fn unknown_model_prices_with_fallback_and_is_estimated() {
711 let mut map = HashMap::new();
712 let mut acc = ModelUsage::default();
713 acc.add(&usage("some-novel-model-xyz", 1_000_000, 0, 0));
714 map.insert("some-novel-model-xyz".to_string(), acc);
715
716 let rows = price_models(&map);
717 assert!(rows[0].pricing_estimated, "fallback pricing is estimated");
718 assert!(rows[0].cost_usd > 0.0);
719 }
720
721 #[test]
726 fn measured_provider_cost_replaces_table_estimate() {
727 const MODEL: &str = "vendor/unlisted-model-20990101";
728 let mut acc = ModelUsage::default();
729 let mut measured = usage(MODEL, 431_600, 22_700, 126_700);
730 measured.provider_cost_usd = Some(0.05);
731 acc.add(&measured);
732
733 let mut map = HashMap::new();
734 map.insert(MODEL.to_string(), acc.clone());
735 let row = &price_models(&map)[0];
736 assert!(
737 (row.cost_usd - 0.05).abs() < 1e-12,
738 "the bill, not the table"
739 );
740 assert!((row.measured_cost_usd - 0.05).abs() < 1e-12);
741 assert_eq!(row.measured_requests, 1);
742 assert!(
743 !row.pricing_estimated,
744 "fully measured spend is not an estimate"
745 );
746
747 acc.add(&usage(MODEL, 10_000, 0, 0));
751 let mut map = HashMap::new();
752 map.insert(MODEL.to_string(), acc);
753 let row = &price_models(&map)[0];
754 assert!(row.cost_usd > 0.05, "unmeasured remainder adds table cost");
755 assert!(row.cost_usd < 0.2, "measured slice must not be re-priced");
756 assert!(
757 row.pricing_estimated,
758 "unmeasured remainder is heuristically priced"
759 );
760 }
761
762 #[test]
763 fn measured_slice_roundtrips_and_defaults_for_legacy_files() {
764 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}}}"#;
765 let p: PersistedUsage = serde_json::from_str(legacy).expect("legacy file loads");
766 assert_eq!(p.models["m"].measured, MeasuredSlice::default());
767
768 let mut p = PersistedUsage::default();
769 let mut acc = ModelUsage::default();
770 let mut m = usage("m", 100, 10, 0);
771 m.provider_cost_usd = Some(0.5);
772 acc.add(&m);
773 p.models.insert("m".into(), acc);
774 let back: PersistedUsage =
775 serde_json::from_str(&serde_json::to_string(&p).unwrap()).unwrap();
776 assert_eq!(back.models["m"].measured.requests, 1);
777 assert!((back.models["m"].measured.cost_usd - 0.5).abs() < 1e-12);
778 }
779
780 #[test]
781 fn dominant_model_picks_highest_token_real_model() {
782 let mut models = HashMap::new();
783 models.insert("claude-haiku-4.5".to_string(), {
784 let mut u = ModelUsage::default();
785 u.add(&usage("claude-haiku-4.5", 100, 100, 0));
786 u
787 });
788 models.insert("claude-opus-4.5".to_string(), {
789 let mut u = ModelUsage::default();
790 u.add(&usage("claude-opus-4.5", 10_000, 10_000, 0));
791 u
792 });
793 models.insert("unknown".to_string(), {
794 let mut u = ModelUsage::default();
795 u.add(&usage("unknown", 999_999, 0, 0));
796 u
797 });
798 let dominant = models
799 .iter()
800 .filter(|(m, _)| m.as_str() != "unknown")
801 .max_by_key(|(_, u)| u.billable_tokens())
802 .map(|(m, _)| m.clone());
803 assert_eq!(dominant.as_deref(), Some("claude-opus-4.5"));
804 }
805
806 #[test]
807 fn empty_model_buckets_as_unknown() {
808 assert_eq!(normalize_key(" "), "unknown");
809 assert_eq!(normalize_key(""), "unknown");
810 assert_eq!(normalize_key("gpt-5.4"), "gpt-5.4");
811 }
812
813 #[test]
814 fn cohort_avg_output_is_mean_per_turn() {
815 let mut c = CohortUsage::default();
816 assert_eq!(c.avg_output(), None, "no observations → None");
817 c.add(&usage("m", 10, 100, 0));
818 c.add(&usage("m", 10, 50, 0));
819 assert_eq!(c.requests, 2);
820 assert_eq!(c.output_tokens, 150);
821 assert!((c.avg_output().unwrap() - 75.0).abs() < f64::EPSILON);
822 }
823
824 #[test]
825 fn persisted_usage_without_cohorts_field_loads() {
826 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}}}"#;
829 let p: PersistedUsage = serde_json::from_str(json).expect("loads legacy file");
830 assert_eq!(p.models.len(), 1);
831 assert!(p.cohorts.is_empty());
832 }
833
834 #[test]
835 fn persisted_usage_roundtrips_cohorts() {
836 let mut p = PersistedUsage::default();
837 p.cohorts.insert(
838 "control".into(),
839 CohortUsage {
840 requests: 3,
841 input_tokens: 30,
842 output_tokens: 300,
843 sum_sq_output: 30_000,
844 },
845 );
846 let json = serde_json::to_string(&p).unwrap();
847 let back: PersistedUsage = serde_json::from_str(&json).unwrap();
848 assert_eq!(back.cohorts.get("control").unwrap().output_tokens, 300);
849 }
850
851 #[test]
852 fn managed_usage_projects_once_and_unmanaged_is_skipped() {
853 let sink = crate::core::ocla::builtin::usage_sink::BuiltinUsageSink::new();
854 let usage = super::super::usage::RealUsage {
855 model: "gpt-5".into(),
856 input_tokens: 100,
857 output_tokens: 40,
858 cache_read_tokens: 20,
859 cache_write_tokens: 5,
860 wire: Some(Box::new(super::super::usage::WireContext {
861 lineage: Some(crate::core::ocla::OclaRequestContext {
862 request_id: "request-1".into(),
863 session_id: "session-1".into(),
864 agent_id: "agent-1".into(),
865 content_ref: "blake3:content".into(),
866 tenant_id: None,
867 trace_id: "tr-unit".into(),
868 }),
869 ..Default::default()
870 })),
871 ..Default::default()
872 };
873
874 project_ocla_usage(&usage, &sink);
875 assert_eq!(sink.record_count(), 1);
876 assert_eq!(sink.total_input_tokens(), 125);
877 assert_eq!(sink.total_output_tokens(), 40);
878
879 project_ocla_usage(
880 &super::super::usage::RealUsage {
881 model: "unmanaged".into(),
882 input_tokens: 999,
883 ..Default::default()
884 },
885 &sink,
886 );
887 assert_eq!(sink.record_count(), 1);
888 }
889}