1use std::collections::HashMap;
16use std::sync::{Mutex, OnceLock};
17
18use serde::Serialize;
19
20use lean_ctx::core::gain::model_pricing::ModelPricing;
21use lean_ctx::core::ocla::{OclaRegistry, UsageSink};
22pub use lean_ctx::core::proxy_usage::{CohortUsage, load_persisted};
23use lean_ctx::core::proxy_usage::{
24 LastModel, MeasuredSlice, ModelUsage, PersistedUsage, usage_path,
25};
26
27fn project_ocla_usage(u: &super::usage::RealUsage, sink: &dyn UsageSink) {
33 use futures::FutureExt;
34 let Ok(Some(record)) = u.to_ocla_usage_record() else {
35 return;
36 };
37 let _ = sink.record_usage(record).now_or_never();
38}
39
40trait AddTurn {
43 fn add(&mut self, u: &super::usage::RealUsage);
44}
45
46impl AddTurn for ModelUsage {
47 fn add(&mut self, u: &super::usage::RealUsage) {
48 if let Some(observation) = &u.response_observation {
49 self.add_observation(observation.clone());
50 }
51 self.requests += 1;
52 self.input_tokens += u.input_tokens;
53 self.output_tokens += u.output_tokens;
54 self.cache_read_tokens += u.cache_read_tokens;
55 self.cache_write_tokens += u.cache_write_tokens;
56 self.reasoning_tokens += u.reasoning_tokens;
57 if let Some(cost) = u.provider_cost_usd {
58 self.measured.merge(&MeasuredSlice {
59 requests: 1,
60 input_tokens: u.input_tokens,
61 output_tokens: u.output_tokens,
62 cache_read_tokens: u.cache_read_tokens,
63 cache_write_tokens: u.cache_write_tokens,
64 cost_usd: cost,
65 });
66 }
67 if let Some(counted) = u
70 .wire
71 .as_deref()
72 .and_then(|w| w.counterfactual.as_ref())
73 .and_then(super::counterfactual::CounterfactualSlot::get)
74 {
75 self.counterfactual_requests += 1;
76 self.counterfactual_input_tokens += counted;
77 self.counterfactual_billed_tokens +=
78 u.input_tokens + u.cache_read_tokens + u.cache_write_tokens;
79 }
80 }
81}
82
83impl AddTurn for CohortUsage {
84 fn add(&mut self, u: &super::usage::RealUsage) {
85 self.requests += 1;
86 self.input_tokens += u.input_tokens;
87 self.output_tokens += u.output_tokens;
88 self.sum_sq_output += u.output_tokens.saturating_mul(u.output_tokens);
89 let prompt = u
90 .input_tokens
91 .saturating_add(u.cache_read_tokens)
92 .saturating_add(u.cache_write_tokens);
93 self.prompt_tokens += prompt;
94 self.sum_sq_prompt += prompt.saturating_mul(prompt);
95 }
96}
97
98#[derive(Debug, Clone, Serialize, PartialEq)]
100pub struct ModelSpend {
101 pub model: String,
102 pub requests: u64,
103 pub input_tokens: u64,
104 pub output_tokens: u64,
105 pub cache_read_tokens: u64,
106 pub cache_write_tokens: u64,
107 pub reasoning_tokens: u64,
108 pub cost_usd: f64,
109 pub measured_cost_usd: f64,
112 pub measured_requests: u64,
114 pub pricing_estimated: bool,
117}
118
119#[must_use]
123pub fn compression_cohort_key(arm: super::holdout::Arm) -> String {
124 format!("compression:{}", arm.as_str())
125}
126
127fn last_model_store() -> &'static Mutex<Option<LastModel>> {
128 static STORE: OnceLock<Mutex<Option<LastModel>>> = OnceLock::new();
129 STORE.get_or_init(|| Mutex::new(None))
130}
131
132const MAX_TRACKED_MODELS: usize = 256;
135
136fn store() -> &'static Mutex<HashMap<String, ModelUsage>> {
137 static STORE: OnceLock<Mutex<HashMap<String, ModelUsage>>> = OnceLock::new();
138 STORE.get_or_init(|| Mutex::new(HashMap::new()))
139}
140
141const MAX_TRACKED_PROJECTS: usize = 512;
143
144fn project_store() -> &'static Mutex<HashMap<String, HashMap<String, ModelUsage>>> {
145 static STORE: OnceLock<Mutex<HashMap<String, HashMap<String, ModelUsage>>>> = OnceLock::new();
146 STORE.get_or_init(|| Mutex::new(HashMap::new()))
147}
148
149fn cohort_store() -> &'static Mutex<HashMap<String, CohortUsage>> {
150 static STORE: OnceLock<Mutex<HashMap<String, CohortUsage>>> = OnceLock::new();
151 STORE.get_or_init(|| Mutex::new(HashMap::new()))
152}
153
154pub fn resume_from_disk() {
158 let Some(persisted) = load_persisted() else {
159 return;
160 };
161 {
162 let mut last = last_model_store()
163 .lock()
164 .unwrap_or_else(std::sync::PoisonError::into_inner);
165 if last.is_none() {
166 last.clone_from(&persisted.last_model);
167 }
168 }
169 let mut map = store()
170 .lock()
171 .unwrap_or_else(std::sync::PoisonError::into_inner);
172 merge_models(&mut map, persisted.models);
173 drop(map);
174 let mut projects = project_store()
175 .lock()
176 .unwrap_or_else(std::sync::PoisonError::into_inner);
177 for (project, models) in persisted.projects {
178 merge_models(projects.entry(project).or_default(), models);
179 }
180 drop(projects);
181 let mut cohorts = cohort_store()
182 .lock()
183 .unwrap_or_else(std::sync::PoisonError::into_inner);
184 for (arm, usage) in persisted.cohorts {
185 cohorts.entry(arm).or_default().merge(&usage);
186 }
187}
188
189fn merge_models(map: &mut HashMap<String, ModelUsage>, models: HashMap<String, ModelUsage>) {
190 for (model, usage) in models {
191 let acc = map.entry(model).or_default();
192 acc.requests += usage.requests;
193 acc.input_tokens += usage.input_tokens;
194 acc.output_tokens += usage.output_tokens;
195 acc.cache_read_tokens += usage.cache_read_tokens;
196 acc.cache_write_tokens += usage.cache_write_tokens;
197 acc.reasoning_tokens += usage.reasoning_tokens;
198 acc.counterfactual_requests += usage.counterfactual_requests;
199 acc.counterfactual_input_tokens += usage.counterfactual_input_tokens;
200 acc.counterfactual_billed_tokens += usage.counterfactual_billed_tokens;
201 acc.measured.merge(&usage.measured);
202 for observation in usage.response_observations {
203 acc.add_observation(observation);
204 }
205 }
206}
207
208pub fn record(u: &super::usage::RealUsage) {
211 super::usage_sink::push(u);
214
215 {
217 let label = u.wire.as_ref().map_or("Unknown", |w| {
218 match w.provider.as_str() {
221 "Anthropic" => "Anthropic",
222 "OpenAI" | "ChatGPT" => "OpenAI",
223 "Gemini" => "Gemini",
224 "Bedrock" => "Bedrock",
225 "Azure" => "Azure",
226 _ => "Unknown",
227 }
228 });
229 let envelope = super::usage_parity::real_usage_to_envelope(u, label);
230 lean_ctx::core::context_kernel::envelope_bridge::record_proxy_envelope(&envelope);
231 }
232
233 project_ocla_usage(u, OclaRegistry::global().usage_sink.as_ref());
237
238 #[cfg(feature = "enterprise")]
244 if let Some(wire) = u.wire.as_deref()
245 && (wire.person.is_some() || wire.project.is_some())
246 {
247 let baseline = lean_ctx::core::config::Config::load()
248 .proxy
249 .baseline
250 .clone();
251 #[allow(clippy::cast_precision_loss)]
252 let cost_usd = if wire.is_local {
253 let billable =
254 u.input_tokens + u.output_tokens + u.cache_read_tokens + u.cache_write_tokens;
255 baseline.effective_local_shadow_rate() / 1_000_000.0 * billable as f64
256 } else if let Some(measured) = u.provider_cost_usd {
257 measured
258 } else {
259 lean_ctx::core::gain::model_pricing::ModelPricing::load()
260 .quote(Some(&u.model))
261 .cost
262 .estimate_usd(
263 u.input_tokens,
264 u.output_tokens,
265 u.cache_write_tokens,
266 u.cache_read_tokens,
267 )
268 };
269 super::policy_gate::record_spend(wire.person.as_deref(), wire.project.as_deref(), cost_usd);
270 }
271
272 if u.cache_read_tokens > 0 {
276 let cost = lean_ctx::core::gain::model_pricing::ModelPricing::load()
277 .quote(Some(&u.model))
278 .cost;
279 #[allow(clippy::cast_precision_loss)]
280 let discount_usd =
281 (cost.input_per_m - cost.cache_read_per_m) / 1_000_000.0 * u.cache_read_tokens as f64;
282 lean_ctx::core::savings_ledger::record_caching_event(
283 &u.model,
284 u.cache_read_tokens,
285 discount_usd,
286 u.wire.as_deref().and_then(|w| w.project_key.as_deref()),
287 );
288 }
289
290 let model = u.model.trim();
293 if !model.is_empty()
294 && model != "unknown"
295 && !lean_ctx::core::model_registry::is_background_model(model)
296 {
297 *last_model_store()
298 .lock()
299 .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(LastModel {
300 model: model.to_string(),
301 at: std::time::SystemTime::now()
302 .duration_since(std::time::UNIX_EPOCH)
303 .unwrap_or_default()
304 .as_secs(),
305 });
306 }
307
308 let key = normalize_key(&u.model);
309 {
310 let mut map = store()
311 .lock()
312 .unwrap_or_else(std::sync::PoisonError::into_inner);
313 let key = if !map.contains_key(&key) && map.len() >= MAX_TRACKED_MODELS {
314 "unknown".to_string()
315 } else {
316 key
317 };
318 map.entry(key.clone()).or_default().add(u);
319 drop(map);
320 if let Some(project) = u.wire.as_deref().and_then(|w| w.project_key.as_deref()) {
321 let mut projects = project_store()
322 .lock()
323 .unwrap_or_else(std::sync::PoisonError::into_inner);
324 if projects.contains_key(project) || projects.len() < MAX_TRACKED_PROJECTS {
325 projects
326 .entry(project.to_string())
327 .or_default()
328 .entry(key)
329 .or_default()
330 .add(u);
331 }
332 }
333 }
334 if u.cohort.output.is_some() || u.cohort.compression.is_some() {
335 let mut cohorts = cohort_store()
336 .lock()
337 .unwrap_or_else(std::sync::PoisonError::into_inner);
338 if let Some(arm) = u.cohort.output {
339 cohorts.entry(arm.as_str().to_string()).or_default().add(u);
340 }
341 if let Some(arm) = u.cohort.compression {
342 cohorts
343 .entry(compression_cohort_key(arm))
344 .or_default()
345 .add(u);
346 }
347 }
348 persist();
349}
350
351#[must_use]
353pub fn cohort_snapshot() -> HashMap<String, CohortUsage> {
354 cohort_store()
355 .lock()
356 .unwrap_or_else(std::sync::PoisonError::into_inner)
357 .clone()
358}
359
360#[must_use]
362pub fn persisted_cohorts() -> HashMap<String, CohortUsage> {
363 load_persisted().map(|p| p.cohorts).unwrap_or_default()
364}
365
366fn normalize_key(model: &str) -> String {
367 let m = model.trim();
368 if m.is_empty() {
369 "unknown".to_string()
370 } else {
371 m.to_string()
372 }
373}
374
375pub fn snapshot() -> Vec<ModelSpend> {
377 let map = store()
378 .lock()
379 .unwrap_or_else(std::sync::PoisonError::into_inner);
380 price_models(&map)
381}
382
383pub fn total_cost_usd() -> f64 {
385 snapshot().iter().map(|m| m.cost_usd).sum()
386}
387
388#[derive(Debug, Clone, Serialize, PartialEq)]
390pub struct VerifiedSavings {
391 pub requests: u64,
393 pub counterfactual_input_tokens: u64,
396 pub billed_input_tokens: u64,
399 pub verified_saved_tokens: i64,
402}
403
404pub fn verified_savings() -> Option<VerifiedSavings> {
409 let map = store()
410 .lock()
411 .unwrap_or_else(std::sync::PoisonError::into_inner);
412 verified_of(&map)
413}
414
415#[allow(clippy::cast_possible_wrap)]
417fn verified_of(map: &HashMap<String, ModelUsage>) -> Option<VerifiedSavings> {
418 let (mut requests, mut counterfactual, mut billed) = (0u64, 0u64, 0u64);
419 for usage in map.values() {
420 requests += usage.counterfactual_requests;
421 counterfactual += usage.counterfactual_input_tokens;
422 billed += usage.counterfactual_billed_tokens;
423 }
424 (requests > 0).then(|| VerifiedSavings {
425 requests,
426 counterfactual_input_tokens: counterfactual,
427 billed_input_tokens: billed,
428 verified_saved_tokens: counterfactual as i64 - billed as i64,
429 })
430}
431
432pub fn price_models(map: &HashMap<String, ModelUsage>) -> Vec<ModelSpend> {
435 let pricing = ModelPricing::load();
436 let mut rows: Vec<ModelSpend> = map
437 .iter()
438 .map(|(model, usage)| price_one(&pricing, model, usage))
439 .collect();
440 rows.sort_by(|a, b| {
441 b.cost_usd
442 .partial_cmp(&a.cost_usd)
443 .unwrap_or(std::cmp::Ordering::Equal)
444 });
445 rows
446}
447
448fn price_one(pricing: &ModelPricing, model: &str, usage: &ModelUsage) -> ModelSpend {
449 let quote = pricing.quote(Some(model));
450 let m = &usage.measured;
453 let derived_cost = quote.cost.estimate_usd(
454 usage.input_tokens.saturating_sub(m.input_tokens),
455 usage.output_tokens.saturating_sub(m.output_tokens),
456 usage
457 .cache_write_tokens
458 .saturating_sub(m.cache_write_tokens),
459 usage.cache_read_tokens.saturating_sub(m.cache_read_tokens),
460 );
461 let derived_requests = usage.requests.saturating_sub(m.requests);
462 ModelSpend {
463 model: model.to_string(),
464 requests: usage.requests,
465 input_tokens: usage.input_tokens,
466 output_tokens: usage.output_tokens,
467 cache_read_tokens: usage.cache_read_tokens,
468 cache_write_tokens: usage.cache_write_tokens,
469 reasoning_tokens: usage.reasoning_tokens,
470 cost_usd: m.cost_usd + derived_cost,
471 measured_cost_usd: m.cost_usd,
472 measured_requests: m.requests,
473 pricing_estimated: quote.match_kind.is_estimated() && derived_requests > 0,
475 }
476}
477
478fn persist() {
480 let Some(path) = usage_path() else {
481 return;
482 };
483 let models = {
484 let map = store()
485 .lock()
486 .unwrap_or_else(std::sync::PoisonError::into_inner);
487 map.clone()
488 };
489 let cohorts = {
490 let map = cohort_store()
491 .lock()
492 .unwrap_or_else(std::sync::PoisonError::into_inner);
493 map.clone()
494 };
495 let last_model = last_model_store()
496 .lock()
497 .unwrap_or_else(std::sync::PoisonError::into_inner)
498 .clone();
499 let projects = project_store()
500 .lock()
501 .unwrap_or_else(std::sync::PoisonError::into_inner)
502 .clone();
503 let payload = PersistedUsage {
504 ts: std::time::SystemTime::now()
505 .duration_since(std::time::UNIX_EPOCH)
506 .unwrap_or_default()
507 .as_secs(),
508 models,
509 cohorts,
510 last_model,
511 projects,
512 };
513 let Ok(json) = serde_json::to_string(&payload) else {
514 return;
515 };
516 let tmp = path.with_extension("json.tmp");
517 if std::fs::write(&tmp, json).is_ok() {
518 let _ = std::fs::rename(&tmp, &path);
519 }
520}
521
522pub fn persisted_snapshot() -> Vec<ModelSpend> {
524 load_persisted()
525 .map(|p| price_models(&p.models))
526 .unwrap_or_default()
527}
528
529pub fn persisted_snapshot_where(keep: impl Fn(&str) -> bool) -> Vec<ModelSpend> {
531 let Some(persisted) = load_persisted() else {
532 return Vec::new();
533 };
534 let mut models: HashMap<String, ModelUsage> = HashMap::new();
535 for (project, usage) in persisted.projects {
536 if keep(&project) {
537 merge_models(&mut models, usage);
538 }
539 }
540 price_models(&models)
541}
542
543pub fn persisted_projects() -> Vec<(String, Vec<ModelSpend>)> {
545 let Some(persisted) = load_persisted() else {
546 return Vec::new();
547 };
548 let mut rows: Vec<(String, Vec<ModelSpend>)> = persisted
549 .projects
550 .iter()
551 .map(|(project, models)| (project.clone(), price_models(models)))
552 .collect();
553 rows.sort_by(|a, b| a.0.cmp(&b.0));
554 rows
555}
556
557#[cfg(test)]
558mod tests {
559 use super::*;
560
561 fn usage(
562 model: &str,
563 input: u64,
564 output: u64,
565 cache_read: u64,
566 ) -> super::super::usage::RealUsage {
567 super::super::usage::RealUsage {
568 model: model.to_string(),
569 input_tokens: input,
570 output_tokens: output,
571 cache_read_tokens: cache_read,
572 ..Default::default()
573 }
574 }
575
576 #[test]
580 fn cohort_merge_keeps_every_running_total() {
581 let mut persisted = CohortUsage::default();
582 for (out, cached) in [(150, 900), (210, 1_200), (180, 0)] {
583 persisted.add(&usage("claude-sonnet-4.5", 1_000, out, cached));
584 }
585 let mut resumed = CohortUsage::default();
586 resumed.merge(&persisted);
587 assert_eq!(resumed, persisted, "every field survives the restart");
588 assert!(resumed.variance_output().is_some_and(|v| v > 0.0));
589
590 let mut live = CohortUsage::default();
592 live.add(&usage("claude-sonnet-4.5", 500, 100, 0));
593 let expected_sum_sq = live.sum_sq_output + persisted.sum_sq_output;
594 live.merge(&persisted);
595 assert_eq!(live.requests, 4);
596 assert_eq!(live.sum_sq_output, expected_sum_sq);
597 assert_eq!(live.prompt_tokens, 500 + persisted.prompt_tokens);
598 }
599
600 #[test]
604 fn spend_is_split_per_project_and_persisted() {
605 let _dir = lean_ctx::core::data_dir::isolated_data_dir();
606 let project = format!("{:016x}", std::process::id() as u64 * 7919 + 17);
607 let other = format!("{:016x}", std::process::id() as u64 * 7919 + 18);
608 let turn = |key: Option<&str>, input| super::super::usage::RealUsage {
609 wire: Some(Box::new(super::super::usage::WireContext {
610 project_key: key.map(str::to_string),
611 ..Default::default()
612 })),
613 ..usage("claude-sonnet-4.5", input, 10, 0)
614 };
615 record(&turn(Some(&project), 1_000));
616 record(&turn(Some(&project), 500));
617 record(&turn(Some(&other), 200));
618 record(&turn(None, 100));
619
620 let mine = persisted_snapshot_where(|key| key == project);
621 assert_eq!(mine.len(), 1);
622 assert_eq!(mine[0].requests, 2);
623 assert_eq!(mine[0].input_tokens, 1_500);
624 let persisted = load_persisted().expect("persisted");
625 let device = &persisted.models["claude-sonnet-4.5"];
626 let split: u64 = persisted_projects()
627 .iter()
628 .filter(|(key, _)| *key == project || *key == other)
629 .flat_map(|(_, models)| models.iter().map(|m| m.input_tokens))
630 .sum();
631 assert_eq!(split, 1_700);
632 assert!(
633 device.input_tokens >= 1_800,
634 "the device total keeps every turn, named or not"
635 );
636 }
637
638 #[test]
643 fn recording_a_turn_leaves_earlier_savings_ledger_entries_unchanged() {
644 use lean_ctx::core::savings_ledger::{self, event::SavingsEvent, store};
645 let _dir = lean_ctx::core::data_dir::isolated_data_dir();
646 let path = store::default_path().expect("ledger path");
647 let legacy: SavingsEvent = serde_json::from_value(serde_json::json!({
648 "ts": "2026-06-01T12:00:00+00:00",
649 "tool": "ctx_read",
650 "model_id": "claude-3.5-sonnet",
651 "tokenizer": "o200k_base",
652 "baseline_tokens": 1000,
653 "actual_tokens": 300,
654 "saved_tokens": 700,
655 "bounce_adjustment": 0,
656 "unit_price_per_m_usd": 3.0,
657 "saved_usd": 0.0021,
658 "repo_hash": "repo",
659 "agent_id": "local",
660 "prev_hash": "",
661 "entry_hash": ""
662 }))
663 .expect("legacy fixture");
664 store::append(&path, legacy).expect("append fixture");
665 record(&usage("claude-sonnet-4.5", 1_000, 50, 400));
666 let before = std::fs::read(&path).expect("ledger bytes");
667 let before_events = savings_ledger::all_events();
668 assert_eq!(before_events.len(), 2, "fixture plus one caching event");
669
670 record(&usage("claude-sonnet-4.5", 2_000, 80, 600));
671
672 let after = std::fs::read(&path).expect("ledger bytes");
673 assert!(after.starts_with(&before), "earlier entries are untouched");
674 let after_events = savings_ledger::all_events();
675 assert_eq!(after_events.len(), 3, "one turn appends one event");
676 for (old, new) in before_events.iter().zip(&after_events) {
677 assert_eq!(old.entry_hash, new.entry_hash);
678 assert_eq!(old.saved_tokens, new.saved_tokens);
679 assert_eq!(old.saved_usd.to_bits(), new.saved_usd.to_bits());
680 }
681 let verified = savings_ledger::verify();
682 assert!(verified.valid, "{verified:?}");
683 assert_eq!(verified.total, 3);
684 }
685
686 #[test]
689 fn counterfactual_pair_recorded_only_when_probe_answered() {
690 let slot = super::super::counterfactual::CounterfactualSlot::new();
691 slot.set(5_000);
692 let with_probe = super::super::usage::RealUsage {
693 wire: Some(Box::new(super::super::usage::WireContext {
694 counterfactual: Some(slot),
695 ..Default::default()
696 })),
697 cache_write_tokens: 200,
698 ..usage("claude-sonnet-4.5", 1_000, 0, 300)
699 };
700 let mut acc = ModelUsage::default();
701 acc.add(&with_probe);
702 assert_eq!(acc.counterfactual_requests, 1);
703 assert_eq!(acc.counterfactual_input_tokens, 5_000);
704 assert_eq!(acc.counterfactual_billed_tokens, 1_000 + 300 + 200);
706
707 let empty = super::super::usage::RealUsage {
710 wire: Some(Box::new(super::super::usage::WireContext {
711 counterfactual: Some(super::super::counterfactual::CounterfactualSlot::new()),
712 ..Default::default()
713 })),
714 ..usage("claude-sonnet-4.5", 1_000, 0, 0)
715 };
716 acc.add(&empty);
717 assert_eq!(acc.requests, 2);
718 assert_eq!(acc.counterfactual_requests, 1, "empty slot adds no pair");
719
720 acc.add(&usage("claude-sonnet-4.5", 10, 0, 0));
722 assert_eq!(acc.counterfactual_requests, 1);
723 }
724
725 #[test]
727 fn verified_of_aggregates_and_reports_signed_savings() {
728 let mut map = HashMap::new();
729 assert!(verified_of(&map).is_none(), "no coverage → None, not zeros");
730
731 map.insert(
732 "claude-sonnet-4.5".to_string(),
733 ModelUsage {
734 counterfactual_requests: 2,
735 counterfactual_input_tokens: 10_000,
736 counterfactual_billed_tokens: 6_000,
737 ..Default::default()
738 },
739 );
740 map.insert(
743 "claude-haiku-4.5".to_string(),
744 ModelUsage {
745 counterfactual_requests: 1,
746 counterfactual_input_tokens: 1_000,
747 counterfactual_billed_tokens: 1_400,
748 ..Default::default()
749 },
750 );
751 let v = verified_of(&map).expect("covered rows present");
752 assert_eq!(v.requests, 3);
753 assert_eq!(v.counterfactual_input_tokens, 11_000);
754 assert_eq!(v.billed_input_tokens, 7_400);
755 assert_eq!(v.verified_saved_tokens, 3_600);
756
757 map.get_mut("claude-sonnet-4.5")
758 .unwrap()
759 .counterfactual_input_tokens = 0;
760 map.get_mut("claude-sonnet-4.5")
761 .unwrap()
762 .counterfactual_billed_tokens = 0;
763 let negative = verified_of(&map).unwrap();
764 assert_eq!(
765 negative.verified_saved_tokens, -400,
766 "a net-negative verified saving is reported, never clamped"
767 );
768 }
769
770 #[test]
773 fn counterfactual_fields_roundtrip_and_default_for_legacy_files() {
774 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}}}"#;
775 let p: PersistedUsage = serde_json::from_str(legacy).expect("legacy file loads");
776 assert_eq!(p.models["m"].counterfactual_requests, 0);
777
778 let mut p = PersistedUsage::default();
779 p.models.insert(
780 "m".into(),
781 ModelUsage {
782 counterfactual_requests: 4,
783 counterfactual_input_tokens: 9_999,
784 counterfactual_billed_tokens: 5_555,
785 ..Default::default()
786 },
787 );
788 let back: PersistedUsage =
789 serde_json::from_str(&serde_json::to_string(&p).unwrap()).unwrap();
790 assert_eq!(back.models["m"].counterfactual_input_tokens, 9_999);
791 assert_eq!(back.models["m"].counterfactual_billed_tokens, 5_555);
792 }
793
794 #[test]
795 fn prices_known_model_with_cache_split() {
796 let mut map = HashMap::new();
797 let mut acc = ModelUsage::default();
798 acc.add(&usage("claude-sonnet-4.5", 1_000_000, 1_000_000, 1_000_000));
799 map.insert("claude-sonnet-4.5".to_string(), acc);
800
801 let rows = price_models(&map);
802 assert_eq!(rows.len(), 1);
803 let row = &rows[0];
804 assert!(
806 (row.cost_usd - 18.30).abs() < 1e-6,
807 "cost was {}",
808 row.cost_usd
809 );
810 assert!(!row.pricing_estimated, "exact model match");
811 assert_eq!(row.requests, 1);
812 }
813
814 #[test]
815 fn unknown_model_prices_with_fallback_and_is_estimated() {
816 let mut map = HashMap::new();
817 let mut acc = ModelUsage::default();
818 acc.add(&usage("some-novel-model-xyz", 1_000_000, 0, 0));
819 map.insert("some-novel-model-xyz".to_string(), acc);
820
821 let rows = price_models(&map);
822 assert!(rows[0].pricing_estimated, "fallback pricing is estimated");
823 assert!(rows[0].cost_usd > 0.0);
824 }
825
826 #[test]
831 fn measured_provider_cost_replaces_table_estimate() {
832 const MODEL: &str = "vendor/unlisted-model-20990101";
833 let mut acc = ModelUsage::default();
834 let mut measured = usage(MODEL, 431_600, 22_700, 126_700);
835 measured.provider_cost_usd = Some(0.05);
836 acc.add(&measured);
837
838 let mut map = HashMap::new();
839 map.insert(MODEL.to_string(), acc.clone());
840 let row = &price_models(&map)[0];
841 assert!(
842 (row.cost_usd - 0.05).abs() < 1e-12,
843 "the bill, not the table"
844 );
845 assert!((row.measured_cost_usd - 0.05).abs() < 1e-12);
846 assert_eq!(row.measured_requests, 1);
847 assert!(
848 !row.pricing_estimated,
849 "fully measured spend is not an estimate"
850 );
851
852 acc.add(&usage(MODEL, 10_000, 0, 0));
856 let mut map = HashMap::new();
857 map.insert(MODEL.to_string(), acc);
858 let row = &price_models(&map)[0];
859 assert!(row.cost_usd > 0.05, "unmeasured remainder adds table cost");
860 assert!(row.cost_usd < 0.2, "measured slice must not be re-priced");
861 assert!(
862 row.pricing_estimated,
863 "unmeasured remainder is heuristically priced"
864 );
865 }
866
867 #[test]
868 fn measured_slice_roundtrips_and_defaults_for_legacy_files() {
869 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}}}"#;
870 let p: PersistedUsage = serde_json::from_str(legacy).expect("legacy file loads");
871 assert_eq!(p.models["m"].measured, MeasuredSlice::default());
872
873 let mut p = PersistedUsage::default();
874 let mut acc = ModelUsage::default();
875 let mut m = usage("m", 100, 10, 0);
876 m.provider_cost_usd = Some(0.5);
877 acc.add(&m);
878 p.models.insert("m".into(), acc);
879 let back: PersistedUsage =
880 serde_json::from_str(&serde_json::to_string(&p).unwrap()).unwrap();
881 assert_eq!(back.models["m"].measured.requests, 1);
882 assert!((back.models["m"].measured.cost_usd - 0.5).abs() < 1e-12);
883 }
884
885 #[test]
886 fn dominant_model_picks_highest_token_real_model() {
887 let mut models = HashMap::new();
888 models.insert("claude-haiku-4.5".to_string(), {
889 let mut u = ModelUsage::default();
890 u.add(&usage("claude-haiku-4.5", 100, 100, 0));
891 u
892 });
893 models.insert("claude-opus-4.5".to_string(), {
894 let mut u = ModelUsage::default();
895 u.add(&usage("claude-opus-4.5", 10_000, 10_000, 0));
896 u
897 });
898 models.insert("unknown".to_string(), {
899 let mut u = ModelUsage::default();
900 u.add(&usage("unknown", 999_999, 0, 0));
901 u
902 });
903 let dominant = models
904 .iter()
905 .filter(|(m, _)| m.as_str() != "unknown")
906 .max_by_key(|(_, u)| u.billable_tokens())
907 .map(|(m, _)| m.clone());
908 assert_eq!(dominant.as_deref(), Some("claude-opus-4.5"));
909 }
910
911 #[test]
912 fn empty_model_buckets_as_unknown() {
913 assert_eq!(normalize_key(" "), "unknown");
914 assert_eq!(normalize_key(""), "unknown");
915 assert_eq!(normalize_key("gpt-5.4"), "gpt-5.4");
916 }
917
918 #[test]
919 fn cohort_avg_output_is_mean_per_turn() {
920 let mut c = CohortUsage::default();
921 assert_eq!(c.avg_output(), None, "no observations → None");
922 c.add(&usage("m", 10, 100, 0));
923 c.add(&usage("m", 10, 50, 0));
924 assert_eq!(c.requests, 2);
925 assert_eq!(c.output_tokens, 150);
926 assert!((c.avg_output().unwrap() - 75.0).abs() < f64::EPSILON);
927 }
928
929 #[test]
930 fn persisted_usage_without_cohorts_field_loads() {
931 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}}}"#;
934 let p: PersistedUsage = serde_json::from_str(json).expect("loads legacy file");
935 assert_eq!(p.models.len(), 1);
936 assert!(p.cohorts.is_empty());
937 }
938
939 #[test]
940 fn persisted_usage_roundtrips_cohorts() {
941 let mut p = PersistedUsage::default();
942 p.cohorts.insert(
943 "control".into(),
944 CohortUsage {
945 requests: 3,
946 input_tokens: 30,
947 output_tokens: 300,
948 sum_sq_output: 30_000,
949 ..CohortUsage::default()
950 },
951 );
952 let json = serde_json::to_string(&p).unwrap();
953 let back: PersistedUsage = serde_json::from_str(&json).unwrap();
954 assert_eq!(back.cohorts.get("control").unwrap().output_tokens, 300);
955 }
956
957 #[test]
958 fn managed_usage_projects_once_and_unmanaged_is_skipped() {
959 let sink = lean_ctx::core::ocla::builtin::usage_sink::BuiltinUsageSink::new();
960 let usage = super::super::usage::RealUsage {
961 model: "gpt-5".into(),
962 input_tokens: 100,
963 output_tokens: 40,
964 cache_read_tokens: 20,
965 cache_write_tokens: 5,
966 wire: Some(Box::new(super::super::usage::WireContext {
967 lineage: Some(lean_ctx::core::ocla::OclaRequestContext {
968 request_id: "request-1".into(),
969 session_id: "session-1".into(),
970 agent_id: "agent-1".into(),
971 content_ref: "blake3:content".into(),
972 tenant_id: None,
973 trace_id: "tr-unit".into(),
974 task_id: None,
975 parent_task_id: None,
976 }),
977 ..Default::default()
978 })),
979 ..Default::default()
980 };
981
982 project_ocla_usage(&usage, &sink);
983 assert_eq!(sink.record_count(), 1);
984 assert_eq!(sink.total_input_tokens(), 125);
985 assert_eq!(sink.total_output_tokens(), 40);
986
987 project_ocla_usage(
988 &super::super::usage::RealUsage {
989 model: "unmanaged".into(),
990 input_tokens: 999,
991 ..Default::default()
992 },
993 &sink,
994 );
995 assert_eq!(sink.record_count(), 1);
996 }
997
998 struct SpyUsage(std::sync::Arc<std::sync::atomic::AtomicUsize>);
1000
1001 impl lean_ctx::core::ocla::traits::OclaService for SpyUsage {
1002 fn capability(&self) -> lean_ctx::core::ocla::types::OclaCapability {
1003 lean_ctx::core::ocla::types::OclaCapability::available(
1004 lean_ctx::core::ocla::types::OclaCapabilityKind::UsageSink,
1005 )
1006 }
1007 }
1008
1009 #[async_trait::async_trait]
1010 impl UsageSink for SpyUsage {
1011 async fn record_usage(
1012 &self,
1013 _usage: lean_ctx::core::ocla::types::UsageRecord,
1014 ) -> lean_ctx::core::ocla::OclaResult<()> {
1015 self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1016 Ok(())
1017 }
1018 }
1019
1020 #[test]
1021 fn proxy_request_projects_measured_usage_to_usage_spy() {
1022 let _dir = lean_ctx::core::data_dir::isolated_data_dir();
1023 let usage_calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
1024 let mut registry = OclaRegistry::with_builtins();
1025 registry.usage_sink = std::sync::Arc::new(SpyUsage(usage_calls.clone()));
1026 let _guard = lean_ctx::core::ocla::registry::with_test_registry(registry);
1027 let usage = super::super::usage::RealUsage {
1028 model: "test-model".into(),
1029 input_tokens: 11,
1030 output_tokens: 7,
1031 wire: Some(Box::new(super::super::usage::WireContext {
1032 lineage: Some(lean_ctx::core::ocla::OclaRequestContext {
1033 request_id: "request-1".into(),
1034 session_id: "session-1".into(),
1035 agent_id: "agent-1".into(),
1036 content_ref: "file:test.rs".into(),
1037 tenant_id: None,
1038 trace_id: "tr-unit".into(),
1039 task_id: None,
1040 parent_task_id: None,
1041 }),
1042 ..Default::default()
1043 })),
1044 ..Default::default()
1045 };
1046
1047 record(&usage);
1048
1049 assert_eq!(usage_calls.load(std::sync::atomic::Ordering::Relaxed), 1);
1050 }
1051}