Skip to main content

lean_ctx/core/
cache.rs

1use md5::{Digest, Md5};
2use std::collections::HashMap;
3use std::time::{Instant, SystemTime};
4
5use super::tokens::count_tokens;
6
7fn normalize_key(path: &str) -> String {
8    crate::core::pathutil::normalize_tool_path(path)
9}
10
11fn max_cache_tokens() -> usize {
12    std::env::var("LEAN_CTX_CACHE_MAX_TOKENS")
13        .ok()
14        .and_then(|v| v.parse().ok())
15        .unwrap_or(500_000)
16}
17
18/// A cached file read: zstd-compressed content, hash, token count, and access metadata.
19#[derive(Clone, Debug)]
20pub struct CacheEntry {
21    compressed_content: Vec<u8>,
22    pub hash: String,
23    pub line_count: usize,
24    pub original_tokens: usize,
25    pub read_count: u32,
26    pub path: String,
27    pub last_access: Instant,
28    pub stored_mtime: Option<SystemTime>,
29    /// Mode-specific compressed outputs (e.g. "map", "signatures") cached to avoid re-parsing.
30    pub compressed_outputs: HashMap<String, String>,
31    /// Whether full (uncompressed) content was already delivered for this hash.
32    /// Prevents cache-stub loops when upgrading from compressed to full mode.
33    pub full_content_delivered: bool,
34    /// Last read mode used for this file (for auto-escalation on edit failure).
35    pub last_mode: String,
36}
37
38const ZSTD_LEVEL: i32 = 3;
39
40fn zstd_compress(data: &str) -> Vec<u8> {
41    zstd::encode_all(data.as_bytes(), ZSTD_LEVEL).unwrap_or_else(|_| data.as_bytes().to_vec())
42}
43
44fn zstd_decompress(data: &[u8]) -> String {
45    zstd::decode_all(data)
46        .ok()
47        .and_then(|v| String::from_utf8(v).ok())
48        .unwrap_or_default()
49}
50
51impl CacheEntry {
52    /// Creates a new entry with zstd-compressed content.
53    pub fn new(
54        content: &str,
55        hash: String,
56        line_count: usize,
57        original_tokens: usize,
58        path: String,
59        stored_mtime: Option<SystemTime>,
60    ) -> Self {
61        let compressed_content = zstd_compress(content);
62        Self {
63            compressed_content,
64            hash,
65            line_count,
66            original_tokens,
67            read_count: 1,
68            path,
69            last_access: Instant::now(),
70            stored_mtime,
71            compressed_outputs: HashMap::new(),
72            full_content_delivered: false,
73            last_mode: String::new(),
74        }
75    }
76
77    /// Decompresses and returns the full file content.
78    pub fn content(&self) -> String {
79        zstd_decompress(&self.compressed_content)
80    }
81
82    /// Replaces the stored content with new zstd-compressed data.
83    pub fn set_content(&mut self, content: &str) {
84        self.compressed_content = zstd_compress(content);
85    }
86
87    /// Approximate RAM usage of the compressed content in bytes.
88    pub fn compressed_size(&self) -> usize {
89        self.compressed_content.len()
90    }
91}
92
93/// Result of a cache store operation, indicating whether it was a hit or new entry.
94#[derive(Debug, Clone)]
95pub struct StoreResult {
96    pub line_count: usize,
97    pub original_tokens: usize,
98    pub read_count: u32,
99    pub was_hit: bool,
100    /// Whether full content was previously delivered for this cache entry.
101    pub full_content_delivered: bool,
102}
103
104impl CacheEntry {
105    /// Computes a legacy eviction score blending recency, frequency, and size.
106    pub fn eviction_score_legacy(&self, now: Instant) -> f64 {
107        let elapsed = now
108            .checked_duration_since(self.last_access)
109            .unwrap_or_default()
110            .as_secs_f64();
111        let recency = 1.0 / (1.0 + elapsed.sqrt());
112        let frequency = (self.read_count as f64 + 1.0).ln();
113        let size_value = (self.original_tokens as f64 + 1.0).ln();
114        recency * 0.4 + frequency * 0.3 + size_value * 0.3
115    }
116
117    pub fn get_compressed(&self, mode_key: &str) -> Option<&String> {
118        self.compressed_outputs.get(mode_key)
119    }
120
121    pub fn set_compressed(&mut self, mode_key: &str, output: String) {
122        const MAX_COMPRESSED_VARIANTS: usize = 3;
123        if self.compressed_outputs.len() >= MAX_COMPRESSED_VARIANTS
124            && !self.compressed_outputs.contains_key(mode_key)
125        {
126            if let Some(oldest_key) = self.compressed_outputs.keys().next().cloned() {
127                self.compressed_outputs.remove(&oldest_key);
128            }
129        }
130        self.compressed_outputs.insert(mode_key.to_string(), output);
131    }
132
133    pub fn mark_full_delivered(&mut self) {
134        self.full_content_delivered = true;
135    }
136}
137
138const RRF_K: f64 = 60.0;
139
140/// Compute Reciprocal Rank Fusion eviction scores for a batch of cache entries.
141/// Each signal (recency, frequency, size) produces an independent ranking.
142/// The final score is the sum of `1/(k + rank)` across all signals.
143/// Higher score = more valuable = keep longer.
144pub fn eviction_scores_rrf(entries: &[(&String, &CacheEntry)], now: Instant) -> Vec<(String, f64)> {
145    if entries.is_empty() {
146        return Vec::new();
147    }
148
149    let n = entries.len();
150
151    let mut recency_order: Vec<usize> = (0..n).collect();
152    recency_order.sort_by(|&a, &b| {
153        let elapsed_a = now
154            .checked_duration_since(entries[a].1.last_access)
155            .unwrap_or_default()
156            .as_secs_f64();
157        let elapsed_b = now
158            .checked_duration_since(entries[b].1.last_access)
159            .unwrap_or_default()
160            .as_secs_f64();
161        elapsed_a
162            .partial_cmp(&elapsed_b)
163            .unwrap_or(std::cmp::Ordering::Equal)
164    });
165
166    let mut frequency_order: Vec<usize> = (0..n).collect();
167    frequency_order.sort_by(|&a, &b| entries[b].1.read_count.cmp(&entries[a].1.read_count));
168
169    let mut size_order: Vec<usize> = (0..n).collect();
170    size_order.sort_by(|&a, &b| {
171        entries[b]
172            .1
173            .original_tokens
174            .cmp(&entries[a].1.original_tokens)
175    });
176
177    let mut recency_ranks = vec![0usize; n];
178    let mut frequency_ranks = vec![0usize; n];
179    let mut size_ranks = vec![0usize; n];
180
181    for (rank, &idx) in recency_order.iter().enumerate() {
182        recency_ranks[idx] = rank;
183    }
184    for (rank, &idx) in frequency_order.iter().enumerate() {
185        frequency_ranks[idx] = rank;
186    }
187    for (rank, &idx) in size_order.iter().enumerate() {
188        size_ranks[idx] = rank;
189    }
190
191    entries
192        .iter()
193        .enumerate()
194        .map(|(i, (path, _))| {
195            let score = 1.0 / (RRF_K + recency_ranks[i] as f64)
196                + 1.0 / (RRF_K + frequency_ranks[i] as f64)
197                + 1.0 / (RRF_K + size_ranks[i] as f64);
198            ((*path).clone(), score)
199        })
200        .collect()
201}
202
203/// Aggregated cache statistics: hits, reads, and token savings.
204#[derive(Debug)]
205pub struct CacheStats {
206    pub total_reads: u64,
207    pub cache_hits: u64,
208    pub total_original_tokens: u64,
209    pub total_sent_tokens: u64,
210    pub files_tracked: usize,
211}
212
213impl CacheStats {
214    /// Returns the cache hit rate as a percentage (0–100).
215    pub fn hit_rate(&self) -> f64 {
216        if self.total_reads == 0 {
217            return 0.0;
218        }
219        (self.cache_hits as f64 / self.total_reads as f64) * 100.0
220    }
221
222    /// Returns the total number of tokens saved by cache hits.
223    pub fn tokens_saved(&self) -> u64 {
224        self.total_original_tokens
225            .saturating_sub(self.total_sent_tokens)
226    }
227
228    /// Returns the savings as a percentage of total original tokens.
229    pub fn savings_percent(&self) -> f64 {
230        if self.total_original_tokens == 0 {
231            return 0.0;
232        }
233        (self.tokens_saved() as f64 / self.total_original_tokens as f64) * 100.0
234    }
235}
236
237/// A block shared across multiple files, identified by its canonical source.
238#[derive(Clone, Debug)]
239pub struct SharedBlock {
240    pub canonical_path: String,
241    pub canonical_ref: String,
242    pub start_line: usize,
243    pub end_line: usize,
244    pub content: String,
245}
246
247/// In-memory file cache with segmented LRU eviction (probationary vs protected),
248/// file references, and cross-file dedup.
249pub struct SessionCache {
250    entries: HashMap<String, CacheEntry>,
251    file_refs: HashMap<String, String>,
252    next_ref: usize,
253    stats: CacheStats,
254    shared_blocks: Vec<SharedBlock>,
255}
256
257impl Default for SessionCache {
258    fn default() -> Self {
259        Self::new()
260    }
261}
262
263impl SessionCache {
264    /// Creates an empty session cache with default stats.
265    pub fn new() -> Self {
266        Self {
267            entries: HashMap::new(),
268            file_refs: HashMap::new(),
269            next_ref: 1,
270            shared_blocks: Vec::new(),
271            stats: CacheStats {
272                total_reads: 0,
273                cache_hits: 0,
274                total_original_tokens: 0,
275                total_sent_tokens: 0,
276                files_tracked: 0,
277            },
278        }
279    }
280
281    /// Returns or assigns a short file reference label (F1, F2, ...) for the given path.
282    pub fn get_file_ref(&mut self, path: &str) -> String {
283        let key = normalize_key(path);
284        if let Some(r) = self.file_refs.get(&key) {
285            return r.clone();
286        }
287        let r = format!("F{}", self.next_ref);
288        self.next_ref += 1;
289        self.file_refs.insert(key, r.clone());
290        r
291    }
292
293    /// Returns the file reference label for a path without assigning a new one.
294    pub fn get_file_ref_readonly(&self, path: &str) -> Option<String> {
295        self.file_refs.get(&normalize_key(path)).cloned()
296    }
297
298    /// Looks up a cached entry by file path.
299    pub fn get(&self, path: &str) -> Option<&CacheEntry> {
300        self.entries.get(&normalize_key(path))
301    }
302
303    /// Mutable lookup of a cached entry by file path.
304    pub fn get_mut(&mut self, path: &str) -> Option<&mut CacheEntry> {
305        self.entries.get_mut(&normalize_key(path))
306    }
307
308    /// Retrieves the full (uncompressed) content for a file path, if cached.
309    /// Used by the CCR (Compress-Cache-Retrieve) mechanism.
310    pub fn get_full_content(&self, path: &str) -> Option<String> {
311        self.entries
312            .get(&normalize_key(path))
313            .map(CacheEntry::content)
314    }
315
316    /// Records a cache hit, updates access stats, and emits a cache-hit event.
317    pub fn record_cache_hit(&mut self, path: &str) -> Option<&CacheEntry> {
318        let key = normalize_key(path);
319        let ref_label = self
320            .file_refs
321            .get(&key)
322            .cloned()
323            .unwrap_or_else(|| "F?".to_string());
324        if let Some(entry) = self.entries.get_mut(&key) {
325            entry.read_count += 1;
326            entry.last_access = Instant::now();
327            self.stats.total_reads += 1;
328            self.stats.cache_hits += 1;
329            self.stats.total_original_tokens += entry.original_tokens as u64;
330            let hit_msg = format!(
331                "{ref_label} cached {}t {}L",
332                entry.read_count, entry.line_count
333            );
334            self.stats.total_sent_tokens += count_tokens(&hit_msg) as u64;
335            crate::core::events::emit_cache_hit(path, entry.original_tokens as u64);
336            Some(entry)
337        } else {
338            None
339        }
340    }
341
342    /// Stores file content in the cache; returns a hit if content hash matches.
343    pub fn store(&mut self, path: &str, content: &str) -> StoreResult {
344        let key = normalize_key(path);
345        let hash = compute_md5(content);
346        let line_count = content.lines().count();
347        let original_tokens = count_tokens(content);
348        let stored_mtime = std::fs::metadata(path).and_then(|m| m.modified()).ok();
349        let now = Instant::now();
350
351        self.stats.total_reads += 1;
352        self.stats.total_original_tokens += original_tokens as u64;
353
354        if let Some(existing) = self.entries.get_mut(&key) {
355            existing.last_access = now;
356            if stored_mtime.is_some() {
357                existing.stored_mtime = stored_mtime;
358            }
359            if existing.hash == hash {
360                existing.read_count += 1;
361                self.stats.cache_hits += 1;
362                let hit_msg = format!(
363                    "{} cached {}t {}L",
364                    self.file_refs.get(&key).unwrap_or(&"F?".to_string()),
365                    existing.read_count,
366                    existing.line_count,
367                );
368                self.stats.total_sent_tokens += count_tokens(&hit_msg) as u64;
369                return StoreResult {
370                    line_count: existing.line_count,
371                    original_tokens: existing.original_tokens,
372                    read_count: existing.read_count,
373                    was_hit: true,
374                    full_content_delivered: existing.full_content_delivered,
375                };
376            }
377            existing.compressed_outputs.clear();
378            existing.set_content(content);
379            existing.hash = hash;
380            existing.line_count = line_count;
381            existing.original_tokens = original_tokens;
382            existing.read_count += 1;
383            existing.full_content_delivered = false;
384            if stored_mtime.is_some() {
385                existing.stored_mtime = stored_mtime;
386            }
387            self.stats.total_sent_tokens += original_tokens as u64;
388            return StoreResult {
389                line_count,
390                original_tokens,
391                read_count: existing.read_count,
392                was_hit: false,
393                full_content_delivered: false,
394            };
395        }
396
397        self.evict_if_needed(original_tokens);
398        self.get_file_ref(&key);
399
400        let entry = CacheEntry::new(
401            content,
402            hash,
403            line_count,
404            original_tokens,
405            key.clone(),
406            stored_mtime,
407        );
408
409        self.entries.insert(key, entry);
410        self.stats.files_tracked += 1;
411        self.stats.total_sent_tokens += original_tokens as u64;
412        StoreResult {
413            line_count,
414            original_tokens,
415            read_count: 1,
416            was_hit: false,
417            full_content_delivered: false,
418        }
419    }
420
421    /// Returns the sum of original token counts across all cached entries.
422    pub fn total_cached_tokens(&self) -> usize {
423        self.entries.values().map(|e| e.original_tokens).sum()
424    }
425
426    /// Evict until cache fits within token budget using segmented LRU:
427    /// probationary (`read_count` ≤ 1, newly inserted / single-touch) first by oldest `last_access`,
428    /// then protected entries (`read_count` > 1) by oldest access.
429    pub fn evict_if_needed(&mut self, incoming_tokens: usize) {
430        let max_tokens = max_cache_tokens();
431        let current = self.total_cached_tokens();
432        if current + incoming_tokens <= max_tokens {
433            return;
434        }
435
436        let mut freed = 0usize;
437        let target = (current + incoming_tokens).saturating_sub(max_tokens);
438
439        let mut probationary: Vec<(String, Instant)> = self
440            .entries
441            .iter()
442            .filter(|(_, e)| e.read_count <= 1)
443            .map(|(p, e)| (p.clone(), e.last_access))
444            .collect();
445        probationary.sort_by_key(|(_, t)| *t);
446
447        let mut protected: Vec<(String, Instant)> = self
448            .entries
449            .iter()
450            .filter(|(_, e)| e.read_count > 1)
451            .map(|(p, e)| (p.clone(), e.last_access))
452            .collect();
453        protected.sort_by_key(|(_, t)| *t);
454
455        for (path, _) in probationary.into_iter().chain(protected) {
456            if freed >= target {
457                break;
458            }
459            if let Some(entry) = self.entries.remove(&path) {
460                freed += entry.original_tokens;
461                self.file_refs.remove(&path);
462            }
463        }
464    }
465
466    /// Returns all cached entries as (path, entry) pairs.
467    pub fn get_all_entries(&self) -> Vec<(&String, &CacheEntry)> {
468        self.entries.iter().collect()
469    }
470
471    /// Returns a reference to the aggregated cache statistics.
472    pub fn get_stats(&self) -> &CacheStats {
473        &self.stats
474    }
475
476    /// Returns the path-to-file-ref mapping (e.g. "/src/main.rs" → "F1").
477    pub fn file_ref_map(&self) -> &HashMap<String, String> {
478        &self.file_refs
479    }
480
481    /// Replaces the cross-file shared blocks used for deduplication.
482    pub fn set_shared_blocks(&mut self, blocks: Vec<SharedBlock>) {
483        self.shared_blocks = blocks;
484    }
485
486    /// Returns the current set of cross-file shared blocks.
487    pub fn get_shared_blocks(&self) -> &[SharedBlock] {
488        &self.shared_blocks
489    }
490
491    /// Replace shared blocks in content with cross-file references.
492    pub fn apply_dedup(&self, path: &str, content: &str) -> Option<String> {
493        if self.shared_blocks.is_empty() {
494            return None;
495        }
496        let refs: Vec<&SharedBlock> = self
497            .shared_blocks
498            .iter()
499            .filter(|b| b.canonical_path != path && content.contains(&b.content))
500            .collect();
501        if refs.is_empty() {
502            return None;
503        }
504        let mut result = content.to_string();
505        for block in refs {
506            result = result.replacen(
507                &block.content,
508                &format!(
509                    "[= {}:{}-{}]",
510                    block.canonical_ref, block.start_line, block.end_line
511                ),
512                1,
513            );
514        }
515        Some(result)
516    }
517
518    /// Removes a file from the cache, forcing a fresh read on next access.
519    pub fn invalidate(&mut self, path: &str) -> bool {
520        self.entries.remove(&normalize_key(path)).is_some()
521    }
522
523    /// Returns a cached compressed output for a given file and mode key.
524    pub fn get_compressed(&self, path: &str, mode_key: &str) -> Option<&String> {
525        self.entries
526            .get(&normalize_key(path))?
527            .get_compressed(mode_key)
528    }
529
530    /// Marks that full (uncompressed) content was delivered for this file.
531    pub fn mark_full_delivered(&mut self, path: &str) {
532        if let Some(entry) = self.entries.get_mut(&normalize_key(path)) {
533            entry.mark_full_delivered();
534        }
535    }
536
537    /// Stores a compressed output for a given file and mode key.
538    pub fn set_compressed(&mut self, path: &str, mode_key: &str, output: String) {
539        if let Some(entry) = self.entries.get_mut(&normalize_key(path)) {
540            entry.set_compressed(mode_key, output);
541        }
542    }
543
544    /// Clears all cached entries, file refs, and resets stats. Returns the number of entries removed.
545    pub fn clear(&mut self) -> usize {
546        let count = self.entries.len();
547        self.entries.clear();
548        self.file_refs.clear();
549        self.shared_blocks.clear();
550        self.next_ref = 1;
551        self.stats = CacheStats {
552            total_reads: 0,
553            cache_hits: 0,
554            total_original_tokens: 0,
555            total_sent_tokens: 0,
556            files_tracked: 0,
557        };
558        count
559    }
560}
561
562pub fn file_mtime(path: &str) -> Option<SystemTime> {
563    std::fs::metadata(path).and_then(|m| m.modified()).ok()
564}
565
566pub fn is_cache_entry_stale(path: &str, cached_mtime: Option<SystemTime>) -> bool {
567    let current = file_mtime(path);
568    match (cached_mtime, current) {
569        (_, None) => false,
570        (None, Some(_)) => true,
571        (Some(cached), Some(current)) => current > cached,
572    }
573}
574
575fn compute_md5(content: &str) -> String {
576    let mut hasher = Md5::new();
577    hasher.update(content.as_bytes());
578    format!("{:x}", hasher.finalize())
579}
580
581#[cfg(test)]
582mod tests {
583    use super::*;
584    use std::time::Duration;
585
586    #[test]
587    fn cache_stores_and_retrieves() {
588        let mut cache = SessionCache::new();
589        let result = cache.store("/test/file.rs", "fn main() {}");
590        assert!(!result.was_hit);
591        assert_eq!(result.line_count, 1);
592        assert!(cache.get("/test/file.rs").is_some());
593    }
594
595    #[test]
596    fn cache_hit_on_same_content() {
597        let mut cache = SessionCache::new();
598        cache.store("/test/file.rs", "content");
599        let result = cache.store("/test/file.rs", "content");
600        assert!(result.was_hit, "same content should be a cache hit");
601    }
602
603    #[test]
604    fn cache_miss_on_changed_content() {
605        let mut cache = SessionCache::new();
606        cache.store("/test/file.rs", "old content");
607        let result = cache.store("/test/file.rs", "new content");
608        assert!(!result.was_hit, "changed content should not be a cache hit");
609    }
610
611    #[test]
612    fn file_refs_are_sequential() {
613        let mut cache = SessionCache::new();
614        assert_eq!(cache.get_file_ref("/a.rs"), "F1");
615        assert_eq!(cache.get_file_ref("/b.rs"), "F2");
616        assert_eq!(cache.get_file_ref("/a.rs"), "F1"); // stable
617    }
618
619    #[test]
620    fn cache_clear_resets_everything() {
621        let mut cache = SessionCache::new();
622        cache.store("/a.rs", "a");
623        cache.store("/b.rs", "b");
624        let count = cache.clear();
625        assert_eq!(count, 2);
626        assert!(cache.get("/a.rs").is_none());
627        assert_eq!(cache.get_file_ref("/c.rs"), "F1"); // refs reset
628    }
629
630    #[test]
631    fn cache_invalidate_removes_entry() {
632        let mut cache = SessionCache::new();
633        cache.store("/test.rs", "test");
634        assert!(cache.invalidate("/test.rs"));
635        assert!(!cache.invalidate("/nonexistent.rs"));
636    }
637
638    #[test]
639    fn cache_stats_track_correctly() {
640        let mut cache = SessionCache::new();
641        cache.store("/a.rs", "hello");
642        cache.store("/a.rs", "hello"); // hit
643        let stats = cache.get_stats();
644        assert_eq!(stats.total_reads, 2);
645        assert_eq!(stats.cache_hits, 1);
646        assert!(stats.hit_rate() > 0.0);
647    }
648
649    #[test]
650    fn md5_is_deterministic() {
651        let h1 = compute_md5("test content");
652        let h2 = compute_md5("test content");
653        assert_eq!(h1, h2);
654        assert_ne!(h1, compute_md5("different"));
655    }
656
657    #[test]
658    fn rrf_eviction_prefers_recent() {
659        let base = Instant::now();
660        std::thread::sleep(std::time::Duration::from_millis(5));
661        let now = Instant::now();
662        let key_a = "a.rs".to_string();
663        let key_b = "b.rs".to_string();
664        let recent = CacheEntry::new("a", "h1".to_string(), 1, 10, "/a.rs".to_string(), None);
665        let old = {
666            let mut e = CacheEntry::new("b", "h2".to_string(), 1, 10, "/b.rs".to_string(), None);
667            e.last_access = base;
668            e
669        };
670        let entries: Vec<(&String, &CacheEntry)> = vec![(&key_a, &recent), (&key_b, &old)];
671        let scores = eviction_scores_rrf(&entries, now);
672        let score_a = scores.iter().find(|(p, _)| p == "a.rs").unwrap().1;
673        let score_b = scores.iter().find(|(p, _)| p == "b.rs").unwrap().1;
674        assert!(
675            score_a > score_b,
676            "recently accessed entries should score higher via RRF"
677        );
678    }
679
680    #[test]
681    fn rrf_eviction_prefers_frequent() {
682        let now = Instant::now();
683        let key_a = "a.rs".to_string();
684        let key_b = "b.rs".to_string();
685        let frequent = {
686            let mut e = CacheEntry::new("a", "h1".to_string(), 1, 10, "/a.rs".to_string(), None);
687            e.read_count = 20;
688            e
689        };
690        let rare = CacheEntry::new("b", "h2".to_string(), 1, 10, "/b.rs".to_string(), None);
691        let entries: Vec<(&String, &CacheEntry)> = vec![(&key_a, &frequent), (&key_b, &rare)];
692        let scores = eviction_scores_rrf(&entries, now);
693        let score_a = scores.iter().find(|(p, _)| p == "a.rs").unwrap().1;
694        let score_b = scores.iter().find(|(p, _)| p == "b.rs").unwrap().1;
695        assert!(
696            score_a > score_b,
697            "frequently accessed entries should score higher via RRF"
698        );
699    }
700
701    #[test]
702    fn evict_if_needed_removes_lowest_score() {
703        std::env::set_var("LEAN_CTX_CACHE_MAX_TOKENS", "50");
704        let mut cache = SessionCache::new();
705        let big_content = "a]".repeat(30); // ~30 tokens
706        cache.store("/old.rs", &big_content);
707        // /old.rs now in cache with ~30 tokens
708
709        let new_content = "b ".repeat(30); // ~30 tokens incoming
710        cache.store("/new.rs", &new_content);
711        // should have evicted /old.rs to make room
712        // (total would be ~60 which exceeds 50)
713
714        // At least one should remain, total should be <= 50
715        assert!(
716            cache.total_cached_tokens() <= 60,
717            "eviction should have kicked in"
718        );
719        std::env::remove_var("LEAN_CTX_CACHE_MAX_TOKENS");
720    }
721
722    #[test]
723    fn stale_detection_flags_newer_file() {
724        let dir = tempfile::tempdir().unwrap();
725        let path = dir.path().join("stale.txt");
726        let p = path.to_string_lossy().to_string();
727
728        std::fs::write(&path, "one").unwrap();
729        let mut cache = SessionCache::new();
730        cache.store(&p, "one");
731
732        let entry = cache.get(&p).unwrap();
733        assert!(!is_cache_entry_stale(&p, entry.stored_mtime));
734
735        // Ensure mtime granularity differences don't make this flaky.
736        std::thread::sleep(Duration::from_secs(1));
737        std::fs::write(&path, "two").unwrap();
738
739        let entry = cache.get(&p).unwrap();
740        assert!(is_cache_entry_stale(&p, entry.stored_mtime));
741    }
742
743    #[test]
744    fn compressed_outputs_cached_and_retrieved() {
745        let mut cache = SessionCache::new();
746        cache.store("/test.rs", "fn main() {}");
747        cache.set_compressed("/test.rs", "map", "compressed map output".to_string());
748        assert_eq!(
749            cache.get_compressed("/test.rs", "map"),
750            Some(&"compressed map output".to_string())
751        );
752        assert_eq!(cache.get_compressed("/test.rs", "signatures"), None);
753    }
754
755    #[test]
756    fn compressed_outputs_cleared_on_content_change() {
757        let mut cache = SessionCache::new();
758        cache.store("/test.rs", "old content");
759        cache.set_compressed("/test.rs", "map", "old map".to_string());
760        assert!(cache.get_compressed("/test.rs", "map").is_some());
761
762        cache.store("/test.rs", "new content");
763        assert_eq!(cache.get_compressed("/test.rs", "map"), None);
764    }
765
766    #[test]
767    fn compressed_outputs_survive_same_content_store() {
768        let mut cache = SessionCache::new();
769        cache.store("/test.rs", "content");
770        cache.set_compressed("/test.rs", "map", "cached map".to_string());
771
772        let result = cache.store("/test.rs", "content");
773        assert!(result.was_hit);
774        assert_eq!(
775            cache.get_compressed("/test.rs", "map"),
776            Some(&"cached map".to_string())
777        );
778    }
779
780    #[test]
781    fn compressed_outputs_cleared_on_invalidate() {
782        let mut cache = SessionCache::new();
783        cache.store("/test.rs", "content");
784        cache.set_compressed("/test.rs", "signatures", "cached sigs".to_string());
785        cache.invalidate("/test.rs");
786        assert_eq!(cache.get_compressed("/test.rs", "signatures"), None);
787    }
788
789    #[test]
790    fn compressed_outputs_cleared_on_clear() {
791        let mut cache = SessionCache::new();
792        cache.store("/a.rs", "a");
793        cache.set_compressed("/a.rs", "map", "map_a".to_string());
794        cache.clear();
795        assert_eq!(cache.get_compressed("/a.rs", "map"), None);
796    }
797}