Skip to main content

lean_ctx/proxy/
usage_meter.rs

1//! Measured per-model spend meter.
2//!
3//! [`record`] aggregates the real provider usage extracted by [`super::usage`]
4//! into per-model token sums, prices them with the shared
5//! `ModelPricing` table, and
6//! persists the totals to `proxy_usage.json` so the dashboard, CLI and the
7//! savings ledger (which run in *other* processes) can read the user's real
8//! provider bill.
9//!
10//! Unlike [`super::metrics`] (which resets per proxy lifetime), this meter is a
11//! lifetime-cumulative spend counter: [`resume_from_disk`] seeds the in-memory
12//! totals on proxy startup so a restart never zeroes the user's measured spend.
13
14use 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
22/// Projects one managed, provider-measured turn into the OCLA usage sink.
23///
24/// This is deliberately one-way and best-effort: the existing meter remains
25/// authoritative for billing/ledger totals, while OCLA receives a validated
26/// payload-free projection only when complete lineage is present.
27fn 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/// Cumulative real token counts for one model. Cost is derived at read time so a
36/// pricing-table change re-values historical usage consistently.
37#[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    /// Requests whose free `count_tokens` probe answered (#701) — the rows the
46    /// verified-savings pair below covers. `serde(default)` keeps pre-#701
47    /// usage files loadable.
48    #[serde(default)]
49    pub counterfactual_requests: u64,
50    /// Provider-counted input tokens the covered requests would have billed
51    /// WITHOUT lean-ctx (sum of probe answers on the original bodies).
52    #[serde(default)]
53    pub counterfactual_input_tokens: u64,
54    /// Input-side tokens those same requests actually billed (input +
55    /// cache read + cache write) — same request, same moment, no confound.
56    #[serde(default)]
57    pub counterfactual_billed_tokens: u64,
58    /// The turns whose response carried the provider's own USD charge
59    /// (OpenRouter usage accounting, #1179). Their tokens are *inside* the
60    /// totals above; pricing subtracts this slice and books its measured USD
61    /// instead of a table estimate. `serde(default)` keeps older files loadable.
62    #[serde(default)]
63    pub measured: MeasuredSlice,
64}
65
66/// Token/cost sums of the turns that reported a measured provider charge.
67#[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    /// Sum of provider-reported USD for those turns — the bill, not a table.
75    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        // Verified-savings pair (#701): only when the probe answered by the
108        // time the billed usage arrived — both sides of the pair or neither.
109        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/// One model's measured, priced spend for `/status` and the dashboard.
128#[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    /// Share of `cost_usd` the provider itself reported (OpenRouter usage
139    /// accounting, #1179) — a bill, not a table estimate.
140    pub measured_cost_usd: f64,
141    /// Requests covered by that provider-reported cost.
142    pub measured_requests: u64,
143    /// True when part of the cost had to be priced from a heuristic/fallback
144    /// match — never set by measured or exactly-priced spend.
145    pub pricing_estimated: bool,
146}
147
148/// Cumulative output-savings cohort totals (#895 Track B). Keyed by arm name
149/// (`"control"` | `"treatment"`); the average output tokens per turn is
150/// `output_tokens / requests`. Only populated while a holdout is active.
151#[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    /// Sum of squared per-turn output tokens, enabling an online sample variance
157    /// (and therefore a confidence interval) without retaining every turn.
158    /// `#[serde(default)]` keeps pre-#895 files loadable.
159    #[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    /// Average output tokens per turn, or `None` with no observations.
172    #[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    /// Unbiased sample variance of per-turn output tokens, or `None` with < 2
183    /// observations. Computed from the running sum / sum-of-squares (clamped at
184    /// 0 to absorb floating-point error on near-constant samples).
185    #[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/// On-disk shape of the measured spend totals.
202#[derive(Debug, Clone, Default, Serialize, Deserialize)]
203pub struct PersistedUsage {
204    pub ts: u64,
205    pub models: HashMap<String, ModelUsage>,
206    /// Output-savings cohort totals (#895). `#[serde(default)]` keeps older
207    /// `proxy_usage.json` files (written before the holdout existed) loadable.
208    #[serde(default)]
209    pub cohorts: HashMap<String, CohortUsage>,
210}
211
212/// Distinct-model bucket cap. `record` keys on the raw response model string, so
213/// overflow folds into "unknown" to keep the map bounded (real model names < ~50).
214const 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
228/// Seeds the in-memory totals from `proxy_usage.json`. Call once on proxy
229/// startup so measured spend is cumulative across restarts. Idempotent-ish: it
230/// merges the persisted totals into whatever is in memory (normally empty).
231pub 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
263/// Records one turn's measured usage against its model bucket (and its
264/// output-savings cohort, when tagged) and persists.
265pub fn record(u: &super::usage::RealUsage) {
266    // Gateway store subscription (enterprise#17): forward the full record to
267    // the installed sink (no-op locally). Never blocks the request path.
268    super::usage_sink::push(u);
269
270    // R33: Feed measured usage into kernel envelope pipeline.
271    {
272        let label = u.wire.as_ref().map_or("Unknown", |w| {
273            // provider_kind_from_label expects a static or owned string;
274            // the label lives long enough for the conversion call.
275            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    // OCLA projection: keep the canonical usage capability on the same single
289    // finalized-turn choke-point without feeding back into billing or ledger
290    // aggregation. Invalid/missing lineage stays explicitly unmanaged.
291    project_ocla_usage(u, OclaRegistry::global().usage_sink.as_ref());
292
293    // Budget windows (enterprise#25): book this turn's measured cost against
294    // the person/day and project/month accumulators the policy gate checks.
295    // The provider's own reported charge wins (#1179); local turns book the
296    // shadow rate — the same valuation the usage store applies — so
297    // local-only budgets stay meaningful.
298    #[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    // Mechanism attribution into the local savings ledger (enterprise#19).
325    // Routing: the gateway served a cheaper model than requested — value the
326    // rate delta on the measured input tokens. Caching: provider prompt-cache
327    // reads billed below the input rate. Both best-effort, never blocking.
328    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/// Live output-savings cohort totals (#895). Empty until a holdout runs.
369#[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/// Cross-process read of the persisted output-savings cohort totals.
378#[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
392/// Live per-model measured spend, priced and sorted by USD descending.
393pub 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
400/// Total measured spend across all models (live in-memory totals).
401pub fn total_cost_usd() -> f64 {
402    snapshot().iter().map(|m| m.cost_usd).sum()
403}
404
405/// Cross-model verified-savings totals (#701) for `/status` and the dashboard.
406#[derive(Debug, Clone, Serialize, PartialEq)]
407pub struct VerifiedSavings {
408    /// Requests covered by a successful `count_tokens` probe.
409    pub requests: u64,
410    /// Provider-counted input tokens those requests would have billed
411    /// without lean-ctx.
412    pub counterfactual_input_tokens: u64,
413    /// Input-side tokens (input + cache read + cache write) they actually
414    /// billed.
415    pub billed_input_tokens: u64,
416    /// `counterfactual - billed`; negative when stub overhead outweighed the
417    /// squeeze — reported honestly, never clamped.
418    pub verified_saved_tokens: i64,
419}
420
421/// Aggregated provider-verified savings across all models (#701), or `None`
422/// until at least one probe-covered request has been recorded. Unlike the
423/// `tokens_saved` estimate (bytes/4), both sides of this pair were counted by
424/// the provider on the same request — receipts, not estimates.
425pub 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/// Pure aggregation behind [`verified_savings`], shared with tests.
433#[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
449/// Prices a model usage map into sorted [`ModelSpend`] rows. Pure: shared by the
450/// in-memory snapshot and the cross-process [`persisted_snapshot`].
451pub 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    // Measured-first (#1179): turns whose response carried the provider's own
468    // USD charge book that exact amount; only the remainder is table-priced.
469    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        // Fully measured spend is never an estimate, whatever the table says.
491        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
501/// Atomically writes the current in-memory totals to disk.
502fn 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
535/// Cross-process read of the persisted measured spend (dashboard / CLI / ledger).
536pub 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
542/// Cross-process priced spend rows, read from disk.
543pub fn persisted_snapshot() -> Vec<ModelSpend> {
544    load_persisted()
545        .map(|p| price_models(&p.models))
546        .unwrap_or_default()
547}
548
549/// The model carrying the most measured tokens (excludes the "unknown" bucket).
550/// Used to value savings against the real dominant model when no explicit model
551/// is configured.
552pub 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    /// #701: the verified pair is recorded only when the probe answered — both
583    /// sides of the pair or neither, never a half row.
584    #[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        // billed side = input + cache read + cache write of the same turn.
601        assert_eq!(acc.counterfactual_billed_tokens, 1_000 + 300 + 200);
602
603        // Empty slot (probe failed / still in flight) → row degrades to the
604        // estimate: no pair recorded, normal usage still counted.
605        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        // No wire context at all (tests / non-forward paths) → no pair.
617        acc.add(&usage("claude-sonnet-4.5", 10, 0, 0));
618        assert_eq!(acc.counterfactual_requests, 1);
619    }
620
621    /// #701: cross-model aggregation and the honest signed difference.
622    #[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        // Stub overhead outweighed the squeeze on this model: billed MORE
637        // than the counterfactual. The total must subtract honestly.
638        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    /// #701: persisted usage files from before the feature load cleanly and
667    /// the new fields round-trip.
668    #[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        // input 3.00 + output 15.00 + cache_read 0.30 (per 1M) = 18.30.
701        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    /// #1179: turns carrying the provider's own charge book that USD; only the
723    /// remaining (unmeasured) turns are table-priced — and a fully measured
724    /// model is never flagged as estimated, even when the table has no entry.
725    /// (Model name chosen to never hit the embedded or live price tables.)
726    #[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        // A second, unmeasured turn on the same model: its tokens are priced
749        // from the table ON TOP of the measured USD, never double-counted
750        // (10k tokens at the $2.50/M blended fallback ≈ $0.025).
751        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        // proxy_usage.json written before #895 has no `cohorts` key; serde(default)
828        // must backfill an empty map so old files stay loadable.
829        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}