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(crate) struct CompressionTotals {
51    pub(crate) input_tokens: u64,
52    pub(crate) 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. Advancing
71    /// turns prices every already-resident result as a cache re-read; results
72    /// emitted in the same turn are not spuriously re-billed against each other.
73    pub(crate) fn record_tool_result_savings(&mut self, saved: u64, turn: u64) {
74        if turn > 0 && self.last_tool_result_turn > 0 && turn > self.last_tool_result_turn {
75            let elapsed = turn - self.last_tool_result_turn;
76            self.reread_tokens_saved = self
77                .reread_tokens_saved
78                .saturating_add(self.active_tool_result_tokens_saved.saturating_mul(elapsed));
79        }
80        if turn > 0 {
81            self.last_tool_result_turn = self.last_tool_result_turn.max(turn);
82            self.active_tool_result_tokens_saved =
83                self.active_tool_result_tokens_saved.saturating_add(saved);
84        }
85        self.first_inject_tokens_saved = self.first_inject_tokens_saved.saturating_add(saved);
86        self.stream_tracked_results = self.stream_tracked_results.saturating_add(1);
87    }
88
89    /// Stream totals through `observed_turn`, including re-reads since the last
90    /// tool result. Legacy stores fall back to measured compressible savings.
91    pub(crate) fn stream_savings(&self, observed_turn: u64) -> (u64, u64) {
92        if self.stream_tracked_results == 0 {
93            return (self.compression_totals().saved_tokens(), 0);
94        }
95        let pending_turns = observed_turn.saturating_sub(self.last_tool_result_turn);
96        let pending = self
97            .active_tool_result_tokens_saved
98            .saturating_mul(pending_turns);
99        (
100            self.first_inject_tokens_saved,
101            self.reread_tokens_saved.saturating_add(pending),
102        )
103    }
104
105    /// Computes effective compression from tagged command rows.
106    pub(crate) fn compression_totals(&self) -> CompressionTotals {
107        self.commands
108            .iter()
109            .filter(|(command, _)| {
110                self.command_classes
111                    .get(*command)
112                    .copied()
113                    .unwrap_or_else(|| classify_command(command))
114                    == TrafficClass::Compressible
115            })
116            .fold(CompressionTotals::default(), |mut totals, (_, stats)| {
117                totals.input_tokens = totals.input_tokens.saturating_add(stats.input_tokens);
118                totals.output_tokens = totals.output_tokens.saturating_add(stats.output_tokens);
119                totals
120            })
121    }
122
123    /// Total reduction across both compressible and passthrough traffic.
124    pub(crate) fn total_reduction_pct(&self) -> f64 {
125        if self.total_input_tokens == 0 {
126            0.0
127        } else {
128            self.total_input_tokens
129                .saturating_sub(self.total_output_tokens) as f64
130                / self.total_input_tokens as f64
131                * 100.0
132        }
133    }
134}
135
136/// Classifies normalized stats keys, with an explicit passthrough default for
137/// control/listing tools so only read, shell, and search output drives the
138/// effective-compression denominator.
139pub(crate) fn classify_command(command: &str) -> TrafficClass {
140    match command {
141        "cli_full" | "cli_raw" | "cli_glob" | "cli_find" | "cli_deps" | "cli_ls"
142        | "ctx_compose" | "ctx_glob" | "ctx_tree" => TrafficClass::Passthrough,
143        c if c.starts_with("cli_") => TrafficClass::Compressible,
144        "ctx_shell" | "ctx_search" | "ctx_semantic_search" => TrafficClass::Compressible,
145        c if c.starts_with("ctx_read")
146            || c.starts_with("ctx_multi_read")
147            || c == "ctx_smart_read"
148            || c == "ctx_git_read"
149            || c == "ctx_url_read" =>
150        {
151            TrafficClass::Compressible
152        }
153        _ => TrafficClass::Passthrough,
154    }
155}
156
157/// Aggregated CEP (Cognitive Efficiency Protocol) metrics across sessions.
158#[derive(Serialize, Deserialize, Clone, Default)]
159pub struct CepStats {
160    pub sessions: u64,
161    pub total_cache_hits: u64,
162    pub total_cache_reads: u64,
163    pub total_tokens_original: u64,
164    pub total_tokens_compressed: u64,
165    pub modes: HashMap<String, u64>,
166    pub scores: Vec<CepSessionSnapshot>,
167    #[serde(default)]
168    pub last_session_pid: Option<u32>,
169    #[serde(default)]
170    pub last_session_original: Option<u64>,
171    #[serde(default)]
172    pub last_session_compressed: Option<u64>,
173    /// Cumulative cache hits/reads observed for the current PID at the last
174    /// snapshot. Used to accumulate *deltas* across repeated snapshots within
175    /// one server process, so `total_cache_hits` keeps tracking cache activity
176    /// after the first checkpoint instead of freezing (#361).
177    #[serde(default)]
178    pub last_session_cache_hits: Option<u64>,
179    #[serde(default)]
180    pub last_session_cache_reads: Option<u64>,
181}
182
183/// Point-in-time snapshot of CEP scores for a single session.
184#[derive(Serialize, Deserialize, Clone)]
185pub struct CepSessionSnapshot {
186    pub timestamp: String,
187    pub score: u32,
188    pub cache_hit_rate: u32,
189    pub mode_diversity: u32,
190    pub compression_rate: u32,
191    pub tool_calls: u64,
192    pub tokens_saved: u64,
193    pub complexity: String,
194}
195
196/// Per-command token statistics: invocation count and input/output totals.
197#[derive(Serialize, Deserialize, Clone, Default, Debug)]
198pub struct CommandStats {
199    pub count: u64,
200    pub input_tokens: u64,
201    pub output_tokens: u64,
202}
203
204/// Daily aggregate: command count and token totals for one calendar day.
205#[derive(Serialize, Deserialize, Clone, Default)]
206pub struct DayStats {
207    pub date: String,
208    pub commands: u64,
209    pub input_tokens: u64,
210    pub output_tokens: u64,
211    /// lean-ctx version active when this day's stats were last recorded.
212    /// Lets `lean-ctx gain` attribute per-day compression changes to a release
213    /// (#307). Empty for days recorded before this field existed.
214    #[serde(default)]
215    pub version: String,
216}
217
218/// High-level token savings summary for display.
219pub struct GainSummary {
220    pub total_saved: u64,
221    pub total_calls: u64,
222}
223
224/// Average LLM pricing per 1M tokens (blended across Claude, GPT, Gemini).
225pub const DEFAULT_INPUT_PRICE_PER_M: f64 = 2.50;
226pub const DEFAULT_OUTPUT_PRICE_PER_M: f64 = 10.0;
227
228/// LLM pricing model for estimating dollar savings from token compression.
229pub struct CostModel {
230    pub model_key: String,
231    pub pricing_match_kind: PricingMatchKind,
232    pub input_price_per_m: f64,
233    pub output_price_per_m: f64,
234    pub avg_verbose_output_per_call: u64,
235    pub avg_concise_output_per_call: u64,
236}
237
238impl Default for CostModel {
239    fn default() -> Self {
240        let pricing = crate::core::gain::model_pricing::ModelPricing::load();
241        let quote = pricing.quote(resolved_gain_model().as_deref());
242        Self::from_quote(quote)
243    }
244}
245
246fn resolved_gain_model() -> Option<String> {
247    std::env::var("LEAN_CTX_MODEL")
248        .or_else(|_| std::env::var("LCTX_MODEL"))
249        .ok()
250        .filter(|s| !s.trim().is_empty())
251        .or_else(|| {
252            crate::core::config::Config::load()
253                .cost
254                .model_for_client("cli")
255        })
256        .or_else(crate::proxy::usage_meter::persisted_dominant_model)
257}
258
259impl CostModel {
260    fn from_quote(quote: ModelQuote) -> Self {
261        Self {
262            model_key: quote.model_key,
263            pricing_match_kind: quote.match_kind,
264            input_price_per_m: quote.cost.input_per_m,
265            output_price_per_m: quote.cost.output_per_m,
266            avg_verbose_output_per_call: 180,
267            avg_concise_output_per_call: 120,
268        }
269    }
270}
271
272/// Detailed cost comparison: with vs. without lean-ctx compression.
273pub struct CostBreakdown {
274    pub input_cost_without: f64,
275    pub input_cost_with: f64,
276    pub output_cost_without: f64,
277    pub output_cost_with: f64,
278    pub total_cost_without: f64,
279    pub total_cost_with: f64,
280    pub total_saved: f64,
281    pub estimated_output_tokens_without: u64,
282    pub estimated_output_tokens_with: u64,
283    pub output_tokens_saved: u64,
284}
285
286impl CostModel {
287    /// Calculates the full cost breakdown from the stats store.
288    pub fn calculate(&self, store: &StatsStore) -> CostBreakdown {
289        let input_cost_without =
290            store.total_input_tokens as f64 / 1_000_000.0 * self.input_price_per_m;
291        let input_cost_with =
292            store.total_output_tokens as f64 / 1_000_000.0 * self.input_price_per_m;
293
294        let input_saved = store
295            .total_input_tokens
296            .saturating_sub(store.total_output_tokens);
297        let compression_rate = if store.total_input_tokens > 0 {
298            input_saved as f64 / store.total_input_tokens as f64
299        } else {
300            0.0
301        };
302        let est_output_without = store.total_commands * self.avg_verbose_output_per_call;
303        let est_output_with = if compression_rate > 0.01 {
304            store.total_commands * self.avg_concise_output_per_call
305        } else {
306            est_output_without
307        };
308        let output_saved = est_output_without.saturating_sub(est_output_with);
309
310        let output_cost_without = est_output_without as f64 / 1_000_000.0 * self.output_price_per_m;
311        let output_cost_with = est_output_with as f64 / 1_000_000.0 * self.output_price_per_m;
312
313        let total_without = input_cost_without + output_cost_without;
314        let total_with = input_cost_with + output_cost_with;
315
316        CostBreakdown {
317            input_cost_without,
318            input_cost_with,
319            output_cost_without,
320            output_cost_with,
321            total_cost_without: total_without,
322            total_cost_with: total_with,
323            total_saved: total_without - total_with,
324            estimated_output_tokens_without: est_output_without,
325            estimated_output_tokens_with: est_output_with,
326            output_tokens_saved: output_saved,
327        }
328    }
329}
330
331#[cfg(test)]
332mod tests {
333    use super::{CompressionTotals, CostModel, StatsStore, TrafficClass, classify_command};
334    use crate::core::gain::model_pricing::PricingMatchKind;
335
336    #[test]
337    fn known_tools_are_classified_by_delivery_contract() {
338        assert_eq!(classify_command("ctx_read"), TrafficClass::Compressible);
339        assert_eq!(classify_command("ctx_shell"), TrafficClass::Compressible);
340        assert_eq!(classify_command("ctx_search"), TrafficClass::Compressible);
341        assert_eq!(classify_command("ctx_compose"), TrafficClass::Passthrough);
342        assert_eq!(classify_command("ctx_glob"), TrafficClass::Passthrough);
343        assert_eq!(classify_command("cli_full"), TrafficClass::Passthrough);
344    }
345
346    #[test]
347    fn cost_model_uses_resolved_model_pricing() {
348        let _lock = crate::core::data_dir::test_env_lock();
349        let old_lean_ctx_model = std::env::var("LEAN_CTX_MODEL").ok();
350        let old_lctx_model = std::env::var("LCTX_MODEL").ok();
351        // SAFETY: test holds exclusive env lock via test_env_lock()
352        unsafe {
353            std::env::set_var("LEAN_CTX_MODEL", "claude-opus-4.5");
354            std::env::remove_var("LCTX_MODEL");
355        }
356
357        let model = CostModel::default();
358
359        assert_eq!(model.model_key, "claude-opus-4.5");
360        assert_eq!(model.pricing_match_kind, PricingMatchKind::Exact);
361        assert_eq!(model.input_price_per_m, 5.0);
362        assert_eq!(model.output_price_per_m, 25.0);
363
364        // SAFETY: test holds exclusive env lock via test_env_lock()
365        unsafe {
366            match old_lean_ctx_model {
367                Some(value) => std::env::set_var("LEAN_CTX_MODEL", value),
368                None => std::env::remove_var("LEAN_CTX_MODEL"),
369            }
370            match old_lctx_model {
371                Some(value) => std::env::set_var("LCTX_MODEL", value),
372                None => std::env::remove_var("LCTX_MODEL"),
373            }
374        }
375    }
376
377    #[test]
378    fn compression_totals_fall_back_for_legacy_stats() {
379        let mut store = StatsStore::default();
380        store.commands.insert(
381            "ctx_read".into(),
382            super::CommandStats {
383                count: 1,
384                input_tokens: 1_000,
385                output_tokens: 400,
386            },
387        );
388        store.commands.insert(
389            "ctx_glob".into(),
390            super::CommandStats {
391                count: 1,
392                input_tokens: 500,
393                output_tokens: 500,
394            },
395        );
396
397        assert_eq!(store.compression_totals().saved_tokens(), 600);
398        assert_eq!(store.compression_totals().compression_pct(), 60.0);
399    }
400
401    #[test]
402    fn explicit_command_tag_overrides_legacy_inference() {
403        let mut store = StatsStore::default();
404        store.commands.insert(
405            "custom_tool".into(),
406            super::CommandStats {
407                count: 1,
408                input_tokens: 100,
409                output_tokens: 25,
410            },
411        );
412        store
413            .command_classes
414            .insert("custom_tool".into(), TrafficClass::Compressible);
415
416        assert_eq!(store.compression_totals().input_tokens, 100);
417        assert_eq!(store.total_reduction_pct(), 0.0);
418    }
419
420    #[test]
421    fn compression_totals_handle_zero_input_without_nan() {
422        let totals = CompressionTotals::default();
423        assert_eq!(totals.saved_tokens(), 0);
424        assert_eq!(totals.compression_pct(), 0.0);
425        assert_eq!(StatsStore::default().total_reduction_pct(), 0.0);
426    }
427
428    #[test]
429    fn first_result_is_first_inject_only() {
430        let mut store = StatsStore::default();
431        store.record_tool_result_savings(1_000, 7);
432        assert_eq!(store.stream_savings(7), (1_000, 0));
433        assert_eq!(store.last_tool_result_turn, 7);
434    }
435
436    #[test]
437    fn next_turn_rereads_all_prior_results() {
438        let mut store = StatsStore::default();
439        store.record_tool_result_savings(1_000, 7);
440        store.record_tool_result_savings(500, 8);
441        assert_eq!(store.stream_savings(8), (1_500, 1_000));
442    }
443
444    #[test]
445    fn parallel_results_on_same_turn_do_not_reread_each_other() {
446        let mut store = StatsStore::default();
447        store.record_tool_result_savings(1_000, 7);
448        store.record_tool_result_savings(500, 7);
449        assert_eq!(store.stream_savings(7), (1_500, 0));
450        assert_eq!(store.stream_savings(8), (1_500, 1_500));
451    }
452
453    #[test]
454    fn skipped_turns_multiply_resident_savings() {
455        let mut store = StatsStore::default();
456        store.record_tool_result_savings(2_000, 3);
457        assert_eq!(store.stream_savings(6), (2_000, 6_000));
458    }
459
460    #[test]
461    fn daemon_free_result_never_guesses_rereads() {
462        let mut store = StatsStore::default();
463        store.record_tool_result_savings(2_000, 0);
464        assert_eq!(store.stream_savings(100), (2_000, 0));
465        assert_eq!(store.active_tool_result_tokens_saved, 0);
466    }
467
468    #[test]
469    fn legacy_stats_fall_back_to_effective_compression() {
470        let mut store = StatsStore::default();
471        store.commands.insert(
472            "ctx_read".into(),
473            super::CommandStats {
474                count: 1,
475                input_tokens: 10_000,
476                output_tokens: 2_500,
477            },
478        );
479        assert_eq!(store.stream_savings(50), (7_500, 0));
480    }
481
482    #[test]
483    fn stream_counters_saturate_instead_of_wrapping() {
484        let mut store = StatsStore::default();
485        store.record_tool_result_savings(u64::MAX, 1);
486        store.record_tool_result_savings(u64::MAX, u64::MAX);
487        assert_eq!(store.first_inject_tokens_saved, u64::MAX);
488        assert_eq!(store.reread_tokens_saved, u64::MAX);
489    }
490}
lean_ctx/core/stats/model.rs

lean_ctx/core/stats/
model.rs
