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    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    // Mechanism attribution into the local savings ledger (enterprise#19).
324    // Routing: the gateway served a cheaper model than requested — value the
325    // rate delta on the measured input tokens. Caching: provider prompt-cache
326    // reads billed below the input rate. Both best-effort, never blocking.
327    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/// Live output-savings cohort totals (#895). Empty until a holdout runs.
368#[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/// Cross-process read of the persisted output-savings cohort totals.
377#[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
391/// Live per-model measured spend, priced and sorted by USD descending.
392pub 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
399/// Total measured spend across all models (live in-memory totals).
400pub fn total_cost_usd() -> f64 {
401    snapshot().iter().map(|m| m.cost_usd).sum()
402}
403
404/// Cross-model verified-savings totals (#701) for `/status` and the dashboard.
405#[derive(Debug, Clone, Serialize, PartialEq)]
406pub struct VerifiedSavings {
407    /// Requests covered by a successful `count_tokens` probe.
408    pub requests: u64,
409    /// Provider-counted input tokens those requests would have billed
410    /// without lean-ctx.
411    pub counterfactual_input_tokens: u64,
412    /// Input-side tokens (input + cache read + cache write) they actually
413    /// billed.
414    pub billed_input_tokens: u64,
415    /// `counterfactual - billed`; negative when stub overhead outweighed the
416    /// squeeze — reported honestly, never clamped.
417    pub verified_saved_tokens: i64,
418}
419
420/// Aggregated provider-verified savings across all models (#701), or `None`
421/// until at least one probe-covered request has been recorded. Unlike the
422/// `tokens_saved` estimate (bytes/4), both sides of this pair were counted by
423/// the provider on the same request — receipts, not estimates.
424pub 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/// Pure aggregation behind [`verified_savings`], shared with tests.
432#[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
448/// Prices a model usage map into sorted [`ModelSpend`] rows. Pure: shared by the
449/// in-memory snapshot and the cross-process [`persisted_snapshot`].
450pub 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    // Measured-first (#1179): turns whose response carried the provider's own
467    // USD charge book that exact amount; only the remainder is table-priced.
468    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        // Fully measured spend is never an estimate, whatever the table says.
490        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
500/// Atomically writes the current in-memory totals to disk.
501fn 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
534/// Cross-process read of the persisted measured spend (dashboard / CLI / ledger).
535pub 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
541/// Cross-process priced spend rows, read from disk.
542pub fn persisted_snapshot() -> Vec<ModelSpend> {
543    load_persisted()
544        .map(|p| price_models(&p.models))
545        .unwrap_or_default()
546}
547
548/// The model carrying the most measured tokens (excludes the "unknown" bucket).
549/// Used to value savings against the real dominant model when no explicit model
550/// is configured.
551pub 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    /// #701: the verified pair is recorded only when the probe answered — both
582    /// sides of the pair or neither, never a half row.
583    #[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        // billed side = input + cache read + cache write of the same turn.
600        assert_eq!(acc.counterfactual_billed_tokens, 1_000 + 300 + 200);
601
602        // Empty slot (probe failed / still in flight) → row degrades to the
603        // estimate: no pair recorded, normal usage still counted.
604        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        // No wire context at all (tests / non-forward paths) → no pair.
616        acc.add(&usage("claude-sonnet-4.5", 10, 0, 0));
617        assert_eq!(acc.counterfactual_requests, 1);
618    }
619
620    /// #701: cross-model aggregation and the honest signed difference.
621    #[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        // Stub overhead outweighed the squeeze on this model: billed MORE
636        // than the counterfactual. The total must subtract honestly.
637        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    /// #701: persisted usage files from before the feature load cleanly and
666    /// the new fields round-trip.
667    #[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        // input 3.00 + output 15.00 + cache_read 0.30 (per 1M) = 18.30.
700        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    /// #1179: turns carrying the provider's own charge book that USD; only the
722    /// remaining (unmeasured) turns are table-priced — and a fully measured
723    /// model is never flagged as estimated, even when the table has no entry.
724    /// (Model name chosen to never hit the embedded or live price tables.)
725    #[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        // A second, unmeasured turn on the same model: its tokens are priced
748        // from the table ON TOP of the measured USD, never double-counted
749        // (10k tokens at the $2.50/M blended fallback ≈ $0.025).
750        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        // proxy_usage.json written before #895 has no `cohorts` key; serde(default)
827        // must backfill an empty map so old files stay loadable.
828        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}