Skip to main content

lean_ctx_proxy/
usage_meter.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Measured per-model spend meter.
3//!
4//! [`record`] aggregates the real provider usage extracted by [`super::usage`]
5//! into per-model token sums, prices them with the shared
6//! `ModelPricing` table, and
7//! persists the totals to `proxy_usage.json` so the dashboard, CLI and the
8//! savings ledger (which run in *other* processes) can read the user's real
9//! provider bill.
10//!
11//! Unlike [`lean_ctx::core::proxy_metrics`] (which resets per proxy lifetime), this meter is a
12//! lifetime-cumulative spend counter: [`resume_from_disk`] seeds the in-memory
13//! totals on proxy startup so a restart never zeroes the user's measured spend.
14
15use std::collections::HashMap;
16use std::sync::{Mutex, OnceLock};
17
18use serde::Serialize;
19
20use lean_ctx::core::gain::model_pricing::ModelPricing;
21use lean_ctx::core::ocla::{OclaRegistry, UsageSink};
22pub use lean_ctx::core::proxy_usage::{CohortUsage, load_persisted};
23use lean_ctx::core::proxy_usage::{
24    LastModel, MeasuredSlice, ModelUsage, PersistedUsage, usage_path,
25};
26
27/// Projects one managed, provider-measured turn into the OCLA usage sink.
28///
29/// This is deliberately one-way and best-effort: the existing meter remains
30/// authoritative for billing/ledger totals, while OCLA receives a validated
31/// payload-free projection only when complete lineage is present.
32fn project_ocla_usage(u: &super::usage::RealUsage, sink: &dyn UsageSink) {
33    use futures::FutureExt;
34    let Ok(Some(record)) = u.to_ocla_usage_record() else {
35        return;
36    };
37    let _ = sink.record_usage(record).now_or_never();
38}
39
40/// Adds one finalized turn to a running total. The totals themselves are the
41/// on-disk record of [`lean_ctx::core::proxy_usage`]; only the meter sees a turn.
42trait AddTurn {
43    fn add(&mut self, u: &super::usage::RealUsage);
44}
45
46impl AddTurn for ModelUsage {
47    fn add(&mut self, u: &super::usage::RealUsage) {
48        if let Some(observation) = &u.response_observation {
49            self.add_observation(observation.clone());
50        }
51        self.requests += 1;
52        self.input_tokens += u.input_tokens;
53        self.output_tokens += u.output_tokens;
54        self.cache_read_tokens += u.cache_read_tokens;
55        self.cache_write_tokens += u.cache_write_tokens;
56        self.reasoning_tokens += u.reasoning_tokens;
57        if let Some(cost) = u.provider_cost_usd {
58            self.measured.merge(&MeasuredSlice {
59                requests: 1,
60                input_tokens: u.input_tokens,
61                output_tokens: u.output_tokens,
62                cache_read_tokens: u.cache_read_tokens,
63                cache_write_tokens: u.cache_write_tokens,
64                cost_usd: cost,
65            });
66        }
67        // Verified-savings pair (#701): only when the probe answered by the
68        // time the billed usage arrived — both sides of the pair or neither.
69        if let Some(counted) = u
70            .wire
71            .as_deref()
72            .and_then(|w| w.counterfactual.as_ref())
73            .and_then(super::counterfactual::CounterfactualSlot::get)
74        {
75            self.counterfactual_requests += 1;
76            self.counterfactual_input_tokens += counted;
77            self.counterfactual_billed_tokens +=
78                u.input_tokens + u.cache_read_tokens + u.cache_write_tokens;
79        }
80    }
81}
82
83impl AddTurn for CohortUsage {
84    fn add(&mut self, u: &super::usage::RealUsage) {
85        self.requests += 1;
86        self.input_tokens += u.input_tokens;
87        self.output_tokens += u.output_tokens;
88        self.sum_sq_output += u.output_tokens.saturating_mul(u.output_tokens);
89        let prompt = u
90            .input_tokens
91            .saturating_add(u.cache_read_tokens)
92            .saturating_add(u.cache_write_tokens);
93        self.prompt_tokens += prompt;
94        self.sum_sq_prompt += prompt.saturating_mul(prompt);
95    }
96}
97
98/// One model's measured, priced spend for `/status` and the dashboard.
99#[derive(Debug, Clone, Serialize, PartialEq)]
100pub struct ModelSpend {
101    pub model: String,
102    pub requests: u64,
103    pub input_tokens: u64,
104    pub output_tokens: u64,
105    pub cache_read_tokens: u64,
106    pub cache_write_tokens: u64,
107    pub reasoning_tokens: u64,
108    pub cost_usd: f64,
109    /// Share of `cost_usd` the provider itself reported (OpenRouter usage
110    /// accounting, #1179) — a bill, not a table estimate.
111    pub measured_cost_usd: f64,
112    /// Requests covered by that provider-reported cost.
113    pub measured_requests: u64,
114    /// True when part of the cost had to be priced from a heuristic/fallback
115    /// match — never set by measured or exactly-priced spend.
116    pub pricing_estimated: bool,
117}
118
119/// Cohort-store key for an input-compression holdout arm (#1905). Prefixed so
120/// it never collides with the output-savings arms stored as `control` /
121/// `treatment`, which keeps existing files and readers unchanged.
122#[must_use]
123pub fn compression_cohort_key(arm: super::holdout::Arm) -> String {
124    format!("compression:{}", arm.as_str())
125}
126
127fn last_model_store() -> &'static Mutex<Option<LastModel>> {
128    static STORE: OnceLock<Mutex<Option<LastModel>>> = OnceLock::new();
129    STORE.get_or_init(|| Mutex::new(None))
130}
131
132/// Distinct-model bucket cap. `record` keys on the raw response model string, so
133/// overflow folds into "unknown" to keep the map bounded (real model names < ~50).
134const MAX_TRACKED_MODELS: usize = 256;
135
136fn store() -> &'static Mutex<HashMap<String, ModelUsage>> {
137    static STORE: OnceLock<Mutex<HashMap<String, ModelUsage>>> = OnceLock::new();
138    STORE.get_or_init(|| Mutex::new(HashMap::new()))
139}
140
141/// Distinct projects kept apart; turns of further projects stay in the device total.
142const MAX_TRACKED_PROJECTS: usize = 512;
143
144fn project_store() -> &'static Mutex<HashMap<String, HashMap<String, ModelUsage>>> {
145    static STORE: OnceLock<Mutex<HashMap<String, HashMap<String, ModelUsage>>>> = OnceLock::new();
146    STORE.get_or_init(|| Mutex::new(HashMap::new()))
147}
148
149fn cohort_store() -> &'static Mutex<HashMap<String, CohortUsage>> {
150    static STORE: OnceLock<Mutex<HashMap<String, CohortUsage>>> = OnceLock::new();
151    STORE.get_or_init(|| Mutex::new(HashMap::new()))
152}
153
154/// Seeds the in-memory totals from `proxy_usage.json`. Call once on proxy
155/// startup so measured spend is cumulative across restarts. Idempotent-ish: it
156/// merges the persisted totals into whatever is in memory (normally empty).
157pub fn resume_from_disk() {
158    let Some(persisted) = load_persisted() else {
159        return;
160    };
161    {
162        let mut last = last_model_store()
163            .lock()
164            .unwrap_or_else(std::sync::PoisonError::into_inner);
165        if last.is_none() {
166            last.clone_from(&persisted.last_model);
167        }
168    }
169    let mut map = store()
170        .lock()
171        .unwrap_or_else(std::sync::PoisonError::into_inner);
172    merge_models(&mut map, persisted.models);
173    drop(map);
174    let mut projects = project_store()
175        .lock()
176        .unwrap_or_else(std::sync::PoisonError::into_inner);
177    for (project, models) in persisted.projects {
178        merge_models(projects.entry(project).or_default(), models);
179    }
180    drop(projects);
181    let mut cohorts = cohort_store()
182        .lock()
183        .unwrap_or_else(std::sync::PoisonError::into_inner);
184    for (arm, usage) in persisted.cohorts {
185        cohorts.entry(arm).or_default().merge(&usage);
186    }
187}
188
189fn merge_models(map: &mut HashMap<String, ModelUsage>, models: HashMap<String, ModelUsage>) {
190    for (model, usage) in models {
191        let acc = map.entry(model).or_default();
192        acc.requests += usage.requests;
193        acc.input_tokens += usage.input_tokens;
194        acc.output_tokens += usage.output_tokens;
195        acc.cache_read_tokens += usage.cache_read_tokens;
196        acc.cache_write_tokens += usage.cache_write_tokens;
197        acc.reasoning_tokens += usage.reasoning_tokens;
198        acc.counterfactual_requests += usage.counterfactual_requests;
199        acc.counterfactual_input_tokens += usage.counterfactual_input_tokens;
200        acc.counterfactual_billed_tokens += usage.counterfactual_billed_tokens;
201        acc.measured.merge(&usage.measured);
202        for observation in usage.response_observations {
203            acc.add_observation(observation);
204        }
205    }
206}
207
208/// Records one turn's measured usage against its model bucket (and its
209/// output-savings cohort, when tagged) and persists.
210pub fn record(u: &super::usage::RealUsage) {
211    // Gateway store subscription (enterprise#17): forward the full record to
212    // the installed sink (no-op locally). Never blocks the request path.
213    super::usage_sink::push(u);
214
215    // R33: Feed measured usage into kernel envelope pipeline.
216    {
217        let label = u.wire.as_ref().map_or("Unknown", |w| {
218            // provider_kind_from_label expects a static or owned string;
219            // the label lives long enough for the conversion call.
220            match w.provider.as_str() {
221                "Anthropic" => "Anthropic",
222                "OpenAI" | "ChatGPT" => "OpenAI",
223                "Gemini" => "Gemini",
224                "Bedrock" => "Bedrock",
225                "Azure" => "Azure",
226                _ => "Unknown",
227            }
228        });
229        let envelope = super::usage_parity::real_usage_to_envelope(u, label);
230        lean_ctx::core::context_kernel::envelope_bridge::record_proxy_envelope(&envelope);
231    }
232
233    // OCLA projection: keep the canonical usage capability on the same single
234    // finalized-turn choke-point without feeding back into billing or ledger
235    // aggregation. Invalid/missing lineage stays explicitly unmanaged.
236    project_ocla_usage(u, OclaRegistry::global().usage_sink.as_ref());
237
238    // Budget windows (enterprise#25): book this turn's measured cost against
239    // the person/day and project/month accumulators the policy gate checks.
240    // The provider's own reported charge wins (#1179); local turns book the
241    // shadow rate — the same valuation the usage store applies — so
242    // local-only budgets stay meaningful.
243    #[cfg(feature = "enterprise")]
244    if let Some(wire) = u.wire.as_deref()
245        && (wire.person.is_some() || wire.project.is_some())
246    {
247        let baseline = lean_ctx::core::config::Config::load()
248            .proxy
249            .baseline
250            .clone();
251        #[allow(clippy::cast_precision_loss)]
252        let cost_usd = if wire.is_local {
253            let billable =
254                u.input_tokens + u.output_tokens + u.cache_read_tokens + u.cache_write_tokens;
255            baseline.effective_local_shadow_rate() / 1_000_000.0 * billable as f64
256        } else if let Some(measured) = u.provider_cost_usd {
257            measured
258        } else {
259            lean_ctx::core::gain::model_pricing::ModelPricing::load()
260                .quote(Some(&u.model))
261                .cost
262                .estimate_usd(
263                    u.input_tokens,
264                    u.output_tokens,
265                    u.cache_write_tokens,
266                    u.cache_read_tokens,
267                )
268        };
269        super::policy_gate::record_spend(wire.person.as_deref(), wire.project.as_deref(), cost_usd);
270    }
271
272    // Mechanism attribution into the local savings ledger (enterprise#19):
273    // provider prompt-cache reads billed below the input rate. Best-effort,
274    // never blocking. A model swap is never booked as a LeanCTX saving.
275    if u.cache_read_tokens > 0 {
276        let cost = lean_ctx::core::gain::model_pricing::ModelPricing::load()
277            .quote(Some(&u.model))
278            .cost;
279        #[allow(clippy::cast_precision_loss)]
280        let discount_usd =
281            (cost.input_per_m - cost.cache_read_per_m) / 1_000_000.0 * u.cache_read_tokens as f64;
282        lean_ctx::core::savings_ledger::record_caching_event(
283            &u.model,
284            u.cache_read_tokens,
285            discount_usd,
286            u.wire.as_deref().and_then(|w| w.project_key.as_deref()),
287        );
288    }
289
290    // Remember the foreground model for attribution in other processes
291    // (savings ledger). Background models never count as "the model in use".
292    let model = u.model.trim();
293    if !model.is_empty()
294        && model != "unknown"
295        && !lean_ctx::core::model_registry::is_background_model(model)
296    {
297        *last_model_store()
298            .lock()
299            .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(LastModel {
300            model: model.to_string(),
301            at: std::time::SystemTime::now()
302                .duration_since(std::time::UNIX_EPOCH)
303                .unwrap_or_default()
304                .as_secs(),
305        });
306    }
307
308    let key = normalize_key(&u.model);
309    {
310        let mut map = store()
311            .lock()
312            .unwrap_or_else(std::sync::PoisonError::into_inner);
313        let key = if !map.contains_key(&key) && map.len() >= MAX_TRACKED_MODELS {
314            "unknown".to_string()
315        } else {
316            key
317        };
318        map.entry(key.clone()).or_default().add(u);
319        drop(map);
320        if let Some(project) = u.wire.as_deref().and_then(|w| w.project_key.as_deref()) {
321            let mut projects = project_store()
322                .lock()
323                .unwrap_or_else(std::sync::PoisonError::into_inner);
324            if projects.contains_key(project) || projects.len() < MAX_TRACKED_PROJECTS {
325                projects
326                    .entry(project.to_string())
327                    .or_default()
328                    .entry(key)
329                    .or_default()
330                    .add(u);
331            }
332        }
333    }
334    if u.cohort.output.is_some() || u.cohort.compression.is_some() {
335        let mut cohorts = cohort_store()
336            .lock()
337            .unwrap_or_else(std::sync::PoisonError::into_inner);
338        if let Some(arm) = u.cohort.output {
339            cohorts.entry(arm.as_str().to_string()).or_default().add(u);
340        }
341        if let Some(arm) = u.cohort.compression {
342            cohorts
343                .entry(compression_cohort_key(arm))
344                .or_default()
345                .add(u);
346        }
347    }
348    persist();
349}
350
351/// Live output-savings cohort totals (#895). Empty until a holdout runs.
352#[must_use]
353pub fn cohort_snapshot() -> HashMap<String, CohortUsage> {
354    cohort_store()
355        .lock()
356        .unwrap_or_else(std::sync::PoisonError::into_inner)
357        .clone()
358}
359
360/// Cross-process read of the persisted output-savings cohort totals.
361#[must_use]
362pub fn persisted_cohorts() -> HashMap<String, CohortUsage> {
363    load_persisted().map(|p| p.cohorts).unwrap_or_default()
364}
365
366fn normalize_key(model: &str) -> String {
367    let m = model.trim();
368    if m.is_empty() {
369        "unknown".to_string()
370    } else {
371        m.to_string()
372    }
373}
374
375/// Live per-model measured spend, priced and sorted by USD descending.
376pub fn snapshot() -> Vec<ModelSpend> {
377    let map = store()
378        .lock()
379        .unwrap_or_else(std::sync::PoisonError::into_inner);
380    price_models(&map)
381}
382
383/// Total measured spend across all models (live in-memory totals).
384pub fn total_cost_usd() -> f64 {
385    snapshot().iter().map(|m| m.cost_usd).sum()
386}
387
388/// Cross-model verified-savings totals (#701) for `/status` and the dashboard.
389#[derive(Debug, Clone, Serialize, PartialEq)]
390pub struct VerifiedSavings {
391    /// Requests covered by a successful `count_tokens` probe.
392    pub requests: u64,
393    /// Provider-counted input tokens those requests would have billed
394    /// without lean-ctx.
395    pub counterfactual_input_tokens: u64,
396    /// Input-side tokens (input + cache read + cache write) they actually
397    /// billed.
398    pub billed_input_tokens: u64,
399    /// `counterfactual - billed`; negative when stub overhead outweighed the
400    /// squeeze — reported honestly, never clamped.
401    pub verified_saved_tokens: i64,
402}
403
404/// Aggregated provider-verified savings across all models (#701), or `None`
405/// until at least one probe-covered request has been recorded. Unlike the
406/// `tokens_saved` estimate (bytes/4), both sides of this pair were counted by
407/// the provider on the same request — receipts, not estimates.
408pub fn verified_savings() -> Option<VerifiedSavings> {
409    let map = store()
410        .lock()
411        .unwrap_or_else(std::sync::PoisonError::into_inner);
412    verified_of(&map)
413}
414
415/// Pure aggregation behind [`verified_savings`], shared with tests.
416#[allow(clippy::cast_possible_wrap)]
417fn verified_of(map: &HashMap<String, ModelUsage>) -> Option<VerifiedSavings> {
418    let (mut requests, mut counterfactual, mut billed) = (0u64, 0u64, 0u64);
419    for usage in map.values() {
420        requests += usage.counterfactual_requests;
421        counterfactual += usage.counterfactual_input_tokens;
422        billed += usage.counterfactual_billed_tokens;
423    }
424    (requests > 0).then(|| VerifiedSavings {
425        requests,
426        counterfactual_input_tokens: counterfactual,
427        billed_input_tokens: billed,
428        verified_saved_tokens: counterfactual as i64 - billed as i64,
429    })
430}
431
432/// Prices a model usage map into sorted [`ModelSpend`] rows. Pure: shared by the
433/// in-memory snapshot and the cross-process [`persisted_snapshot`].
434pub fn price_models(map: &HashMap<String, ModelUsage>) -> Vec<ModelSpend> {
435    let pricing = ModelPricing::load();
436    let mut rows: Vec<ModelSpend> = map
437        .iter()
438        .map(|(model, usage)| price_one(&pricing, model, usage))
439        .collect();
440    rows.sort_by(|a, b| {
441        b.cost_usd
442            .partial_cmp(&a.cost_usd)
443            .unwrap_or(std::cmp::Ordering::Equal)
444    });
445    rows
446}
447
448fn price_one(pricing: &ModelPricing, model: &str, usage: &ModelUsage) -> ModelSpend {
449    let quote = pricing.quote(Some(model));
450    // Measured-first (#1179): turns whose response carried the provider's own
451    // USD charge book that exact amount; only the remainder is table-priced.
452    let m = &usage.measured;
453    let derived_cost = quote.cost.estimate_usd(
454        usage.input_tokens.saturating_sub(m.input_tokens),
455        usage.output_tokens.saturating_sub(m.output_tokens),
456        usage
457            .cache_write_tokens
458            .saturating_sub(m.cache_write_tokens),
459        usage.cache_read_tokens.saturating_sub(m.cache_read_tokens),
460    );
461    let derived_requests = usage.requests.saturating_sub(m.requests);
462    ModelSpend {
463        model: model.to_string(),
464        requests: usage.requests,
465        input_tokens: usage.input_tokens,
466        output_tokens: usage.output_tokens,
467        cache_read_tokens: usage.cache_read_tokens,
468        cache_write_tokens: usage.cache_write_tokens,
469        reasoning_tokens: usage.reasoning_tokens,
470        cost_usd: m.cost_usd + derived_cost,
471        measured_cost_usd: m.cost_usd,
472        measured_requests: m.requests,
473        // Fully measured spend is never an estimate, whatever the table says.
474        pricing_estimated: quote.match_kind.is_estimated() && derived_requests > 0,
475    }
476}
477
478/// Atomically writes the current in-memory totals to disk.
479fn persist() {
480    let Some(path) = usage_path() else {
481        return;
482    };
483    let models = {
484        let map = store()
485            .lock()
486            .unwrap_or_else(std::sync::PoisonError::into_inner);
487        map.clone()
488    };
489    let cohorts = {
490        let map = cohort_store()
491            .lock()
492            .unwrap_or_else(std::sync::PoisonError::into_inner);
493        map.clone()
494    };
495    let last_model = last_model_store()
496        .lock()
497        .unwrap_or_else(std::sync::PoisonError::into_inner)
498        .clone();
499    let projects = project_store()
500        .lock()
501        .unwrap_or_else(std::sync::PoisonError::into_inner)
502        .clone();
503    let payload = PersistedUsage {
504        ts: std::time::SystemTime::now()
505            .duration_since(std::time::UNIX_EPOCH)
506            .unwrap_or_default()
507            .as_secs(),
508        models,
509        cohorts,
510        last_model,
511        projects,
512    };
513    let Ok(json) = serde_json::to_string(&payload) else {
514        return;
515    };
516    let tmp = path.with_extension("json.tmp");
517    if std::fs::write(&tmp, json).is_ok() {
518        let _ = std::fs::rename(&tmp, &path);
519    }
520}
521
522/// Cross-process priced spend rows, read from disk.
523pub fn persisted_snapshot() -> Vec<ModelSpend> {
524    load_persisted()
525        .map(|p| price_models(&p.models))
526        .unwrap_or_default()
527}
528
529/// Priced spend rows of the projects whose key `keep` accepts, summed per model.
530pub fn persisted_snapshot_where(keep: impl Fn(&str) -> bool) -> Vec<ModelSpend> {
531    let Some(persisted) = load_persisted() else {
532        return Vec::new();
533    };
534    let mut models: HashMap<String, ModelUsage> = HashMap::new();
535    for (project, usage) in persisted.projects {
536        if keep(&project) {
537            merge_models(&mut models, usage);
538        }
539    }
540    price_models(&models)
541}
542
543/// Priced spend per project key (the device view's split).
544pub fn persisted_projects() -> Vec<(String, Vec<ModelSpend>)> {
545    let Some(persisted) = load_persisted() else {
546        return Vec::new();
547    };
548    let mut rows: Vec<(String, Vec<ModelSpend>)> = persisted
549        .projects
550        .iter()
551        .map(|(project, models)| (project.clone(), price_models(models)))
552        .collect();
553    rows.sort_by(|a, b| a.0.cmp(&b.0));
554    rows
555}
556
557#[cfg(test)]
558mod tests {
559    use super::*;
560
561    fn usage(
562        model: &str,
563        input: u64,
564        output: u64,
565        cache_read: u64,
566    ) -> super::super::usage::RealUsage {
567        super::super::usage::RealUsage {
568            model: model.to_string(),
569            input_tokens: input,
570            output_tokens: output,
571            cache_read_tokens: cache_read,
572            ..Default::default()
573        }
574    }
575
576    /// #895: resuming from disk must carry the sums of squares and prompt
577    /// totals; dropping them zeroed the variance after every proxy restart and
578    /// turned the CI into a zero-width "measurement".
579    #[test]
580    fn cohort_merge_keeps_every_running_total() {
581        let mut persisted = CohortUsage::default();
582        for (out, cached) in [(150, 900), (210, 1_200), (180, 0)] {
583            persisted.add(&usage("claude-sonnet-4.5", 1_000, out, cached));
584        }
585        let mut resumed = CohortUsage::default();
586        resumed.merge(&persisted);
587        assert_eq!(resumed, persisted, "every field survives the restart");
588        assert!(resumed.variance_output().is_some_and(|v| v > 0.0));
589
590        // Merging into live in-memory totals adds field by field.
591        let mut live = CohortUsage::default();
592        live.add(&usage("claude-sonnet-4.5", 500, 100, 0));
593        let expected_sum_sq = live.sum_sq_output + persisted.sum_sq_output;
594        live.merge(&persisted);
595        assert_eq!(live.requests, 4);
596        assert_eq!(live.sum_sq_output, expected_sum_sq);
597        assert_eq!(live.prompt_tokens, 500 + persisted.prompt_tokens);
598    }
599
600    /// A turn whose request named a project is booked to the device total and
601    /// to that project; one that named none to the device total only. Both
602    /// survive the file round trip.
603    #[test]
604    fn spend_is_split_per_project_and_persisted() {
605        let _dir = lean_ctx::core::data_dir::isolated_data_dir();
606        let project = format!("{:016x}", std::process::id() as u64 * 7919 + 17);
607        let other = format!("{:016x}", std::process::id() as u64 * 7919 + 18);
608        let turn = |key: Option<&str>, input| super::super::usage::RealUsage {
609            wire: Some(Box::new(super::super::usage::WireContext {
610                project_key: key.map(str::to_string),
611                ..Default::default()
612            })),
613            ..usage("claude-sonnet-4.5", input, 10, 0)
614        };
615        record(&turn(Some(&project), 1_000));
616        record(&turn(Some(&project), 500));
617        record(&turn(Some(&other), 200));
618        record(&turn(None, 100));
619
620        let mine = persisted_snapshot_where(|key| key == project);
621        assert_eq!(mine.len(), 1);
622        assert_eq!(mine[0].requests, 2);
623        assert_eq!(mine[0].input_tokens, 1_500);
624        let persisted = load_persisted().expect("persisted");
625        let device = &persisted.models["claude-sonnet-4.5"];
626        let split: u64 = persisted_projects()
627            .iter()
628            .filter(|(key, _)| *key == project || *key == other)
629            .flat_map(|(_, models)| models.iter().map(|m| m.input_tokens))
630            .sum();
631        assert_eq!(split, 1_700);
632        assert!(
633            device.input_tokens >= 1_800,
634            "the device total keeps every turn, named or not"
635        );
636    }
637
638    /// 3.12.1: counting each proxied request once in the kernel evidence must
639    /// not rewrite what the savings ledger already holds. A pre-existing legacy
640    /// entry and the caching event of an earlier turn keep their exact bytes and
641    /// still verify; the next turn appends exactly one event.
642    #[test]
643    fn recording_a_turn_leaves_earlier_savings_ledger_entries_unchanged() {
644        use lean_ctx::core::savings_ledger::{self, event::SavingsEvent, store};
645        let _dir = lean_ctx::core::data_dir::isolated_data_dir();
646        let path = store::default_path().expect("ledger path");
647        let legacy: SavingsEvent = serde_json::from_value(serde_json::json!({
648            "ts": "2026-06-01T12:00:00+00:00",
649            "tool": "ctx_read",
650            "model_id": "claude-3.5-sonnet",
651            "tokenizer": "o200k_base",
652            "baseline_tokens": 1000,
653            "actual_tokens": 300,
654            "saved_tokens": 700,
655            "bounce_adjustment": 0,
656            "unit_price_per_m_usd": 3.0,
657            "saved_usd": 0.0021,
658            "repo_hash": "repo",
659            "agent_id": "local",
660            "prev_hash": "",
661            "entry_hash": ""
662        }))
663        .expect("legacy fixture");
664        store::append(&path, legacy).expect("append fixture");
665        record(&usage("claude-sonnet-4.5", 1_000, 50, 400));
666        let before = std::fs::read(&path).expect("ledger bytes");
667        let before_events = savings_ledger::all_events();
668        assert_eq!(before_events.len(), 2, "fixture plus one caching event");
669
670        record(&usage("claude-sonnet-4.5", 2_000, 80, 600));
671
672        let after = std::fs::read(&path).expect("ledger bytes");
673        assert!(after.starts_with(&before), "earlier entries are untouched");
674        let after_events = savings_ledger::all_events();
675        assert_eq!(after_events.len(), 3, "one turn appends one event");
676        for (old, new) in before_events.iter().zip(&after_events) {
677            assert_eq!(old.entry_hash, new.entry_hash);
678            assert_eq!(old.saved_tokens, new.saved_tokens);
679            assert_eq!(old.saved_usd.to_bits(), new.saved_usd.to_bits());
680        }
681        let verified = savings_ledger::verify();
682        assert!(verified.valid, "{verified:?}");
683        assert_eq!(verified.total, 3);
684    }
685
686    /// #701: the verified pair is recorded only when the probe answered — both
687    /// sides of the pair or neither, never a half row.
688    #[test]
689    fn counterfactual_pair_recorded_only_when_probe_answered() {
690        let slot = super::super::counterfactual::CounterfactualSlot::new();
691        slot.set(5_000);
692        let with_probe = super::super::usage::RealUsage {
693            wire: Some(Box::new(super::super::usage::WireContext {
694                counterfactual: Some(slot),
695                ..Default::default()
696            })),
697            cache_write_tokens: 200,
698            ..usage("claude-sonnet-4.5", 1_000, 0, 300)
699        };
700        let mut acc = ModelUsage::default();
701        acc.add(&with_probe);
702        assert_eq!(acc.counterfactual_requests, 1);
703        assert_eq!(acc.counterfactual_input_tokens, 5_000);
704        // billed side = input + cache read + cache write of the same turn.
705        assert_eq!(acc.counterfactual_billed_tokens, 1_000 + 300 + 200);
706
707        // Empty slot (probe failed / still in flight) → row degrades to the
708        // estimate: no pair recorded, normal usage still counted.
709        let empty = super::super::usage::RealUsage {
710            wire: Some(Box::new(super::super::usage::WireContext {
711                counterfactual: Some(super::super::counterfactual::CounterfactualSlot::new()),
712                ..Default::default()
713            })),
714            ..usage("claude-sonnet-4.5", 1_000, 0, 0)
715        };
716        acc.add(&empty);
717        assert_eq!(acc.requests, 2);
718        assert_eq!(acc.counterfactual_requests, 1, "empty slot adds no pair");
719
720        // No wire context at all (tests / non-forward paths) → no pair.
721        acc.add(&usage("claude-sonnet-4.5", 10, 0, 0));
722        assert_eq!(acc.counterfactual_requests, 1);
723    }
724
725    /// #701: cross-model aggregation and the honest signed difference.
726    #[test]
727    fn verified_of_aggregates_and_reports_signed_savings() {
728        let mut map = HashMap::new();
729        assert!(verified_of(&map).is_none(), "no coverage → None, not zeros");
730
731        map.insert(
732            "claude-sonnet-4.5".to_string(),
733            ModelUsage {
734                counterfactual_requests: 2,
735                counterfactual_input_tokens: 10_000,
736                counterfactual_billed_tokens: 6_000,
737                ..Default::default()
738            },
739        );
740        // Stub overhead outweighed the squeeze on this model: billed MORE
741        // than the counterfactual. The total must subtract honestly.
742        map.insert(
743            "claude-haiku-4.5".to_string(),
744            ModelUsage {
745                counterfactual_requests: 1,
746                counterfactual_input_tokens: 1_000,
747                counterfactual_billed_tokens: 1_400,
748                ..Default::default()
749            },
750        );
751        let v = verified_of(&map).expect("covered rows present");
752        assert_eq!(v.requests, 3);
753        assert_eq!(v.counterfactual_input_tokens, 11_000);
754        assert_eq!(v.billed_input_tokens, 7_400);
755        assert_eq!(v.verified_saved_tokens, 3_600);
756
757        map.get_mut("claude-sonnet-4.5")
758            .unwrap()
759            .counterfactual_input_tokens = 0;
760        map.get_mut("claude-sonnet-4.5")
761            .unwrap()
762            .counterfactual_billed_tokens = 0;
763        let negative = verified_of(&map).unwrap();
764        assert_eq!(
765            negative.verified_saved_tokens, -400,
766            "a net-negative verified saving is reported, never clamped"
767        );
768    }
769
770    /// #701: persisted usage files from before the feature load cleanly and
771    /// the new fields round-trip.
772    #[test]
773    fn counterfactual_fields_roundtrip_and_default_for_legacy_files() {
774        let legacy = r#"{"ts":1,"models":{"m":{"requests":1,"input_tokens":10,"output_tokens":5,"cache_read_tokens":0,"cache_write_tokens":0,"reasoning_tokens":0}}}"#;
775        let p: PersistedUsage = serde_json::from_str(legacy).expect("legacy file loads");
776        assert_eq!(p.models["m"].counterfactual_requests, 0);
777
778        let mut p = PersistedUsage::default();
779        p.models.insert(
780            "m".into(),
781            ModelUsage {
782                counterfactual_requests: 4,
783                counterfactual_input_tokens: 9_999,
784                counterfactual_billed_tokens: 5_555,
785                ..Default::default()
786            },
787        );
788        let back: PersistedUsage =
789            serde_json::from_str(&serde_json::to_string(&p).unwrap()).unwrap();
790        assert_eq!(back.models["m"].counterfactual_input_tokens, 9_999);
791        assert_eq!(back.models["m"].counterfactual_billed_tokens, 5_555);
792    }
793
794    #[test]
795    fn prices_known_model_with_cache_split() {
796        let mut map = HashMap::new();
797        let mut acc = ModelUsage::default();
798        acc.add(&usage("claude-sonnet-4.5", 1_000_000, 1_000_000, 1_000_000));
799        map.insert("claude-sonnet-4.5".to_string(), acc);
800
801        let rows = price_models(&map);
802        assert_eq!(rows.len(), 1);
803        let row = &rows[0];
804        // input 3.00 + output 15.00 + cache_read 0.30 (per 1M) = 18.30.
805        assert!(
806            (row.cost_usd - 18.30).abs() < 1e-6,
807            "cost was {}",
808            row.cost_usd
809        );
810        assert!(!row.pricing_estimated, "exact model match");
811        assert_eq!(row.requests, 1);
812    }
813
814    #[test]
815    fn unknown_model_prices_with_fallback_and_is_estimated() {
816        let mut map = HashMap::new();
817        let mut acc = ModelUsage::default();
818        acc.add(&usage("some-novel-model-xyz", 1_000_000, 0, 0));
819        map.insert("some-novel-model-xyz".to_string(), acc);
820
821        let rows = price_models(&map);
822        assert!(rows[0].pricing_estimated, "fallback pricing is estimated");
823        assert!(rows[0].cost_usd > 0.0);
824    }
825
826    /// #1179: turns carrying the provider's own charge book that USD; only the
827    /// remaining (unmeasured) turns are table-priced — and a fully measured
828    /// model is never flagged as estimated, even when the table has no entry.
829    /// (Model name chosen to never hit the embedded or live price tables.)
830    #[test]
831    fn measured_provider_cost_replaces_table_estimate() {
832        const MODEL: &str = "vendor/unlisted-model-20990101";
833        let mut acc = ModelUsage::default();
834        let mut measured = usage(MODEL, 431_600, 22_700, 126_700);
835        measured.provider_cost_usd = Some(0.05);
836        acc.add(&measured);
837
838        let mut map = HashMap::new();
839        map.insert(MODEL.to_string(), acc.clone());
840        let row = &price_models(&map)[0];
841        assert!(
842            (row.cost_usd - 0.05).abs() < 1e-12,
843            "the bill, not the table"
844        );
845        assert!((row.measured_cost_usd - 0.05).abs() < 1e-12);
846        assert_eq!(row.measured_requests, 1);
847        assert!(
848            !row.pricing_estimated,
849            "fully measured spend is not an estimate"
850        );
851
852        // A second, unmeasured turn on the same model: its tokens are priced
853        // from the table ON TOP of the measured USD, never double-counted
854        // (10k tokens at the $2.50/M blended fallback ≈ $0.025).
855        acc.add(&usage(MODEL, 10_000, 0, 0));
856        let mut map = HashMap::new();
857        map.insert(MODEL.to_string(), acc);
858        let row = &price_models(&map)[0];
859        assert!(row.cost_usd > 0.05, "unmeasured remainder adds table cost");
860        assert!(row.cost_usd < 0.2, "measured slice must not be re-priced");
861        assert!(
862            row.pricing_estimated,
863            "unmeasured remainder is heuristically priced"
864        );
865    }
866
867    #[test]
868    fn measured_slice_roundtrips_and_defaults_for_legacy_files() {
869        let legacy = r#"{"ts":1,"models":{"m":{"requests":1,"input_tokens":10,"output_tokens":5,"cache_read_tokens":0,"cache_write_tokens":0,"reasoning_tokens":0}}}"#;
870        let p: PersistedUsage = serde_json::from_str(legacy).expect("legacy file loads");
871        assert_eq!(p.models["m"].measured, MeasuredSlice::default());
872
873        let mut p = PersistedUsage::default();
874        let mut acc = ModelUsage::default();
875        let mut m = usage("m", 100, 10, 0);
876        m.provider_cost_usd = Some(0.5);
877        acc.add(&m);
878        p.models.insert("m".into(), acc);
879        let back: PersistedUsage =
880            serde_json::from_str(&serde_json::to_string(&p).unwrap()).unwrap();
881        assert_eq!(back.models["m"].measured.requests, 1);
882        assert!((back.models["m"].measured.cost_usd - 0.5).abs() < 1e-12);
883    }
884
885    #[test]
886    fn dominant_model_picks_highest_token_real_model() {
887        let mut models = HashMap::new();
888        models.insert("claude-haiku-4.5".to_string(), {
889            let mut u = ModelUsage::default();
890            u.add(&usage("claude-haiku-4.5", 100, 100, 0));
891            u
892        });
893        models.insert("claude-opus-4.5".to_string(), {
894            let mut u = ModelUsage::default();
895            u.add(&usage("claude-opus-4.5", 10_000, 10_000, 0));
896            u
897        });
898        models.insert("unknown".to_string(), {
899            let mut u = ModelUsage::default();
900            u.add(&usage("unknown", 999_999, 0, 0));
901            u
902        });
903        let dominant = models
904            .iter()
905            .filter(|(m, _)| m.as_str() != "unknown")
906            .max_by_key(|(_, u)| u.billable_tokens())
907            .map(|(m, _)| m.clone());
908        assert_eq!(dominant.as_deref(), Some("claude-opus-4.5"));
909    }
910
911    #[test]
912    fn empty_model_buckets_as_unknown() {
913        assert_eq!(normalize_key("  "), "unknown");
914        assert_eq!(normalize_key(""), "unknown");
915        assert_eq!(normalize_key("gpt-5.4"), "gpt-5.4");
916    }
917
918    #[test]
919    fn cohort_avg_output_is_mean_per_turn() {
920        let mut c = CohortUsage::default();
921        assert_eq!(c.avg_output(), None, "no observations → None");
922        c.add(&usage("m", 10, 100, 0));
923        c.add(&usage("m", 10, 50, 0));
924        assert_eq!(c.requests, 2);
925        assert_eq!(c.output_tokens, 150);
926        assert!((c.avg_output().unwrap() - 75.0).abs() < f64::EPSILON);
927    }
928
929    #[test]
930    fn persisted_usage_without_cohorts_field_loads() {
931        // proxy_usage.json written before #895 has no `cohorts` key; serde(default)
932        // must backfill an empty map so old files stay loadable.
933        let json = r#"{"ts":1,"models":{"gpt-5.4":{"requests":1,"input_tokens":10,"output_tokens":5,"cache_read_tokens":0,"cache_write_tokens":0,"reasoning_tokens":0}}}"#;
934        let p: PersistedUsage = serde_json::from_str(json).expect("loads legacy file");
935        assert_eq!(p.models.len(), 1);
936        assert!(p.cohorts.is_empty());
937    }
938
939    #[test]
940    fn persisted_usage_roundtrips_cohorts() {
941        let mut p = PersistedUsage::default();
942        p.cohorts.insert(
943            "control".into(),
944            CohortUsage {
945                requests: 3,
946                input_tokens: 30,
947                output_tokens: 300,
948                sum_sq_output: 30_000,
949                ..CohortUsage::default()
950            },
951        );
952        let json = serde_json::to_string(&p).unwrap();
953        let back: PersistedUsage = serde_json::from_str(&json).unwrap();
954        assert_eq!(back.cohorts.get("control").unwrap().output_tokens, 300);
955    }
956
957    #[test]
958    fn managed_usage_projects_once_and_unmanaged_is_skipped() {
959        let sink = lean_ctx::core::ocla::builtin::usage_sink::BuiltinUsageSink::new();
960        let usage = super::super::usage::RealUsage {
961            model: "gpt-5".into(),
962            input_tokens: 100,
963            output_tokens: 40,
964            cache_read_tokens: 20,
965            cache_write_tokens: 5,
966            wire: Some(Box::new(super::super::usage::WireContext {
967                lineage: Some(lean_ctx::core::ocla::OclaRequestContext {
968                    request_id: "request-1".into(),
969                    session_id: "session-1".into(),
970                    agent_id: "agent-1".into(),
971                    content_ref: "blake3:content".into(),
972                    tenant_id: None,
973                    trace_id: "tr-unit".into(),
974                    task_id: None,
975                    parent_task_id: None,
976                }),
977                ..Default::default()
978            })),
979            ..Default::default()
980        };
981
982        project_ocla_usage(&usage, &sink);
983        assert_eq!(sink.record_count(), 1);
984        assert_eq!(sink.total_input_tokens(), 125);
985        assert_eq!(sink.total_output_tokens(), 40);
986
987        project_ocla_usage(
988            &super::super::usage::RealUsage {
989                model: "unmanaged".into(),
990                input_tokens: 999,
991                ..Default::default()
992            },
993            &sink,
994        );
995        assert_eq!(sink.record_count(), 1);
996    }
997
998    /// Counts the records the global registry's usage sink receives.
999    struct SpyUsage(std::sync::Arc<std::sync::atomic::AtomicUsize>);
1000
1001    impl lean_ctx::core::ocla::traits::OclaService for SpyUsage {
1002        fn capability(&self) -> lean_ctx::core::ocla::types::OclaCapability {
1003            lean_ctx::core::ocla::types::OclaCapability::available(
1004                lean_ctx::core::ocla::types::OclaCapabilityKind::UsageSink,
1005            )
1006        }
1007    }
1008
1009    #[async_trait::async_trait]
1010    impl UsageSink for SpyUsage {
1011        async fn record_usage(
1012            &self,
1013            _usage: lean_ctx::core::ocla::types::UsageRecord,
1014        ) -> lean_ctx::core::ocla::OclaResult<()> {
1015            self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1016            Ok(())
1017        }
1018    }
1019
1020    #[test]
1021    fn proxy_request_projects_measured_usage_to_usage_spy() {
1022        let _dir = lean_ctx::core::data_dir::isolated_data_dir();
1023        let usage_calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
1024        let mut registry = OclaRegistry::with_builtins();
1025        registry.usage_sink = std::sync::Arc::new(SpyUsage(usage_calls.clone()));
1026        let _guard = lean_ctx::core::ocla::registry::with_test_registry(registry);
1027        let usage = super::super::usage::RealUsage {
1028            model: "test-model".into(),
1029            input_tokens: 11,
1030            output_tokens: 7,
1031            wire: Some(Box::new(super::super::usage::WireContext {
1032                lineage: Some(lean_ctx::core::ocla::OclaRequestContext {
1033                    request_id: "request-1".into(),
1034                    session_id: "session-1".into(),
1035                    agent_id: "agent-1".into(),
1036                    content_ref: "file:test.rs".into(),
1037                    tenant_id: None,
1038                    trace_id: "tr-unit".into(),
1039                    task_id: None,
1040                    parent_task_id: None,
1041                }),
1042                ..Default::default()
1043            })),
1044            ..Default::default()
1045        };
1046
1047        record(&usage);
1048
1049        assert_eq!(usage_calls.load(std::sync::atomic::Ordering::Relaxed), 1);
1050    }
1051}