1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3
4use crate::core::gain::model_pricing::{ModelQuote, PricingMatchKind};
5
6#[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 #[serde(default)]
21 pub command_classes: HashMap<String, TrafficClass>,
22 #[serde(default)]
24 pub first_inject_tokens_saved: u64,
25 #[serde(default)]
27 pub reread_tokens_saved: u64,
28 #[serde(default)]
30 pub active_tool_result_tokens_saved: u64,
31 #[serde(default)]
33 pub last_tool_result_turn: u64,
34 #[serde(default)]
37 pub stream_tracked_results: u64,
38}
39
40#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
42#[serde(rename_all = "snake_case")]
43pub enum TrafficClass {
44 Compressible,
45 Passthrough,
46}
47
48#[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 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 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 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 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
134pub 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#[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 #[serde(default)]
176 pub last_session_cache_hits: Option<u64>,
177 #[serde(default)]
178 pub last_session_cache_reads: Option<u64>,
179}
180
181#[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#[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#[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 #[serde(default)]
213 pub version: String,
214 #[serde(default)]
218 pub compressible_input_tokens: u64,
219 #[serde(default)]
220 pub compressible_output_tokens: u64,
221}
222
223pub struct GainSummary {
225 pub total_saved: u64,
226 pub total_calls: u64,
227}
228
229pub const DEFAULT_INPUT_PRICE_PER_M: f64 = 2.50;
231pub const DEFAULT_OUTPUT_PRICE_PER_M: f64 = 10.0;
232
233pub 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
277pub 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 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 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; assert_eq!(store.stream_savings(7), (1_000, 300));
452 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 assert_eq!(store.reread_tokens_saved, 0);
494 }
495}