Skip to main content

lean_ctx/core/stats/
model.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3
4use crate::core::gain::model_pricing::{ModelQuote, PricingMatchKind};
5
6/// Persistent store for all-time token savings, command stats, and daily history.
7#[derive(Serialize, Deserialize, Default, Clone)]
8pub struct StatsStore {
9    pub total_commands: u64,
10    pub total_input_tokens: u64,
11    pub total_output_tokens: u64,
12    pub first_use: Option<String>,
13    pub last_use: Option<String>,
14    pub commands: HashMap<String, CommandStats>,
15    pub daily: Vec<DayStats>,
16    #[serde(default)]
17    pub cep: CepStats,
18    /// Delivery classification recorded for each command. Older stats files do
19    /// not have this map; callers infer classifications from the command key.
20    #[serde(default)]
21    pub command_classes: HashMap<String, TrafficClass>,
22    /// Savings when a compressed tool result first enters provider context.
23    #[serde(default)]
24    pub first_inject_tokens_saved: u64,
25    /// Savings from previously injected tool results carried into later turns.
26    #[serde(default)]
27    pub reread_tokens_saved: u64,
28    /// Savings still resident in the active transcript and eligible for re-read.
29    #[serde(default)]
30    pub active_tool_result_tokens_saved: u64,
31    /// Last provider turn for which active tool-result re-reads were accrued.
32    #[serde(default)]
33    pub last_tool_result_turn: u64,
34    /// Number of results recorded with stream-aware accounting. Zero identifies
35    /// legacy stats files whose aggregate savings need a display-time fallback.
36    #[serde(default)]
37    pub stream_tracked_results: u64,
38}
39
40/// Whether a recorded command's output is controlled by compression.
41#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
42#[serde(rename_all = "snake_case")]
43pub enum TrafficClass {
44    Compressible,
45    Passthrough,
46}
47
48/// Token totals for traffic that lean-ctx can compress.
49#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
50pub struct CompressionTotals {
51    pub input_tokens: u64,
52    pub output_tokens: u64,
53}
54
55impl CompressionTotals {
56    pub(crate) fn saved_tokens(self) -> u64 {
57        self.input_tokens.saturating_sub(self.output_tokens)
58    }
59
60    pub(crate) fn compression_pct(self) -> f64 {
61        if self.input_tokens == 0 {
62            0.0
63        } else {
64            self.saved_tokens() as f64 / self.input_tokens as f64 * 100.0
65        }
66    }
67}
68
69impl StatsStore {
70    /// Records one tool result at the provider turn that caused it.
71    ///
72    /// #1283: this used to price every already-resident result as a re-read on
73    /// each turn advance (`reread += active * elapsed_turns`). Per session that
74    /// model is defensible — a compressed result resident in context saves its
75    /// delta on every later turn — but this store is *global and persistent*:
76    /// `active_tool_result_tokens_saved` accumulated across months and all
77    /// concurrent sessions (measured: 932M), and `last_tool_result_turn` mixed
78    /// unrelated per-session counters (measured: 36k). The product exploded
79    /// `reread_tokens_saved` to ~119 trillion — 50,000x every input token ever
80    /// processed. Persistent counters now hold measured savings only; re-read
81    /// savings come exclusively from the real cache-hit path (`record_reread`).
82    pub(crate) fn record_tool_result_savings(&mut self, saved: u64, turn: u64) {
83        if turn > 0 {
84            self.last_tool_result_turn = self.last_tool_result_turn.max(turn);
85            self.active_tool_result_tokens_saved =
86                self.active_tool_result_tokens_saved.saturating_add(saved);
87        }
88        self.first_inject_tokens_saved = self.first_inject_tokens_saved.saturating_add(saved);
89        self.stream_tracked_results = self.stream_tracked_results.saturating_add(1);
90    }
91
92    /// Stream totals: measured first-inject savings and measured re-read
93    /// savings. Legacy stores fall back to measured compressible savings.
94    /// No turn-based extrapolation (#1283) — the parameter remains for call
95    /// sites that report the observation turn.
96    pub(crate) fn stream_savings(&self, _observed_turn: u64) -> (u64, u64) {
97        if self.stream_tracked_results == 0 {
98            return (self.compression_totals().saved_tokens(), 0);
99        }
100        (self.first_inject_tokens_saved, self.reread_tokens_saved)
101    }
102
103    /// Computes effective compression from tagged command rows.
104    pub fn compression_totals(&self) -> CompressionTotals {
105        self.commands
106            .iter()
107            .filter(|(command, _)| {
108                self.command_classes
109                    .get(*command)
110                    .copied()
111                    .unwrap_or_else(|| classify_command(command))
112                    == TrafficClass::Compressible
113            })
114            .fold(CompressionTotals::default(), |mut totals, (_, stats)| {
115                totals.input_tokens = totals.input_tokens.saturating_add(stats.input_tokens);
116                totals.output_tokens = totals.output_tokens.saturating_add(stats.output_tokens);
117                totals
118            })
119    }
120
121    /// Total reduction across both compressible and passthrough traffic.
122    pub(crate) fn total_reduction_pct(&self) -> f64 {
123        if self.total_input_tokens == 0 {
124            0.0
125        } else {
126            self.total_input_tokens
127                .saturating_sub(self.total_output_tokens) as f64
128                / self.total_input_tokens as f64
129                * 100.0
130        }
131    }
132}
133
134/// Classifies normalized stats keys, with an explicit passthrough default for
135/// control/listing tools so only read, shell, and search output drives the
136/// effective-compression denominator.
137pub fn classify_command(command: &str) -> TrafficClass {
138    match command {
139        "cli_glob" | "cli_find" | "cli_deps" | "cli_ls" | "ctx_compose" | "ctx_glob"
140        | "ctx_tree" => TrafficClass::Passthrough,
141        c if c.starts_with("cli_") => TrafficClass::Compressible,
142        "ctx_shell" | "ctx_search" | "ctx_semantic_search" => TrafficClass::Compressible,
143        c if c.starts_with("ctx_read")
144            || c.starts_with("ctx_multi_read")
145            || c == "ctx_smart_read"
146            || c == "ctx_git_read"
147            || c == "ctx_url_read" =>
148        {
149            TrafficClass::Compressible
150        }
151        _ => TrafficClass::Passthrough,
152    }
153}
154
155/// Aggregated CEP (Cognitive Efficiency Protocol) metrics across sessions.
156#[derive(Serialize, Deserialize, Clone, Default)]
157pub struct CepStats {
158    pub sessions: u64,
159    pub total_cache_hits: u64,
160    pub total_cache_reads: u64,
161    pub total_tokens_original: u64,
162    pub total_tokens_compressed: u64,
163    pub modes: HashMap<String, u64>,
164    pub scores: Vec<CepSessionSnapshot>,
165    #[serde(default)]
166    pub last_session_pid: Option<u32>,
167    #[serde(default)]
168    pub last_session_original: Option<u64>,
169    #[serde(default)]
170    pub last_session_compressed: Option<u64>,
171    /// Cumulative cache hits/reads observed for the current PID at the last
172    /// snapshot. Used to accumulate *deltas* across repeated snapshots within
173    /// one server process, so `total_cache_hits` keeps tracking cache activity
174    /// after the first checkpoint instead of freezing (#361).
175    #[serde(default)]
176    pub last_session_cache_hits: Option<u64>,
177    #[serde(default)]
178    pub last_session_cache_reads: Option<u64>,
179}
180
181/// Point-in-time snapshot of CEP scores for a single session.
182#[derive(Serialize, Deserialize, Clone)]
183pub struct CepSessionSnapshot {
184    pub timestamp: String,
185    pub score: u32,
186    pub cache_hit_rate: u32,
187    pub mode_diversity: u32,
188    pub compression_rate: u32,
189    pub tool_calls: u64,
190    pub tokens_saved: u64,
191    pub complexity: String,
192}
193
194/// Per-command token statistics: invocation count and input/output totals.
195#[derive(Serialize, Deserialize, Clone, Default, Debug)]
196pub struct CommandStats {
197    pub count: u64,
198    pub input_tokens: u64,
199    pub output_tokens: u64,
200}
201
202/// Daily aggregate: command count and token totals for one calendar day.
203#[derive(Serialize, Deserialize, Clone, Default)]
204pub struct DayStats {
205    pub date: String,
206    pub commands: u64,
207    pub input_tokens: u64,
208    pub output_tokens: u64,
209    /// lean-ctx version active when this day's stats were last recorded.
210    /// Lets `lean-ctx gain` attribute per-day compression changes to a release
211    /// (#307). Empty for days recorded before this field existed.
212    #[serde(default)]
213    pub version: String,
214    /// Tokens from operations classified as `Compressible` (file reads, shell
215    /// output). Excludes terminal polling, meta-tools, and stream-tracking to
216    /// give the dashboard an undiluted compression rate.
217    #[serde(default)]
218    pub compressible_input_tokens: u64,
219    #[serde(default)]
220    pub compressible_output_tokens: u64,
221}
222
223/// High-level token savings summary for display.
224pub struct GainSummary {
225    pub total_saved: u64,
226    pub total_calls: u64,
227}
228
229/// Average LLM pricing per 1M tokens (blended across Claude, GPT, Gemini).
230pub const DEFAULT_INPUT_PRICE_PER_M: f64 = 2.50;
231pub const DEFAULT_OUTPUT_PRICE_PER_M: f64 = 10.0;
232
233/// LLM pricing model for estimating dollar savings from token compression.
234pub struct CostModel {
235    pub model_key: String,
236    pub pricing_match_kind: PricingMatchKind,
237    pub input_price_per_m: f64,
238    pub output_price_per_m: f64,
239    pub avg_verbose_output_per_call: u64,
240    pub avg_concise_output_per_call: u64,
241}
242
243impl Default for CostModel {
244    fn default() -> Self {
245        let pricing = crate::core::gain::model_pricing::ModelPricing::load();
246        let quote = pricing.quote(resolved_gain_model().as_deref());
247        Self::from_quote(quote)
248    }
249}
250
251fn resolved_gain_model() -> Option<String> {
252    std::env::var("LEAN_CTX_MODEL")
253        .or_else(|_| std::env::var("LCTX_MODEL"))
254        .ok()
255        .filter(|s| !s.trim().is_empty())
256        .or_else(|| {
257            crate::core::config::Config::load()
258                .cost
259                .model_for_client("cli")
260        })
261        .or_else(crate::core::proxy_usage::persisted_dominant_model)
262}
263
264impl CostModel {
265    fn from_quote(quote: ModelQuote) -> Self {
266        Self {
267            model_key: quote.model_key,
268            pricing_match_kind: quote.match_kind,
269            input_price_per_m: quote.cost.input_per_m,
270            output_price_per_m: quote.cost.output_per_m,
271            avg_verbose_output_per_call: 180,
272            avg_concise_output_per_call: 120,
273        }
274    }
275}
276
277/// Detailed cost comparison: with vs. without lean-ctx compression.
278pub struct CostBreakdown {
279    pub input_cost_without: f64,
280    pub input_cost_with: f64,
281    pub output_cost_without: f64,
282    pub output_cost_with: f64,
283    pub total_cost_without: f64,
284    pub total_cost_with: f64,
285    pub total_saved: f64,
286    pub estimated_output_tokens_without: u64,
287    pub estimated_output_tokens_with: u64,
288    pub output_tokens_saved: u64,
289}
290
291impl CostModel {
292    /// Calculates the full cost breakdown from the stats store.
293    pub fn calculate(&self, store: &StatsStore) -> CostBreakdown {
294        let input_cost_without =
295            store.total_input_tokens as f64 / 1_000_000.0 * self.input_price_per_m;
296        let input_cost_with =
297            store.total_output_tokens as f64 / 1_000_000.0 * self.input_price_per_m;
298
299        let input_saved = store
300            .total_input_tokens
301            .saturating_sub(store.total_output_tokens);
302        let compression_rate = if store.total_input_tokens > 0 {
303            input_saved as f64 / store.total_input_tokens as f64
304        } else {
305            0.0
306        };
307        let est_output_without = store.total_commands * self.avg_verbose_output_per_call;
308        let est_output_with = if compression_rate > 0.01 {
309            store.total_commands * self.avg_concise_output_per_call
310        } else {
311            est_output_without
312        };
313        let output_saved = est_output_without.saturating_sub(est_output_with);
314
315        let output_cost_without = est_output_without as f64 / 1_000_000.0 * self.output_price_per_m;
316        let output_cost_with = est_output_with as f64 / 1_000_000.0 * self.output_price_per_m;
317
318        let total_without = input_cost_without + output_cost_without;
319        let total_with = input_cost_with + output_cost_with;
320
321        CostBreakdown {
322            input_cost_without,
323            input_cost_with,
324            output_cost_without,
325            output_cost_with,
326            total_cost_without: total_without,
327            total_cost_with: total_with,
328            total_saved: total_without - total_with,
329            estimated_output_tokens_without: est_output_without,
330            estimated_output_tokens_with: est_output_with,
331            output_tokens_saved: output_saved,
332        }
333    }
334}
335
336#[cfg(test)]
337mod tests {
338    use super::{CompressionTotals, CostModel, StatsStore, TrafficClass, classify_command};
339    use crate::core::gain::model_pricing::PricingMatchKind;
340
341    #[test]
342    fn known_tools_are_classified_by_delivery_contract() {
343        assert_eq!(classify_command("ctx_read"), TrafficClass::Compressible);
344        assert_eq!(classify_command("ctx_shell"), TrafficClass::Compressible);
345        assert_eq!(classify_command("ctx_search"), TrafficClass::Compressible);
346        assert_eq!(classify_command("ctx_compose"), TrafficClass::Passthrough);
347        assert_eq!(classify_command("ctx_glob"), TrafficClass::Passthrough);
348        assert_eq!(classify_command("cli_full"), TrafficClass::Compressible);
349    }
350
351    #[test]
352    fn cost_model_uses_resolved_model_pricing() {
353        let _lock = crate::core::data_dir::test_env_lock();
354        let old_lean_ctx_model = std::env::var("LEAN_CTX_MODEL").ok();
355        let old_lctx_model = std::env::var("LCTX_MODEL").ok();
356        crate::test_env::set_var("LEAN_CTX_MODEL", "claude-opus-4.5");
357        crate::test_env::remove_var("LCTX_MODEL");
358
359        let model = CostModel::default();
360
361        assert_eq!(model.model_key, "claude-opus-4.5");
362        assert_eq!(model.pricing_match_kind, PricingMatchKind::Exact);
363        assert_eq!(model.input_price_per_m, 5.0);
364        assert_eq!(model.output_price_per_m, 25.0);
365
366        match old_lean_ctx_model {
367            Some(value) => crate::test_env::set_var("LEAN_CTX_MODEL", value),
368            None => crate::test_env::remove_var("LEAN_CTX_MODEL"),
369        }
370        match old_lctx_model {
371            Some(value) => crate::test_env::set_var("LCTX_MODEL", value),
372            None => crate::test_env::remove_var("LCTX_MODEL"),
373        }
374    }
375
376    #[test]
377    fn compression_totals_fall_back_for_legacy_stats() {
378        let mut store = StatsStore::default();
379        store.commands.insert(
380            "ctx_read".into(),
381            super::CommandStats {
382                count: 1,
383                input_tokens: 1_000,
384                output_tokens: 400,
385            },
386        );
387        store.commands.insert(
388            "ctx_glob".into(),
389            super::CommandStats {
390                count: 1,
391                input_tokens: 500,
392                output_tokens: 500,
393            },
394        );
395
396        assert_eq!(store.compression_totals().saved_tokens(), 600);
397        assert_eq!(store.compression_totals().compression_pct(), 60.0);
398    }
399
400    #[test]
401    fn explicit_command_tag_overrides_legacy_inference() {
402        let mut store = StatsStore::default();
403        store.commands.insert(
404            "custom_tool".into(),
405            super::CommandStats {
406                count: 1,
407                input_tokens: 100,
408                output_tokens: 25,
409            },
410        );
411        store
412            .command_classes
413            .insert("custom_tool".into(), TrafficClass::Compressible);
414
415        assert_eq!(store.compression_totals().input_tokens, 100);
416        assert_eq!(store.total_reduction_pct(), 0.0);
417    }
418
419    #[test]
420    fn compression_totals_handle_zero_input_without_nan() {
421        let totals = CompressionTotals::default();
422        assert_eq!(totals.saved_tokens(), 0);
423        assert_eq!(totals.compression_pct(), 0.0);
424        assert_eq!(StatsStore::default().total_reduction_pct(), 0.0);
425    }
426
427    #[test]
428    fn first_result_is_first_inject_only() {
429        let mut store = StatsStore::default();
430        store.record_tool_result_savings(1_000, 7);
431        assert_eq!(store.stream_savings(7), (1_000, 0));
432        assert_eq!(store.last_tool_result_turn, 7);
433    }
434
435    #[test]
436    fn turn_advance_never_extrapolates_rereads() {
437        // #1283: re-read savings are measured (record_reread), never derived
438        // from turn arithmetic on the lifetime-cumulative active set.
439        let mut store = StatsStore::default();
440        store.record_tool_result_savings(1_000, 7);
441        store.record_tool_result_savings(500, 8);
442        assert_eq!(store.stream_savings(8), (1_500, 0));
443        assert_eq!(store.reread_tokens_saved, 0);
444    }
445
446    #[test]
447    fn measured_rereads_pass_through_unscaled() {
448        let mut store = StatsStore::default();
449        store.record_tool_result_savings(1_000, 7);
450        store.reread_tokens_saved = 300; // as recorded by the cache-hit path
451        assert_eq!(store.stream_savings(7), (1_000, 300));
452        // A later observation turn must not inflate the measured value.
453        assert_eq!(store.stream_savings(5_000), (1_000, 300));
454    }
455
456    #[test]
457    fn skipped_turns_do_not_multiply_resident_savings() {
458        let mut store = StatsStore::default();
459        store.record_tool_result_savings(2_000, 3);
460        assert_eq!(store.stream_savings(6), (2_000, 0));
461    }
462
463    #[test]
464    fn daemon_free_result_never_guesses_rereads() {
465        let mut store = StatsStore::default();
466        store.record_tool_result_savings(2_000, 0);
467        assert_eq!(store.stream_savings(100), (2_000, 0));
468        assert_eq!(store.active_tool_result_tokens_saved, 0);
469    }
470
471    #[test]
472    fn legacy_stats_fall_back_to_effective_compression() {
473        let mut store = StatsStore::default();
474        store.commands.insert(
475            "ctx_read".into(),
476            super::CommandStats {
477                count: 1,
478                input_tokens: 10_000,
479                output_tokens: 2_500,
480            },
481        );
482        assert_eq!(store.stream_savings(50), (7_500, 0));
483    }
484
485    #[test]
486    fn stream_counters_saturate_instead_of_wrapping() {
487        let mut store = StatsStore::default();
488        store.record_tool_result_savings(u64::MAX, 1);
489        store.record_tool_result_savings(u64::MAX, u64::MAX);
490        assert_eq!(store.first_inject_tokens_saved, u64::MAX);
491        assert_eq!(store.active_tool_result_tokens_saved, u64::MAX);
492        // #1283: no extrapolation — reread stays at its measured value.
493        assert_eq!(store.reread_tokens_saved, 0);
494    }
495}