Skip to main content

lean_ctx/core/
cache.rs

1use md5::{Digest, Md5};
2use std::collections::HashMap;
3use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
4use std::sync::OnceLock;
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        {
179            if let Some(oldest_key) = self.compressed_outputs.keys().next().cloned() {
180                self.compressed_outputs.remove(&oldest_key);
181            }
182        }
183        self.compressed_outputs.insert(mode_key.to_string(), output);
184    }
185
186    pub fn mark_full_delivered(&mut self) {
187        self.full_content_delivered = true;
188    }
189}
190
191const RRF_K: f64 = 60.0;
192
193/// Compute Reciprocal Rank Fusion eviction scores for a batch of cache entries.
194/// Each signal (recency, frequency, size) produces an independent ranking.
195/// The final score is the sum of `1/(k + rank)` across all signals.
196/// Higher score = more valuable = keep longer.
197pub fn eviction_scores_rrf(entries: &[(&String, &CacheEntry)], now: Instant) -> Vec<(String, f64)> {
198    if entries.is_empty() {
199        return Vec::new();
200    }
201
202    let n = entries.len();
203
204    let mut recency_order: Vec<usize> = (0..n).collect();
205    recency_order.sort_by(|&a, &b| {
206        let elapsed_a = now
207            .checked_duration_since(entries[a].1.last_access())
208            .unwrap_or_default()
209            .as_secs_f64();
210        let elapsed_b = now
211            .checked_duration_since(entries[b].1.last_access())
212            .unwrap_or_default()
213            .as_secs_f64();
214        elapsed_a
215            .partial_cmp(&elapsed_b)
216            .unwrap_or(std::cmp::Ordering::Equal)
217    });
218
219    let mut frequency_order: Vec<usize> = (0..n).collect();
220    frequency_order.sort_by(|&a, &b| entries[b].1.read_count().cmp(&entries[a].1.read_count()));
221
222    let mut size_order: Vec<usize> = (0..n).collect();
223    size_order.sort_by(|&a, &b| {
224        entries[b]
225            .1
226            .original_tokens
227            .cmp(&entries[a].1.original_tokens)
228    });
229
230    let mut recency_ranks = vec![0usize; n];
231    let mut frequency_ranks = vec![0usize; n];
232    let mut size_ranks = vec![0usize; n];
233
234    for (rank, &idx) in recency_order.iter().enumerate() {
235        recency_ranks[idx] = rank;
236    }
237    for (rank, &idx) in frequency_order.iter().enumerate() {
238        frequency_ranks[idx] = rank;
239    }
240    for (rank, &idx) in size_order.iter().enumerate() {
241        size_ranks[idx] = rank;
242    }
243
244    entries
245        .iter()
246        .enumerate()
247        .map(|(i, (path, _))| {
248            let score = 1.0 / (RRF_K + recency_ranks[i] as f64)
249                + 1.0 / (RRF_K + frequency_ranks[i] as f64)
250                + 1.0 / (RRF_K + size_ranks[i] as f64);
251            ((*path).clone(), score)
252        })
253        .collect()
254}
255
256/// Aggregated cache statistics: hits, reads, and token savings.
257///
258/// Counters are atomic so they can be updated on the read-locked cache-hit
259/// fast path without taking a write lock.
260#[derive(Debug, Default)]
261pub struct CacheStats {
262    total_reads: AtomicU64,
263    cache_hits: AtomicU64,
264    total_original_tokens: AtomicU64,
265    total_sent_tokens: AtomicU64,
266    files_tracked: AtomicU64,
267}
268
269impl CacheStats {
270    /// Total number of read operations recorded.
271    pub fn total_reads(&self) -> u64 {
272        self.total_reads.load(Ordering::Relaxed)
273    }
274
275    /// Total number of cache hits recorded.
276    pub fn cache_hits(&self) -> u64 {
277        self.cache_hits.load(Ordering::Relaxed)
278    }
279
280    /// Sum of original (uncompressed) token counts across all reads.
281    pub fn total_original_tokens(&self) -> u64 {
282        self.total_original_tokens.load(Ordering::Relaxed)
283    }
284
285    /// Sum of tokens actually sent to the model.
286    pub fn total_sent_tokens(&self) -> u64 {
287        self.total_sent_tokens.load(Ordering::Relaxed)
288    }
289
290    /// Number of distinct files currently tracked.
291    pub fn files_tracked(&self) -> u64 {
292        self.files_tracked.load(Ordering::Relaxed)
293    }
294
295    /// Returns the cache hit rate as a percentage (0–100).
296    pub fn hit_rate(&self) -> f64 {
297        let total = self.total_reads();
298        if total == 0 {
299            return 0.0;
300        }
301        (self.cache_hits() as f64 / total as f64) * 100.0
302    }
303
304    /// Returns the total number of tokens saved by cache hits.
305    pub fn tokens_saved(&self) -> u64 {
306        self.total_original_tokens()
307            .saturating_sub(self.total_sent_tokens())
308    }
309
310    /// Returns the savings as a percentage of total original tokens.
311    pub fn savings_percent(&self) -> f64 {
312        let original = self.total_original_tokens();
313        if original == 0 {
314            return 0.0;
315        }
316        (self.tokens_saved() as f64 / original as f64) * 100.0
317    }
318}
319
320/// A block shared across multiple files, identified by its canonical source.
321#[derive(Clone, Debug)]
322pub struct SharedBlock {
323    pub canonical_path: String,
324    pub canonical_ref: String,
325    pub start_line: usize,
326    pub end_line: usize,
327    pub content: String,
328}
329
330/// In-memory file cache with segmented LRU eviction (probationary vs protected),
331/// file references, and cross-file dedup.
332pub struct SessionCache {
333    entries: HashMap<String, CacheEntry>,
334    file_refs: HashMap<String, String>,
335    next_ref: usize,
336    stats: CacheStats,
337    shared_blocks: Vec<SharedBlock>,
338}
339
340impl Default for SessionCache {
341    fn default() -> Self {
342        Self::new()
343    }
344}
345
346impl SessionCache {
347    /// Creates an empty session cache with default stats.
348    pub fn new() -> Self {
349        Self {
350            entries: HashMap::new(),
351            file_refs: HashMap::new(),
352            next_ref: 1,
353            shared_blocks: Vec::new(),
354            stats: CacheStats::default(),
355        }
356    }
357
358    /// Returns or assigns a short file reference label (F1, F2, ...) for the given path.
359    pub fn get_file_ref(&mut self, path: &str) -> String {
360        let key = normalize_key(path);
361        if let Some(r) = self.file_refs.get(&key) {
362            return r.clone();
363        }
364        let r = format!("F{}", self.next_ref);
365        self.next_ref += 1;
366        self.file_refs.insert(key, r.clone());
367        r
368    }
369
370    /// Returns the file reference label for a path without assigning a new one.
371    pub fn get_file_ref_readonly(&self, path: &str) -> Option<String> {
372        self.file_refs.get(&normalize_key(path)).cloned()
373    }
374
375    /// Looks up a cached entry by file path.
376    pub fn get(&self, path: &str) -> Option<&CacheEntry> {
377        self.entries.get(&normalize_key(path))
378    }
379
380    /// Mutable lookup of a cached entry by file path.
381    pub fn get_mut(&mut self, path: &str) -> Option<&mut CacheEntry> {
382        self.entries.get_mut(&normalize_key(path))
383    }
384
385    /// Retrieves the full (uncompressed) content for a file path, if cached.
386    /// Used by the CCR (Compress-Cache-Retrieve) mechanism.
387    pub fn get_full_content(&self, path: &str) -> Option<String> {
388        self.entries
389            .get(&normalize_key(path))
390            .and_then(CacheEntry::content)
391    }
392
393    /// Records a cache hit, updates access stats, and emits a cache-hit event.
394    ///
395    /// Takes `&self`: the hit counters use interior-mutable atomics, so this
396    /// runs under a shared (read) lock and lets parallel reads of different
397    /// files proceed concurrently instead of serializing on a write lock.
398    pub fn record_cache_hit(&self, path: &str) -> Option<&CacheEntry> {
399        let key = normalize_key(path);
400        let ref_label = self
401            .file_refs
402            .get(&key)
403            .cloned()
404            .unwrap_or_else(|| "F?".to_string());
405        let entry = self.entries.get(&key)?;
406        let new_count = entry.bump_read_count();
407        entry.touch();
408        self.stats.total_reads.fetch_add(1, Ordering::Relaxed);
409        self.stats.cache_hits.fetch_add(1, Ordering::Relaxed);
410        self.stats
411            .total_original_tokens
412            .fetch_add(entry.original_tokens as u64, Ordering::Relaxed);
413        let hit_msg = format!("{ref_label} cached {new_count}t {}L", entry.line_count);
414        self.stats
415            .total_sent_tokens
416            .fetch_add(count_tokens(&hit_msg) as u64, Ordering::Relaxed);
417        crate::core::events::emit_cache_hit(path, entry.original_tokens as u64);
418        Some(entry)
419    }
420
421    /// Stores file content in the cache; returns a hit if content hash matches.
422    pub fn store(&mut self, path: &str, content: &str) -> StoreResult {
423        let key = normalize_key(path);
424        let hash = compute_md5(content);
425        let line_count = content.lines().count();
426        let original_tokens = count_tokens(content);
427        let stored_mtime = std::fs::metadata(path).and_then(|m| m.modified()).ok();
428        let now = Instant::now();
429
430        self.stats.total_reads.fetch_add(1, Ordering::Relaxed);
431        self.stats
432            .total_original_tokens
433            .fetch_add(original_tokens as u64, Ordering::Relaxed);
434
435        if let Some(existing) = self.entries.get_mut(&key) {
436            existing.set_last_access(now);
437            if stored_mtime.is_some() {
438                existing.stored_mtime = stored_mtime;
439            }
440            if existing.hash == hash {
441                let new_count = existing.bump_read_count();
442                self.stats.cache_hits.fetch_add(1, Ordering::Relaxed);
443                let hit_msg = format!(
444                    "{} cached {new_count}t {}L",
445                    self.file_refs.get(&key).unwrap_or(&"F?".to_string()),
446                    existing.line_count,
447                );
448                self.stats
449                    .total_sent_tokens
450                    .fetch_add(count_tokens(&hit_msg) as u64, Ordering::Relaxed);
451                return StoreResult {
452                    line_count: existing.line_count,
453                    original_tokens: existing.original_tokens,
454                    read_count: new_count,
455                    was_hit: true,
456                    full_content_delivered: existing.full_content_delivered,
457                };
458            }
459            existing.compressed_outputs.clear();
460            existing.set_content(content);
461            existing.hash = hash;
462            existing.line_count = line_count;
463            existing.original_tokens = original_tokens;
464            let new_count = existing.bump_read_count();
465            existing.full_content_delivered = false;
466            if stored_mtime.is_some() {
467                existing.stored_mtime = stored_mtime;
468            }
469            self.stats
470                .total_sent_tokens
471                .fetch_add(original_tokens as u64, Ordering::Relaxed);
472            return StoreResult {
473                line_count,
474                original_tokens,
475                read_count: new_count,
476                was_hit: false,
477                full_content_delivered: false,
478            };
479        }
480
481        self.evict_if_needed(original_tokens);
482        self.get_file_ref(&key);
483
484        let entry = CacheEntry::new(
485            content,
486            hash,
487            line_count,
488            original_tokens,
489            key.clone(),
490            stored_mtime,
491        );
492
493        self.entries.insert(key, entry);
494        self.stats.files_tracked.fetch_add(1, Ordering::Relaxed);
495        self.stats
496            .total_sent_tokens
497            .fetch_add(original_tokens as u64, Ordering::Relaxed);
498        StoreResult {
499            line_count,
500            original_tokens,
501            read_count: 1,
502            was_hit: false,
503            full_content_delivered: false,
504        }
505    }
506
507    /// Returns the sum of original token counts across all cached entries.
508    pub fn total_cached_tokens(&self) -> usize {
509        self.entries.values().map(|e| e.original_tokens).sum()
510    }
511
512    /// Evict until cache fits within token budget using RRF (Reciprocal Rank Fusion).
513    /// Combines recency, frequency, and size signals to evict least-valuable entries first.
514    pub fn evict_if_needed(&mut self, incoming_tokens: usize) {
515        let max_tokens = max_cache_tokens();
516        let current = self.total_cached_tokens();
517        if current + incoming_tokens <= max_tokens {
518            return;
519        }
520
521        let now = Instant::now();
522        let all: Vec<(&String, &CacheEntry)> = self.entries.iter().collect();
523        let mut scores = eviction_scores_rrf(&all, now);
524        // Sort ascending: lowest RRF score = least valuable = evict first
525        scores.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
526
527        let mut freed = 0usize;
528        let target = (current + incoming_tokens).saturating_sub(max_tokens);
529
530        for (path, _score) in &scores {
531            if freed >= target {
532                break;
533            }
534            if let Some(entry) = self.entries.remove(path) {
535                freed += entry.original_tokens;
536                self.file_refs.remove(path);
537            }
538        }
539    }
540
541    /// Returns all cached entries as (path, entry) pairs.
542    pub fn get_all_entries(&self) -> Vec<(&String, &CacheEntry)> {
543        self.entries.iter().collect()
544    }
545
546    /// Returns a reference to the aggregated cache statistics.
547    pub fn get_stats(&self) -> &CacheStats {
548        &self.stats
549    }
550
551    /// Returns the path-to-file-ref mapping (e.g. "/src/main.rs" → "F1").
552    pub fn file_ref_map(&self) -> &HashMap<String, String> {
553        &self.file_refs
554    }
555
556    /// Replaces the cross-file shared blocks used for deduplication.
557    pub fn set_shared_blocks(&mut self, blocks: Vec<SharedBlock>) {
558        self.shared_blocks = blocks;
559    }
560
561    /// Returns the current set of cross-file shared blocks.
562    pub fn get_shared_blocks(&self) -> &[SharedBlock] {
563        &self.shared_blocks
564    }
565
566    /// Replace shared blocks in content with cross-file references.
567    pub fn apply_dedup(&self, path: &str, content: &str) -> Option<String> {
568        if self.shared_blocks.is_empty() {
569            return None;
570        }
571        let refs: Vec<&SharedBlock> = self
572            .shared_blocks
573            .iter()
574            .filter(|b| b.canonical_path != path && content.contains(&b.content))
575            .collect();
576        if refs.is_empty() {
577            return None;
578        }
579        let mut result = content.to_string();
580        for block in refs {
581            result = result.replacen(
582                &block.content,
583                &format!(
584                    "[= {}:{}-{}]",
585                    block.canonical_ref, block.start_line, block.end_line
586                ),
587                1,
588            );
589        }
590        Some(result)
591    }
592
593    /// Removes a file from the cache, forcing a fresh read on next access.
594    pub fn invalidate(&mut self, path: &str) -> bool {
595        self.entries.remove(&normalize_key(path)).is_some()
596    }
597
598    /// Returns a cached compressed output for a given file and mode key.
599    pub fn get_compressed(&self, path: &str, mode_key: &str) -> Option<&String> {
600        self.entries
601            .get(&normalize_key(path))?
602            .get_compressed(mode_key)
603    }
604
605    /// Marks that full (uncompressed) content was delivered for this file.
606    pub fn mark_full_delivered(&mut self, path: &str) {
607        if let Some(entry) = self.entries.get_mut(&normalize_key(path)) {
608            entry.mark_full_delivered();
609        }
610    }
611
612    /// Stores a compressed output for a given file and mode key.
613    pub fn set_compressed(&mut self, path: &str, mode_key: &str, output: String) {
614        if let Some(entry) = self.entries.get_mut(&normalize_key(path)) {
615            entry.set_compressed(mode_key, output);
616        }
617    }
618
619    /// Resets `full_content_delivered` for all entries without removing them.
620    /// Used after host context compaction — forces re-delivery on next read
621    /// while preserving compressed content and file refs.
622    pub fn reset_delivery_flags(&mut self) -> usize {
623        let mut count = 0;
624        for entry in self.entries.values_mut() {
625            if entry.full_content_delivered {
626                entry.full_content_delivered = false;
627                count += 1;
628            }
629        }
630        count
631    }
632
633    /// Returns whether full content was previously delivered for this path.
634    pub fn is_full_delivered(&self, path: &str) -> bool {
635        self.entries
636            .get(&normalize_key(path))
637            .is_some_and(|e| e.full_content_delivered)
638    }
639
640    /// Removes all compressed output variants (map, signatures, etc.) from every entry,
641    /// keeping the full zstd-compressed content intact. Returns the number of entries trimmed.
642    pub fn trim_compressed_outputs(&mut self) -> usize {
643        let mut trimmed = 0;
644        for entry in self.entries.values_mut() {
645            if !entry.compressed_outputs.is_empty() {
646                entry.compressed_outputs.clear();
647                trimmed += 1;
648            }
649        }
650        trimmed
651    }
652
653    /// Evicts all entries that have been read at most once (probationary).
654    /// Returns the number of entries removed.
655    pub fn evict_probationary(&mut self) -> usize {
656        let to_remove: Vec<String> = self
657            .entries
658            .iter()
659            .filter(|(_, e)| e.read_count() <= 1)
660            .map(|(k, _)| k.clone())
661            .collect();
662        let count = to_remove.len();
663        for key in &to_remove {
664            self.entries.remove(key);
665            self.file_refs.remove(key);
666        }
667        count
668    }
669
670    /// Evicts entries via RRF scoring until total tokens are at or below `target_tokens`.
671    pub fn evict_to_budget(&mut self, target_tokens: usize) {
672        let current = self.total_cached_tokens();
673        if current <= target_tokens {
674            return;
675        }
676        let now = Instant::now();
677        let all: Vec<(&String, &CacheEntry)> = self.entries.iter().collect();
678        let mut scores = eviction_scores_rrf(&all, now);
679        scores.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
680
681        let mut freed = 0usize;
682        let target_free = current.saturating_sub(target_tokens);
683        for (path, _score) in &scores {
684            if freed >= target_free {
685                break;
686            }
687            if let Some(entry) = self.entries.remove(path) {
688                freed += entry.original_tokens;
689                self.file_refs.remove(path);
690            }
691        }
692    }
693
694    /// Estimates the approximate heap memory usage in bytes.
695    pub fn approximate_bytes(&self) -> usize {
696        let entries_bytes: usize = self
697            .entries
698            .values()
699            .map(|e| {
700                e.compressed_content.len()
701                    + e.hash.len()
702                    + e.path.len()
703                    + e.compressed_outputs
704                        .iter()
705                        .map(|(k, v)| k.len() + v.len())
706                        .sum::<usize>()
707                    + 128 // fixed overhead per entry
708            })
709            .sum();
710        let refs_bytes: usize = self.file_refs.iter().map(|(k, v)| k.len() + v.len()).sum();
711        let blocks_bytes: usize = self
712            .shared_blocks
713            .iter()
714            .map(|b| b.canonical_path.len() + b.canonical_ref.len() + b.content.len() + 32)
715            .sum();
716        entries_bytes + refs_bytes + blocks_bytes
717    }
718
719    const MAX_SHARED_BLOCKS: usize = 100;
720
721    /// Trims shared blocks to a maximum count, keeping the most recent.
722    pub fn trim_shared_blocks(&mut self) {
723        if self.shared_blocks.len() > Self::MAX_SHARED_BLOCKS {
724            let excess = self.shared_blocks.len() - Self::MAX_SHARED_BLOCKS;
725            self.shared_blocks.drain(..excess);
726        }
727    }
728
729    /// Clears all cached entries, file refs, and resets stats. Returns the number of entries removed.
730    pub fn clear(&mut self) -> usize {
731        let count = self.entries.len();
732        self.entries.clear();
733        self.file_refs.clear();
734        self.shared_blocks.clear();
735        self.next_ref = 1;
736        self.stats = CacheStats::default();
737        count
738    }
739}
740
741pub fn file_mtime(path: &str) -> Option<SystemTime> {
742    std::fs::metadata(path).and_then(|m| m.modified()).ok()
743}
744
745pub fn is_cache_entry_stale(path: &str, cached_mtime: Option<SystemTime>) -> bool {
746    let current = file_mtime(path);
747    match (cached_mtime, current) {
748        // Both unavailable (e.g. WSL DrvFS): can't tell → assume fresh (conservative).
749        (None, None) => false,
750        // One side missing: metadata changed or appeared/disappeared → stale.
751        (Some(_), None) | (None, Some(_)) => true,
752        (Some(cached), Some(current)) => current > cached,
753    }
754}
755
756fn compute_md5(content: &str) -> String {
757    let mut hasher = Md5::new();
758    hasher.update(content.as_bytes());
759    format!("{:x}", hasher.finalize())
760}
761
762#[cfg(test)]
763mod tests {
764    use super::*;
765    use std::time::Duration;
766
767    #[test]
768    fn cache_stores_and_retrieves() {
769        let mut cache = SessionCache::new();
770        let result = cache.store("/test/file.rs", "fn main() {}");
771        assert!(!result.was_hit);
772        assert_eq!(result.line_count, 1);
773        assert!(cache.get("/test/file.rs").is_some());
774    }
775
776    #[test]
777    fn cache_hit_on_same_content() {
778        let mut cache = SessionCache::new();
779        cache.store("/test/file.rs", "content");
780        let result = cache.store("/test/file.rs", "content");
781        assert!(result.was_hit, "same content should be a cache hit");
782    }
783
784    #[test]
785    fn cache_miss_on_changed_content() {
786        let mut cache = SessionCache::new();
787        cache.store("/test/file.rs", "old content");
788        let result = cache.store("/test/file.rs", "new content");
789        assert!(!result.was_hit, "changed content should not be a cache hit");
790    }
791
792    #[test]
793    fn file_refs_are_sequential() {
794        let mut cache = SessionCache::new();
795        assert_eq!(cache.get_file_ref("/a.rs"), "F1");
796        assert_eq!(cache.get_file_ref("/b.rs"), "F2");
797        assert_eq!(cache.get_file_ref("/a.rs"), "F1"); // stable
798    }
799
800    #[test]
801    fn cache_clear_resets_everything() {
802        let mut cache = SessionCache::new();
803        cache.store("/a.rs", "a");
804        cache.store("/b.rs", "b");
805        let count = cache.clear();
806        assert_eq!(count, 2);
807        assert!(cache.get("/a.rs").is_none());
808        assert_eq!(cache.get_file_ref("/c.rs"), "F1"); // refs reset
809    }
810
811    #[test]
812    fn cache_invalidate_removes_entry() {
813        let mut cache = SessionCache::new();
814        cache.store("/test.rs", "test");
815        assert!(cache.invalidate("/test.rs"));
816        assert!(!cache.invalidate("/nonexistent.rs"));
817    }
818
819    #[test]
820    fn cache_stats_track_correctly() {
821        let mut cache = SessionCache::new();
822        cache.store("/a.rs", "hello");
823        cache.store("/a.rs", "hello"); // hit
824        let stats = cache.get_stats();
825        assert_eq!(stats.total_reads(), 2);
826        assert_eq!(stats.cache_hits(), 1);
827        assert!(stats.hit_rate() > 0.0);
828    }
829
830    #[test]
831    fn record_cache_hit_works_through_shared_ref() {
832        let mut cache = SessionCache::new();
833        cache.store("/x.rs", "hello world");
834        // &self path: a cache hit can be recorded without a write lock.
835        let shared: &SessionCache = &cache;
836        assert!(shared.record_cache_hit("/x.rs").is_some());
837        assert!(shared.record_cache_hit("/x.rs").is_some());
838        // store=1 + two hits => read_count 3, cache_hits 2.
839        assert_eq!(cache.get("/x.rs").unwrap().read_count(), 3);
840        assert_eq!(cache.get_stats().cache_hits(), 2);
841    }
842
843    #[test]
844    fn concurrent_cache_hits_are_lossless() {
845        use std::sync::Arc;
846        let mut cache = SessionCache::new();
847        cache.store("/a.rs", "a");
848        cache.store("/b.rs", "b");
849        // Shared (no RwLock): proves SessionCache is Sync and hit recording is
850        // lock-free and atomic — the whole point of the read-mostly refactor.
851        let cache = Arc::new(cache);
852        let threads = 8;
853        let iters = 1_000;
854        let handles: Vec<_> = (0..threads)
855            .map(|_| {
856                let c = Arc::clone(&cache);
857                std::thread::spawn(move || {
858                    for _ in 0..iters {
859                        c.record_cache_hit("/a.rs");
860                        c.record_cache_hit("/b.rs");
861                    }
862                })
863            })
864            .collect();
865        for h in handles {
866            h.join().unwrap();
867        }
868        let total = (threads * iters) as u64;
869        assert_eq!(cache.get_stats().cache_hits(), total * 2);
870        assert_eq!(cache.get("/a.rs").unwrap().read_count(), 1 + total as u32);
871        assert_eq!(cache.get("/b.rs").unwrap().read_count(), 1 + total as u32);
872    }
873
874    #[test]
875    fn md5_is_deterministic() {
876        let h1 = compute_md5("test content");
877        let h2 = compute_md5("test content");
878        assert_eq!(h1, h2);
879        assert_ne!(h1, compute_md5("different"));
880    }
881
882    #[test]
883    fn rrf_eviction_prefers_recent() {
884        let key_a = "a.rs".to_string();
885        let key_b = "b.rs".to_string();
886        // Construct entries first so the global instant base is initialized,
887        // then assign access times relative to a post-init reference.
888        let recent = CacheEntry::new("a", "h1".to_string(), 1, 10, "/a.rs".to_string(), None);
889        let old = CacheEntry::new("b", "h2".to_string(), 1, 10, "/b.rs".to_string(), None);
890        let t_old = Instant::now();
891        std::thread::sleep(std::time::Duration::from_millis(10));
892        let t_recent = Instant::now();
893        old.set_last_access(t_old);
894        recent.set_last_access(t_recent);
895        let now = Instant::now();
896        let entries: Vec<(&String, &CacheEntry)> = vec![(&key_a, &recent), (&key_b, &old)];
897        let scores = eviction_scores_rrf(&entries, now);
898        let score_a = scores.iter().find(|(p, _)| p == "a.rs").unwrap().1;
899        let score_b = scores.iter().find(|(p, _)| p == "b.rs").unwrap().1;
900        assert!(
901            score_a > score_b,
902            "recently accessed entries should score higher via RRF"
903        );
904    }
905
906    #[test]
907    fn rrf_eviction_prefers_frequent() {
908        let now = Instant::now();
909        let key_a = "a.rs".to_string();
910        let key_b = "b.rs".to_string();
911        let frequent = {
912            let e = CacheEntry::new("a", "h1".to_string(), 1, 10, "/a.rs".to_string(), None);
913            e.set_read_count(20);
914            e
915        };
916        let rare = CacheEntry::new("b", "h2".to_string(), 1, 10, "/b.rs".to_string(), None);
917        let entries: Vec<(&String, &CacheEntry)> = vec![(&key_a, &frequent), (&key_b, &rare)];
918        let scores = eviction_scores_rrf(&entries, now);
919        let score_a = scores.iter().find(|(p, _)| p == "a.rs").unwrap().1;
920        let score_b = scores.iter().find(|(p, _)| p == "b.rs").unwrap().1;
921        assert!(
922            score_a > score_b,
923            "frequently accessed entries should score higher via RRF"
924        );
925    }
926
927    #[test]
928    fn evict_if_needed_removes_lowest_score() {
929        std::env::set_var("LEAN_CTX_CACHE_MAX_TOKENS", "50");
930        let mut cache = SessionCache::new();
931        let big_content = "a]".repeat(30); // ~30 tokens
932        cache.store("/old.rs", &big_content);
933        // /old.rs now in cache with ~30 tokens
934
935        let new_content = "b ".repeat(30); // ~30 tokens incoming
936        cache.store("/new.rs", &new_content);
937        // should have evicted /old.rs to make room
938        // (total would be ~60 which exceeds 50)
939
940        // At least one should remain, total should be <= 50
941        assert!(
942            cache.total_cached_tokens() <= 60,
943            "eviction should have kicked in"
944        );
945        std::env::remove_var("LEAN_CTX_CACHE_MAX_TOKENS");
946    }
947
948    #[test]
949    fn stale_detection_flags_newer_file() {
950        let dir = tempfile::tempdir().unwrap();
951        let path = dir.path().join("stale.txt");
952        let p = path.to_string_lossy().to_string();
953
954        std::fs::write(&path, "one").unwrap();
955        let mut cache = SessionCache::new();
956        cache.store(&p, "one");
957
958        let entry = cache.get(&p).unwrap();
959        assert!(!is_cache_entry_stale(&p, entry.stored_mtime));
960
961        // Ensure mtime granularity differences don't make this flaky.
962        std::thread::sleep(Duration::from_secs(1));
963        std::fs::write(&path, "two").unwrap();
964
965        let entry = cache.get(&p).unwrap();
966        assert!(is_cache_entry_stale(&p, entry.stored_mtime));
967    }
968
969    #[test]
970    fn compressed_outputs_cached_and_retrieved() {
971        let mut cache = SessionCache::new();
972        cache.store("/test.rs", "fn main() {}");
973        cache.set_compressed("/test.rs", "map", "compressed map output".to_string());
974        assert_eq!(
975            cache.get_compressed("/test.rs", "map"),
976            Some(&"compressed map output".to_string())
977        );
978        assert_eq!(cache.get_compressed("/test.rs", "signatures"), None);
979    }
980
981    #[test]
982    fn compressed_outputs_cleared_on_content_change() {
983        let mut cache = SessionCache::new();
984        cache.store("/test.rs", "old content");
985        cache.set_compressed("/test.rs", "map", "old map".to_string());
986        assert!(cache.get_compressed("/test.rs", "map").is_some());
987
988        cache.store("/test.rs", "new content");
989        assert_eq!(cache.get_compressed("/test.rs", "map"), None);
990    }
991
992    #[test]
993    fn compressed_outputs_survive_same_content_store() {
994        let mut cache = SessionCache::new();
995        cache.store("/test.rs", "content");
996        cache.set_compressed("/test.rs", "map", "cached map".to_string());
997
998        let result = cache.store("/test.rs", "content");
999        assert!(result.was_hit);
1000        assert_eq!(
1001            cache.get_compressed("/test.rs", "map"),
1002            Some(&"cached map".to_string())
1003        );
1004    }
1005
1006    #[test]
1007    fn compressed_outputs_cleared_on_invalidate() {
1008        let mut cache = SessionCache::new();
1009        cache.store("/test.rs", "content");
1010        cache.set_compressed("/test.rs", "signatures", "cached sigs".to_string());
1011        cache.invalidate("/test.rs");
1012        assert_eq!(cache.get_compressed("/test.rs", "signatures"), None);
1013    }
1014
1015    #[test]
1016    fn compressed_outputs_cleared_on_clear() {
1017        let mut cache = SessionCache::new();
1018        cache.store("/a.rs", "a");
1019        cache.set_compressed("/a.rs", "map", "map_a".to_string());
1020        cache.clear();
1021        assert_eq!(cache.get_compressed("/a.rs", "map"), None);
1022    }
1023}