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;
20
21/// Cumulative real token counts for one model. Cost is derived at read time so a
22/// pricing-table change re-values historical usage consistently.
23#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
24pub struct ModelUsage {
25    pub requests: u64,
26    pub input_tokens: u64,
27    pub output_tokens: u64,
28    pub cache_read_tokens: u64,
29    pub cache_write_tokens: u64,
30    pub reasoning_tokens: u64,
31    /// Requests whose free `count_tokens` probe answered (#701) — the rows the
32    /// verified-savings pair below covers. `serde(default)` keeps pre-#701
33    /// usage files loadable.
34    #[serde(default)]
35    pub counterfactual_requests: u64,
36    /// Provider-counted input tokens the covered requests would have billed
37    /// WITHOUT lean-ctx (sum of probe answers on the original bodies).
38    #[serde(default)]
39    pub counterfactual_input_tokens: u64,
40    /// Input-side tokens those same requests actually billed (input +
41    /// cache read + cache write) — same request, same moment, no confound.
42    #[serde(default)]
43    pub counterfactual_billed_tokens: u64,
44    /// The turns whose response carried the provider's own USD charge
45    /// (OpenRouter usage accounting, #1179). Their tokens are *inside* the
46    /// totals above; pricing subtracts this slice and books its measured USD
47    /// instead of a table estimate. `serde(default)` keeps older files loadable.
48    #[serde(default)]
49    pub measured: MeasuredSlice,
50}
51
52/// Token/cost sums of the turns that reported a measured provider charge.
53#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq)]
54pub struct MeasuredSlice {
55    pub requests: u64,
56    pub input_tokens: u64,
57    pub output_tokens: u64,
58    pub cache_read_tokens: u64,
59    pub cache_write_tokens: u64,
60    /// Sum of provider-reported USD for those turns — the bill, not a table.
61    pub cost_usd: f64,
62}
63
64impl MeasuredSlice {
65    fn merge(&mut self, other: &Self) {
66        self.requests += other.requests;
67        self.input_tokens += other.input_tokens;
68        self.output_tokens += other.output_tokens;
69        self.cache_read_tokens += other.cache_read_tokens;
70        self.cache_write_tokens += other.cache_write_tokens;
71        self.cost_usd += other.cost_usd;
72    }
73}
74
75impl ModelUsage {
76    fn add(&mut self, u: &super::usage::RealUsage) {
77        self.requests += 1;
78        self.input_tokens += u.input_tokens;
79        self.output_tokens += u.output_tokens;
80        self.cache_read_tokens += u.cache_read_tokens;
81        self.cache_write_tokens += u.cache_write_tokens;
82        self.reasoning_tokens += u.reasoning_tokens;
83        if let Some(cost) = u.provider_cost_usd {
84            self.measured.merge(&MeasuredSlice {
85                requests: 1,
86                input_tokens: u.input_tokens,
87                output_tokens: u.output_tokens,
88                cache_read_tokens: u.cache_read_tokens,
89                cache_write_tokens: u.cache_write_tokens,
90                cost_usd: cost,
91            });
92        }
93        // Verified-savings pair (#701): only when the probe answered by the
94        // time the billed usage arrived — both sides of the pair or neither.
95        if let Some(counted) = u
96            .wire
97            .as_deref()
98            .and_then(|w| w.counterfactual.as_ref())
99            .and_then(super::counterfactual::CounterfactualSlot::get)
100        {
101            self.counterfactual_requests += 1;
102            self.counterfactual_input_tokens += counted;
103            self.counterfactual_billed_tokens +=
104                u.input_tokens + u.cache_read_tokens + u.cache_write_tokens;
105        }
106    }
107
108    fn billable_tokens(&self) -> u64 {
109        self.input_tokens + self.output_tokens + self.cache_read_tokens + self.cache_write_tokens
110    }
111}
112
113/// One model's measured, priced spend for `/status` and the dashboard.
114#[derive(Debug, Clone, Serialize, PartialEq)]
115pub struct ModelSpend {
116    pub model: String,
117    pub requests: u64,
118    pub input_tokens: u64,
119    pub output_tokens: u64,
120    pub cache_read_tokens: u64,
121    pub cache_write_tokens: u64,
122    pub reasoning_tokens: u64,
123    pub cost_usd: f64,
124    /// Share of `cost_usd` the provider itself reported (OpenRouter usage
125    /// accounting, #1179) — a bill, not a table estimate.
126    pub measured_cost_usd: f64,
127    /// Requests covered by that provider-reported cost.
128    pub measured_requests: u64,
129    /// True when part of the cost had to be priced from a heuristic/fallback
130    /// match — never set by measured or exactly-priced spend.
131    pub pricing_estimated: bool,
132}
133
134/// Cumulative output-savings cohort totals (#895 Track B). Keyed by arm name
135/// (`"control"` | `"treatment"`); the average output tokens per turn is
136/// `output_tokens / requests`. Only populated while a holdout is active.
137#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
138pub struct CohortUsage {
139    pub requests: u64,
140    pub input_tokens: u64,
141    pub output_tokens: u64,
142    /// Sum of squared per-turn output tokens, enabling an online sample variance
143    /// (and therefore a confidence interval) without retaining every turn.
144    /// `#[serde(default)]` keeps pre-#895 files loadable.
145    #[serde(default)]
146    pub sum_sq_output: u64,
147}
148
149impl CohortUsage {
150    fn add(&mut self, u: &super::usage::RealUsage) {
151        self.requests += 1;
152        self.input_tokens += u.input_tokens;
153        self.output_tokens += u.output_tokens;
154        self.sum_sq_output += u.output_tokens.saturating_mul(u.output_tokens);
155    }
156
157    /// Average output tokens per turn, or `None` with no observations.
158    #[must_use]
159    pub fn avg_output(&self) -> Option<f64> {
160        if self.requests == 0 {
161            None
162        } else {
163            #[allow(clippy::cast_precision_loss)]
164            Some(self.output_tokens as f64 / self.requests as f64)
165        }
166    }
167
168    /// Unbiased sample variance of per-turn output tokens, or `None` with < 2
169    /// observations. Computed from the running sum / sum-of-squares (clamped at
170    /// 0 to absorb floating-point error on near-constant samples).
171    #[must_use]
172    pub fn variance_output(&self) -> Option<f64> {
173        if self.requests < 2 {
174            return None;
175        }
176        #[allow(clippy::cast_precision_loss)]
177        let n = self.requests as f64;
178        #[allow(clippy::cast_precision_loss)]
179        let sum = self.output_tokens as f64;
180        #[allow(clippy::cast_precision_loss)]
181        let sum_sq = self.sum_sq_output as f64;
182        let var = (sum_sq - sum * sum / n) / (n - 1.0);
183        Some(var.max(0.0))
184    }
185}
186
187/// On-disk shape of the measured spend totals.
188#[derive(Debug, Clone, Default, Serialize, Deserialize)]
189pub struct PersistedUsage {
190    pub ts: u64,
191    pub models: HashMap<String, ModelUsage>,
192    /// Output-savings cohort totals (#895). `#[serde(default)]` keeps older
193    /// `proxy_usage.json` files (written before the holdout existed) loadable.
194    #[serde(default)]
195    pub cohorts: HashMap<String, CohortUsage>,
196}
197
198/// Distinct-model bucket cap. `record` keys on the raw response model string, so
199/// overflow folds into "unknown" to keep the map bounded (real model names < ~50).
200const MAX_TRACKED_MODELS: usize = 256;
201
202const PROXY_USAGE_FILE: &str = "proxy_usage.json";
203
204fn store() -> &'static Mutex<HashMap<String, ModelUsage>> {
205    static STORE: OnceLock<Mutex<HashMap<String, ModelUsage>>> = OnceLock::new();
206    STORE.get_or_init(|| Mutex::new(HashMap::new()))
207}
208
209fn cohort_store() -> &'static Mutex<HashMap<String, CohortUsage>> {
210    static STORE: OnceLock<Mutex<HashMap<String, CohortUsage>>> = OnceLock::new();
211    STORE.get_or_init(|| Mutex::new(HashMap::new()))
212}
213
214/// Seeds the in-memory totals from `proxy_usage.json`. Call once on proxy
215/// startup so measured spend is cumulative across restarts. Idempotent-ish: it
216/// merges the persisted totals into whatever is in memory (normally empty).
217pub fn resume_from_disk() {
218    let Some(persisted) = load_persisted() else {
219        return;
220    };
221    let mut map = store()
222        .lock()
223        .unwrap_or_else(std::sync::PoisonError::into_inner);
224    for (model, usage) in persisted.models {
225        let acc = map.entry(model).or_default();
226        acc.requests += usage.requests;
227        acc.input_tokens += usage.input_tokens;
228        acc.output_tokens += usage.output_tokens;
229        acc.cache_read_tokens += usage.cache_read_tokens;
230        acc.cache_write_tokens += usage.cache_write_tokens;
231        acc.reasoning_tokens += usage.reasoning_tokens;
232        acc.counterfactual_requests += usage.counterfactual_requests;
233        acc.counterfactual_input_tokens += usage.counterfactual_input_tokens;
234        acc.counterfactual_billed_tokens += usage.counterfactual_billed_tokens;
235        acc.measured.merge(&usage.measured);
236    }
237    drop(map);
238    let mut cohorts = cohort_store()
239        .lock()
240        .unwrap_or_else(std::sync::PoisonError::into_inner);
241    for (arm, usage) in persisted.cohorts {
242        let acc = cohorts.entry(arm).or_default();
243        acc.requests += usage.requests;
244        acc.input_tokens += usage.input_tokens;
245        acc.output_tokens += usage.output_tokens;
246    }
247}
248
249/// Records one turn's measured usage against its model bucket (and its
250/// output-savings cohort, when tagged) and persists.
251pub fn record(u: &super::usage::RealUsage) {
252    // Gateway store subscription (enterprise#17): forward the full record to
253    // the installed sink (no-op locally). Never blocks the request path.
254    super::usage_sink::push(u);
255
256    // Budget windows (enterprise#25): book this turn's measured cost against
257    // the person/day and project/month accumulators the policy gate checks.
258    // The provider's own reported charge wins (#1179); local turns book the
259    // shadow rate — the same valuation the usage store applies — so
260    // local-only budgets stay meaningful.
261    if let Some(wire) = u.wire.as_deref()
262        && (wire.person.is_some() || wire.project.is_some())
263    {
264        let baseline = crate::core::config::Config::load().proxy.baseline.clone();
265        #[allow(clippy::cast_precision_loss)]
266        let cost_usd = if wire.is_local {
267            let billable =
268                u.input_tokens + u.output_tokens + u.cache_read_tokens + u.cache_write_tokens;
269            baseline.effective_local_shadow_rate() / 1_000_000.0 * billable as f64
270        } else if let Some(measured) = u.provider_cost_usd {
271            measured
272        } else {
273            crate::core::gain::model_pricing::ModelPricing::load()
274                .quote(Some(&u.model))
275                .cost
276                .estimate_usd(
277                    u.input_tokens,
278                    u.output_tokens,
279                    u.cache_write_tokens,
280                    u.cache_read_tokens,
281                )
282        };
283        super::policy_gate::record_spend(wire.person.as_deref(), wire.project.as_deref(), cost_usd);
284    }
285
286    // Mechanism attribution into the local savings ledger (enterprise#19).
287    // Routing: the gateway served a cheaper model than requested — value the
288    // rate delta on the measured input tokens. Caching: provider prompt-cache
289    // reads billed below the input rate. Both best-effort, never blocking.
290    if let Some(wire) = u.wire.as_deref()
291        && let Some(routed_from) = wire.routed_from.as_deref()
292    {
293        crate::core::savings_ledger::record_routing_event(routed_from, &u.model, u.input_tokens);
294    }
295    if u.cache_read_tokens > 0 {
296        let cost = crate::core::gain::model_pricing::ModelPricing::load()
297            .quote(Some(&u.model))
298            .cost;
299        #[allow(clippy::cast_precision_loss)]
300        let discount_usd =
301            (cost.input_per_m - cost.cache_read_per_m) / 1_000_000.0 * u.cache_read_tokens as f64;
302        crate::core::savings_ledger::record_caching_event(
303            &u.model,
304            u.cache_read_tokens,
305            discount_usd,
306        );
307    }
308
309    let key = normalize_key(&u.model);
310    {
311        let mut map = store()
312            .lock()
313            .unwrap_or_else(std::sync::PoisonError::into_inner);
314        let key = if !map.contains_key(&key) && map.len() >= MAX_TRACKED_MODELS {
315            "unknown".to_string()
316        } else {
317            key
318        };
319        map.entry(key).or_default().add(u);
320    }
321    if let Some(arm) = u.cohort {
322        let mut cohorts = cohort_store()
323            .lock()
324            .unwrap_or_else(std::sync::PoisonError::into_inner);
325        cohorts.entry(arm.as_str().to_string()).or_default().add(u);
326    }
327    persist();
328}
329
330/// Live output-savings cohort totals (#895). Empty until a holdout runs.
331#[must_use]
332pub fn cohort_snapshot() -> HashMap<String, CohortUsage> {
333    cohort_store()
334        .lock()
335        .unwrap_or_else(std::sync::PoisonError::into_inner)
336        .clone()
337}
338
339/// Cross-process read of the persisted output-savings cohort totals.
340#[must_use]
341pub fn persisted_cohorts() -> HashMap<String, CohortUsage> {
342    load_persisted().map(|p| p.cohorts).unwrap_or_default()
343}
344
345fn normalize_key(model: &str) -> String {
346    let m = model.trim();
347    if m.is_empty() {
348        "unknown".to_string()
349    } else {
350        m.to_string()
351    }
352}
353
354/// Live per-model measured spend, priced and sorted by USD descending.
355pub fn snapshot() -> Vec<ModelSpend> {
356    let map = store()
357        .lock()
358        .unwrap_or_else(std::sync::PoisonError::into_inner);
359    price_models(&map)
360}
361
362/// Total measured spend across all models (live in-memory totals).
363pub fn total_cost_usd() -> f64 {
364    snapshot().iter().map(|m| m.cost_usd).sum()
365}
366
367/// Cross-model verified-savings totals (#701) for `/status` and the dashboard.
368#[derive(Debug, Clone, Serialize, PartialEq)]
369pub struct VerifiedSavings {
370    /// Requests covered by a successful `count_tokens` probe.
371    pub requests: u64,
372    /// Provider-counted input tokens those requests would have billed
373    /// without lean-ctx.
374    pub counterfactual_input_tokens: u64,
375    /// Input-side tokens (input + cache read + cache write) they actually
376    /// billed.
377    pub billed_input_tokens: u64,
378    /// `counterfactual - billed`; negative when stub overhead outweighed the
379    /// squeeze — reported honestly, never clamped.
380    pub verified_saved_tokens: i64,
381}
382
383/// Aggregated provider-verified savings across all models (#701), or `None`
384/// until at least one probe-covered request has been recorded. Unlike the
385/// `tokens_saved` estimate (bytes/4), both sides of this pair were counted by
386/// the provider on the same request — receipts, not estimates.
387pub fn verified_savings() -> Option<VerifiedSavings> {
388    let map = store()
389        .lock()
390        .unwrap_or_else(std::sync::PoisonError::into_inner);
391    verified_of(&map)
392}
393
394/// Pure aggregation behind [`verified_savings`], shared with tests.
395#[allow(clippy::cast_possible_wrap)]
396fn verified_of(map: &HashMap<String, ModelUsage>) -> Option<VerifiedSavings> {
397    let (mut requests, mut counterfactual, mut billed) = (0u64, 0u64, 0u64);
398    for usage in map.values() {
399        requests += usage.counterfactual_requests;
400        counterfactual += usage.counterfactual_input_tokens;
401        billed += usage.counterfactual_billed_tokens;
402    }
403    (requests > 0).then(|| VerifiedSavings {
404        requests,
405        counterfactual_input_tokens: counterfactual,
406        billed_input_tokens: billed,
407        verified_saved_tokens: counterfactual as i64 - billed as i64,
408    })
409}
410
411/// Prices a model usage map into sorted [`ModelSpend`] rows. Pure: shared by the
412/// in-memory snapshot and the cross-process [`persisted_snapshot`].
413pub fn price_models(map: &HashMap<String, ModelUsage>) -> Vec<ModelSpend> {
414    let pricing = ModelPricing::load();
415    let mut rows: Vec<ModelSpend> = map
416        .iter()
417        .map(|(model, usage)| price_one(&pricing, model, usage))
418        .collect();
419    rows.sort_by(|a, b| {
420        b.cost_usd
421            .partial_cmp(&a.cost_usd)
422            .unwrap_or(std::cmp::Ordering::Equal)
423    });
424    rows
425}
426
427fn price_one(pricing: &ModelPricing, model: &str, usage: &ModelUsage) -> ModelSpend {
428    let quote = pricing.quote(Some(model));
429    // Measured-first (#1179): turns whose response carried the provider's own
430    // USD charge book that exact amount; only the remainder is table-priced.
431    let m = &usage.measured;
432    let derived_cost = quote.cost.estimate_usd(
433        usage.input_tokens.saturating_sub(m.input_tokens),
434        usage.output_tokens.saturating_sub(m.output_tokens),
435        usage
436            .cache_write_tokens
437            .saturating_sub(m.cache_write_tokens),
438        usage.cache_read_tokens.saturating_sub(m.cache_read_tokens),
439    );
440    let derived_requests = usage.requests.saturating_sub(m.requests);
441    ModelSpend {
442        model: model.to_string(),
443        requests: usage.requests,
444        input_tokens: usage.input_tokens,
445        output_tokens: usage.output_tokens,
446        cache_read_tokens: usage.cache_read_tokens,
447        cache_write_tokens: usage.cache_write_tokens,
448        reasoning_tokens: usage.reasoning_tokens,
449        cost_usd: m.cost_usd + derived_cost,
450        measured_cost_usd: m.cost_usd,
451        measured_requests: m.requests,
452        // Fully measured spend is never an estimate, whatever the table says.
453        pricing_estimated: quote.match_kind.is_estimated() && derived_requests > 0,
454    }
455}
456
457fn usage_path() -> Option<std::path::PathBuf> {
458    crate::core::data_dir::lean_ctx_data_dir()
459        .ok()
460        .map(|d| d.join(PROXY_USAGE_FILE))
461}
462
463/// Atomically writes the current in-memory totals to disk.
464fn persist() {
465    let Some(path) = usage_path() else {
466        return;
467    };
468    let models = {
469        let map = store()
470            .lock()
471            .unwrap_or_else(std::sync::PoisonError::into_inner);
472        map.clone()
473    };
474    let cohorts = {
475        let map = cohort_store()
476            .lock()
477            .unwrap_or_else(std::sync::PoisonError::into_inner);
478        map.clone()
479    };
480    let payload = PersistedUsage {
481        ts: std::time::SystemTime::now()
482            .duration_since(std::time::UNIX_EPOCH)
483            .unwrap_or_default()
484            .as_secs(),
485        models,
486        cohorts,
487    };
488    let Ok(json) = serde_json::to_string(&payload) else {
489        return;
490    };
491    let tmp = path.with_extension("json.tmp");
492    if std::fs::write(&tmp, json).is_ok() {
493        let _ = std::fs::rename(&tmp, &path);
494    }
495}
496
497/// Cross-process read of the persisted measured spend (dashboard / CLI / ledger).
498pub fn load_persisted() -> Option<PersistedUsage> {
499    let path = usage_path()?;
500    let data = std::fs::read_to_string(path).ok()?;
501    serde_json::from_str(&data).ok()
502}
503
504/// Cross-process priced spend rows, read from disk.
505pub fn persisted_snapshot() -> Vec<ModelSpend> {
506    load_persisted()
507        .map(|p| price_models(&p.models))
508        .unwrap_or_default()
509}
510
511/// The model carrying the most measured tokens (excludes the "unknown" bucket).
512/// Used to value savings against the real dominant model when no explicit model
513/// is configured.
514pub fn persisted_dominant_model() -> Option<String> {
515    let persisted = load_persisted()?;
516    persisted
517        .models
518        .iter()
519        .filter(|(m, _)| m.as_str() != "unknown" && !m.trim().is_empty())
520        .max_by_key(|(_, u)| u.billable_tokens())
521        .filter(|(_, u)| u.billable_tokens() > 0)
522        .map(|(m, _)| m.clone())
523}
524
525#[cfg(test)]
526mod tests {
527    use super::*;
528
529    fn usage(
530        model: &str,
531        input: u64,
532        output: u64,
533        cache_read: u64,
534    ) -> super::super::usage::RealUsage {
535        super::super::usage::RealUsage {
536            model: model.to_string(),
537            input_tokens: input,
538            output_tokens: output,
539            cache_read_tokens: cache_read,
540            ..Default::default()
541        }
542    }
543
544    /// #701: the verified pair is recorded only when the probe answered — both
545    /// sides of the pair or neither, never a half row.
546    #[test]
547    fn counterfactual_pair_recorded_only_when_probe_answered() {
548        let slot = super::super::counterfactual::CounterfactualSlot::new();
549        slot.set(5_000);
550        let with_probe = super::super::usage::RealUsage {
551            wire: Some(Box::new(super::super::usage::WireContext {
552                counterfactual: Some(slot),
553                ..Default::default()
554            })),
555            cache_write_tokens: 200,
556            ..usage("claude-sonnet-4.5", 1_000, 0, 300)
557        };
558        let mut acc = ModelUsage::default();
559        acc.add(&with_probe);
560        assert_eq!(acc.counterfactual_requests, 1);
561        assert_eq!(acc.counterfactual_input_tokens, 5_000);
562        // billed side = input + cache read + cache write of the same turn.
563        assert_eq!(acc.counterfactual_billed_tokens, 1_000 + 300 + 200);
564
565        // Empty slot (probe failed / still in flight) → row degrades to the
566        // estimate: no pair recorded, normal usage still counted.
567        let empty = super::super::usage::RealUsage {
568            wire: Some(Box::new(super::super::usage::WireContext {
569                counterfactual: Some(super::super::counterfactual::CounterfactualSlot::new()),
570                ..Default::default()
571            })),
572            ..usage("claude-sonnet-4.5", 1_000, 0, 0)
573        };
574        acc.add(&empty);
575        assert_eq!(acc.requests, 2);
576        assert_eq!(acc.counterfactual_requests, 1, "empty slot adds no pair");
577
578        // No wire context at all (tests / non-forward paths) → no pair.
579        acc.add(&usage("claude-sonnet-4.5", 10, 0, 0));
580        assert_eq!(acc.counterfactual_requests, 1);
581    }
582
583    /// #701: cross-model aggregation and the honest signed difference.
584    #[test]
585    fn verified_of_aggregates_and_reports_signed_savings() {
586        let mut map = HashMap::new();
587        assert!(verified_of(&map).is_none(), "no coverage → None, not zeros");
588
589        map.insert(
590            "claude-sonnet-4.5".to_string(),
591            ModelUsage {
592                counterfactual_requests: 2,
593                counterfactual_input_tokens: 10_000,
594                counterfactual_billed_tokens: 6_000,
595                ..Default::default()
596            },
597        );
598        // Stub overhead outweighed the squeeze on this model: billed MORE
599        // than the counterfactual. The total must subtract honestly.
600        map.insert(
601            "claude-haiku-4.5".to_string(),
602            ModelUsage {
603                counterfactual_requests: 1,
604                counterfactual_input_tokens: 1_000,
605                counterfactual_billed_tokens: 1_400,
606                ..Default::default()
607            },
608        );
609        let v = verified_of(&map).expect("covered rows present");
610        assert_eq!(v.requests, 3);
611        assert_eq!(v.counterfactual_input_tokens, 11_000);
612        assert_eq!(v.billed_input_tokens, 7_400);
613        assert_eq!(v.verified_saved_tokens, 3_600);
614
615        map.get_mut("claude-sonnet-4.5")
616            .unwrap()
617            .counterfactual_input_tokens = 0;
618        map.get_mut("claude-sonnet-4.5")
619            .unwrap()
620            .counterfactual_billed_tokens = 0;
621        let negative = verified_of(&map).unwrap();
622        assert_eq!(
623            negative.verified_saved_tokens, -400,
624            "a net-negative verified saving is reported, never clamped"
625        );
626    }
627
628    /// #701: persisted usage files from before the feature load cleanly and
629    /// the new fields round-trip.
630    #[test]
631    fn counterfactual_fields_roundtrip_and_default_for_legacy_files() {
632        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}}}"#;
633        let p: PersistedUsage = serde_json::from_str(legacy).expect("legacy file loads");
634        assert_eq!(p.models["m"].counterfactual_requests, 0);
635
636        let mut p = PersistedUsage::default();
637        p.models.insert(
638            "m".into(),
639            ModelUsage {
640                counterfactual_requests: 4,
641                counterfactual_input_tokens: 9_999,
642                counterfactual_billed_tokens: 5_555,
643                ..Default::default()
644            },
645        );
646        let back: PersistedUsage =
647            serde_json::from_str(&serde_json::to_string(&p).unwrap()).unwrap();
648        assert_eq!(back.models["m"].counterfactual_input_tokens, 9_999);
649        assert_eq!(back.models["m"].counterfactual_billed_tokens, 5_555);
650    }
651
652    #[test]
653    fn prices_known_model_with_cache_split() {
654        let mut map = HashMap::new();
655        let mut acc = ModelUsage::default();
656        acc.add(&usage("claude-sonnet-4.5", 1_000_000, 1_000_000, 1_000_000));
657        map.insert("claude-sonnet-4.5".to_string(), acc);
658
659        let rows = price_models(&map);
660        assert_eq!(rows.len(), 1);
661        let row = &rows[0];
662        // input 3.00 + output 15.00 + cache_read 0.30 (per 1M) = 18.30.
663        assert!(
664            (row.cost_usd - 18.30).abs() < 1e-6,
665            "cost was {}",
666            row.cost_usd
667        );
668        assert!(!row.pricing_estimated, "exact model match");
669        assert_eq!(row.requests, 1);
670    }
671
672    #[test]
673    fn unknown_model_prices_with_fallback_and_is_estimated() {
674        let mut map = HashMap::new();
675        let mut acc = ModelUsage::default();
676        acc.add(&usage("some-novel-model-xyz", 1_000_000, 0, 0));
677        map.insert("some-novel-model-xyz".to_string(), acc);
678
679        let rows = price_models(&map);
680        assert!(rows[0].pricing_estimated, "fallback pricing is estimated");
681        assert!(rows[0].cost_usd > 0.0);
682    }
683
684    /// #1179: turns carrying the provider's own charge book that USD; only the
685    /// remaining (unmeasured) turns are table-priced — and a fully measured
686    /// model is never flagged as estimated, even when the table has no entry.
687    /// (Model name chosen to never hit the embedded or live price tables.)
688    #[test]
689    fn measured_provider_cost_replaces_table_estimate() {
690        const MODEL: &str = "vendor/unlisted-model-20990101";
691        let mut acc = ModelUsage::default();
692        let mut measured = usage(MODEL, 431_600, 22_700, 126_700);
693        measured.provider_cost_usd = Some(0.05);
694        acc.add(&measured);
695
696        let mut map = HashMap::new();
697        map.insert(MODEL.to_string(), acc.clone());
698        let row = &price_models(&map)[0];
699        assert!(
700            (row.cost_usd - 0.05).abs() < 1e-12,
701            "the bill, not the table"
702        );
703        assert!((row.measured_cost_usd - 0.05).abs() < 1e-12);
704        assert_eq!(row.measured_requests, 1);
705        assert!(
706            !row.pricing_estimated,
707            "fully measured spend is not an estimate"
708        );
709
710        // A second, unmeasured turn on the same model: its tokens are priced
711        // from the table ON TOP of the measured USD, never double-counted
712        // (10k tokens at the $2.50/M blended fallback ≈ $0.025).
713        acc.add(&usage(MODEL, 10_000, 0, 0));
714        let mut map = HashMap::new();
715        map.insert(MODEL.to_string(), acc);
716        let row = &price_models(&map)[0];
717        assert!(row.cost_usd > 0.05, "unmeasured remainder adds table cost");
718        assert!(row.cost_usd < 0.2, "measured slice must not be re-priced");
719        assert!(
720            row.pricing_estimated,
721            "unmeasured remainder is heuristically priced"
722        );
723    }
724
725    #[test]
726    fn measured_slice_roundtrips_and_defaults_for_legacy_files() {
727        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}}}"#;
728        let p: PersistedUsage = serde_json::from_str(legacy).expect("legacy file loads");
729        assert_eq!(p.models["m"].measured, MeasuredSlice::default());
730
731        let mut p = PersistedUsage::default();
732        let mut acc = ModelUsage::default();
733        let mut m = usage("m", 100, 10, 0);
734        m.provider_cost_usd = Some(0.5);
735        acc.add(&m);
736        p.models.insert("m".into(), acc);
737        let back: PersistedUsage =
738            serde_json::from_str(&serde_json::to_string(&p).unwrap()).unwrap();
739        assert_eq!(back.models["m"].measured.requests, 1);
740        assert!((back.models["m"].measured.cost_usd - 0.5).abs() < 1e-12);
741    }
742
743    #[test]
744    fn dominant_model_picks_highest_token_real_model() {
745        let mut models = HashMap::new();
746        models.insert("claude-haiku-4.5".to_string(), {
747            let mut u = ModelUsage::default();
748            u.add(&usage("claude-haiku-4.5", 100, 100, 0));
749            u
750        });
751        models.insert("claude-opus-4.5".to_string(), {
752            let mut u = ModelUsage::default();
753            u.add(&usage("claude-opus-4.5", 10_000, 10_000, 0));
754            u
755        });
756        models.insert("unknown".to_string(), {
757            let mut u = ModelUsage::default();
758            u.add(&usage("unknown", 999_999, 0, 0));
759            u
760        });
761        let dominant = models
762            .iter()
763            .filter(|(m, _)| m.as_str() != "unknown")
764            .max_by_key(|(_, u)| u.billable_tokens())
765            .map(|(m, _)| m.clone());
766        assert_eq!(dominant.as_deref(), Some("claude-opus-4.5"));
767    }
768
769    #[test]
770    fn empty_model_buckets_as_unknown() {
771        assert_eq!(normalize_key("  "), "unknown");
772        assert_eq!(normalize_key(""), "unknown");
773        assert_eq!(normalize_key("gpt-5.4"), "gpt-5.4");
774    }
775
776    #[test]
777    fn cohort_avg_output_is_mean_per_turn() {
778        let mut c = CohortUsage::default();
779        assert_eq!(c.avg_output(), None, "no observations → None");
780        c.add(&usage("m", 10, 100, 0));
781        c.add(&usage("m", 10, 50, 0));
782        assert_eq!(c.requests, 2);
783        assert_eq!(c.output_tokens, 150);
784        assert!((c.avg_output().unwrap() - 75.0).abs() < f64::EPSILON);
785    }
786
787    #[test]
788    fn persisted_usage_without_cohorts_field_loads() {
789        // proxy_usage.json written before #895 has no `cohorts` key; serde(default)
790        // must backfill an empty map so old files stay loadable.
791        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}}}"#;
792        let p: PersistedUsage = serde_json::from_str(json).expect("loads legacy file");
793        assert_eq!(p.models.len(), 1);
794        assert!(p.cohorts.is_empty());
795    }
796
797    #[test]
798    fn persisted_usage_roundtrips_cohorts() {
799        let mut p = PersistedUsage::default();
800        p.cohorts.insert(
801            "control".into(),
802            CohortUsage {
803                requests: 3,
804                input_tokens: 30,
805                output_tokens: 300,
806                sum_sq_output: 30_000,
807            },
808        );
809        let json = serde_json::to_string(&p).unwrap();
810        let back: PersistedUsage = serde_json::from_str(&json).unwrap();
811        assert_eq!(back.cohorts.get("control").unwrap().output_tokens, 300);
812    }
813}