1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5use super::context_field::{
6 ContextItemId, ContextKind, ContextState, Provenance, ViewCosts, ViewKind,
7};
8
9const DEFAULT_CONTEXT_WINDOW: usize = 128_000;
10
11fn ledger_path(agent_id: &str) -> Result<std::path::PathBuf, String> {
12 let dir = crate::core::data_dir::lean_ctx_data_dir()?;
13 if agent_id == "default" {
14 Ok(dir.join("context_ledger.json"))
15 } else {
16 let ledger_dir = dir.join("ledger");
17 let safe_id: String = agent_id
18 .chars()
19 .map(|c| {
20 if c.is_alphanumeric() || c == '-' || c == '_' {
21 c
22 } else {
23 '_'
24 }
25 })
26 .collect();
27 Ok(ledger_dir.join(format!("{safe_id}.json")))
28 }
29}
30
31fn atomic_write_json(path: &std::path::Path, data: &str) {
32 let _ = crate::config_io::write_atomic(path, data);
33}
34
35#[cfg(unix)]
38fn acquire_ledger_lock(path: &std::path::Path) -> Option<std::fs::File> {
39 use std::os::unix::io::AsRawFd;
40 let lock_path = path.with_extension("json.lock");
41 let file = std::fs::OpenOptions::new()
42 .create(true)
43 .write(true)
44 .truncate(false)
45 .open(&lock_path)
46 .ok()?;
47 let fd = file.as_raw_fd();
48 let ret = unsafe { libc::flock(fd, libc::LOCK_EX | libc::LOCK_NB) };
49 if ret != 0 {
50 use std::time::{Duration, Instant};
52 let deadline = Instant::now() + Duration::from_secs(2);
53 loop {
54 std::thread::sleep(Duration::from_millis(50));
55 let ret = unsafe { libc::flock(fd, libc::LOCK_EX | libc::LOCK_NB) };
56 if ret == 0 {
57 break;
58 }
59 if Instant::now() >= deadline {
60 return None;
61 }
62 }
63 }
64 Some(file)
65}
66
67#[cfg(not(unix))]
68fn acquire_ledger_lock(_path: &std::path::Path) -> Option<std::fs::File> {
69 None
70}
71
72#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct ContextLedger {
74 pub window_size: usize,
75 pub entries: Vec<LedgerEntry>,
76 pub total_tokens_sent: usize,
77 pub total_tokens_saved: usize,
78 #[serde(skip)]
79 last_flush: Option<std::time::Instant>,
80}
81
82#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct LedgerEntry {
84 pub path: String,
85 pub mode: String,
86 pub original_tokens: usize,
87 pub sent_tokens: usize,
88 pub timestamp: i64,
89 #[serde(default)]
90 pub id: Option<ContextItemId>,
91 #[serde(default)]
92 pub kind: Option<ContextKind>,
93 #[serde(default)]
94 pub source_hash: Option<String>,
95 #[serde(default)]
96 pub state: Option<ContextState>,
97 #[serde(default)]
98 pub phi: Option<f64>,
99 #[serde(default)]
100 pub view_costs: Option<ViewCosts>,
101 #[serde(default)]
102 pub active_view: Option<ViewKind>,
103 #[serde(default)]
104 pub provenance: Option<Provenance>,
105}
106
107#[derive(Debug, Clone)]
108pub struct ContextPressure {
109 pub utilization: f64,
110 pub remaining_tokens: usize,
111 pub entries_count: usize,
112 pub recommendation: PressureAction,
113}
114
115#[derive(Debug, Clone, Copy, PartialEq, Eq)]
116pub enum PressureAction {
117 NoAction,
118 SuggestCompression,
119 ForceCompression,
120 EvictLeastRelevant,
121}
122
123impl ContextLedger {
124 pub fn new() -> Self {
125 Self {
126 window_size: DEFAULT_CONTEXT_WINDOW,
127 entries: Vec::new(),
128 total_tokens_sent: 0,
129 total_tokens_saved: 0,
130 last_flush: None,
131 }
132 }
133
134 pub fn with_window_size(size: usize) -> Self {
135 Self {
136 window_size: size,
137 entries: Vec::new(),
138 total_tokens_sent: 0,
139 total_tokens_saved: 0,
140 last_flush: None,
141 }
142 }
143
144 pub fn record(&mut self, path: &str, mode: &str, original_tokens: usize, sent_tokens: usize) {
145 self.record_with_task(path, mode, original_tokens, sent_tokens, None);
146 }
147
148 pub fn record_with_task(
149 &mut self,
150 path: &str,
151 mode: &str,
152 original_tokens: usize,
153 sent_tokens: usize,
154 task: Option<&str>,
155 ) {
156 let path = crate::core::pathutil::normalize_tool_path(path);
157 let item_id = ContextItemId::from_file(&path);
158
159 let phi =
160 Self::compute_real_phi(&path, sent_tokens, original_tokens, self.window_size, task);
161
162 if let Some(existing) = self.entries.iter_mut().find(|e| e.path == path) {
163 self.total_tokens_sent -= existing.sent_tokens;
164 self.total_tokens_saved -= existing
165 .original_tokens
166 .saturating_sub(existing.sent_tokens);
167 existing.mode = mode.to_string();
168 existing.original_tokens = original_tokens;
169 existing.sent_tokens = sent_tokens;
170 existing.timestamp = chrono::Utc::now().timestamp();
171 existing.active_view = Some(ViewKind::parse(mode));
172 if existing.id.is_none() {
173 existing.id = Some(item_id);
174 }
175 if existing.state.is_none() || existing.state == Some(ContextState::Candidate) {
176 existing.state = Some(ContextState::Included);
177 }
178 if existing.phi.is_none() {
179 existing.phi = Some(phi);
180 }
181 } else {
182 self.entries.push(LedgerEntry {
183 path: path.clone(),
184 mode: mode.to_string(),
185 original_tokens,
186 sent_tokens,
187 timestamp: chrono::Utc::now().timestamp(),
188 id: Some(item_id),
189 kind: Some(ContextKind::File),
190 source_hash: None,
191 state: Some(ContextState::Included),
192 phi: Some(phi),
193 view_costs: Some(ViewCosts::from_full_tokens(original_tokens)),
194 active_view: Some(ViewKind::parse(mode)),
195 provenance: None,
196 });
197 }
198 self.total_tokens_sent += sent_tokens;
199 self.total_tokens_saved += original_tokens.saturating_sub(sent_tokens);
200 }
201
202 fn compute_real_phi(
203 path: &str,
204 sent_tokens: usize,
205 original_tokens: usize,
206 window_size: usize,
207 task: Option<&str>,
208 ) -> f64 {
209 use crate::core::context_field::{compute_signals_for_path, ContextField};
210
211 let (signals, _costs) =
212 compute_signals_for_path(path, task, None, window_size, original_tokens);
213 let phi = ContextField::new().compute_phi(&signals);
214 if phi > 0.0 {
215 return phi;
216 }
217
218 Self::compute_lightweight_phi(sent_tokens, window_size)
219 }
220
221 fn compute_lightweight_phi(sent_tokens: usize, window_size: usize) -> f64 {
222 use crate::core::context_field::{ContextField, FieldSignals};
223 let token_cost_norm = if window_size > 0 {
224 (sent_tokens as f64 / window_size as f64).min(1.0)
225 } else {
226 0.0
227 };
228 let signals = FieldSignals {
229 relevance: 1.0,
230 surprise: 0.5,
231 graph_proximity: 0.0,
232 history_signal: 0.0,
233 token_cost_norm,
234 redundancy: 0.0,
235 };
236 ContextField::new().compute_phi(&signals)
237 }
238
239 pub fn upsert(
241 &mut self,
242 path: &str,
243 mode: &str,
244 original_tokens: usize,
245 sent_tokens: usize,
246 source_hash: Option<&str>,
247 kind: ContextKind,
248 provenance: Option<Provenance>,
249 ) {
250 self.record(path, mode, original_tokens, sent_tokens);
251 if let Some(entry) = self.entries.iter_mut().find(|e| e.path == path) {
252 entry.kind = Some(kind);
253 if let Some(h) = source_hash {
254 if entry.source_hash.as_deref() != Some(h) {
255 if entry.source_hash.is_some() {
256 entry.state = Some(ContextState::Stale);
257 }
258 entry.source_hash = Some(h.to_string());
259 }
260 }
261 if let Some(prov) = provenance {
262 entry.provenance = Some(prov);
263 }
264 }
265 }
266
267 pub fn update_phi(&mut self, path: &str, phi: f64) {
269 if let Some(entry) = self.entries.iter_mut().find(|e| e.path == path) {
270 entry.phi = Some(phi);
271 }
272 }
273
274 pub fn set_state(&mut self, path: &str, state: ContextState) {
276 if let Some(entry) = self.entries.iter_mut().find(|e| e.path == path) {
277 entry.state = Some(state);
278 }
279 }
280
281 pub fn find_by_id(&self, id: &ContextItemId) -> Option<&LedgerEntry> {
283 self.entries.iter().find(|e| e.id.as_ref() == Some(id))
284 }
285
286 pub fn items_by_state(&self, state: ContextState) -> Vec<&LedgerEntry> {
288 self.entries
289 .iter()
290 .filter(|e| e.state == Some(state))
291 .collect()
292 }
293
294 pub fn eviction_candidates_by_phi(&self, keep_count: usize) -> Vec<String> {
297 if self.entries.len() <= keep_count {
298 return Vec::new();
299 }
300 let mut sorted = self.entries.clone();
301 sorted.sort_by(|a, b| {
302 let a_phi = a.phi.unwrap_or(0.0);
303 let b_phi = b.phi.unwrap_or(0.0);
304 a_phi
305 .partial_cmp(&b_phi)
306 .unwrap_or(std::cmp::Ordering::Equal)
307 .then_with(|| a.timestamp.cmp(&b.timestamp))
308 });
309 sorted
310 .iter()
311 .filter(|e| e.state != Some(ContextState::Pinned))
312 .take(self.entries.len() - keep_count)
313 .map(|e| e.path.clone())
314 .collect()
315 }
316
317 pub fn mark_stale_by_hash(&mut self, path: &str, new_hash: &str) {
319 if let Some(entry) = self.entries.iter_mut().find(|e| e.path == path) {
320 if let Some(ref old_hash) = entry.source_hash {
321 if old_hash != new_hash {
322 entry.state = Some(ContextState::Stale);
323 entry.source_hash = Some(new_hash.to_string());
324 }
325 }
326 }
327 }
328
329 pub fn pressure(&self) -> ContextPressure {
330 let utilization = self.total_tokens_sent as f64 / self.window_size as f64;
331
332 let pinned_count = self
333 .entries
334 .iter()
335 .filter(|e| e.state == Some(ContextState::Pinned))
336 .count();
337 let stale_count = self
338 .entries
339 .iter()
340 .filter(|e| e.state == Some(ContextState::Stale))
341 .count();
342 let pinned_pressure = pinned_count as f64 * 0.02;
343 let stale_penalty = stale_count as f64 * 0.01;
344 let effective_utilization = (utilization + pinned_pressure + stale_penalty).min(1.0);
345
346 let effective_used = (effective_utilization * self.window_size as f64).round() as usize;
347 let remaining = self.window_size.saturating_sub(effective_used);
348
349 let recommendation = if effective_utilization > 0.9 {
350 PressureAction::EvictLeastRelevant
351 } else if effective_utilization > 0.75 {
352 PressureAction::ForceCompression
353 } else if effective_utilization > 0.5 {
354 PressureAction::SuggestCompression
355 } else {
356 PressureAction::NoAction
357 };
358
359 ContextPressure {
360 utilization: effective_utilization,
361 remaining_tokens: remaining,
362 entries_count: self.entries.len(),
363 recommendation,
364 }
365 }
366
367 pub fn compression_ratio(&self) -> f64 {
368 let total_original: usize = self.entries.iter().map(|e| e.original_tokens).sum();
369 if total_original == 0 {
370 return 1.0;
371 }
372 self.total_tokens_sent as f64 / total_original as f64
373 }
374
375 pub fn files_by_token_cost(&self) -> Vec<(String, usize)> {
376 let mut costs: Vec<(String, usize)> = self
377 .entries
378 .iter()
379 .map(|e| (e.path.clone(), e.sent_tokens))
380 .collect();
381 costs.sort_by_key(|b| std::cmp::Reverse(b.1));
382 costs
383 }
384
385 pub fn mode_distribution(&self) -> HashMap<String, usize> {
386 let mut dist: HashMap<String, usize> = HashMap::new();
387 for entry in &self.entries {
388 *dist.entry(entry.mode.clone()).or_insert(0) += 1;
389 }
390 dist
391 }
392
393 pub fn eviction_candidates(&self, keep_count: usize) -> Vec<String> {
394 if self.entries.len() <= keep_count {
395 return Vec::new();
396 }
397 let mut sorted = self.entries.clone();
398 sorted.sort_by_key(|e| e.timestamp);
399 sorted
400 .iter()
401 .take(self.entries.len() - keep_count)
402 .map(|e| e.path.clone())
403 .collect()
404 }
405
406 pub fn remove(&mut self, path: &str) -> bool {
407 if let Some(idx) = self.entries.iter().position(|e| e.path == path) {
408 let entry = &self.entries[idx];
409 self.total_tokens_sent = self.total_tokens_sent.saturating_sub(entry.sent_tokens);
410 self.total_tokens_saved = self
411 .total_tokens_saved
412 .saturating_sub(entry.original_tokens.saturating_sub(entry.sent_tokens));
413 self.entries.remove(idx);
414 true
415 } else {
416 false
417 }
418 }
419
420 pub fn reset(&mut self) {
422 let pinned_count = self
423 .entries
424 .iter()
425 .filter(|e| e.state == Some(ContextState::Pinned))
426 .count();
427 self.entries.clear();
428 self.total_tokens_sent = 0;
429 self.total_tokens_saved = 0;
430 if pinned_count > 0 {
431 tracing::info!("{pinned_count} pinned entries were also cleared");
432 }
433 }
434
435 pub fn evict_paths(&mut self, paths: &[&str]) -> usize {
438 let mut removed = 0;
439 for path in paths {
440 let normalized = crate::core::pathutil::normalize_tool_path(path);
441 if self.remove(&normalized) {
442 removed += 1;
443 }
444 }
445 removed
446 }
447
448 pub fn save(&self) {
449 self.save_for_agent("default");
450 }
451
452 pub fn save_debounced(&mut self) {
455 let now = std::time::Instant::now();
456 if let Some(last) = self.last_flush {
457 if now.duration_since(last) < std::time::Duration::from_secs(3) {
458 return;
459 }
460 }
461 self.save();
462 self.last_flush = Some(now);
463 }
464
465 pub fn save_for_agent(&self, agent_id: &str) {
466 if let Ok(path) = ledger_path(agent_id) {
467 if let Some(parent) = path.parent() {
468 let _ = std::fs::create_dir_all(parent);
469 }
470 let _lock = acquire_ledger_lock(&path);
471 if let Ok(json) = serde_json::to_string(self) {
472 atomic_write_json(&path, &json);
473 }
474 }
475 }
476
477 const MAX_LEDGER_ENTRIES: usize = 200;
478 const STALE_AGE_SECS: i64 = 7 * 24 * 3600;
479
480 pub fn prune(&mut self) -> usize {
481 let before = self.entries.len();
482 let now = chrono::Utc::now().timestamp();
483
484 for entry in &mut self.entries {
485 if let Some(phi) = entry.phi {
486 let hours_since = ((now - entry.timestamp) as f64 / 3600.0).max(0.0);
487 let decayed = phi * 0.95_f64.powf(hours_since);
488 entry.phi = Some(decayed.max(0.0));
489 }
490 }
491
492 self.entries
493 .retain(|e| !(e.mode == "error" && e.original_tokens == 0));
494
495 self.entries.retain(|e| {
496 let age = now - e.timestamp;
497 let phi = e.phi.unwrap_or(0.0);
498 !(age > Self::STALE_AGE_SECS && phi < 0.1)
499 });
500
501 let mut seen = std::collections::HashSet::new();
502 self.entries.sort_by_key(|e| std::cmp::Reverse(e.timestamp));
503 self.entries.retain(|e| {
504 let key = crate::core::pathutil::normalize_tool_path(&e.path);
505 seen.insert(key)
506 });
507
508 if self.entries.len() > Self::MAX_LEDGER_ENTRIES {
509 self.entries.sort_by(|a, b| {
510 let pa = a.phi.unwrap_or(0.0);
511 let pb = b.phi.unwrap_or(0.0);
512 pb.partial_cmp(&pa).unwrap_or(std::cmp::Ordering::Equal)
513 });
514 self.entries.truncate(Self::MAX_LEDGER_ENTRIES);
515 }
516
517 self.rebuild_totals();
518 before - self.entries.len()
519 }
520
521 fn rebuild_totals(&mut self) {
522 self.total_tokens_sent = self.entries.iter().map(|e| e.sent_tokens).sum();
523 self.total_tokens_saved = self
524 .entries
525 .iter()
526 .map(|e| e.original_tokens.saturating_sub(e.sent_tokens))
527 .sum();
528 }
529
530 pub fn load() -> Self {
531 Self::load_for_agent("default")
532 }
533
534 pub fn load_for_agent(agent_id: &str) -> Self {
535 let mut ledger: Self = ledger_path(agent_id)
536 .ok()
537 .and_then(|p| {
538 let _lock = acquire_ledger_lock(&p);
539 std::fs::read_to_string(p).ok()
540 })
541 .and_then(|s| serde_json::from_str(&s).ok())
542 .unwrap_or_default();
543 if let Some((_model, window)) = crate::hook_handlers::load_detected_model() {
544 ledger.window_size = window;
545 }
546 let pruned = ledger.prune();
547 if pruned > 0 {
548 ledger.save_for_agent(agent_id);
549 }
550 ledger
551 }
552
553 pub fn format_summary(&self) -> String {
554 let pressure = self.pressure();
555 format!(
556 "CTX: {}/{} tokens ({:.0}%), {} files, ratio {:.2}, action: {:?}",
557 self.total_tokens_sent,
558 self.window_size,
559 pressure.utilization * 100.0,
560 self.entries.len(),
561 self.compression_ratio(),
562 pressure.recommendation,
563 )
564 }
565
566 pub fn adjusted_total_saved(&self) -> isize {
567 if let Ok(bt) = crate::core::bounce_tracker::global().lock() {
568 bt.adjusted_savings(self.total_tokens_saved)
569 } else {
570 self.total_tokens_saved as isize
571 }
572 }
573}
574
575#[derive(Debug, Clone)]
576pub struct ReinjectionAction {
577 pub path: String,
578 pub current_mode: String,
579 pub new_mode: String,
580 pub tokens_freed: usize,
581}
582
583#[derive(Debug, Clone)]
584pub struct ReinjectionPlan {
585 pub actions: Vec<ReinjectionAction>,
586 pub total_tokens_freed: usize,
587 pub new_utilization: f64,
588}
589
590impl ContextLedger {
591 pub fn reinjection_plan(
592 &self,
593 intent: &super::intent_engine::StructuredIntent,
594 target_utilization: f64,
595 ) -> ReinjectionPlan {
596 let current_util = self.total_tokens_sent as f64 / self.window_size as f64;
597 if current_util <= target_utilization {
598 return ReinjectionPlan {
599 actions: Vec::new(),
600 total_tokens_freed: 0,
601 new_utilization: current_util,
602 };
603 }
604
605 let tokens_to_free =
606 self.total_tokens_sent - (self.window_size as f64 * target_utilization) as usize;
607
608 let target_set: std::collections::HashSet<&str> = intent
609 .targets
610 .iter()
611 .map(std::string::String::as_str)
612 .collect();
613
614 let mut candidates: Vec<(usize, &LedgerEntry)> = self
615 .entries
616 .iter()
617 .enumerate()
618 .filter(|(_, e)| !target_set.iter().any(|t| e.path.contains(t)))
619 .collect();
620
621 candidates.sort_by(|a, b| {
622 let a_phi = a.1.phi.unwrap_or(0.0);
623 let b_phi = b.1.phi.unwrap_or(0.0);
624 a_phi
625 .partial_cmp(&b_phi)
626 .unwrap_or_else(|| a.1.timestamp.cmp(&b.1.timestamp))
627 });
628
629 let mut actions = Vec::new();
630 let mut freed = 0usize;
631
632 for (_, entry) in &candidates {
633 if freed >= tokens_to_free {
634 break;
635 }
636 if let Some((new_mode, new_tokens)) = downgrade_mode(&entry.mode, entry.sent_tokens) {
637 let saving = entry.sent_tokens.saturating_sub(new_tokens);
638 if saving > 0 {
639 actions.push(ReinjectionAction {
640 path: entry.path.clone(),
641 current_mode: entry.mode.clone(),
642 new_mode,
643 tokens_freed: saving,
644 });
645 freed += saving;
646 }
647 }
648 }
649
650 let new_sent = self.total_tokens_sent.saturating_sub(freed);
651 let new_utilization = new_sent as f64 / self.window_size as f64;
652
653 ReinjectionPlan {
654 actions,
655 total_tokens_freed: freed,
656 new_utilization,
657 }
658 }
659}
660
661fn downgrade_mode(current_mode: &str, current_tokens: usize) -> Option<(String, usize)> {
662 match current_mode {
663 "full" => Some(("signatures".to_string(), current_tokens / 5)),
664 "aggressive" => Some(("signatures".to_string(), current_tokens / 3)),
665 "signatures" => Some(("map".to_string(), current_tokens / 2)),
666 "map" => Some(("reference".to_string(), current_tokens / 4)),
667 _ => None,
668 }
669}
670
671impl Default for ContextLedger {
672 fn default() -> Self {
673 Self::new()
674 }
675}
676
677#[cfg(test)]
678mod tests {
679 use super::*;
680
681 #[test]
682 fn new_ledger_is_empty() {
683 let ledger = ContextLedger::new();
684 assert_eq!(ledger.total_tokens_sent, 0);
685 assert_eq!(ledger.entries.len(), 0);
686 assert_eq!(ledger.pressure().recommendation, PressureAction::NoAction);
687 }
688
689 #[test]
690 fn record_tracks_tokens() {
691 let mut ledger = ContextLedger::with_window_size(10000);
692 ledger.record("src/main.rs", "full", 500, 500);
693 ledger.record("src/lib.rs", "signatures", 1000, 200);
694 assert_eq!(ledger.total_tokens_sent, 700);
695 assert_eq!(ledger.total_tokens_saved, 800);
696 assert_eq!(ledger.entries.len(), 2);
697 }
698
699 #[test]
700 fn record_updates_existing_entry() {
701 let mut ledger = ContextLedger::with_window_size(10000);
702 ledger.record("src/main.rs", "full", 500, 500);
703 ledger.record("src/main.rs", "signatures", 500, 100);
704 assert_eq!(ledger.entries.len(), 1);
705 assert_eq!(ledger.total_tokens_sent, 100);
706 assert_eq!(ledger.total_tokens_saved, 400);
707 }
708
709 #[test]
710 fn pressure_escalates() {
711 let mut ledger = ContextLedger::with_window_size(1000);
712 ledger.record("a.rs", "full", 600, 600);
713 assert_eq!(
714 ledger.pressure().recommendation,
715 PressureAction::SuggestCompression
716 );
717 ledger.record("b.rs", "full", 200, 200);
718 assert_eq!(
719 ledger.pressure().recommendation,
720 PressureAction::ForceCompression
721 );
722 ledger.record("c.rs", "full", 150, 150);
723 assert_eq!(
724 ledger.pressure().recommendation,
725 PressureAction::EvictLeastRelevant
726 );
727 }
728
729 #[test]
730 fn compression_ratio_accurate() {
731 let mut ledger = ContextLedger::with_window_size(10000);
732 ledger.record("a.rs", "full", 1000, 1000);
733 ledger.record("b.rs", "signatures", 1000, 200);
734 let ratio = ledger.compression_ratio();
735 assert!((ratio - 0.6).abs() < 0.01);
736 }
737
738 #[test]
739 fn eviction_returns_oldest() {
740 let mut ledger = ContextLedger::with_window_size(10000);
741 ledger.record("old.rs", "full", 100, 100);
742 std::thread::sleep(std::time::Duration::from_millis(10));
743 ledger.record("new.rs", "full", 100, 100);
744 let candidates = ledger.eviction_candidates(1);
745 assert_eq!(candidates, vec!["old.rs"]);
746 }
747
748 #[test]
749 fn remove_updates_totals() {
750 let mut ledger = ContextLedger::with_window_size(10000);
751 ledger.record("a.rs", "full", 500, 500);
752 ledger.record("b.rs", "full", 300, 300);
753 assert!(ledger.remove("a.rs"));
754 assert_eq!(ledger.total_tokens_sent, 300);
755 assert_eq!(ledger.entries.len(), 1);
756 assert!(!ledger.remove("nonexistent.rs"));
757 }
758
759 #[test]
760 fn reset_clears_everything() {
761 let mut ledger = ContextLedger::with_window_size(10000);
762 ledger.record("a.rs", "full", 500, 500);
763 ledger.record("b.rs", "full", 300, 300);
764 ledger.reset();
765 assert_eq!(ledger.entries.len(), 0);
766 assert_eq!(ledger.total_tokens_sent, 0);
767 assert_eq!(ledger.total_tokens_saved, 0);
768 assert_eq!(ledger.pressure().recommendation, PressureAction::NoAction);
769 }
770
771 #[test]
772 fn evict_paths_removes_matching() {
773 let mut ledger = ContextLedger::with_window_size(10000);
774 ledger.record("a.rs", "full", 500, 500);
775 ledger.record("b.rs", "full", 300, 300);
776 ledger.record("c.rs", "full", 200, 200);
777 let removed = ledger.evict_paths(&["a.rs", "c.rs", "nonexistent.rs"]);
778 assert_eq!(removed, 2);
779 assert_eq!(ledger.entries.len(), 1);
780 assert_eq!(ledger.entries[0].path, "b.rs");
781 assert_eq!(ledger.total_tokens_sent, 300);
782 }
783
784 #[test]
785 fn mode_distribution_counts() {
786 let mut ledger = ContextLedger::new();
787 ledger.record("a.rs", "full", 100, 100);
788 ledger.record("b.rs", "signatures", 100, 50);
789 ledger.record("c.rs", "full", 100, 100);
790 let dist = ledger.mode_distribution();
791 assert_eq!(dist.get("full"), Some(&2));
792 assert_eq!(dist.get("signatures"), Some(&1));
793 }
794
795 #[test]
796 fn format_summary_includes_key_info() {
797 let mut ledger = ContextLedger::with_window_size(10000);
798 ledger.record("a.rs", "full", 500, 500);
799 let summary = ledger.format_summary();
800 assert!(summary.contains("500/10000"));
801 assert!(summary.contains("1 files"));
802 }
803
804 #[test]
805 fn reinjection_no_action_when_low_pressure() {
806 use crate::core::intent_engine::StructuredIntent;
807
808 let mut ledger = ContextLedger::with_window_size(10000);
809 ledger.record("a.rs", "full", 100, 100);
810 let intent = StructuredIntent::from_query("fix bug in a.rs");
811 let plan = ledger.reinjection_plan(&intent, 0.7);
812 assert!(plan.actions.is_empty());
813 assert_eq!(plan.total_tokens_freed, 0);
814 }
815
816 #[test]
817 fn reinjection_downgrades_non_target_files() {
818 use crate::core::intent_engine::StructuredIntent;
819
820 let mut ledger = ContextLedger::with_window_size(1000);
821 ledger.record("src/target.rs", "full", 400, 400);
822 std::thread::sleep(std::time::Duration::from_millis(10));
823 ledger.record("src/other.rs", "full", 400, 400);
824 std::thread::sleep(std::time::Duration::from_millis(10));
825 ledger.record("src/utils.rs", "full", 200, 200);
826
827 let intent = StructuredIntent::from_query("fix bug in target.rs");
828 let plan = ledger.reinjection_plan(&intent, 0.5);
829
830 assert!(!plan.actions.is_empty());
831 assert!(
832 plan.actions.iter().all(|a| !a.path.contains("target")),
833 "should not downgrade target file"
834 );
835 assert!(plan.total_tokens_freed > 0);
836 }
837
838 #[test]
839 fn reinjection_preserves_targets() {
840 use crate::core::intent_engine::StructuredIntent;
841
842 let mut ledger = ContextLedger::with_window_size(1000);
843 ledger.record("src/auth.rs", "full", 900, 900);
844 let intent = StructuredIntent::from_query("fix bug in auth.rs");
845 let plan = ledger.reinjection_plan(&intent, 0.5);
846 assert!(
847 plan.actions.is_empty(),
848 "should not downgrade target files even under pressure"
849 );
850 }
851
852 #[test]
853 fn downgrade_mode_chain() {
854 assert_eq!(
855 downgrade_mode("full", 1000),
856 Some(("signatures".to_string(), 200))
857 );
858 assert_eq!(
859 downgrade_mode("signatures", 200),
860 Some(("map".to_string(), 100))
861 );
862 assert_eq!(
863 downgrade_mode("map", 100),
864 Some(("reference".to_string(), 25))
865 );
866 assert_eq!(downgrade_mode("reference", 25), None);
867 }
868
869 #[test]
870 fn record_assigns_item_id() {
871 let mut ledger = ContextLedger::new();
872 ledger.record("src/main.rs", "full", 500, 500);
873 let entry = &ledger.entries[0];
874 assert!(entry.id.is_some());
875 assert_eq!(entry.id.as_ref().unwrap().as_str(), "file:src/main.rs");
876 }
877
878 #[test]
879 fn record_sets_state_to_included() {
880 let mut ledger = ContextLedger::new();
881 ledger.record("src/main.rs", "full", 500, 500);
882 assert_eq!(
883 ledger.entries[0].state,
884 Some(crate::core::context_field::ContextState::Included)
885 );
886 }
887
888 #[test]
889 fn record_generates_view_costs() {
890 let mut ledger = ContextLedger::new();
891 ledger.record("src/main.rs", "full", 5000, 5000);
892 let vc = ledger.entries[0].view_costs.as_ref().unwrap();
893 assert_eq!(vc.get(&crate::core::context_field::ViewKind::Full), 5000);
894 assert_eq!(
895 vc.get(&crate::core::context_field::ViewKind::Signatures),
896 1000
897 );
898 }
899
900 #[test]
901 fn update_phi_works() {
902 let mut ledger = ContextLedger::new();
903 ledger.record("a.rs", "full", 100, 100);
904 ledger.update_phi("a.rs", 0.85);
905 assert_eq!(ledger.entries[0].phi, Some(0.85));
906 }
907
908 #[test]
909 fn set_state_works() {
910 let mut ledger = ContextLedger::new();
911 ledger.record("a.rs", "full", 100, 100);
912 ledger.set_state("a.rs", crate::core::context_field::ContextState::Pinned);
913 assert_eq!(
914 ledger.entries[0].state,
915 Some(crate::core::context_field::ContextState::Pinned)
916 );
917 }
918
919 #[test]
920 fn items_by_state_filters() {
921 let mut ledger = ContextLedger::new();
922 ledger.record("a.rs", "full", 100, 100);
923 ledger.record("b.rs", "full", 100, 100);
924 ledger.set_state("b.rs", crate::core::context_field::ContextState::Excluded);
925 let included = ledger.items_by_state(crate::core::context_field::ContextState::Included);
926 assert_eq!(included.len(), 1);
927 assert_eq!(included[0].path, "a.rs");
928 }
929
930 #[test]
931 fn eviction_by_phi_prefers_low_phi() {
932 let mut ledger = ContextLedger::with_window_size(10000);
933 ledger.record("high.rs", "full", 100, 100);
934 ledger.update_phi("high.rs", 0.9);
935 ledger.record("low.rs", "full", 100, 100);
936 ledger.update_phi("low.rs", 0.1);
937 let candidates = ledger.eviction_candidates_by_phi(1);
938 assert_eq!(candidates, vec!["low.rs"]);
939 }
940
941 #[test]
942 fn eviction_by_phi_skips_pinned() {
943 let mut ledger = ContextLedger::with_window_size(10000);
944 ledger.record("pinned.rs", "full", 100, 100);
945 ledger.update_phi("pinned.rs", 0.01);
946 ledger.set_state(
947 "pinned.rs",
948 crate::core::context_field::ContextState::Pinned,
949 );
950 ledger.record("normal.rs", "full", 100, 100);
951 ledger.update_phi("normal.rs", 0.5);
952 let candidates = ledger.eviction_candidates_by_phi(1);
953 assert_eq!(candidates, vec!["normal.rs"]);
954 }
955
956 #[test]
957 fn mark_stale_by_hash_detects_change() {
958 let mut ledger = ContextLedger::new();
959 ledger.record("a.rs", "full", 100, 100);
960 ledger.entries[0].source_hash = Some("hash_v1".to_string());
961 ledger.mark_stale_by_hash("a.rs", "hash_v2");
962 assert_eq!(
963 ledger.entries[0].state,
964 Some(crate::core::context_field::ContextState::Stale)
965 );
966 }
967
968 #[test]
969 fn find_by_id_works() {
970 let mut ledger = ContextLedger::new();
971 ledger.record("src/lib.rs", "full", 100, 100);
972 let id = crate::core::context_field::ContextItemId::from_file("src/lib.rs");
973 assert!(ledger.find_by_id(&id).is_some());
974 }
975
976 #[test]
977 fn upsert_sets_source_hash_and_kind() {
978 let mut ledger = ContextLedger::new();
979 ledger.upsert(
980 "src/main.rs",
981 "full",
982 500,
983 500,
984 Some("sha256_abc"),
985 crate::core::context_field::ContextKind::File,
986 None,
987 );
988 let entry = &ledger.entries[0];
989 assert_eq!(entry.source_hash.as_deref(), Some("sha256_abc"));
990 assert_eq!(
991 entry.kind,
992 Some(crate::core::context_field::ContextKind::File)
993 );
994 }
995}