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, PricingMatchKind};
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}
32
33impl ModelUsage {
34    fn add(&mut self, u: &super::usage::RealUsage) {
35        self.requests += 1;
36        self.input_tokens += u.input_tokens;
37        self.output_tokens += u.output_tokens;
38        self.cache_read_tokens += u.cache_read_tokens;
39        self.cache_write_tokens += u.cache_write_tokens;
40        self.reasoning_tokens += u.reasoning_tokens;
41    }
42
43    fn billable_tokens(&self) -> u64 {
44        self.input_tokens + self.output_tokens + self.cache_read_tokens + self.cache_write_tokens
45    }
46}
47
48/// One model's measured, priced spend for `/status` and the dashboard.
49#[derive(Debug, Clone, Serialize, PartialEq)]
50pub struct ModelSpend {
51    pub model: String,
52    pub requests: u64,
53    pub input_tokens: u64,
54    pub output_tokens: u64,
55    pub cache_read_tokens: u64,
56    pub cache_write_tokens: u64,
57    pub reasoning_tokens: u64,
58    pub cost_usd: f64,
59    /// True when pricing came from a heuristic/fallback match, not an exact one.
60    pub pricing_estimated: bool,
61}
62
63/// On-disk shape of the measured spend totals.
64#[derive(Debug, Clone, Default, Serialize, Deserialize)]
65pub struct PersistedUsage {
66    pub ts: u64,
67    pub models: HashMap<String, ModelUsage>,
68}
69
70/// Distinct-model bucket cap. `record` keys on the raw response model string, so
71/// overflow folds into "unknown" to keep the map bounded (real model names < ~50).
72const MAX_TRACKED_MODELS: usize = 256;
73
74const PROXY_USAGE_FILE: &str = "proxy_usage.json";
75
76fn store() -> &'static Mutex<HashMap<String, ModelUsage>> {
77    static STORE: OnceLock<Mutex<HashMap<String, ModelUsage>>> = OnceLock::new();
78    STORE.get_or_init(|| Mutex::new(HashMap::new()))
79}
80
81/// Seeds the in-memory totals from `proxy_usage.json`. Call once on proxy
82/// startup so measured spend is cumulative across restarts. Idempotent-ish: it
83/// merges the persisted totals into whatever is in memory (normally empty).
84pub fn resume_from_disk() {
85    let Some(persisted) = load_persisted() else {
86        return;
87    };
88    let mut map = store()
89        .lock()
90        .unwrap_or_else(std::sync::PoisonError::into_inner);
91    for (model, usage) in persisted.models {
92        let acc = map.entry(model).or_default();
93        acc.requests += usage.requests;
94        acc.input_tokens += usage.input_tokens;
95        acc.output_tokens += usage.output_tokens;
96        acc.cache_read_tokens += usage.cache_read_tokens;
97        acc.cache_write_tokens += usage.cache_write_tokens;
98        acc.reasoning_tokens += usage.reasoning_tokens;
99    }
100}
101
102/// Records one turn's measured usage against its model bucket and persists.
103pub fn record(u: &super::usage::RealUsage) {
104    let key = normalize_key(&u.model);
105    {
106        let mut map = store()
107            .lock()
108            .unwrap_or_else(std::sync::PoisonError::into_inner);
109        let key = if !map.contains_key(&key) && map.len() >= MAX_TRACKED_MODELS {
110            "unknown".to_string()
111        } else {
112            key
113        };
114        map.entry(key).or_default().add(u);
115    }
116    persist();
117}
118
119fn normalize_key(model: &str) -> String {
120    let m = model.trim();
121    if m.is_empty() {
122        "unknown".to_string()
123    } else {
124        m.to_string()
125    }
126}
127
128/// Live per-model measured spend, priced and sorted by USD descending.
129pub fn snapshot() -> Vec<ModelSpend> {
130    let map = store()
131        .lock()
132        .unwrap_or_else(std::sync::PoisonError::into_inner);
133    price_models(&map)
134}
135
136/// Total measured spend across all models (live in-memory totals).
137pub fn total_cost_usd() -> f64 {
138    snapshot().iter().map(|m| m.cost_usd).sum()
139}
140
141/// Prices a model usage map into sorted [`ModelSpend`] rows. Pure: shared by the
142/// in-memory snapshot and the cross-process [`persisted_snapshot`].
143pub fn price_models(map: &HashMap<String, ModelUsage>) -> Vec<ModelSpend> {
144    let pricing = ModelPricing::load();
145    let mut rows: Vec<ModelSpend> = map
146        .iter()
147        .map(|(model, usage)| price_one(&pricing, model, usage))
148        .collect();
149    rows.sort_by(|a, b| {
150        b.cost_usd
151            .partial_cmp(&a.cost_usd)
152            .unwrap_or(std::cmp::Ordering::Equal)
153    });
154    rows
155}
156
157fn price_one(pricing: &ModelPricing, model: &str, usage: &ModelUsage) -> ModelSpend {
158    let quote = pricing.quote(Some(model));
159    let cost_usd = quote.cost.estimate_usd(
160        usage.input_tokens,
161        usage.output_tokens,
162        usage.cache_write_tokens,
163        usage.cache_read_tokens,
164    );
165    ModelSpend {
166        model: model.to_string(),
167        requests: usage.requests,
168        input_tokens: usage.input_tokens,
169        output_tokens: usage.output_tokens,
170        cache_read_tokens: usage.cache_read_tokens,
171        cache_write_tokens: usage.cache_write_tokens,
172        reasoning_tokens: usage.reasoning_tokens,
173        cost_usd,
174        pricing_estimated: !matches!(quote.match_kind, PricingMatchKind::Exact),
175    }
176}
177
178fn usage_path() -> Option<std::path::PathBuf> {
179    crate::core::data_dir::lean_ctx_data_dir()
180        .ok()
181        .map(|d| d.join(PROXY_USAGE_FILE))
182}
183
184/// Atomically writes the current in-memory totals to disk.
185fn persist() {
186    let Some(path) = usage_path() else {
187        return;
188    };
189    let models = {
190        let map = store()
191            .lock()
192            .unwrap_or_else(std::sync::PoisonError::into_inner);
193        map.clone()
194    };
195    let payload = PersistedUsage {
196        ts: std::time::SystemTime::now()
197            .duration_since(std::time::UNIX_EPOCH)
198            .unwrap_or_default()
199            .as_secs(),
200        models,
201    };
202    let Ok(json) = serde_json::to_string(&payload) else {
203        return;
204    };
205    let tmp = path.with_extension("json.tmp");
206    if std::fs::write(&tmp, json).is_ok() {
207        let _ = std::fs::rename(&tmp, &path);
208    }
209}
210
211/// Cross-process read of the persisted measured spend (dashboard / CLI / ledger).
212pub fn load_persisted() -> Option<PersistedUsage> {
213    let path = usage_path()?;
214    let data = std::fs::read_to_string(path).ok()?;
215    serde_json::from_str(&data).ok()
216}
217
218/// Cross-process priced spend rows, read from disk.
219pub fn persisted_snapshot() -> Vec<ModelSpend> {
220    load_persisted()
221        .map(|p| price_models(&p.models))
222        .unwrap_or_default()
223}
224
225/// The model carrying the most measured tokens (excludes the "unknown" bucket).
226/// Used to value savings against the real dominant model when no explicit model
227/// is configured.
228pub fn persisted_dominant_model() -> Option<String> {
229    let persisted = load_persisted()?;
230    persisted
231        .models
232        .iter()
233        .filter(|(m, _)| m.as_str() != "unknown" && !m.trim().is_empty())
234        .max_by_key(|(_, u)| u.billable_tokens())
235        .filter(|(_, u)| u.billable_tokens() > 0)
236        .map(|(m, _)| m.clone())
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242
243    fn usage(
244        model: &str,
245        input: u64,
246        output: u64,
247        cache_read: u64,
248    ) -> super::super::usage::RealUsage {
249        super::super::usage::RealUsage {
250            model: model.to_string(),
251            input_tokens: input,
252            output_tokens: output,
253            cache_read_tokens: cache_read,
254            cache_write_tokens: 0,
255            reasoning_tokens: 0,
256        }
257    }
258
259    #[test]
260    fn prices_known_model_with_cache_split() {
261        let mut map = HashMap::new();
262        let mut acc = ModelUsage::default();
263        acc.add(&usage("claude-sonnet-4.5", 1_000_000, 1_000_000, 1_000_000));
264        map.insert("claude-sonnet-4.5".to_string(), acc);
265
266        let rows = price_models(&map);
267        assert_eq!(rows.len(), 1);
268        let row = &rows[0];
269        // input 3.00 + output 15.00 + cache_read 0.30 (per 1M) = 18.30.
270        assert!(
271            (row.cost_usd - 18.30).abs() < 1e-6,
272            "cost was {}",
273            row.cost_usd
274        );
275        assert!(!row.pricing_estimated, "exact model match");
276        assert_eq!(row.requests, 1);
277    }
278
279    #[test]
280    fn unknown_model_prices_with_fallback_and_is_estimated() {
281        let mut map = HashMap::new();
282        let mut acc = ModelUsage::default();
283        acc.add(&usage("some-novel-model-xyz", 1_000_000, 0, 0));
284        map.insert("some-novel-model-xyz".to_string(), acc);
285
286        let rows = price_models(&map);
287        assert!(rows[0].pricing_estimated, "fallback pricing is estimated");
288        assert!(rows[0].cost_usd > 0.0);
289    }
290
291    #[test]
292    fn dominant_model_picks_highest_token_real_model() {
293        let mut models = HashMap::new();
294        models.insert("claude-haiku-4.5".to_string(), {
295            let mut u = ModelUsage::default();
296            u.add(&usage("claude-haiku-4.5", 100, 100, 0));
297            u
298        });
299        models.insert("claude-opus-4.5".to_string(), {
300            let mut u = ModelUsage::default();
301            u.add(&usage("claude-opus-4.5", 10_000, 10_000, 0));
302            u
303        });
304        models.insert("unknown".to_string(), {
305            let mut u = ModelUsage::default();
306            u.add(&usage("unknown", 999_999, 0, 0));
307            u
308        });
309        let dominant = models
310            .iter()
311            .filter(|(m, _)| m.as_str() != "unknown")
312            .max_by_key(|(_, u)| u.billable_tokens())
313            .map(|(m, _)| m.clone());
314        assert_eq!(dominant.as_deref(), Some("claude-opus-4.5"));
315    }
316
317    #[test]
318    fn empty_model_buckets_as_unknown() {
319        assert_eq!(normalize_key("  "), "unknown");
320        assert_eq!(normalize_key(""), "unknown");
321        assert_eq!(normalize_key("gpt-5.4"), "gpt-5.4");
322    }
323}
lean_ctx/proxy/usage_meter.rs

lean_ctx/proxy/
usage_meter.rs
