Skip to main content

lean_ctx/core/gain/
mod.rs

1pub mod bridge_status;
2pub mod gain_score;
3pub mod model_pricing;
4pub mod task_classifier;
5
6use serde::{Deserialize, Serialize};
7
8use crate::core::a2a::cost_attribution::CostStore;
9use crate::core::gain::gain_score::GainScore;
10use crate::core::gain::model_pricing::{ModelPricing, ModelQuote};
11use crate::core::gain::task_classifier::{TaskCategory, TaskClassifier};
12use crate::core::heatmap::HeatMap;
13use crate::core::stats::StatsStore;
14
15#[derive(Clone)]
16pub struct GainEngine {
17    pub stats: StatsStore,
18    pub costs: CostStore,
19    pub heatmap: HeatMap,
20    pub pricing: ModelPricing,
21    pub events: Vec<crate::core::events::LeanCtxEvent>,
22    pub session: Option<crate::core::session::SessionState>,
23}
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct GainSummary {
27    pub model: ModelQuote,
28    pub total_commands: u64,
29    pub input_tokens: u64,
30    pub output_tokens: u64,
31    pub tokens_saved: u64,
32    pub gain_rate_pct: f64,
33    /// Fixed per-turn context lean-ctx injects (tool schemas + server
34    /// instructions + rules block). On a provider WITHOUT prompt caching this
35    /// rides — and is re-billed on — every turn, so the net bill impact is
36    /// `tokens_saved − injected_overhead_tokens_per_turn × turns`. Note that
37    /// `tokens_saved` / `gain_rate_pct` are measured against lean-ctx-touched
38    /// traffic (their denominator), not the full provider bill (GitHub #361).
39    #[serde(default)]
40    pub injected_overhead_tokens_per_turn: u64,
41    pub avoided_usd: f64,
42    /// Estimated grid energy avoided (Wh) by keeping `tokens_saved` out of context.
43    pub energy_wh: f64,
44    /// Estimated CO₂-equivalent avoided (grams), derived from `energy_wh`.
45    pub co2_grams: f64,
46    pub tool_spend_usd: f64,
47    pub roi: Option<f64>,
48    pub score: GainScore,
49    #[serde(skip_serializing_if = "Option::is_none", default)]
50    pub daemon_hint: Option<String>,
51}
52
53#[derive(Debug, Clone, Serialize, Deserialize)]
54pub struct TaskGainRow {
55    pub category: TaskCategory,
56    pub commands: u64,
57    pub tokens_saved: u64,
58    pub tool_calls: u64,
59    pub tool_spend_usd: f64,
60}
61
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct FileGainRow {
64    pub path: String,
65    pub access_count: u32,
66    pub tokens_saved: u64,
67    pub compression_pct: f32,
68}
69
70impl GainEngine {
71    pub fn load() -> Self {
72        Self {
73            stats: crate::core::stats::load(),
74            costs: crate::core::a2a::cost_attribution::CostStore::load(),
75            heatmap: crate::core::heatmap::HeatMap::load(),
76            pricing: ModelPricing::load(),
77            events: crate::core::events::load_events_from_file(500),
78            session: crate::core::session::SessionState::load_latest(),
79        }
80    }
81
82    pub fn summary(&self, model: Option<&str>) -> GainSummary {
83        let quote = self.pricing.quote(model);
84        let tokens_saved = self
85            .stats
86            .total_input_tokens
87            .saturating_sub(self.stats.total_output_tokens);
88        let gain_rate_pct = if self.stats.total_input_tokens > 0 {
89            tokens_saved as f64 / self.stats.total_input_tokens as f64 * 100.0
90        } else {
91            0.0
92        };
93        let avoided_usd = quote.cost.estimate_usd(tokens_saved, 0, 0, 0);
94        let tool_spend_usd = self.costs.total_cost().max(0.0);
95        let roi = if tool_spend_usd > 0.0 {
96            Some(avoided_usd / tool_spend_usd)
97        } else {
98            None
99        };
100        let score = GainScore::compute(&self.stats, &self.costs, &self.pricing, model);
101        #[cfg(unix)]
102        let daemon_hint = if crate::daemon::is_daemon_running() {
103            None
104        } else {
105            Some(
106                "daemon not running — stats tracked locally (lean-ctx serve -d for full tracking)"
107                    .to_string(),
108            )
109        };
110        #[cfg(not(unix))]
111        let daemon_hint: Option<String> = None;
112        let injected_overhead_tokens_per_turn =
113            crate::core::context_overhead::ContextOverhead::cached().total_tokens() as u64;
114        GainSummary {
115            model: quote,
116            total_commands: self.stats.total_commands,
117            input_tokens: self.stats.total_input_tokens,
118            output_tokens: self.stats.total_output_tokens,
119            tokens_saved,
120            gain_rate_pct,
121            injected_overhead_tokens_per_turn,
122            avoided_usd,
123            energy_wh: crate::core::energy::wh_for_tokens(tokens_saved),
124            co2_grams: crate::core::energy::co2_grams_for_tokens(tokens_saved),
125            tool_spend_usd,
126            roi,
127            score,
128            daemon_hint,
129        }
130    }
131
132    pub fn gain_score(&self, model: Option<&str>) -> GainScore {
133        GainScore::compute(&self.stats, &self.costs, &self.pricing, model)
134    }
135
136    pub fn task_breakdown(&self) -> Vec<TaskGainRow> {
137        use std::collections::HashMap;
138
139        let mut by_cat: HashMap<TaskCategory, TaskGainRow> = HashMap::new();
140
141        for (cmd_key, st) in &self.stats.commands {
142            let cat = TaskClassifier::classify_command_key(cmd_key);
143            let row = by_cat.entry(cat).or_insert(TaskGainRow {
144                category: cat,
145                commands: 0,
146                tokens_saved: 0,
147                tool_calls: 0,
148                tool_spend_usd: 0.0,
149            });
150            row.commands += st.count;
151            row.tokens_saved += st.input_tokens.saturating_sub(st.output_tokens);
152        }
153
154        for (tool, tc) in &self.costs.tools {
155            let cat = TaskClassifier::classify_tool(tool);
156            let row = by_cat.entry(cat).or_insert(TaskGainRow {
157                category: cat,
158                commands: 0,
159                tokens_saved: 0,
160                tool_calls: 0,
161                tool_spend_usd: 0.0,
162            });
163            row.tool_calls += tc.total_calls;
164            row.tool_spend_usd += tc.cost_usd;
165        }
166
167        let mut out: Vec<TaskGainRow> = by_cat.into_values().collect();
168        out.sort_by_key(|x| std::cmp::Reverse(x.tokens_saved));
169        out
170    }
171
172    pub fn heatmap_gains(&self, limit: usize) -> Vec<FileGainRow> {
173        let mut items: Vec<_> = self.heatmap.entries.values().collect();
174        items.sort_by_key(|x| std::cmp::Reverse(x.total_tokens_saved));
175        items.truncate(limit);
176        items
177            .into_iter()
178            .map(|e| FileGainRow {
179                path: e.path.clone(),
180                access_count: e.access_count,
181                tokens_saved: e.total_tokens_saved,
182                compression_pct: e.avg_compression_ratio * 100.0,
183            })
184            .collect()
185    }
186}