Skip to main content

lean_ctx/core/
cache.rs

1use md5::{Digest, Md5};
2use std::collections::HashMap;
3use std::sync::OnceLock;
4use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
5use std::time::{Duration, Instant, SystemTime};
6
7use super::tokens::count_tokens;
8
9/// Process-global monotonic base for encoding `Instant`s into an `AtomicU64`.
10/// Stored as milliseconds since this base, which is sufficient resolution for
11/// LRU/RRF eviction recency while allowing lock-free access on cache hits.
12fn instant_base() -> Instant {
13    static BASE: OnceLock<Instant> = OnceLock::new();
14    *BASE.get_or_init(Instant::now)
15}
16
17fn encode_instant(i: Instant) -> u64 {
18    i.saturating_duration_since(instant_base()).as_millis() as u64
19}
20
21fn decode_instant(ms: u64) -> Instant {
22    instant_base() + Duration::from_millis(ms)
23}
24
25fn normalize_key(path: &str) -> String {
26    crate::core::pathutil::normalize_tool_path(path)
27}
28
29fn max_cache_tokens() -> usize {
30    std::env::var("LEAN_CTX_CACHE_MAX_TOKENS")
31        .ok()
32        .and_then(|v| v.parse().ok())
33        .unwrap_or(500_000)
34}
35
36/// A cached file read: zstd-compressed content, hash, token count, and access metadata.
37///
38/// `read_count` and `last_access` use interior mutability (atomics) so cache
39/// hits can be recorded under a shared (read) lock — parallel reads of distinct
40/// files no longer serialize on a global write lock.
41#[derive(Debug)]
42pub struct CacheEntry {
43    compressed_content: Vec<u8>,
44    pub hash: String,
45    pub line_count: usize,
46    pub original_tokens: usize,
47    read_count: AtomicU32,
48    pub path: String,
49    last_access: AtomicU64,
50    pub stored_mtime: Option<SystemTime>,
51    /// Mode-specific compressed outputs (e.g. "map", "signatures") cached to avoid re-parsing.
52    pub compressed_outputs: HashMap<String, String>,
53    /// Whether full (uncompressed) content was already delivered for this hash.
54    /// Prevents cache-stub loops when upgrading from compressed to full mode.
55    pub full_content_delivered: bool,
56    /// Last read mode used for this file (for auto-escalation on edit failure).
57    pub last_mode: String,
58}
59
60const ZSTD_LEVEL: i32 = 3;
61
62fn zstd_compress(data: &str) -> Vec<u8> {
63    zstd::encode_all(data.as_bytes(), ZSTD_LEVEL).unwrap_or_else(|_| data.as_bytes().to_vec())
64}
65
66fn zstd_decompress(data: &[u8]) -> Option<String> {
67    zstd::decode_all(data)
68        .ok()
69        .and_then(|v| String::from_utf8(v).ok())
70}
71
72impl CacheEntry {
73    /// Creates a new entry with zstd-compressed content.
74    pub fn new(
75        content: &str,
76        hash: String,
77        line_count: usize,
78        original_tokens: usize,
79        path: String,
80        stored_mtime: Option<SystemTime>,
81    ) -> Self {
82        let compressed_content = zstd_compress(content);
83        Self {
84            compressed_content,
85            hash,
86            line_count,
87            original_tokens,
88            read_count: AtomicU32::new(1),
89            path,
90            last_access: AtomicU64::new(encode_instant(Instant::now())),
91            stored_mtime,
92            compressed_outputs: HashMap::new(),
93            full_content_delivered: false,
94            last_mode: String::new(),
95        }
96    }
97
98    /// Current read count (lock-free).
99    pub fn read_count(&self) -> u32 {
100        self.read_count.load(Ordering::Relaxed)
101    }
102
103    /// Atomically increments the read count and returns the new value (lock-free).
104    pub fn bump_read_count(&self) -> u32 {
105        self.read_count.fetch_add(1, Ordering::Relaxed) + 1
106    }
107
108    /// Overwrites the read count (used by `store` and tests).
109    pub fn set_read_count(&self, n: u32) {
110        self.read_count.store(n, Ordering::Relaxed);
111    }
112
113    /// Last access time, decoded from the atomic millisecond offset.
114    pub fn last_access(&self) -> Instant {
115        decode_instant(self.last_access.load(Ordering::Relaxed))
116    }
117
118    /// Marks the entry as accessed now (lock-free).
119    pub fn touch(&self) {
120        self.last_access
121            .store(encode_instant(Instant::now()), Ordering::Relaxed);
122    }
123
124    /// Overwrites the last-access time (used by tests and eviction setup).
125    pub fn set_last_access(&self, when: Instant) {
126        self.last_access
127            .store(encode_instant(when), Ordering::Relaxed);
128    }
129
130    /// Decompresses and returns the full file content.
131    pub fn content(&self) -> Option<String> {
132        zstd_decompress(&self.compressed_content)
133    }
134
135    /// Replaces the stored content with new zstd-compressed data.
136    pub fn set_content(&mut self, content: &str) {
137        self.compressed_content = zstd_compress(content);
138    }
139
140    /// Approximate RAM usage of the compressed content in bytes.
141    pub fn compressed_size(&self) -> usize {
142        self.compressed_content.len()
143    }
144}
145
146/// Result of a cache store operation, indicating whether it was a hit or new entry.
147#[derive(Debug, Clone)]
148pub struct StoreResult {
149    pub line_count: usize,
150    pub original_tokens: usize,
151    pub read_count: u32,
152    pub was_hit: bool,
153    /// Whether full content was previously delivered for this cache entry.
154    pub full_content_delivered: bool,
155}
156
157impl CacheEntry {
158    /// Computes a legacy eviction score blending recency, frequency, and size.
159    pub fn eviction_score_legacy(&self, now: Instant) -> f64 {
160        let elapsed = now
161            .checked_duration_since(self.last_access())
162            .unwrap_or_default()
163            .as_secs_f64();
164        let recency = 1.0 / (1.0 + elapsed.sqrt());
165        let frequency = (self.read_count() as f64 + 1.0).ln();
166        let size_value = (self.original_tokens as f64 + 1.0).ln();
167        recency * 0.4 + frequency * 0.3 + size_value * 0.3
168    }
169
170    pub fn get_compressed(&self, mode_key: &str) -> Option<&String> {
171        self.compressed_outputs.get(mode_key)
172    }
173
174    pub fn set_compressed(&mut self, mode_key: &str, output: String) {
175        const MAX_COMPRESSED_VARIANTS: usize = 3;
176        if self.compressed_outputs.len() >= MAX_COMPRESSED_VARIANTS
177            && !self.compressed_outputs.contains_key(mode_key)
178            && let Some(oldest_key) = self.compressed_outputs.keys().next().cloned()
179        {
180            self.compressed_outputs.remove(&oldest_key);
181        }
182        self.compressed_outputs.insert(mode_key.to_string(), output);
183    }
184
185    pub fn mark_full_delivered(&mut self) {
186        self.full_content_delivered = true;
187    }
188}
189
190const RRF_K: f64 = 60.0;
191
192/// Compute Reciprocal Rank Fusion eviction scores for a batch of cache entries.
193/// Each signal (recency, frequency, size) produces an independent ranking.
194/// The final score is the sum of `1/(k + rank)` across all signals.
195/// Higher score = more valuable = keep longer.
196pub fn eviction_scores_rrf(entries: &[(&String, &CacheEntry)], now: Instant) -> Vec<(String, f64)> {
197    if entries.is_empty() {
198        return Vec::new();
199    }
200
201    let n = entries.len();
202
203    let mut recency_order: Vec<usize> = (0..n).collect();
204    recency_order.sort_by(|&a, &b| {
205        let elapsed_a = now
206            .checked_duration_since(entries[a].1.last_access())
207            .unwrap_or_default()
208            .as_secs_f64();
209        let elapsed_b = now
210            .checked_duration_since(entries[b].1.last_access())
211            .unwrap_or_default()
212            .as_secs_f64();
213        elapsed_a
214            .partial_cmp(&elapsed_b)
215            .unwrap_or(std::cmp::Ordering::Equal)
216    });
217
218    let mut frequency_order: Vec<usize> = (0..n).collect();
219    frequency_order.sort_by(|&a, &b| entries[b].1.read_count().cmp(&entries[a].1.read_count()));
220
221    let mut size_order: Vec<usize> = (0..n).collect();
222    size_order.sort_by(|&a, &b| {
223        entries[b]
224            .1
225            .original_tokens
226            .cmp(&entries[a].1.original_tokens)
227    });
228
229    let mut recency_ranks = vec![0usize; n];
230    let mut frequency_ranks = vec![0usize; n];
231    let mut size_ranks = vec![0usize; n];
232
233    for (rank, &idx) in recency_order.iter().enumerate() {
234        recency_ranks[idx] = rank;
235    }
236    for (rank, &idx) in frequency_order.iter().enumerate() {
237        frequency_ranks[idx] = rank;
238    }
239    for (rank, &idx) in size_order.iter().enumerate() {
240        size_ranks[idx] = rank;
241    }
242
243    entries
244        .iter()
245        .enumerate()
246        .map(|(i, (path, _))| {
247            let score = 1.0 / (RRF_K + recency_ranks[i] as f64)
248                + 1.0 / (RRF_K + frequency_ranks[i] as f64)
249                + 1.0 / (RRF_K + size_ranks[i] as f64);
250            ((*path).clone(), score)
251        })
252        .collect()
253}
254
255/// Aggregated cache statistics: hits, reads, and token savings.
256///
257/// Counters are atomic so they can be updated on the read-locked cache-hit
258/// fast path without taking a write lock.
259#[derive(Debug, Default)]
260pub struct CacheStats {
261    total_reads: AtomicU64,
262    cache_hits: AtomicU64,
263    total_original_tokens: AtomicU64,
264    total_sent_tokens: AtomicU64,
265    files_tracked: AtomicU64,
266}
267
268impl CacheStats {
269    /// Total number of read operations recorded.
270    pub fn total_reads(&self) -> u64 {
271        self.total_reads.load(Ordering::Relaxed)
272    }
273
274    /// Total number of cache hits recorded.
275    pub fn cache_hits(&self) -> u64 {
276        self.cache_hits.load(Ordering::Relaxed)
277    }
278
279    /// Sum of original (uncompressed) token counts across all reads.
280    pub fn total_original_tokens(&self) -> u64 {
281        self.total_original_tokens.load(Ordering::Relaxed)
282    }
283
284    /// Sum of tokens actually sent to the model.
285    pub fn total_sent_tokens(&self) -> u64 {
286        self.total_sent_tokens.load(Ordering::Relaxed)
287    }
288
289    /// Number of distinct files currently tracked.
290    pub fn files_tracked(&self) -> u64 {
291        self.files_tracked.load(Ordering::Relaxed)
292    }
293
294    /// Returns the cache hit rate as a percentage (0–100).
295    pub fn hit_rate(&self) -> f64 {
296        let total = self.total_reads();
297        if total == 0 {
298            return 0.0;
299        }
300        (self.cache_hits() as f64 / total as f64) * 100.0
301    }
302
303    /// Returns the total number of tokens saved by cache hits.
304    pub fn tokens_saved(&self) -> u64 {
305        self.total_original_tokens()
306            .saturating_sub(self.total_sent_tokens())
307    }
308
309    /// Returns the savings as a percentage of total original tokens.
310    pub fn savings_percent(&self) -> f64 {
311        let original = self.total_original_tokens();
312        if original == 0 {
313            return 0.0;
314        }
315        (self.tokens_saved() as f64 / original as f64) * 100.0
316    }
317}
318
319/// A block shared across multiple files, identified by its canonical source.
320#[derive(Clone, Debug)]
321pub struct SharedBlock {
322    pub canonical_path: String,
323    pub canonical_ref: String,
324    pub start_line: usize,
325    pub end_line: usize,
326    pub content: String,
327}
328
329/// In-memory file cache with segmented LRU eviction (probationary vs protected),
330/// file references, and cross-file dedup.
331pub struct SessionCache {
332    entries: HashMap<String, CacheEntry>,
333    file_refs: HashMap<String, String>,
334    next_ref: usize,
335    stats: CacheStats,
336    shared_blocks: Vec<SharedBlock>,
337}
338
339impl Default for SessionCache {
340    fn default() -> Self {
341        Self::new()
342    }
343}
344
345impl SessionCache {
346    /// Creates an empty session cache with default stats.
347    pub fn new() -> Self {
348        Self {
349            entries: HashMap::new(),
350            file_refs: HashMap::new(),
351            next_ref: 1,
352            shared_blocks: Vec::new(),
353            stats: CacheStats::default(),
354        }
355    }
356
357    /// Returns or assigns a short file reference label (F1, F2, ...) for the given path.
358    pub fn get_file_ref(&mut self, path: &str) -> String {
359        let key = normalize_key(path);
360        if let Some(r) = self.file_refs.get(&key) {
361            return r.clone();
362        }
363        let r = format!("F{}", self.next_ref);
364        self.next_ref += 1;
365        self.file_refs.insert(key, r.clone());
366        r
367    }
368
369    /// Returns the file reference label for a path without assigning a new one.
370    pub fn get_file_ref_readonly(&self, path: &str) -> Option<String> {
371        self.file_refs.get(&normalize_key(path)).cloned()
372    }
373
374    /// Looks up a cached entry by file path.
375    pub fn get(&self, path: &str) -> Option<&CacheEntry> {
376        self.entries.get(&normalize_key(path))
377    }
378
379    /// Mutable lookup of a cached entry by file path.
380    pub fn get_mut(&mut self, path: &str) -> Option<&mut CacheEntry> {
381        self.entries.get_mut(&normalize_key(path))
382    }
383
384    /// Retrieves the full (uncompressed) content for a file path, if cached.
385    /// Used by the CCR (Compress-Cache-Retrieve) mechanism.
386    pub fn get_full_content(&self, path: &str) -> Option<String> {
387        self.entries
388            .get(&normalize_key(path))
389            .and_then(CacheEntry::content)
390    }
391
392    /// Records a cache hit, updates access stats, and emits a cache-hit event.
393    ///
394    /// Takes `&self`: the hit counters use interior-mutable atomics, so this
395    /// runs under a shared (read) lock and lets parallel reads of different
396    /// files proceed concurrently instead of serializing on a write lock.
397    pub fn record_cache_hit(&self, path: &str) -> Option<&CacheEntry> {
398        let key = normalize_key(path);
399        let ref_label = self
400            .file_refs
401            .get(&key)
402            .cloned()
403            .unwrap_or_else(|| "F?".to_string());
404        let entry = self.entries.get(&key)?;
405        let new_count = entry.bump_read_count();
406        entry.touch();
407        self.stats.total_reads.fetch_add(1, Ordering::Relaxed);
408        self.stats.cache_hits.fetch_add(1, Ordering::Relaxed);
409        self.stats
410            .total_original_tokens
411            .fetch_add(entry.original_tokens as u64, Ordering::Relaxed);
412        let hit_msg = format!("{ref_label} cached {new_count}t {}L", entry.line_count);
413        self.stats
414            .total_sent_tokens
415            .fetch_add(count_tokens(&hit_msg) as u64, Ordering::Relaxed);
416        crate::core::events::emit_cache_hit(path, entry.original_tokens as u64);
417        Some(entry)
418    }
419
420    /// Stores file content in the cache; returns a hit if content hash matches.
421    pub fn store(&mut self, path: &str, content: &str) -> StoreResult {
422        let key = normalize_key(path);
423        let hash = compute_md5(content);
424        let line_count = content.lines().count();
425        let original_tokens = count_tokens(content);
426        let stored_mtime = std::fs::metadata(path).and_then(|m| m.modified()).ok();
427        let now = Instant::now();
428
429        self.stats.total_reads.fetch_add(1, Ordering::Relaxed);
430        self.stats
431            .total_original_tokens
432            .fetch_add(original_tokens as u64, Ordering::Relaxed);
433
434        if let Some(existing) = self.entries.get_mut(&key) {
435            existing.set_last_access(now);
436            if stored_mtime.is_some() {
437                existing.stored_mtime = stored_mtime;
438            }
439            if existing.hash == hash {
440                let new_count = existing.bump_read_count();
441                self.stats.cache_hits.fetch_add(1, Ordering::Relaxed);
442                let hit_msg = format!(
443                    "{} cached {new_count}t {}L",
444                    self.file_refs.get(&key).unwrap_or(&"F?".to_string()),
445                    existing.line_count,
446                );
447                self.stats
448                    .total_sent_tokens
449                    .fetch_add(count_tokens(&hit_msg) as u64, Ordering::Relaxed);
450                return StoreResult {
451                    line_count: existing.line_count,
452                    original_tokens: existing.original_tokens,
453                    read_count: new_count,
454                    was_hit: true,
455                    full_content_delivered: existing.full_content_delivered,
456                };
457            }
458            existing.compressed_outputs.clear();
459            existing.set_content(content);
460            existing.hash = hash;
461            existing.line_count = line_count;
462            existing.original_tokens = original_tokens;
463            let new_count = existing.bump_read_count();
464            existing.full_content_delivered = false;
465            if stored_mtime.is_some() {
466                existing.stored_mtime = stored_mtime;
467            }
468            self.stats
469                .total_sent_tokens
470                .fetch_add(original_tokens as u64, Ordering::Relaxed);
471            return StoreResult {
472                line_count,
473                original_tokens,
474                read_count: new_count,
475                was_hit: false,
476                full_content_delivered: false,
477            };
478        }
479
480        self.evict_if_needed(original_tokens);
481        self.get_file_ref(&key);
482
483        let entry = CacheEntry::new(
484            content,
485            hash,
486            line_count,
487            original_tokens,
488            key.clone(),
489            stored_mtime,
490        );
491
492        self.entries.insert(key, entry);
493        self.stats.files_tracked.fetch_add(1, Ordering::Relaxed);
494        self.stats
495            .total_sent_tokens
496            .fetch_add(original_tokens as u64, Ordering::Relaxed);
497        StoreResult {
498            line_count,
499            original_tokens,
500            read_count: 1,
501            was_hit: false,
502            full_content_delivered: false,
503        }
504    }
505
506    /// Returns the sum of original token counts across all cached entries.
507    pub fn total_cached_tokens(&self) -> usize {
508        self.entries.values().map(|e| e.original_tokens).sum()
509    }
510
511    /// Evict until cache fits within token budget using RRF (Reciprocal Rank Fusion).
512    /// Combines recency, frequency, and size signals to evict least-valuable entries first.
513    pub fn evict_if_needed(&mut self, incoming_tokens: usize) {
514        let max_tokens = max_cache_tokens();
515        let current = self.total_cached_tokens();
516        if current + incoming_tokens <= max_tokens {
517            return;
518        }
519
520        let now = Instant::now();
521        let all: Vec<(&String, &CacheEntry)> = self.entries.iter().collect();
522        let mut scores = eviction_scores_rrf(&all, now);
523        // Sort ascending: lowest RRF score = least valuable = evict first
524        scores.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
525
526        let mut freed = 0usize;
527        let target = (current + incoming_tokens).saturating_sub(max_tokens);
528
529        for (path, _score) in &scores {
530            if freed >= target {
531                break;
532            }
533            if let Some(entry) = self.entries.remove(path) {
534                freed += entry.original_tokens;
535                self.file_refs.remove(path);
536            }
537        }
538    }
539
540    /// Returns all cached entries as (path, entry) pairs.
541    pub fn get_all_entries(&self) -> Vec<(&String, &CacheEntry)> {
542        self.entries.iter().collect()
543    }
544
545    /// Returns a reference to the aggregated cache statistics.
546    pub fn get_stats(&self) -> &CacheStats {
547        &self.stats
548    }
549
550    /// Returns the path-to-file-ref mapping (e.g. "/src/main.rs" → "F1").
551    pub fn file_ref_map(&self) -> &HashMap<String, String> {
552        &self.file_refs
553    }
554
555    /// Replaces the cross-file shared blocks used for deduplication.
556    pub fn set_shared_blocks(&mut self, blocks: Vec<SharedBlock>) {
557        self.shared_blocks = blocks;
558    }
559
560    /// Returns the current set of cross-file shared blocks.
561    pub fn get_shared_blocks(&self) -> &[SharedBlock] {
562        &self.shared_blocks
563    }
564
565    /// Replace shared blocks in content with cross-file references.
566    pub fn apply_dedup(&self, path: &str, content: &str) -> Option<String> {
567        if self.shared_blocks.is_empty() {
568            return None;
569        }
570        let refs: Vec<&SharedBlock> = self
571            .shared_blocks
572            .iter()
573            .filter(|b| b.canonical_path != path && content.contains(&b.content))
574            .collect();
575        if refs.is_empty() {
576            return None;
577        }
578        let mut result = content.to_string();
579        for block in refs {
580            result = result.replacen(
581                &block.content,
582                &format!(
583                    "[= {}:{}-{}]",
584                    block.canonical_ref, block.start_line, block.end_line
585                ),
586                1,
587            );
588        }
589        Some(result)
590    }
591
592    /// Removes a file from the cache, forcing a fresh read on next access.
593    pub fn invalidate(&mut self, path: &str) -> bool {
594        self.entries.remove(&normalize_key(path)).is_some()
595    }
596
597    /// Returns a cached compressed output for a given file and mode key.
598    pub fn get_compressed(&self, path: &str, mode_key: &str) -> Option<&String> {
599        self.entries
600            .get(&normalize_key(path))?
601            .get_compressed(mode_key)
602    }
603
604    /// Marks that full (uncompressed) content was delivered for this file.
605    pub fn mark_full_delivered(&mut self, path: &str) {
606        if let Some(entry) = self.entries.get_mut(&normalize_key(path)) {
607            entry.mark_full_delivered();
608        }
609    }
610
611    /// Stores a compressed output for a given file and mode key.
612    pub fn set_compressed(&mut self, path: &str, mode_key: &str, output: String) {
613        if let Some(entry) = self.entries.get_mut(&normalize_key(path)) {
614            entry.set_compressed(mode_key, output);
615        }
616    }
617
618    /// Resets `full_content_delivered` for all entries without removing them.
619    /// Used after host context compaction — forces re-delivery on next read
620    /// while preserving compressed content and file refs.
621    pub fn reset_delivery_flags(&mut self) -> usize {
622        let mut count = 0;
623        for entry in self.entries.values_mut() {
624            if entry.full_content_delivered {
625                entry.full_content_delivered = false;
626                count += 1;
627            }
628        }
629        count
630    }
631
632    /// Returns whether full content was previously delivered for this path.
633    pub fn is_full_delivered(&self, path: &str) -> bool {
634        self.entries
635            .get(&normalize_key(path))
636            .is_some_and(|e| e.full_content_delivered)
637    }
638
639    /// Removes all compressed output variants (map, signatures, etc.) from every entry,
640    /// keeping the full zstd-compressed content intact. Returns the number of entries trimmed.
641    pub fn trim_compressed_outputs(&mut self) -> usize {
642        let mut trimmed = 0;
643        for entry in self.entries.values_mut() {
644            if !entry.compressed_outputs.is_empty() {
645                entry.compressed_outputs.clear();
646                trimmed += 1;
647            }
648        }
649        trimmed
650    }
651
652    /// Evicts all entries that have been read at most once (probationary).
653    /// Returns the number of entries removed.
654    pub fn evict_probationary(&mut self) -> usize {
655        let to_remove: Vec<String> = self
656            .entries
657            .iter()
658            .filter(|(_, e)| e.read_count() <= 1)
659            .map(|(k, _)| k.clone())
660            .collect();
661        let count = to_remove.len();
662        for key in &to_remove {
663            self.entries.remove(key);
664            self.file_refs.remove(key);
665        }
666        count
667    }
668
669    /// Evicts entries via RRF scoring until total tokens are at or below `target_tokens`.
670    pub fn evict_to_budget(&mut self, target_tokens: usize) {
671        let current = self.total_cached_tokens();
672        if current <= target_tokens {
673            return;
674        }
675        let now = Instant::now();
676        let all: Vec<(&String, &CacheEntry)> = self.entries.iter().collect();
677        let mut scores = eviction_scores_rrf(&all, now);
678        scores.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
679
680        let mut freed = 0usize;
681        let target_free = current.saturating_sub(target_tokens);
682        for (path, _score) in &scores {
683            if freed >= target_free {
684                break;
685            }
686            if let Some(entry) = self.entries.remove(path) {
687                freed += entry.original_tokens;
688                self.file_refs.remove(path);
689            }
690        }
691    }
692
693    /// Estimates the approximate heap memory usage in bytes.
694    pub fn approximate_bytes(&self) -> usize {
695        let entries_bytes: usize = self
696            .entries
697            .values()
698            .map(|e| {
699                e.compressed_content.len()
700                    + e.hash.len()
701                    + e.path.len()
702                    + e.compressed_outputs
703                        .iter()
704                        .map(|(k, v)| k.len() + v.len())
705                        .sum::<usize>()
706                    + 128 // fixed overhead per entry
707            })
708            .sum();
709        let refs_bytes: usize = self.file_refs.iter().map(|(k, v)| k.len() + v.len()).sum();
710        let blocks_bytes: usize = self
711            .shared_blocks
712            .iter()
713            .map(|b| b.canonical_path.len() + b.canonical_ref.len() + b.content.len() + 32)
714            .sum();
715        entries_bytes + refs_bytes + blocks_bytes
716    }
717
718    const MAX_SHARED_BLOCKS: usize = 100;
719
720    /// Trims shared blocks to a maximum count, keeping the most recent.
721    pub fn trim_shared_blocks(&mut self) {
722        if self.shared_blocks.len() > Self::MAX_SHARED_BLOCKS {
723            let excess = self.shared_blocks.len() - Self::MAX_SHARED_BLOCKS;
724            self.shared_blocks.drain(..excess);
725        }
726    }
727
728    /// Clears all cached entries, file refs, and resets stats. Returns the number of entries removed.
729    pub fn clear(&mut self) -> usize {
730        let count = self.entries.len();
731        self.entries.clear();
732        self.file_refs.clear();
733        self.shared_blocks.clear();
734        self.next_ref = 1;
735        self.stats = CacheStats::default();
736        count
737    }
738}
739
740pub fn file_mtime(path: &str) -> Option<SystemTime> {
741    std::fs::metadata(path).and_then(|m| m.modified()).ok()
742}
743
744pub fn is_cache_entry_stale(path: &str, cached_mtime: Option<SystemTime>) -> bool {
745    let current = file_mtime(path);
746    match (cached_mtime, current) {
747        // Both unavailable (e.g. WSL DrvFS): can't tell → assume fresh (conservative).
748        (None, None) => false,
749        // One side missing: metadata changed or appeared/disappeared → stale.
750        (Some(_), None) | (None, Some(_)) => true,
751        // `!=`, not `>`: a *backward* mtime (git checkout, touch -t, snapshot
752        // restore) is just as much a content change as a forward one.
753        (Some(cached), Some(current)) => current != cached,
754    }
755}
756
757/// Files larger than this are not content-hashed for stub verification; the
758/// mtime check alone decides. Keeps the stub fast-path O(small-file-read).
759const VERIFY_HASH_CAP_BYTES: u64 = 8 * 1024 * 1024;
760
761fn cache_verify_enabled() -> bool {
762    std::env::var("LEAN_CTX_CACHE_VERIFY").map_or(true, |v| v != "0")
763}
764
765/// Staleness with content verification: like [`is_cache_entry_stale`], but when
766/// the mtime claims "unchanged", additionally compares the md5 of the on-disk
767/// content against the cached hash.
768///
769/// mtime alone cannot be trusted for *correctness*: same-second writes are
770/// invisible on coarse-granularity filesystems (HFS+ 1s, FAT 2s) and mtimes can
771/// be restored by tools. Serving an `[unchanged]` stub for changed content
772/// would silently mislead the agent — the worst failure mode a context layer
773/// can have. The extra disk read costs microseconds for typical source files;
774/// the stub's token savings are unaffected. Opt out: `LEAN_CTX_CACHE_VERIFY=0`.
775///
776/// Note: entries whose stored content differs from disk by design (e.g. secret
777/// redaction) hash differently and therefore never serve stubs — conservative
778/// and correct.
779pub fn is_cache_entry_stale_verified(
780    path: &str,
781    cached_mtime: Option<SystemTime>,
782    cached_hash: &str,
783) -> bool {
784    if is_cache_entry_stale(path, cached_mtime) {
785        return true;
786    }
787    if cached_hash.is_empty() || !cache_verify_enabled() {
788        return false;
789    }
790    let Ok(meta) = std::fs::metadata(path) else {
791        // Can't stat → never serve a stub on top of it.
792        return true;
793    };
794    if meta.len() > VERIFY_HASH_CAP_BYTES {
795        return false;
796    }
797    match std::fs::read(path) {
798        // Hash the same view of the bytes that `store()` hashed (lossy UTF-8).
799        Ok(bytes) => compute_md5(&String::from_utf8_lossy(&bytes)) != cached_hash,
800        Err(_) => true,
801    }
802}
803
804fn compute_md5(content: &str) -> String {
805    let mut hasher = Md5::new();
806    hasher.update(content.as_bytes());
807    format!("{:x}", hasher.finalize())
808}
809
810#[cfg(test)]
811mod tests {
812    use super::*;
813    use std::time::Duration;
814
815    #[test]
816    fn cache_stores_and_retrieves() {
817        let mut cache = SessionCache::new();
818        let result = cache.store("/test/file.rs", "fn main() {}");
819        assert!(!result.was_hit);
820        assert_eq!(result.line_count, 1);
821        assert!(cache.get("/test/file.rs").is_some());
822    }
823
824    #[test]
825    fn cache_hit_on_same_content() {
826        let mut cache = SessionCache::new();
827        cache.store("/test/file.rs", "content");
828        let result = cache.store("/test/file.rs", "content");
829        assert!(result.was_hit, "same content should be a cache hit");
830    }
831
832    #[test]
833    fn cache_miss_on_changed_content() {
834        let mut cache = SessionCache::new();
835        cache.store("/test/file.rs", "old content");
836        let result = cache.store("/test/file.rs", "new content");
837        assert!(!result.was_hit, "changed content should not be a cache hit");
838    }
839
840    #[test]
841    fn file_refs_are_sequential() {
842        let mut cache = SessionCache::new();
843        assert_eq!(cache.get_file_ref("/a.rs"), "F1");
844        assert_eq!(cache.get_file_ref("/b.rs"), "F2");
845        assert_eq!(cache.get_file_ref("/a.rs"), "F1"); // stable
846    }
847
848    #[test]
849    fn cache_clear_resets_everything() {
850        let mut cache = SessionCache::new();
851        cache.store("/a.rs", "a");
852        cache.store("/b.rs", "b");
853        let count = cache.clear();
854        assert_eq!(count, 2);
855        assert!(cache.get("/a.rs").is_none());
856        assert_eq!(cache.get_file_ref("/c.rs"), "F1"); // refs reset
857    }
858
859    #[test]
860    fn cache_invalidate_removes_entry() {
861        let mut cache = SessionCache::new();
862        cache.store("/test.rs", "test");
863        assert!(cache.invalidate("/test.rs"));
864        assert!(!cache.invalidate("/nonexistent.rs"));
865    }
866
867    #[test]
868    fn cache_stats_track_correctly() {
869        let mut cache = SessionCache::new();
870        cache.store("/a.rs", "hello");
871        cache.store("/a.rs", "hello"); // hit
872        let stats = cache.get_stats();
873        assert_eq!(stats.total_reads(), 2);
874        assert_eq!(stats.cache_hits(), 1);
875        assert!(stats.hit_rate() > 0.0);
876    }
877
878    #[test]
879    fn record_cache_hit_works_through_shared_ref() {
880        let mut cache = SessionCache::new();
881        cache.store("/x.rs", "hello world");
882        // &self path: a cache hit can be recorded without a write lock.
883        let shared: &SessionCache = &cache;
884        assert!(shared.record_cache_hit("/x.rs").is_some());
885        assert!(shared.record_cache_hit("/x.rs").is_some());
886        // store=1 + two hits => read_count 3, cache_hits 2.
887        assert_eq!(cache.get("/x.rs").unwrap().read_count(), 3);
888        assert_eq!(cache.get_stats().cache_hits(), 2);
889    }
890
891    #[test]
892    fn concurrent_cache_hits_are_lossless() {
893        use std::sync::Arc;
894        let mut cache = SessionCache::new();
895        cache.store("/a.rs", "a");
896        cache.store("/b.rs", "b");
897        // Shared (no RwLock): proves SessionCache is Sync and hit recording is
898        // lock-free and atomic — the whole point of the read-mostly refactor.
899        let cache = Arc::new(cache);
900        let threads = 8;
901        let iters = 1_000;
902        let handles: Vec<_> = (0..threads)
903            .map(|_| {
904                let c = Arc::clone(&cache);
905                std::thread::spawn(move || {
906                    for _ in 0..iters {
907                        c.record_cache_hit("/a.rs");
908                        c.record_cache_hit("/b.rs");
909                    }
910                })
911            })
912            .collect();
913        for h in handles {
914            h.join().unwrap();
915        }
916        let total = (threads * iters) as u64;
917        assert_eq!(cache.get_stats().cache_hits(), total * 2);
918        assert_eq!(cache.get("/a.rs").unwrap().read_count(), 1 + total as u32);
919        assert_eq!(cache.get("/b.rs").unwrap().read_count(), 1 + total as u32);
920    }
921
922    #[test]
923    fn md5_is_deterministic() {
924        let h1 = compute_md5("test content");
925        let h2 = compute_md5("test content");
926        assert_eq!(h1, h2);
927        assert_ne!(h1, compute_md5("different"));
928    }
929
930    #[test]
931    fn rrf_eviction_prefers_recent() {
932        let key_a = "a.rs".to_string();
933        let key_b = "b.rs".to_string();
934        // Construct entries first so the global instant base is initialized,
935        // then assign access times relative to a post-init reference.
936        let recent = CacheEntry::new("a", "h1".to_string(), 1, 10, "/a.rs".to_string(), None);
937        let old = CacheEntry::new("b", "h2".to_string(), 1, 10, "/b.rs".to_string(), None);
938        let t_old = Instant::now();
939        std::thread::sleep(std::time::Duration::from_millis(10));
940        let t_recent = Instant::now();
941        old.set_last_access(t_old);
942        recent.set_last_access(t_recent);
943        let now = Instant::now();
944        let entries: Vec<(&String, &CacheEntry)> = vec![(&key_a, &recent), (&key_b, &old)];
945        let scores = eviction_scores_rrf(&entries, now);
946        let score_a = scores.iter().find(|(p, _)| p == "a.rs").unwrap().1;
947        let score_b = scores.iter().find(|(p, _)| p == "b.rs").unwrap().1;
948        assert!(
949            score_a > score_b,
950            "recently accessed entries should score higher via RRF"
951        );
952    }
953
954    #[test]
955    fn rrf_eviction_prefers_frequent() {
956        let now = Instant::now();
957        let key_a = "a.rs".to_string();
958        let key_b = "b.rs".to_string();
959        let frequent = {
960            let e = CacheEntry::new("a", "h1".to_string(), 1, 10, "/a.rs".to_string(), None);
961            e.set_read_count(20);
962            e
963        };
964        let rare = CacheEntry::new("b", "h2".to_string(), 1, 10, "/b.rs".to_string(), None);
965        let entries: Vec<(&String, &CacheEntry)> = vec![(&key_a, &frequent), (&key_b, &rare)];
966        let scores = eviction_scores_rrf(&entries, now);
967        let score_a = scores.iter().find(|(p, _)| p == "a.rs").unwrap().1;
968        let score_b = scores.iter().find(|(p, _)| p == "b.rs").unwrap().1;
969        assert!(
970            score_a > score_b,
971            "frequently accessed entries should score higher via RRF"
972        );
973    }
974
975    #[test]
976    fn evict_if_needed_removes_lowest_score() {
977        crate::test_env::set_var("LEAN_CTX_CACHE_MAX_TOKENS", "50");
978        let mut cache = SessionCache::new();
979        let big_content = "a]".repeat(30); // ~30 tokens
980        cache.store("/old.rs", &big_content);
981        // /old.rs now in cache with ~30 tokens
982
983        let new_content = "b ".repeat(30); // ~30 tokens incoming
984        cache.store("/new.rs", &new_content);
985        // should have evicted /old.rs to make room
986        // (total would be ~60 which exceeds 50)
987
988        // At least one should remain, total should be <= 50
989        assert!(
990            cache.total_cached_tokens() <= 60,
991            "eviction should have kicked in"
992        );
993        crate::test_env::remove_var("LEAN_CTX_CACHE_MAX_TOKENS");
994    }
995
996    #[test]
997    fn stale_detection_flags_newer_file() {
998        let dir = tempfile::tempdir().unwrap();
999        let path = dir.path().join("stale.txt");
1000        let p = path.to_string_lossy().to_string();
1001
1002        std::fs::write(&path, "one").unwrap();
1003        let mut cache = SessionCache::new();
1004        cache.store(&p, "one");
1005
1006        let entry = cache.get(&p).unwrap();
1007        assert!(!is_cache_entry_stale(&p, entry.stored_mtime));
1008
1009        // Ensure mtime granularity differences don't make this flaky.
1010        std::thread::sleep(Duration::from_secs(1));
1011        std::fs::write(&path, "two").unwrap();
1012
1013        let entry = cache.get(&p).unwrap();
1014        assert!(is_cache_entry_stale(&p, entry.stored_mtime));
1015    }
1016
1017    // P0-7 (#419): a *backward* mtime (git checkout, touch -t) is a change.
1018    #[test]
1019    fn stale_detection_flags_backward_mtime() {
1020        let dir = tempfile::tempdir().unwrap();
1021        let path = dir.path().join("backward.txt");
1022        let p = path.to_string_lossy().to_string();
1023
1024        std::fs::write(&path, "one").unwrap();
1025        let mut cache = SessionCache::new();
1026        cache.store(&p, "one");
1027        let entry_mtime = cache.get(&p).unwrap().stored_mtime;
1028        assert!(!is_cache_entry_stale(&p, entry_mtime));
1029
1030        // Simulate `git checkout` of an older version: content + older mtime.
1031        std::fs::write(&path, "zero").unwrap();
1032        let f = std::fs::OpenOptions::new().write(true).open(&path).unwrap();
1033        f.set_modified(SystemTime::now() - Duration::from_hours(1))
1034            .unwrap();
1035        drop(f);
1036
1037        assert!(
1038            is_cache_entry_stale(&p, entry_mtime),
1039            "older mtime must read as stale"
1040        );
1041    }
1042
1043    // P0-7 (#419): identical mtime with different content (same-second write,
1044    // restored timestamps) is caught by the content-hash verification.
1045    #[test]
1046    fn verified_staleness_catches_same_mtime_content_change() {
1047        let dir = tempfile::tempdir().unwrap();
1048        let path = dir.path().join("sneaky.txt");
1049        let p = path.to_string_lossy().to_string();
1050
1051        std::fs::write(&path, "one").unwrap();
1052        let original_mtime = std::fs::metadata(&path).unwrap().modified().unwrap();
1053        let mut cache = SessionCache::new();
1054        cache.store(&p, "one");
1055        let (mtime, hash) = {
1056            let e = cache.get(&p).unwrap();
1057            (e.stored_mtime, e.hash.clone())
1058        };
1059
1060        // Unchanged file: both checks agree it is fresh.
1061        assert!(!is_cache_entry_stale_verified(&p, mtime, &hash));
1062
1063        // Change the content but restore the exact original mtime.
1064        std::fs::write(&path, "two").unwrap();
1065        let f = std::fs::OpenOptions::new().write(true).open(&path).unwrap();
1066        f.set_modified(original_mtime).unwrap();
1067        drop(f);
1068
1069        assert!(
1070            !is_cache_entry_stale(&p, mtime),
1071            "test premise: the mtime check alone is fooled"
1072        );
1073        assert!(
1074            is_cache_entry_stale_verified(&p, mtime, &hash),
1075            "hash verification must catch the change"
1076        );
1077    }
1078
1079    #[test]
1080    fn verified_staleness_flags_unreadable_file() {
1081        let mut cache = SessionCache::new();
1082        cache.store("/nonexistent/file.rs", "content");
1083        let (mtime, hash) = {
1084            let e = cache.get("/nonexistent/file.rs").unwrap();
1085            (e.stored_mtime, e.hash.clone())
1086        };
1087        assert!(is_cache_entry_stale_verified(
1088            "/nonexistent/file.rs",
1089            mtime,
1090            &hash
1091        ));
1092    }
1093
1094    #[test]
1095    fn compressed_outputs_cached_and_retrieved() {
1096        let mut cache = SessionCache::new();
1097        cache.store("/test.rs", "fn main() {}");
1098        cache.set_compressed("/test.rs", "map", "compressed map output".to_string());
1099        assert_eq!(
1100            cache.get_compressed("/test.rs", "map"),
1101            Some(&"compressed map output".to_string())
1102        );
1103        assert_eq!(cache.get_compressed("/test.rs", "signatures"), None);
1104    }
1105
1106    #[test]
1107    fn compressed_outputs_cleared_on_content_change() {
1108        let mut cache = SessionCache::new();
1109        cache.store("/test.rs", "old content");
1110        cache.set_compressed("/test.rs", "map", "old map".to_string());
1111        assert!(cache.get_compressed("/test.rs", "map").is_some());
1112
1113        cache.store("/test.rs", "new content");
1114        assert_eq!(cache.get_compressed("/test.rs", "map"), None);
1115    }
1116
1117    #[test]
1118    fn compressed_outputs_survive_same_content_store() {
1119        let mut cache = SessionCache::new();
1120        cache.store("/test.rs", "content");
1121        cache.set_compressed("/test.rs", "map", "cached map".to_string());
1122
1123        let result = cache.store("/test.rs", "content");
1124        assert!(result.was_hit);
1125        assert_eq!(
1126            cache.get_compressed("/test.rs", "map"),
1127            Some(&"cached map".to_string())
1128        );
1129    }
1130
1131    #[test]
1132    fn compressed_outputs_cleared_on_invalidate() {
1133        let mut cache = SessionCache::new();
1134        cache.store("/test.rs", "content");
1135        cache.set_compressed("/test.rs", "signatures", "cached sigs".to_string());
1136        cache.invalidate("/test.rs");
1137        assert_eq!(cache.get_compressed("/test.rs", "signatures"), None);
1138    }
1139
1140    #[test]
1141    fn compressed_outputs_cleared_on_clear() {
1142        let mut cache = SessionCache::new();
1143        cache.store("/a.rs", "a");
1144        cache.set_compressed("/a.rs", "map", "map_a".to_string());
1145        cache.clear();
1146        assert_eq!(cache.get_compressed("/a.rs", "map"), None);
1147    }
1148}