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 #[cfg(feature = "enterprise")]
299 if let Some(wire) = u.wire.as_deref()
300 && (wire.person.is_some() || wire.project.is_some())
301 {
302 let baseline = crate::core::config::Config::load().proxy.baseline.clone();
303 #[allow(clippy::cast_precision_loss)]
304 let cost_usd = if wire.is_local {
305 let billable =
306 u.input_tokens + u.output_tokens + u.cache_read_tokens + u.cache_write_tokens;
307 baseline.effective_local_shadow_rate() / 1_000_000.0 * billable as f64
308 } else if let Some(measured) = u.provider_cost_usd {
309 measured
310 } else {
311 crate::core::gain::model_pricing::ModelPricing::load()
312 .quote(Some(&u.model))
313 .cost
314 .estimate_usd(
315 u.input_tokens,
316 u.output_tokens,
317 u.cache_write_tokens,
318 u.cache_read_tokens,
319 )
320 };
321 super::policy_gate::record_spend(wire.person.as_deref(), wire.project.as_deref(), cost_usd);
322 }
323
324 if let Some(wire) = u.wire.as_deref()
329 && let Some(routed_from) = wire.routed_from.as_deref()
330 {
331 crate::core::savings_ledger::record_routing_event(routed_from, &u.model, u.input_tokens);
332 }
333 if u.cache_read_tokens > 0 {
334 let cost = crate::core::gain::model_pricing::ModelPricing::load()
335 .quote(Some(&u.model))
336 .cost;
337 #[allow(clippy::cast_precision_loss)]
338 let discount_usd =
339 (cost.input_per_m - cost.cache_read_per_m) / 1_000_000.0 * u.cache_read_tokens as f64;
340 crate::core::savings_ledger::record_caching_event(
341 &u.model,
342 u.cache_read_tokens,
343 discount_usd,
344 );
345 }
346
347 let key = normalize_key(&u.model);
348 {
349 let mut map = store()
350 .lock()
351 .unwrap_or_else(std::sync::PoisonError::into_inner);
352 let key = if !map.contains_key(&key) && map.len() >= MAX_TRACKED_MODELS {
353 "unknown".to_string()
354 } else {
355 key
356 };
357 map.entry(key).or_default().add(u);
358 }
359 if let Some(arm) = u.cohort {
360 let mut cohorts = cohort_store()
361 .lock()
362 .unwrap_or_else(std::sync::PoisonError::into_inner);
363 cohorts.entry(arm.as_str().to_string()).or_default().add(u);
364 }
365 persist();
366}
367
368#[must_use]
370pub fn cohort_snapshot() -> HashMap<String, CohortUsage> {
371 cohort_store()
372 .lock()
373 .unwrap_or_else(std::sync::PoisonError::into_inner)
374 .clone()
375}
376
377#[must_use]
379pub fn persisted_cohorts() -> HashMap<String, CohortUsage> {
380 load_persisted().map(|p| p.cohorts).unwrap_or_default()
381}
382
383fn normalize_key(model: &str) -> String {
384 let m = model.trim();
385 if m.is_empty() {
386 "unknown".to_string()
387 } else {
388 m.to_string()
389 }
390}
391
392pub fn snapshot() -> Vec<ModelSpend> {
394 let map = store()
395 .lock()
396 .unwrap_or_else(std::sync::PoisonError::into_inner);
397 price_models(&map)
398}
399
400pub fn total_cost_usd() -> f64 {
402 snapshot().iter().map(|m| m.cost_usd).sum()
403}
404
405#[derive(Debug, Clone, Serialize, PartialEq)]
407pub struct VerifiedSavings {
408 pub requests: u64,
410 pub counterfactual_input_tokens: u64,
413 pub billed_input_tokens: u64,
416 pub verified_saved_tokens: i64,
419}
420
421pub fn verified_savings() -> Option<VerifiedSavings> {
426 let map = store()
427 .lock()
428 .unwrap_or_else(std::sync::PoisonError::into_inner);
429 verified_of(&map)
430}
431
432#[allow(clippy::cast_possible_wrap)]
434fn verified_of(map: &HashMap<String, ModelUsage>) -> Option<VerifiedSavings> {
435 let (mut requests, mut counterfactual, mut billed) = (0u64, 0u64, 0u64);
436 for usage in map.values() {
437 requests += usage.counterfactual_requests;
438 counterfactual += usage.counterfactual_input_tokens;
439 billed += usage.counterfactual_billed_tokens;
440 }
441 (requests > 0).then(|| VerifiedSavings {
442 requests,
443 counterfactual_input_tokens: counterfactual,
444 billed_input_tokens: billed,
445 verified_saved_tokens: counterfactual as i64 - billed as i64,
446 })
447}
448
449pub fn price_models(map: &HashMap<String, ModelUsage>) -> Vec<ModelSpend> {
452 let pricing = ModelPricing::load();
453 let mut rows: Vec<ModelSpend> = map
454 .iter()
455 .map(|(model, usage)| price_one(&pricing, model, usage))
456 .collect();
457 rows.sort_by(|a, b| {
458 b.cost_usd
459 .partial_cmp(&a.cost_usd)
460 .unwrap_or(std::cmp::Ordering::Equal)
461 });
462 rows
463}
464
465fn price_one(pricing: &ModelPricing, model: &str, usage: &ModelUsage) -> ModelSpend {
466 let quote = pricing.quote(Some(model));
467 let m = &usage.measured;
470 let derived_cost = quote.cost.estimate_usd(
471 usage.input_tokens.saturating_sub(m.input_tokens),
472 usage.output_tokens.saturating_sub(m.output_tokens),
473 usage
474 .cache_write_tokens
475 .saturating_sub(m.cache_write_tokens),
476 usage.cache_read_tokens.saturating_sub(m.cache_read_tokens),
477 );
478 let derived_requests = usage.requests.saturating_sub(m.requests);
479 ModelSpend {
480 model: model.to_string(),
481 requests: usage.requests,
482 input_tokens: usage.input_tokens,
483 output_tokens: usage.output_tokens,
484 cache_read_tokens: usage.cache_read_tokens,
485 cache_write_tokens: usage.cache_write_tokens,
486 reasoning_tokens: usage.reasoning_tokens,
487 cost_usd: m.cost_usd + derived_cost,
488 measured_cost_usd: m.cost_usd,
489 measured_requests: m.requests,
490 pricing_estimated: quote.match_kind.is_estimated() && derived_requests > 0,
492 }
493}
494
495fn usage_path() -> Option<std::path::PathBuf> {
496 crate::core::data_dir::lean_ctx_data_dir()
497 .ok()
498 .map(|d| d.join(PROXY_USAGE_FILE))
499}
500
501fn persist() {
503 let Some(path) = usage_path() else {
504 return;
505 };
506 let models = {
507 let map = store()
508 .lock()
509 .unwrap_or_else(std::sync::PoisonError::into_inner);
510 map.clone()
511 };
512 let cohorts = {
513 let map = cohort_store()
514 .lock()
515 .unwrap_or_else(std::sync::PoisonError::into_inner);
516 map.clone()
517 };
518 let payload = PersistedUsage {
519 ts: std::time::SystemTime::now()
520 .duration_since(std::time::UNIX_EPOCH)
521 .unwrap_or_default()
522 .as_secs(),
523 models,
524 cohorts,
525 };
526 let Ok(json) = serde_json::to_string(&payload) else {
527 return;
528 };
529 let tmp = path.with_extension("json.tmp");
530 if std::fs::write(&tmp, json).is_ok() {
531 let _ = std::fs::rename(&tmp, &path);
532 }
533}
534
535pub fn load_persisted() -> Option<PersistedUsage> {
537 let path = usage_path()?;
538 let data = std::fs::read_to_string(path).ok()?;
539 serde_json::from_str(&data).ok()
540}
541
542pub fn persisted_snapshot() -> Vec<ModelSpend> {
544 load_persisted()
545 .map(|p| price_models(&p.models))
546 .unwrap_or_default()
547}
548
549pub fn persisted_dominant_model() -> Option<String> {
553 let persisted = load_persisted()?;
554 persisted
555 .models
556 .iter()
557 .filter(|(m, _)| m.as_str() != "unknown" && !m.trim().is_empty())
558 .max_by_key(|(_, u)| u.billable_tokens())
559 .filter(|(_, u)| u.billable_tokens() > 0)
560 .map(|(m, _)| m.clone())
561}
562
563#[cfg(test)]
564mod tests {
565 use super::*;
566
567 fn usage(
568 model: &str,
569 input: u64,
570 output: u64,
571 cache_read: u64,
572 ) -> super::super::usage::RealUsage {
573 super::super::usage::RealUsage {
574 model: model.to_string(),
575 input_tokens: input,
576 output_tokens: output,
577 cache_read_tokens: cache_read,
578 ..Default::default()
579 }
580 }
581
582 #[test]
585 fn counterfactual_pair_recorded_only_when_probe_answered() {
586 let slot = super::super::counterfactual::CounterfactualSlot::new();
587 slot.set(5_000);
588 let with_probe = super::super::usage::RealUsage {
589 wire: Some(Box::new(super::super::usage::WireContext {
590 counterfactual: Some(slot),
591 ..Default::default()
592 })),
593 cache_write_tokens: 200,
594 ..usage("claude-sonnet-4.5", 1_000, 0, 300)
595 };
596 let mut acc = ModelUsage::default();
597 acc.add(&with_probe);
598 assert_eq!(acc.counterfactual_requests, 1);
599 assert_eq!(acc.counterfactual_input_tokens, 5_000);
600 assert_eq!(acc.counterfactual_billed_tokens, 1_000 + 300 + 200);
602
603 let empty = super::super::usage::RealUsage {
606 wire: Some(Box::new(super::super::usage::WireContext {
607 counterfactual: Some(super::super::counterfactual::CounterfactualSlot::new()),
608 ..Default::default()
609 })),
610 ..usage("claude-sonnet-4.5", 1_000, 0, 0)
611 };
612 acc.add(&empty);
613 assert_eq!(acc.requests, 2);
614 assert_eq!(acc.counterfactual_requests, 1, "empty slot adds no pair");
615
616 acc.add(&usage("claude-sonnet-4.5", 10, 0, 0));
618 assert_eq!(acc.counterfactual_requests, 1);
619 }
620
621 #[test]
623 fn verified_of_aggregates_and_reports_signed_savings() {
624 let mut map = HashMap::new();
625 assert!(verified_of(&map).is_none(), "no coverage → None, not zeros");
626
627 map.insert(
628 "claude-sonnet-4.5".to_string(),
629 ModelUsage {
630 counterfactual_requests: 2,
631 counterfactual_input_tokens: 10_000,
632 counterfactual_billed_tokens: 6_000,
633 ..Default::default()
634 },
635 );
636 map.insert(
639 "claude-haiku-4.5".to_string(),
640 ModelUsage {
641 counterfactual_requests: 1,
642 counterfactual_input_tokens: 1_000,
643 counterfactual_billed_tokens: 1_400,
644 ..Default::default()
645 },
646 );
647 let v = verified_of(&map).expect("covered rows present");
648 assert_eq!(v.requests, 3);
649 assert_eq!(v.counterfactual_input_tokens, 11_000);
650 assert_eq!(v.billed_input_tokens, 7_400);
651 assert_eq!(v.verified_saved_tokens, 3_600);
652
653 map.get_mut("claude-sonnet-4.5")
654 .unwrap()
655 .counterfactual_input_tokens = 0;
656 map.get_mut("claude-sonnet-4.5")
657 .unwrap()
658 .counterfactual_billed_tokens = 0;
659 let negative = verified_of(&map).unwrap();
660 assert_eq!(
661 negative.verified_saved_tokens, -400,
662 "a net-negative verified saving is reported, never clamped"
663 );
664 }
665
666 #[test]
669 fn counterfactual_fields_roundtrip_and_default_for_legacy_files() {
670 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}}}"#;
671 let p: PersistedUsage = serde_json::from_str(legacy).expect("legacy file loads");
672 assert_eq!(p.models["m"].counterfactual_requests, 0);
673
674 let mut p = PersistedUsage::default();
675 p.models.insert(
676 "m".into(),
677 ModelUsage {
678 counterfactual_requests: 4,
679 counterfactual_input_tokens: 9_999,
680 counterfactual_billed_tokens: 5_555,
681 ..Default::default()
682 },
683 );
684 let back: PersistedUsage =
685 serde_json::from_str(&serde_json::to_string(&p).unwrap()).unwrap();
686 assert_eq!(back.models["m"].counterfactual_input_tokens, 9_999);
687 assert_eq!(back.models["m"].counterfactual_billed_tokens, 5_555);
688 }
689
690 #[test]
691 fn prices_known_model_with_cache_split() {
692 let mut map = HashMap::new();
693 let mut acc = ModelUsage::default();
694 acc.add(&usage("claude-sonnet-4.5", 1_000_000, 1_000_000, 1_000_000));
695 map.insert("claude-sonnet-4.5".to_string(), acc);
696
697 let rows = price_models(&map);
698 assert_eq!(rows.len(), 1);
699 let row = &rows[0];
700 assert!(
702 (row.cost_usd - 18.30).abs() < 1e-6,
703 "cost was {}",
704 row.cost_usd
705 );
706 assert!(!row.pricing_estimated, "exact model match");
707 assert_eq!(row.requests, 1);
708 }
709
710 #[test]
711 fn unknown_model_prices_with_fallback_and_is_estimated() {
712 let mut map = HashMap::new();
713 let mut acc = ModelUsage::default();
714 acc.add(&usage("some-novel-model-xyz", 1_000_000, 0, 0));
715 map.insert("some-novel-model-xyz".to_string(), acc);
716
717 let rows = price_models(&map);
718 assert!(rows[0].pricing_estimated, "fallback pricing is estimated");
719 assert!(rows[0].cost_usd > 0.0);
720 }
721
722 #[test]
727 fn measured_provider_cost_replaces_table_estimate() {
728 const MODEL: &str = "vendor/unlisted-model-20990101";
729 let mut acc = ModelUsage::default();
730 let mut measured = usage(MODEL, 431_600, 22_700, 126_700);
731 measured.provider_cost_usd = Some(0.05);
732 acc.add(&measured);
733
734 let mut map = HashMap::new();
735 map.insert(MODEL.to_string(), acc.clone());
736 let row = &price_models(&map)[0];
737 assert!(
738 (row.cost_usd - 0.05).abs() < 1e-12,
739 "the bill, not the table"
740 );
741 assert!((row.measured_cost_usd - 0.05).abs() < 1e-12);
742 assert_eq!(row.measured_requests, 1);
743 assert!(
744 !row.pricing_estimated,
745 "fully measured spend is not an estimate"
746 );
747
748 acc.add(&usage(MODEL, 10_000, 0, 0));
752 let mut map = HashMap::new();
753 map.insert(MODEL.to_string(), acc);
754 let row = &price_models(&map)[0];
755 assert!(row.cost_usd > 0.05, "unmeasured remainder adds table cost");
756 assert!(row.cost_usd < 0.2, "measured slice must not be re-priced");
757 assert!(
758 row.pricing_estimated,
759 "unmeasured remainder is heuristically priced"
760 );
761 }
762
763 #[test]
764 fn measured_slice_roundtrips_and_defaults_for_legacy_files() {
765 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}}}"#;
766 let p: PersistedUsage = serde_json::from_str(legacy).expect("legacy file loads");
767 assert_eq!(p.models["m"].measured, MeasuredSlice::default());
768
769 let mut p = PersistedUsage::default();
770 let mut acc = ModelUsage::default();
771 let mut m = usage("m", 100, 10, 0);
772 m.provider_cost_usd = Some(0.5);
773 acc.add(&m);
774 p.models.insert("m".into(), acc);
775 let back: PersistedUsage =
776 serde_json::from_str(&serde_json::to_string(&p).unwrap()).unwrap();
777 assert_eq!(back.models["m"].measured.requests, 1);
778 assert!((back.models["m"].measured.cost_usd - 0.5).abs() < 1e-12);
779 }
780
781 #[test]
782 fn dominant_model_picks_highest_token_real_model() {
783 let mut models = HashMap::new();
784 models.insert("claude-haiku-4.5".to_string(), {
785 let mut u = ModelUsage::default();
786 u.add(&usage("claude-haiku-4.5", 100, 100, 0));
787 u
788 });
789 models.insert("claude-opus-4.5".to_string(), {
790 let mut u = ModelUsage::default();
791 u.add(&usage("claude-opus-4.5", 10_000, 10_000, 0));
792 u
793 });
794 models.insert("unknown".to_string(), {
795 let mut u = ModelUsage::default();
796 u.add(&usage("unknown", 999_999, 0, 0));
797 u
798 });
799 let dominant = models
800 .iter()
801 .filter(|(m, _)| m.as_str() != "unknown")
802 .max_by_key(|(_, u)| u.billable_tokens())
803 .map(|(m, _)| m.clone());
804 assert_eq!(dominant.as_deref(), Some("claude-opus-4.5"));
805 }
806
807 #[test]
808 fn empty_model_buckets_as_unknown() {
809 assert_eq!(normalize_key(" "), "unknown");
810 assert_eq!(normalize_key(""), "unknown");
811 assert_eq!(normalize_key("gpt-5.4"), "gpt-5.4");
812 }
813
814 #[test]
815 fn cohort_avg_output_is_mean_per_turn() {
816 let mut c = CohortUsage::default();
817 assert_eq!(c.avg_output(), None, "no observations → None");
818 c.add(&usage("m", 10, 100, 0));
819 c.add(&usage("m", 10, 50, 0));
820 assert_eq!(c.requests, 2);
821 assert_eq!(c.output_tokens, 150);
822 assert!((c.avg_output().unwrap() - 75.0).abs() < f64::EPSILON);
823 }
824
825 #[test]
826 fn persisted_usage_without_cohorts_field_loads() {
827 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}}}"#;
830 let p: PersistedUsage = serde_json::from_str(json).expect("loads legacy file");
831 assert_eq!(p.models.len(), 1);
832 assert!(p.cohorts.is_empty());
833 }
834
835 #[test]
836 fn persisted_usage_roundtrips_cohorts() {
837 let mut p = PersistedUsage::default();
838 p.cohorts.insert(
839 "control".into(),
840 CohortUsage {
841 requests: 3,
842 input_tokens: 30,
843 output_tokens: 300,
844 sum_sq_output: 30_000,
845 },
846 );
847 let json = serde_json::to_string(&p).unwrap();
848 let back: PersistedUsage = serde_json::from_str(&json).unwrap();
849 assert_eq!(back.cohorts.get("control").unwrap().output_tokens, 300);
850 }
851
852 #[test]
853 fn managed_usage_projects_once_and_unmanaged_is_skipped() {
854 let sink = crate::core::ocla::builtin::usage_sink::BuiltinUsageSink::new();
855 let usage = super::super::usage::RealUsage {
856 model: "gpt-5".into(),
857 input_tokens: 100,
858 output_tokens: 40,
859 cache_read_tokens: 20,
860 cache_write_tokens: 5,
861 wire: Some(Box::new(super::super::usage::WireContext {
862 lineage: Some(crate::core::ocla::OclaRequestContext {
863 request_id: "request-1".into(),
864 session_id: "session-1".into(),
865 agent_id: "agent-1".into(),
866 content_ref: "blake3:content".into(),
867 tenant_id: None,
868 trace_id: "tr-unit".into(),
869 }),
870 ..Default::default()
871 })),
872 ..Default::default()
873 };
874
875 project_ocla_usage(&usage, &sink);
876 assert_eq!(sink.record_count(), 1);
877 assert_eq!(sink.total_input_tokens(), 125);
878 assert_eq!(sink.total_output_tokens(), 40);
879
880 project_ocla_usage(
881 &super::super::usage::RealUsage {
882 model: "unmanaged".into(),
883 input_tokens: 999,
884 ..Default::default()
885 },
886 &sink,
887 );
888 assert_eq!(sink.record_count(), 1);
889 }
890}