lean_ctx/proxy/
cost.rs

1//! Per-model proxy savings accounting.
2//!
3//! Headroom's per-model cost breakdown was the one metric a tester found clearer
4//! than lean-ctx's single flat number. This module buckets request-side savings
5//! by model and prices them with the shared [`ModelPricing`] table so `/status`
6//! can report estimated USD avoided per model.
7//!
8//! Honesty contract: token counts here are request-side *estimates* (the bytes
9//! the proxy removed before forwarding). They deliberately do NOT try to model
10//! re-reads the agent may perform later, so figures are conservative and
11//! labelled `estimated` in the output.
12
13use std::collections::HashMap;
14use std::sync::{Mutex, OnceLock};
15
16use serde::Serialize;
17
18use crate::core::gain::model_pricing::{ModelPricing, PricingMatchKind};
19
20#[derive(Default, Clone)]
21struct ModelAccum {
22    requests: u64,
23    tokens_saved: u64,
24    bytes_original: u64,
25    bytes_compressed: u64,
26}
27
28/// One model's aggregated, priced savings for the `/status` endpoint.
29#[derive(Debug, Clone, Serialize, PartialEq)]
30pub struct ModelStat {
31    pub model: String,
32    pub requests: u64,
33    pub tokens_saved: u64,
34    pub usd_saved: f64,
35    /// True when the price came from a fallback/heuristic match, not an exact one.
36    pub pricing_estimated: bool,
37}
38
39/// Cap on distinct model buckets. `record` sees the raw request model string, so an
40/// arbitrary-model client could otherwise grow the map without bound; overflow folds
41/// into "unknown" (real model names number < ~50, so this is generous).
42const MAX_TRACKED_MODELS: usize = 256;
43
44fn store() -> &'static Mutex<HashMap<String, ModelAccum>> {
45    static STORE: OnceLock<Mutex<HashMap<String, ModelAccum>>> = OnceLock::new();
46    STORE.get_or_init(|| Mutex::new(HashMap::new()))
47}
48
49/// Records one request's request-side savings against its model bucket.
50///
51/// `model` is taken from the request body (`None`/empty buckets under
52/// `"unknown"`). Recording never blocks request handling on poisoning.
53pub fn record(model: Option<&str>, tokens_saved: u64, bytes_original: u64, bytes_compressed: u64) {
54    let key = model
55        .map(str::trim)
56        .filter(|m| !m.is_empty())
57        .unwrap_or("unknown")
58        .to_string();
59
60    let mut map = store()
61        .lock()
62        .unwrap_or_else(std::sync::PoisonError::into_inner);
63    // Bound distinct model buckets. `record` takes the raw request model string (it does
64    // NOT pass through normalize_model), so a client sending arbitrary model names could
65    // otherwise grow this map unbounded. Fold overflow into "unknown" so aggregate totals
66    // stay exact — only per-model granularity for rare/novel names is lost.
67    let key = if !map.contains_key(&key) && map.len() >= MAX_TRACKED_MODELS {
68        "unknown".to_string()
69    } else {
70        key
71    };
72    let acc = map.entry(key).or_default();
73    acc.requests += 1;
74    acc.tokens_saved += tokens_saved;
75    acc.bytes_original += bytes_original;
76    acc.bytes_compressed += bytes_compressed;
77}
78
79/// Returns per-model stats, priced and sorted by USD saved (descending).
80pub fn snapshot() -> Vec<ModelStat> {
81    let pricing = ModelPricing::load();
82    let map = store()
83        .lock()
84        .unwrap_or_else(std::sync::PoisonError::into_inner);
85
86    let mut stats: Vec<ModelStat> = map
87        .iter()
88        .map(|(model, acc)| {
89            let quote = pricing.quote(Some(model));
90            // Compression removes *input* tokens, so price against the input rate.
91            let usd_saved = acc.tokens_saved as f64 / 1_000_000.0 * quote.cost.input_per_m;
92            let pricing_estimated = !matches!(quote.match_kind, PricingMatchKind::Exact);
93            ModelStat {
94                model: model.clone(),
95                requests: acc.requests,
96                tokens_saved: acc.tokens_saved,
97                usd_saved,
98                pricing_estimated,
99            }
100        })
101        .collect();
102
103    stats.sort_by(|a, b| {
104        b.usd_saved
105            .partial_cmp(&a.usd_saved)
106            .unwrap_or(std::cmp::Ordering::Equal)
107    });
108    stats
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn unknown_model_buckets_and_prices_without_panic() {
117        record(None, 1000, 4000, 0);
118        record(Some("   "), 500, 2000, 0);
119        let stats = snapshot();
120        let unknown = stats.iter().find(|s| s.model == "unknown");
121        assert!(
122            unknown.is_some(),
123            "blank/None models bucket under 'unknown'"
124        );
125        assert!(unknown.unwrap().requests >= 2);
126    }
127
128    #[test]
129    fn known_model_yields_positive_usd() {
130        record(
131            Some("claude-opus-4-8-zzz-cost-test"),
132            2_000_000,
133            8_000_000,
134            100,
135        );
136        let stats = snapshot();
137        let row = stats
138            .iter()
139            .find(|s| s.model.contains("opus-4-8-zzz-cost-test"))
140            .expect("recorded model present");
141        // Fallback pricing still produces a finite, non-negative estimate.
142        assert!(row.usd_saved >= 0.0 && row.usd_saved.is_finite());
143        assert_eq!(row.tokens_saved, 2_000_000);
144    }
145}
lean_ctx/proxy/cost.rs

lean_ctx/proxy/
cost.rs
