Skip to main content

lean_ctx/core/
context_ledger.rs

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