lean_ctx_proxy/
cost.rs

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

lean_ctx_proxy/
cost.rs
