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