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            // Lexical key only: entries were normalized when written, and the
515            // full variant would `realpath` every persisted path — the daemon
516            // runs this at boot (ContextLedger::load → prune) and stat-ing
517            // stored paths under ~/Documents from a launchd process pops the
518            // macOS TCC prompt (#356).
519            let key = crate::core::pathutil::normalize_tool_path_lexical(&e.path);
520            seen.insert(key)
521        });
522
523        if self.entries.len() > Self::MAX_LEDGER_ENTRIES {
524            self.entries.sort_by(|a, b| {
525                let pa = a.phi.unwrap_or(0.0);
526                let pb = b.phi.unwrap_or(0.0);
527                pb.partial_cmp(&pa).unwrap_or(std::cmp::Ordering::Equal)
528            });
529            self.entries.truncate(Self::MAX_LEDGER_ENTRIES);
530        }
531
532        self.rebuild_totals();
533        before - self.entries.len()
534    }
535
536    fn rebuild_totals(&mut self) {
537        self.total_tokens_sent = self.entries.iter().map(|e| e.sent_tokens).sum();
538        self.total_tokens_saved = self
539            .entries
540            .iter()
541            .map(|e| e.original_tokens.saturating_sub(e.sent_tokens))
542            .sum();
543    }
544
545    pub fn load() -> Self {
546        Self::load_for_agent("default")
547    }
548
549    pub fn load_for_agent(agent_id: &str) -> Self {
550        let mut ledger: Self = ledger_path(agent_id)
551            .ok()
552            .and_then(|p| {
553                let _lock = acquire_ledger_lock(&p);
554                std::fs::read_to_string(p).ok()
555            })
556            .and_then(|s| serde_json::from_str(&s).ok())
557            .unwrap_or_default();
558        if let Some((_model, window)) = crate::hook_handlers::load_detected_model() {
559            ledger.window_size = window;
560        }
561        let pruned = ledger.prune();
562        if pruned > 0 {
563            ledger.save_for_agent(agent_id);
564        }
565        ledger
566    }
567
568    pub fn format_summary(&self) -> String {
569        let pressure = self.pressure();
570        format!(
571            "CTX: {}/{} tokens ({:.0}%), {} files, ratio {:.2}, action: {:?}",
572            self.total_tokens_sent,
573            self.window_size,
574            pressure.utilization * 100.0,
575            self.entries.len(),
576            self.compression_ratio(),
577            pressure.recommendation,
578        )
579    }
580
581    pub fn adjusted_total_saved(&self) -> isize {
582        if let Ok(bt) = crate::core::bounce_tracker::global().lock() {
583            bt.adjusted_savings(self.total_tokens_saved)
584        } else {
585            self.total_tokens_saved as isize
586        }
587    }
588}
589
590#[derive(Debug, Clone)]
591pub struct ReinjectionAction {
592    pub path: String,
593    pub current_mode: String,
594    pub new_mode: String,
595    pub tokens_freed: usize,
596}
597
598#[derive(Debug, Clone)]
599pub struct ReinjectionPlan {
600    pub actions: Vec<ReinjectionAction>,
601    pub total_tokens_freed: usize,
602    pub new_utilization: f64,
603}
604
605impl ContextLedger {
606    pub fn reinjection_plan(
607        &self,
608        intent: &super::intent_engine::StructuredIntent,
609        target_utilization: f64,
610    ) -> ReinjectionPlan {
611        let current_util = self.total_tokens_sent as f64 / self.window_size as f64;
612        if current_util <= target_utilization {
613            return ReinjectionPlan {
614                actions: Vec::new(),
615                total_tokens_freed: 0,
616                new_utilization: current_util,
617            };
618        }
619
620        let tokens_to_free =
621            self.total_tokens_sent - (self.window_size as f64 * target_utilization) as usize;
622
623        let target_set: std::collections::HashSet<&str> = intent
624            .targets
625            .iter()
626            .map(std::string::String::as_str)
627            .collect();
628
629        let mut candidates: Vec<(usize, &LedgerEntry)> = self
630            .entries
631            .iter()
632            .enumerate()
633            .filter(|(_, e)| !target_set.iter().any(|t| e.path.contains(t)))
634            .collect();
635
636        candidates.sort_by(|a, b| {
637            let a_phi = a.1.phi.unwrap_or(0.0);
638            let b_phi = b.1.phi.unwrap_or(0.0);
639            a_phi
640                .partial_cmp(&b_phi)
641                .unwrap_or_else(|| a.1.timestamp.cmp(&b.1.timestamp))
642        });
643
644        let mut actions = Vec::new();
645        let mut freed = 0usize;
646
647        for (_, entry) in &candidates {
648            if freed >= tokens_to_free {
649                break;
650            }
651            if let Some((new_mode, new_tokens)) = downgrade_mode(&entry.mode, entry.sent_tokens) {
652                let saving = entry.sent_tokens.saturating_sub(new_tokens);
653                if saving > 0 {
654                    actions.push(ReinjectionAction {
655                        path: entry.path.clone(),
656                        current_mode: entry.mode.clone(),
657                        new_mode,
658                        tokens_freed: saving,
659                    });
660                    freed += saving;
661                }
662            }
663        }
664
665        let new_sent = self.total_tokens_sent.saturating_sub(freed);
666        let new_utilization = new_sent as f64 / self.window_size as f64;
667
668        ReinjectionPlan {
669            actions,
670            total_tokens_freed: freed,
671            new_utilization,
672        }
673    }
674}
675
676fn downgrade_mode(current_mode: &str, current_tokens: usize) -> Option<(String, usize)> {
677    match current_mode {
678        "full" => Some(("signatures".to_string(), current_tokens / 5)),
679        "aggressive" => Some(("signatures".to_string(), current_tokens / 3)),
680        "signatures" => Some(("map".to_string(), current_tokens / 2)),
681        "map" => Some(("reference".to_string(), current_tokens / 4)),
682        _ => None,
683    }
684}
685
686impl Default for ContextLedger {
687    fn default() -> Self {
688        Self::new()
689    }
690}
691
692#[cfg(test)]
693mod tests {
694    use super::*;
695
696    #[test]
697    fn new_ledger_is_empty() {
698        let ledger = ContextLedger::new();
699        assert_eq!(ledger.total_tokens_sent, 0);
700        assert_eq!(ledger.entries.len(), 0);
701        assert_eq!(ledger.pressure().recommendation, PressureAction::NoAction);
702    }
703
704    #[test]
705    fn record_tracks_tokens() {
706        let mut ledger = ContextLedger::with_window_size(10000);
707        ledger.record("src/main.rs", "full", 500, 500);
708        ledger.record("src/lib.rs", "signatures", 1000, 200);
709        assert_eq!(ledger.total_tokens_sent, 700);
710        assert_eq!(ledger.total_tokens_saved, 800);
711        assert_eq!(ledger.entries.len(), 2);
712    }
713
714    #[test]
715    fn record_updates_existing_entry() {
716        let mut ledger = ContextLedger::with_window_size(10000);
717        ledger.record("src/main.rs", "full", 500, 500);
718        ledger.record("src/main.rs", "signatures", 500, 100);
719        assert_eq!(ledger.entries.len(), 1);
720        assert_eq!(ledger.total_tokens_sent, 100);
721        assert_eq!(ledger.total_tokens_saved, 400);
722    }
723
724    #[test]
725    fn access_count_tracks_rereads() {
726        let mut ledger = ContextLedger::with_window_size(10000);
727        ledger.record("src/main.rs", "full", 500, 500);
728        assert_eq!(ledger.entries[0].access_count, 1);
729        ledger.record("src/main.rs", "signatures", 500, 100);
730        ledger.record("src/main.rs", "map", 500, 50);
731        assert_eq!(ledger.entries[0].access_count, 3);
732        // A different file starts its own count.
733        ledger.record("src/other.rs", "full", 200, 200);
734        let other = ledger.entries.iter().find(|e| e.path == "src/other.rs");
735        assert_eq!(other.map(|e| e.access_count), Some(1));
736    }
737
738    #[test]
739    fn pressure_escalates() {
740        let mut ledger = ContextLedger::with_window_size(1000);
741        ledger.record("a.rs", "full", 600, 600);
742        assert_eq!(
743            ledger.pressure().recommendation,
744            PressureAction::SuggestCompression
745        );
746        ledger.record("b.rs", "full", 200, 200);
747        assert_eq!(
748            ledger.pressure().recommendation,
749            PressureAction::ForceCompression
750        );
751        ledger.record("c.rs", "full", 150, 150);
752        assert_eq!(
753            ledger.pressure().recommendation,
754            PressureAction::EvictLeastRelevant
755        );
756    }
757
758    #[test]
759    fn compression_ratio_accurate() {
760        let mut ledger = ContextLedger::with_window_size(10000);
761        ledger.record("a.rs", "full", 1000, 1000);
762        ledger.record("b.rs", "signatures", 1000, 200);
763        let ratio = ledger.compression_ratio();
764        assert!((ratio - 0.6).abs() < 0.01);
765    }
766
767    #[test]
768    fn eviction_returns_oldest() {
769        let mut ledger = ContextLedger::with_window_size(10000);
770        ledger.record("old.rs", "full", 100, 100);
771        std::thread::sleep(std::time::Duration::from_millis(10));
772        ledger.record("new.rs", "full", 100, 100);
773        let candidates = ledger.eviction_candidates(1);
774        assert_eq!(candidates, vec!["old.rs"]);
775    }
776
777    #[test]
778    fn remove_updates_totals() {
779        let mut ledger = ContextLedger::with_window_size(10000);
780        ledger.record("a.rs", "full", 500, 500);
781        ledger.record("b.rs", "full", 300, 300);
782        assert!(ledger.remove("a.rs"));
783        assert_eq!(ledger.total_tokens_sent, 300);
784        assert_eq!(ledger.entries.len(), 1);
785        assert!(!ledger.remove("nonexistent.rs"));
786    }
787
788    #[test]
789    fn reset_clears_everything() {
790        let mut ledger = ContextLedger::with_window_size(10000);
791        ledger.record("a.rs", "full", 500, 500);
792        ledger.record("b.rs", "full", 300, 300);
793        ledger.reset();
794        assert_eq!(ledger.entries.len(), 0);
795        assert_eq!(ledger.total_tokens_sent, 0);
796        assert_eq!(ledger.total_tokens_saved, 0);
797        assert_eq!(ledger.pressure().recommendation, PressureAction::NoAction);
798    }
799
800    #[test]
801    fn evict_paths_removes_matching() {
802        let mut ledger = ContextLedger::with_window_size(10000);
803        ledger.record("a.rs", "full", 500, 500);
804        ledger.record("b.rs", "full", 300, 300);
805        ledger.record("c.rs", "full", 200, 200);
806        let removed = ledger.evict_paths(&["a.rs", "c.rs", "nonexistent.rs"]);
807        assert_eq!(removed, 2);
808        assert_eq!(ledger.entries.len(), 1);
809        assert_eq!(ledger.entries[0].path, "b.rs");
810        assert_eq!(ledger.total_tokens_sent, 300);
811    }
812
813    #[test]
814    fn mode_distribution_counts() {
815        let mut ledger = ContextLedger::new();
816        ledger.record("a.rs", "full", 100, 100);
817        ledger.record("b.rs", "signatures", 100, 50);
818        ledger.record("c.rs", "full", 100, 100);
819        let dist = ledger.mode_distribution();
820        assert_eq!(dist.get("full"), Some(&2));
821        assert_eq!(dist.get("signatures"), Some(&1));
822    }
823
824    #[test]
825    fn format_summary_includes_key_info() {
826        let mut ledger = ContextLedger::with_window_size(10000);
827        ledger.record("a.rs", "full", 500, 500);
828        let summary = ledger.format_summary();
829        assert!(summary.contains("500/10000"));
830        assert!(summary.contains("1 files"));
831    }
832
833    #[test]
834    fn reinjection_no_action_when_low_pressure() {
835        use crate::core::intent_engine::StructuredIntent;
836
837        let mut ledger = ContextLedger::with_window_size(10000);
838        ledger.record("a.rs", "full", 100, 100);
839        let intent = StructuredIntent::from_query("fix bug in a.rs");
840        let plan = ledger.reinjection_plan(&intent, 0.7);
841        assert!(plan.actions.is_empty());
842        assert_eq!(plan.total_tokens_freed, 0);
843    }
844
845    #[test]
846    fn reinjection_downgrades_non_target_files() {
847        use crate::core::intent_engine::StructuredIntent;
848
849        let mut ledger = ContextLedger::with_window_size(1000);
850        ledger.record("src/target.rs", "full", 400, 400);
851        std::thread::sleep(std::time::Duration::from_millis(10));
852        ledger.record("src/other.rs", "full", 400, 400);
853        std::thread::sleep(std::time::Duration::from_millis(10));
854        ledger.record("src/utils.rs", "full", 200, 200);
855
856        let intent = StructuredIntent::from_query("fix bug in target.rs");
857        let plan = ledger.reinjection_plan(&intent, 0.5);
858
859        assert!(!plan.actions.is_empty());
860        assert!(
861            plan.actions.iter().all(|a| !a.path.contains("target")),
862            "should not downgrade target file"
863        );
864        assert!(plan.total_tokens_freed > 0);
865    }
866
867    #[test]
868    fn reinjection_preserves_targets() {
869        use crate::core::intent_engine::StructuredIntent;
870
871        let mut ledger = ContextLedger::with_window_size(1000);
872        ledger.record("src/auth.rs", "full", 900, 900);
873        let intent = StructuredIntent::from_query("fix bug in auth.rs");
874        let plan = ledger.reinjection_plan(&intent, 0.5);
875        assert!(
876            plan.actions.is_empty(),
877            "should not downgrade target files even under pressure"
878        );
879    }
880
881    #[test]
882    fn downgrade_mode_chain() {
883        assert_eq!(
884            downgrade_mode("full", 1000),
885            Some(("signatures".to_string(), 200))
886        );
887        assert_eq!(
888            downgrade_mode("signatures", 200),
889            Some(("map".to_string(), 100))
890        );
891        assert_eq!(
892            downgrade_mode("map", 100),
893            Some(("reference".to_string(), 25))
894        );
895        assert_eq!(downgrade_mode("reference", 25), None);
896    }
897
898    #[test]
899    fn record_assigns_item_id() {
900        let mut ledger = ContextLedger::new();
901        ledger.record("src/main.rs", "full", 500, 500);
902        let entry = &ledger.entries[0];
903        assert!(entry.id.is_some());
904        assert_eq!(entry.id.as_ref().unwrap().as_str(), "file:src/main.rs");
905    }
906
907    #[test]
908    fn record_sets_state_to_included() {
909        let mut ledger = ContextLedger::new();
910        ledger.record("src/main.rs", "full", 500, 500);
911        assert_eq!(
912            ledger.entries[0].state,
913            Some(crate::core::context_field::ContextState::Included)
914        );
915    }
916
917    #[test]
918    fn record_generates_view_costs() {
919        let mut ledger = ContextLedger::new();
920        ledger.record("src/main.rs", "full", 5000, 5000);
921        let vc = ledger.entries[0].view_costs.as_ref().unwrap();
922        assert_eq!(vc.get(&crate::core::context_field::ViewKind::Full), 5000);
923        assert_eq!(
924            vc.get(&crate::core::context_field::ViewKind::Signatures),
925            1000
926        );
927    }
928
929    #[test]
930    fn update_phi_works() {
931        let mut ledger = ContextLedger::new();
932        ledger.record("a.rs", "full", 100, 100);
933        ledger.update_phi("a.rs", 0.85);
934        assert_eq!(ledger.entries[0].phi, Some(0.85));
935    }
936
937    #[test]
938    fn set_state_works() {
939        let mut ledger = ContextLedger::new();
940        ledger.record("a.rs", "full", 100, 100);
941        ledger.set_state("a.rs", crate::core::context_field::ContextState::Pinned);
942        assert_eq!(
943            ledger.entries[0].state,
944            Some(crate::core::context_field::ContextState::Pinned)
945        );
946    }
947
948    #[test]
949    fn items_by_state_filters() {
950        let mut ledger = ContextLedger::new();
951        ledger.record("a.rs", "full", 100, 100);
952        ledger.record("b.rs", "full", 100, 100);
953        ledger.set_state("b.rs", crate::core::context_field::ContextState::Excluded);
954        let included = ledger.items_by_state(crate::core::context_field::ContextState::Included);
955        assert_eq!(included.len(), 1);
956        assert_eq!(included[0].path, "a.rs");
957    }
958
959    #[test]
960    fn eviction_by_phi_prefers_low_phi() {
961        let mut ledger = ContextLedger::with_window_size(10000);
962        ledger.record("high.rs", "full", 100, 100);
963        ledger.update_phi("high.rs", 0.9);
964        ledger.record("low.rs", "full", 100, 100);
965        ledger.update_phi("low.rs", 0.1);
966        let candidates = ledger.eviction_candidates_by_phi(1);
967        assert_eq!(candidates, vec!["low.rs"]);
968    }
969
970    #[test]
971    fn eviction_by_phi_skips_pinned() {
972        let mut ledger = ContextLedger::with_window_size(10000);
973        ledger.record("pinned.rs", "full", 100, 100);
974        ledger.update_phi("pinned.rs", 0.01);
975        ledger.set_state(
976            "pinned.rs",
977            crate::core::context_field::ContextState::Pinned,
978        );
979        ledger.record("normal.rs", "full", 100, 100);
980        ledger.update_phi("normal.rs", 0.5);
981        let candidates = ledger.eviction_candidates_by_phi(1);
982        assert_eq!(candidates, vec!["normal.rs"]);
983    }
984
985    #[test]
986    fn mark_stale_by_hash_detects_change() {
987        let mut ledger = ContextLedger::new();
988        ledger.record("a.rs", "full", 100, 100);
989        ledger.entries[0].source_hash = Some("hash_v1".to_string());
990        ledger.mark_stale_by_hash("a.rs", "hash_v2");
991        assert_eq!(
992            ledger.entries[0].state,
993            Some(crate::core::context_field::ContextState::Stale)
994        );
995    }
996
997    #[test]
998    fn find_by_id_works() {
999        let mut ledger = ContextLedger::new();
1000        ledger.record("src/lib.rs", "full", 100, 100);
1001        let id = crate::core::context_field::ContextItemId::from_file("src/lib.rs");
1002        assert!(ledger.find_by_id(&id).is_some());
1003    }
1004
1005    #[test]
1006    fn upsert_sets_source_hash_and_kind() {
1007        let mut ledger = ContextLedger::new();
1008        ledger.upsert(
1009            "src/main.rs",
1010            "full",
1011            500,
1012            500,
1013            Some("sha256_abc"),
1014            crate::core::context_field::ContextKind::File,
1015            None,
1016        );
1017        let entry = &ledger.entries[0];
1018        assert_eq!(entry.source_hash.as_deref(), Some("sha256_abc"));
1019        assert_eq!(
1020            entry.kind,
1021            Some(crate::core::context_field::ContextKind::File)
1022        );
1023    }
1024}