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> =
218 intent.targets.iter().map(|t| t.as_str()).collect();
219
220 let mut candidates: Vec<(usize, &LedgerEntry)> = self
221 .entries
222 .iter()
223 .enumerate()
224 .filter(|(_, e)| !target_set.iter().any(|t| e.path.contains(t)))
225 .collect();
226
227 candidates.sort_by(|a, b| {
228 let a_age = a.1.timestamp;
229 let b_age = b.1.timestamp;
230 a_age.cmp(&b_age)
231 });
232
233 let mut actions = Vec::new();
234 let mut freed = 0usize;
235
236 for (_, entry) in &candidates {
237 if freed >= tokens_to_free {
238 break;
239 }
240 if let Some((new_mode, new_tokens)) = downgrade_mode(&entry.mode, entry.sent_tokens) {
241 let saving = entry.sent_tokens.saturating_sub(new_tokens);
242 if saving > 0 {
243 actions.push(ReinjectionAction {
244 path: entry.path.clone(),
245 current_mode: entry.mode.clone(),
246 new_mode,
247 tokens_freed: saving,
248 });
249 freed += saving;
250 }
251 }
252 }
253
254 let new_sent = self.total_tokens_sent.saturating_sub(freed);
255 let new_utilization = new_sent as f64 / self.window_size as f64;
256
257 ReinjectionPlan {
258 actions,
259 total_tokens_freed: freed,
260 new_utilization,
261 }
262 }
263}
264
265fn downgrade_mode(current_mode: &str, current_tokens: usize) -> Option<(String, usize)> {
266 match current_mode {
267 "full" => Some(("signatures".to_string(), current_tokens / 5)),
268 "aggressive" => Some(("signatures".to_string(), current_tokens / 3)),
269 "signatures" => Some(("map".to_string(), current_tokens / 2)),
270 "map" => Some(("reference".to_string(), current_tokens / 4)),
271 _ => None,
272 }
273}
274
275impl Default for ContextLedger {
276 fn default() -> Self {
277 Self::new()
278 }
279}
280
281#[cfg(test)]
282mod tests {
283 use super::*;
284
285 #[test]
286 fn new_ledger_is_empty() {
287 let ledger = ContextLedger::new();
288 assert_eq!(ledger.total_tokens_sent, 0);
289 assert_eq!(ledger.entries.len(), 0);
290 assert_eq!(ledger.pressure().recommendation, PressureAction::NoAction);
291 }
292
293 #[test]
294 fn record_tracks_tokens() {
295 let mut ledger = ContextLedger::with_window_size(10000);
296 ledger.record("src/main.rs", "full", 500, 500);
297 ledger.record("src/lib.rs", "signatures", 1000, 200);
298 assert_eq!(ledger.total_tokens_sent, 700);
299 assert_eq!(ledger.total_tokens_saved, 800);
300 assert_eq!(ledger.entries.len(), 2);
301 }
302
303 #[test]
304 fn record_updates_existing_entry() {
305 let mut ledger = ContextLedger::with_window_size(10000);
306 ledger.record("src/main.rs", "full", 500, 500);
307 ledger.record("src/main.rs", "signatures", 500, 100);
308 assert_eq!(ledger.entries.len(), 1);
309 assert_eq!(ledger.total_tokens_sent, 100);
310 assert_eq!(ledger.total_tokens_saved, 400);
311 }
312
313 #[test]
314 fn pressure_escalates() {
315 let mut ledger = ContextLedger::with_window_size(1000);
316 ledger.record("a.rs", "full", 600, 600);
317 assert_eq!(
318 ledger.pressure().recommendation,
319 PressureAction::SuggestCompression
320 );
321 ledger.record("b.rs", "full", 200, 200);
322 assert_eq!(
323 ledger.pressure().recommendation,
324 PressureAction::ForceCompression
325 );
326 ledger.record("c.rs", "full", 150, 150);
327 assert_eq!(
328 ledger.pressure().recommendation,
329 PressureAction::EvictLeastRelevant
330 );
331 }
332
333 #[test]
334 fn compression_ratio_accurate() {
335 let mut ledger = ContextLedger::with_window_size(10000);
336 ledger.record("a.rs", "full", 1000, 1000);
337 ledger.record("b.rs", "signatures", 1000, 200);
338 let ratio = ledger.compression_ratio();
339 assert!((ratio - 0.6).abs() < 0.01);
340 }
341
342 #[test]
343 fn eviction_returns_oldest() {
344 let mut ledger = ContextLedger::with_window_size(10000);
345 ledger.record("old.rs", "full", 100, 100);
346 std::thread::sleep(std::time::Duration::from_millis(10));
347 ledger.record("new.rs", "full", 100, 100);
348 let candidates = ledger.eviction_candidates(1);
349 assert_eq!(candidates, vec!["old.rs"]);
350 }
351
352 #[test]
353 fn remove_updates_totals() {
354 let mut ledger = ContextLedger::with_window_size(10000);
355 ledger.record("a.rs", "full", 500, 500);
356 ledger.record("b.rs", "full", 300, 300);
357 ledger.remove("a.rs");
358 assert_eq!(ledger.total_tokens_sent, 300);
359 assert_eq!(ledger.entries.len(), 1);
360 }
361
362 #[test]
363 fn mode_distribution_counts() {
364 let mut ledger = ContextLedger::new();
365 ledger.record("a.rs", "full", 100, 100);
366 ledger.record("b.rs", "signatures", 100, 50);
367 ledger.record("c.rs", "full", 100, 100);
368 let dist = ledger.mode_distribution();
369 assert_eq!(dist.get("full"), Some(&2));
370 assert_eq!(dist.get("signatures"), Some(&1));
371 }
372
373 #[test]
374 fn format_summary_includes_key_info() {
375 let mut ledger = ContextLedger::with_window_size(10000);
376 ledger.record("a.rs", "full", 500, 500);
377 let summary = ledger.format_summary();
378 assert!(summary.contains("500/10000"));
379 assert!(summary.contains("1 files"));
380 }
381
382 #[test]
383 fn reinjection_no_action_when_low_pressure() {
384 use crate::core::intent_engine::StructuredIntent;
385
386 let mut ledger = ContextLedger::with_window_size(10000);
387 ledger.record("a.rs", "full", 100, 100);
388 let intent = StructuredIntent::from_query("fix bug in a.rs");
389 let plan = ledger.reinjection_plan(&intent, 0.7);
390 assert!(plan.actions.is_empty());
391 assert_eq!(plan.total_tokens_freed, 0);
392 }
393
394 #[test]
395 fn reinjection_downgrades_non_target_files() {
396 use crate::core::intent_engine::StructuredIntent;
397
398 let mut ledger = ContextLedger::with_window_size(1000);
399 ledger.record("src/target.rs", "full", 400, 400);
400 std::thread::sleep(std::time::Duration::from_millis(10));
401 ledger.record("src/other.rs", "full", 400, 400);
402 std::thread::sleep(std::time::Duration::from_millis(10));
403 ledger.record("src/utils.rs", "full", 200, 200);
404
405 let intent = StructuredIntent::from_query("fix bug in target.rs");
406 let plan = ledger.reinjection_plan(&intent, 0.5);
407
408 assert!(!plan.actions.is_empty());
409 assert!(
410 plan.actions.iter().all(|a| !a.path.contains("target")),
411 "should not downgrade target file"
412 );
413 assert!(plan.total_tokens_freed > 0);
414 }
415
416 #[test]
417 fn reinjection_preserves_targets() {
418 use crate::core::intent_engine::StructuredIntent;
419
420 let mut ledger = ContextLedger::with_window_size(1000);
421 ledger.record("src/auth.rs", "full", 900, 900);
422 let intent = StructuredIntent::from_query("fix bug in auth.rs");
423 let plan = ledger.reinjection_plan(&intent, 0.5);
424 assert!(
425 plan.actions.is_empty(),
426 "should not downgrade target files even under pressure"
427 );
428 }
429
430 #[test]
431 fn downgrade_mode_chain() {
432 assert_eq!(
433 downgrade_mode("full", 1000),
434 Some(("signatures".to_string(), 200))
435 );
436 assert_eq!(
437 downgrade_mode("signatures", 200),
438 Some(("map".to_string(), 100))
439 );
440 assert_eq!(
441 downgrade_mode("map", 100),
442 Some(("reference".to_string(), 25))
443 );
444 assert_eq!(downgrade_mode("reference", 25), None);
445 }
446}