1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5const DEFAULT_CONTEXT_WINDOW: usize = 128_000;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct ContextLedger {
9 pub window_size: usize,
10 pub entries: Vec<LedgerEntry>,
11 pub total_tokens_sent: usize,
12 pub total_tokens_saved: usize,
13}
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct LedgerEntry {
17 pub path: String,
18 pub mode: String,
19 pub original_tokens: usize,
20 pub sent_tokens: usize,
21 pub timestamp: i64,
22}
23
24#[derive(Debug, Clone)]
25pub struct ContextPressure {
26 pub utilization: f64,
27 pub remaining_tokens: usize,
28 pub entries_count: usize,
29 pub recommendation: PressureAction,
30}
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum PressureAction {
34 NoAction,
35 SuggestCompression,
36 ForceCompression,
37 EvictLeastRelevant,
38}
39
40impl ContextLedger {
41 pub fn new() -> Self {
42 Self {
43 window_size: DEFAULT_CONTEXT_WINDOW,
44 entries: Vec::new(),
45 total_tokens_sent: 0,
46 total_tokens_saved: 0,
47 }
48 }
49
50 pub fn with_window_size(size: usize) -> Self {
51 Self {
52 window_size: size,
53 entries: Vec::new(),
54 total_tokens_sent: 0,
55 total_tokens_saved: 0,
56 }
57 }
58
59 pub fn record(&mut self, path: &str, mode: &str, original_tokens: usize, sent_tokens: usize) {
60 if let Some(existing) = self.entries.iter_mut().find(|e| e.path == path) {
61 self.total_tokens_sent -= existing.sent_tokens;
62 self.total_tokens_saved -= existing
63 .original_tokens
64 .saturating_sub(existing.sent_tokens);
65 existing.mode = mode.to_string();
66 existing.original_tokens = original_tokens;
67 existing.sent_tokens = sent_tokens;
68 existing.timestamp = chrono::Utc::now().timestamp();
69 } else {
70 self.entries.push(LedgerEntry {
71 path: path.to_string(),
72 mode: mode.to_string(),
73 original_tokens,
74 sent_tokens,
75 timestamp: chrono::Utc::now().timestamp(),
76 });
77 }
78 self.total_tokens_sent += sent_tokens;
79 self.total_tokens_saved += original_tokens.saturating_sub(sent_tokens);
80 }
81
82 pub fn pressure(&self) -> ContextPressure {
83 let utilization = self.total_tokens_sent as f64 / self.window_size as f64;
84 let remaining = self.window_size.saturating_sub(self.total_tokens_sent);
85
86 let recommendation = if utilization > 0.9 {
87 PressureAction::EvictLeastRelevant
88 } else if utilization > 0.75 {
89 PressureAction::ForceCompression
90 } else if utilization > 0.5 {
91 PressureAction::SuggestCompression
92 } else {
93 PressureAction::NoAction
94 };
95
96 ContextPressure {
97 utilization,
98 remaining_tokens: remaining,
99 entries_count: self.entries.len(),
100 recommendation,
101 }
102 }
103
104 pub fn compression_ratio(&self) -> f64 {
105 let total_original: usize = self.entries.iter().map(|e| e.original_tokens).sum();
106 if total_original == 0 {
107 return 1.0;
108 }
109 self.total_tokens_sent as f64 / total_original as f64
110 }
111
112 pub fn files_by_token_cost(&self) -> Vec<(String, usize)> {
113 let mut costs: Vec<(String, usize)> = self
114 .entries
115 .iter()
116 .map(|e| (e.path.clone(), e.sent_tokens))
117 .collect();
118 costs.sort_by_key(|b| std::cmp::Reverse(b.1));
119 costs
120 }
121
122 pub fn mode_distribution(&self) -> HashMap<String, usize> {
123 let mut dist: HashMap<String, usize> = HashMap::new();
124 for entry in &self.entries {
125 *dist.entry(entry.mode.clone()).or_insert(0) += 1;
126 }
127 dist
128 }
129
130 pub fn eviction_candidates(&self, keep_count: usize) -> Vec<String> {
131 if self.entries.len() <= keep_count {
132 return Vec::new();
133 }
134 let mut sorted = self.entries.clone();
135 sorted.sort_by_key(|e| e.timestamp);
136 sorted
137 .iter()
138 .take(self.entries.len() - keep_count)
139 .map(|e| e.path.clone())
140 .collect()
141 }
142
143 pub fn remove(&mut self, path: &str) {
144 if let Some(idx) = self.entries.iter().position(|e| e.path == path) {
145 let entry = &self.entries[idx];
146 self.total_tokens_sent -= entry.sent_tokens;
147 self.total_tokens_saved -= entry.original_tokens.saturating_sub(entry.sent_tokens);
148 self.entries.remove(idx);
149 }
150 }
151
152 pub fn save(&self) {
153 if let Ok(dir) = crate::core::data_dir::lean_ctx_data_dir() {
154 let path = dir.join("context_ledger.json");
155 if let Ok(json) = serde_json::to_string(self) {
156 let _ = std::fs::write(path, json);
157 }
158 }
159 }
160
161 pub fn load() -> Self {
162 crate::core::data_dir::lean_ctx_data_dir()
163 .ok()
164 .map(|d| d.join("context_ledger.json"))
165 .and_then(|p| std::fs::read_to_string(p).ok())
166 .and_then(|s| serde_json::from_str(&s).ok())
167 .unwrap_or_default()
168 }
169
170 pub fn format_summary(&self) -> String {
171 let pressure = self.pressure();
172 format!(
173 "CTX: {}/{} tokens ({:.0}%), {} files, ratio {:.2}, action: {:?}",
174 self.total_tokens_sent,
175 self.window_size,
176 pressure.utilization * 100.0,
177 self.entries.len(),
178 self.compression_ratio(),
179 pressure.recommendation,
180 )
181 }
182}
183
184#[derive(Debug, Clone)]
185pub struct ReinjectionAction {
186 pub path: String,
187 pub current_mode: String,
188 pub new_mode: String,
189 pub tokens_freed: usize,
190}
191
192#[derive(Debug, Clone)]
193pub struct ReinjectionPlan {
194 pub actions: Vec<ReinjectionAction>,
195 pub total_tokens_freed: usize,
196 pub new_utilization: f64,
197}
198
199impl ContextLedger {
200 pub fn reinjection_plan(
201 &self,
202 intent: &super::intent_engine::StructuredIntent,
203 target_utilization: f64,
204 ) -> ReinjectionPlan {
205 let current_util = self.total_tokens_sent as f64 / self.window_size as f64;
206 if current_util <= target_utilization {
207 return ReinjectionPlan {
208 actions: Vec::new(),
209 total_tokens_freed: 0,
210 new_utilization: current_util,
211 };
212 }
213
214 let tokens_to_free =
215 self.total_tokens_sent - (self.window_size as f64 * target_utilization) as usize;
216
217 let target_set: std::collections::HashSet<&str> = intent
218 .targets
219 .iter()
220 .map(std::string::String::as_str)
221 .collect();
222
223 let mut candidates: Vec<(usize, &LedgerEntry)> = self
224 .entries
225 .iter()
226 .enumerate()
227 .filter(|(_, e)| !target_set.iter().any(|t| e.path.contains(t)))
228 .collect();
229
230 candidates.sort_by(|a, b| {
231 let a_age = a.1.timestamp;
232 let b_age = b.1.timestamp;
233 a_age.cmp(&b_age)
234 });
235
236 let mut actions = Vec::new();
237 let mut freed = 0usize;
238
239 for (_, entry) in &candidates {
240 if freed >= tokens_to_free {
241 break;
242 }
243 if let Some((new_mode, new_tokens)) = downgrade_mode(&entry.mode, entry.sent_tokens) {
244 let saving = entry.sent_tokens.saturating_sub(new_tokens);
245 if saving > 0 {
246 actions.push(ReinjectionAction {
247 path: entry.path.clone(),
248 current_mode: entry.mode.clone(),
249 new_mode,
250 tokens_freed: saving,
251 });
252 freed += saving;
253 }
254 }
255 }
256
257 let new_sent = self.total_tokens_sent.saturating_sub(freed);
258 let new_utilization = new_sent as f64 / self.window_size as f64;
259
260 ReinjectionPlan {
261 actions,
262 total_tokens_freed: freed,
263 new_utilization,
264 }
265 }
266}
267
268fn downgrade_mode(current_mode: &str, current_tokens: usize) -> Option<(String, usize)> {
269 match current_mode {
270 "full" => Some(("signatures".to_string(), current_tokens / 5)),
271 "aggressive" => Some(("signatures".to_string(), current_tokens / 3)),
272 "signatures" => Some(("map".to_string(), current_tokens / 2)),
273 "map" => Some(("reference".to_string(), current_tokens / 4)),
274 _ => None,
275 }
276}
277
278impl Default for ContextLedger {
279 fn default() -> Self {
280 Self::new()
281 }
282}
283
284#[cfg(test)]
285mod tests {
286 use super::*;
287
288 #[test]
289 fn new_ledger_is_empty() {
290 let ledger = ContextLedger::new();
291 assert_eq!(ledger.total_tokens_sent, 0);
292 assert_eq!(ledger.entries.len(), 0);
293 assert_eq!(ledger.pressure().recommendation, PressureAction::NoAction);
294 }
295
296 #[test]
297 fn record_tracks_tokens() {
298 let mut ledger = ContextLedger::with_window_size(10000);
299 ledger.record("src/main.rs", "full", 500, 500);
300 ledger.record("src/lib.rs", "signatures", 1000, 200);
301 assert_eq!(ledger.total_tokens_sent, 700);
302 assert_eq!(ledger.total_tokens_saved, 800);
303 assert_eq!(ledger.entries.len(), 2);
304 }
305
306 #[test]
307 fn record_updates_existing_entry() {
308 let mut ledger = ContextLedger::with_window_size(10000);
309 ledger.record("src/main.rs", "full", 500, 500);
310 ledger.record("src/main.rs", "signatures", 500, 100);
311 assert_eq!(ledger.entries.len(), 1);
312 assert_eq!(ledger.total_tokens_sent, 100);
313 assert_eq!(ledger.total_tokens_saved, 400);
314 }
315
316 #[test]
317 fn pressure_escalates() {
318 let mut ledger = ContextLedger::with_window_size(1000);
319 ledger.record("a.rs", "full", 600, 600);
320 assert_eq!(
321 ledger.pressure().recommendation,
322 PressureAction::SuggestCompression
323 );
324 ledger.record("b.rs", "full", 200, 200);
325 assert_eq!(
326 ledger.pressure().recommendation,
327 PressureAction::ForceCompression
328 );
329 ledger.record("c.rs", "full", 150, 150);
330 assert_eq!(
331 ledger.pressure().recommendation,
332 PressureAction::EvictLeastRelevant
333 );
334 }
335
336 #[test]
337 fn compression_ratio_accurate() {
338 let mut ledger = ContextLedger::with_window_size(10000);
339 ledger.record("a.rs", "full", 1000, 1000);
340 ledger.record("b.rs", "signatures", 1000, 200);
341 let ratio = ledger.compression_ratio();
342 assert!((ratio - 0.6).abs() < 0.01);
343 }
344
345 #[test]
346 fn eviction_returns_oldest() {
347 let mut ledger = ContextLedger::with_window_size(10000);
348 ledger.record("old.rs", "full", 100, 100);
349 std::thread::sleep(std::time::Duration::from_millis(10));
350 ledger.record("new.rs", "full", 100, 100);
351 let candidates = ledger.eviction_candidates(1);
352 assert_eq!(candidates, vec!["old.rs"]);
353 }
354
355 #[test]
356 fn remove_updates_totals() {
357 let mut ledger = ContextLedger::with_window_size(10000);
358 ledger.record("a.rs", "full", 500, 500);
359 ledger.record("b.rs", "full", 300, 300);
360 ledger.remove("a.rs");
361 assert_eq!(ledger.total_tokens_sent, 300);
362 assert_eq!(ledger.entries.len(), 1);
363 }
364
365 #[test]
366 fn mode_distribution_counts() {
367 let mut ledger = ContextLedger::new();
368 ledger.record("a.rs", "full", 100, 100);
369 ledger.record("b.rs", "signatures", 100, 50);
370 ledger.record("c.rs", "full", 100, 100);
371 let dist = ledger.mode_distribution();
372 assert_eq!(dist.get("full"), Some(&2));
373 assert_eq!(dist.get("signatures"), Some(&1));
374 }
375
376 #[test]
377 fn format_summary_includes_key_info() {
378 let mut ledger = ContextLedger::with_window_size(10000);
379 ledger.record("a.rs", "full", 500, 500);
380 let summary = ledger.format_summary();
381 assert!(summary.contains("500/10000"));
382 assert!(summary.contains("1 files"));
383 }
384
385 #[test]
386 fn reinjection_no_action_when_low_pressure() {
387 use crate::core::intent_engine::StructuredIntent;
388
389 let mut ledger = ContextLedger::with_window_size(10000);
390 ledger.record("a.rs", "full", 100, 100);
391 let intent = StructuredIntent::from_query("fix bug in a.rs");
392 let plan = ledger.reinjection_plan(&intent, 0.7);
393 assert!(plan.actions.is_empty());
394 assert_eq!(plan.total_tokens_freed, 0);
395 }
396
397 #[test]
398 fn reinjection_downgrades_non_target_files() {
399 use crate::core::intent_engine::StructuredIntent;
400
401 let mut ledger = ContextLedger::with_window_size(1000);
402 ledger.record("src/target.rs", "full", 400, 400);
403 std::thread::sleep(std::time::Duration::from_millis(10));
404 ledger.record("src/other.rs", "full", 400, 400);
405 std::thread::sleep(std::time::Duration::from_millis(10));
406 ledger.record("src/utils.rs", "full", 200, 200);
407
408 let intent = StructuredIntent::from_query("fix bug in target.rs");
409 let plan = ledger.reinjection_plan(&intent, 0.5);
410
411 assert!(!plan.actions.is_empty());
412 assert!(
413 plan.actions.iter().all(|a| !a.path.contains("target")),
414 "should not downgrade target file"
415 );
416 assert!(plan.total_tokens_freed > 0);
417 }
418
419 #[test]
420 fn reinjection_preserves_targets() {
421 use crate::core::intent_engine::StructuredIntent;
422
423 let mut ledger = ContextLedger::with_window_size(1000);
424 ledger.record("src/auth.rs", "full", 900, 900);
425 let intent = StructuredIntent::from_query("fix bug in auth.rs");
426 let plan = ledger.reinjection_plan(&intent, 0.5);
427 assert!(
428 plan.actions.is_empty(),
429 "should not downgrade target files even under pressure"
430 );
431 }
432
433 #[test]
434 fn downgrade_mode_chain() {
435 assert_eq!(
436 downgrade_mode("full", 1000),
437 Some(("signatures".to_string(), 200))
438 );
439 assert_eq!(
440 downgrade_mode("signatures", 200),
441 Some(("map".to_string(), 100))
442 );
443 assert_eq!(
444 downgrade_mode("map", 100),
445 Some(("reference".to_string(), 25))
446 );
447 assert_eq!(downgrade_mode("reference", 25), None);
448 }
449}