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
9fn 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#[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 pub compressed_outputs: HashMap<String, String>,
53 pub full_content_delivered: bool,
56 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 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 pub fn read_count(&self) -> u32 {
100 self.read_count.load(Ordering::Relaxed)
101 }
102
103 pub fn bump_read_count(&self) -> u32 {
105 self.read_count.fetch_add(1, Ordering::Relaxed) + 1
106 }
107
108 pub fn set_read_count(&self, n: u32) {
110 self.read_count.store(n, Ordering::Relaxed);
111 }
112
113 pub fn last_access(&self) -> Instant {
115 decode_instant(self.last_access.load(Ordering::Relaxed))
116 }
117
118 pub fn touch(&self) {
120 self.last_access
121 .store(encode_instant(Instant::now()), Ordering::Relaxed);
122 }
123
124 pub fn set_last_access(&self, when: Instant) {
126 self.last_access
127 .store(encode_instant(when), Ordering::Relaxed);
128 }
129
130 pub fn content(&self) -> Option<String> {
132 zstd_decompress(&self.compressed_content)
133 }
134
135 pub fn set_content(&mut self, content: &str) {
137 self.compressed_content = zstd_compress(content);
138 }
139
140 pub fn compressed_size(&self) -> usize {
142 self.compressed_content.len()
143 }
144}
145
146#[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 pub full_content_delivered: bool,
155}
156
157impl CacheEntry {
158 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
192pub 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#[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 pub fn total_reads(&self) -> u64 {
271 self.total_reads.load(Ordering::Relaxed)
272 }
273
274 pub fn cache_hits(&self) -> u64 {
276 self.cache_hits.load(Ordering::Relaxed)
277 }
278
279 pub fn total_original_tokens(&self) -> u64 {
281 self.total_original_tokens.load(Ordering::Relaxed)
282 }
283
284 pub fn total_sent_tokens(&self) -> u64 {
286 self.total_sent_tokens.load(Ordering::Relaxed)
287 }
288
289 pub fn files_tracked(&self) -> u64 {
291 self.files_tracked.load(Ordering::Relaxed)
292 }
293
294 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 pub fn tokens_saved(&self) -> u64 {
305 self.total_original_tokens()
306 .saturating_sub(self.total_sent_tokens())
307 }
308
309 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#[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
329pub 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 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 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 pub fn get_file_ref_readonly(&self, path: &str) -> Option<String> {
371 self.file_refs.get(&normalize_key(path)).cloned()
372 }
373
374 pub fn get(&self, path: &str) -> Option<&CacheEntry> {
376 self.entries.get(&normalize_key(path))
377 }
378
379 pub fn get_mut(&mut self, path: &str) -> Option<&mut CacheEntry> {
381 self.entries.get_mut(&normalize_key(path))
382 }
383
384 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 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 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 pub fn total_cached_tokens(&self) -> usize {
508 self.entries.values().map(|e| e.original_tokens).sum()
509 }
510
511 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 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 pub fn get_all_entries(&self) -> Vec<(&String, &CacheEntry)> {
542 self.entries.iter().collect()
543 }
544
545 pub fn get_stats(&self) -> &CacheStats {
547 &self.stats
548 }
549
550 pub fn file_ref_map(&self) -> &HashMap<String, String> {
552 &self.file_refs
553 }
554
555 pub fn set_shared_blocks(&mut self, blocks: Vec<SharedBlock>) {
557 self.shared_blocks = blocks;
558 }
559
560 pub fn get_shared_blocks(&self) -> &[SharedBlock] {
562 &self.shared_blocks
563 }
564
565 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 pub fn invalidate(&mut self, path: &str) -> bool {
594 self.entries.remove(&normalize_key(path)).is_some()
595 }
596
597 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 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 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 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 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 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 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 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 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 })
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 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 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 (None, None) => false,
749 (Some(_), None) | (None, Some(_)) => true,
751 (Some(cached), Some(current)) => current != cached,
754 }
755}
756
757const 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
765pub 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 return true;
793 };
794 if meta.len() > VERIFY_HASH_CAP_BYTES {
795 return false;
796 }
797 match std::fs::read(path) {
798 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"); }
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"); }
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"); 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 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 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 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 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); cache.store("/old.rs", &big_content);
981 let new_content = "b ".repeat(30); cache.store("/new.rs", &new_content);
985 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 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 #[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 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 #[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 assert!(!is_cache_entry_stale_verified(&p, mtime, &hash));
1062
1063 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}