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
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 {
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
193pub 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#[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 pub fn total_reads(&self) -> u64 {
272 self.total_reads.load(Ordering::Relaxed)
273 }
274
275 pub fn cache_hits(&self) -> u64 {
277 self.cache_hits.load(Ordering::Relaxed)
278 }
279
280 pub fn total_original_tokens(&self) -> u64 {
282 self.total_original_tokens.load(Ordering::Relaxed)
283 }
284
285 pub fn total_sent_tokens(&self) -> u64 {
287 self.total_sent_tokens.load(Ordering::Relaxed)
288 }
289
290 pub fn files_tracked(&self) -> u64 {
292 self.files_tracked.load(Ordering::Relaxed)
293 }
294
295 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 pub fn tokens_saved(&self) -> u64 {
306 self.total_original_tokens()
307 .saturating_sub(self.total_sent_tokens())
308 }
309
310 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#[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
330pub 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 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 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 pub fn get_file_ref_readonly(&self, path: &str) -> Option<String> {
372 self.file_refs.get(&normalize_key(path)).cloned()
373 }
374
375 pub fn get(&self, path: &str) -> Option<&CacheEntry> {
377 self.entries.get(&normalize_key(path))
378 }
379
380 pub fn get_mut(&mut self, path: &str) -> Option<&mut CacheEntry> {
382 self.entries.get_mut(&normalize_key(path))
383 }
384
385 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 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 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 pub fn total_cached_tokens(&self) -> usize {
509 self.entries.values().map(|e| e.original_tokens).sum()
510 }
511
512 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 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 pub fn get_all_entries(&self) -> Vec<(&String, &CacheEntry)> {
543 self.entries.iter().collect()
544 }
545
546 pub fn get_stats(&self) -> &CacheStats {
548 &self.stats
549 }
550
551 pub fn file_ref_map(&self) -> &HashMap<String, String> {
553 &self.file_refs
554 }
555
556 pub fn set_shared_blocks(&mut self, blocks: Vec<SharedBlock>) {
558 self.shared_blocks = blocks;
559 }
560
561 pub fn get_shared_blocks(&self) -> &[SharedBlock] {
563 &self.shared_blocks
564 }
565
566 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 pub fn invalidate(&mut self, path: &str) -> bool {
595 self.entries.remove(&normalize_key(path)).is_some()
596 }
597
598 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 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 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 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 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 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 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 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 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 })
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 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 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 (None, None) => false,
750 (Some(_), None) | (None, Some(_)) => true,
752 (Some(cached), Some(current)) => current != cached,
755 }
756}
757
758const VERIFY_HASH_CAP_BYTES: u64 = 8 * 1024 * 1024;
761
762fn cache_verify_enabled() -> bool {
763 std::env::var("LEAN_CTX_CACHE_VERIFY").map_or(true, |v| v != "0")
764}
765
766pub fn is_cache_entry_stale_verified(
781 path: &str,
782 cached_mtime: Option<SystemTime>,
783 cached_hash: &str,
784) -> bool {
785 if is_cache_entry_stale(path, cached_mtime) {
786 return true;
787 }
788 if cached_hash.is_empty() || !cache_verify_enabled() {
789 return false;
790 }
791 let Ok(meta) = std::fs::metadata(path) else {
792 return true;
794 };
795 if meta.len() > VERIFY_HASH_CAP_BYTES {
796 return false;
797 }
798 match std::fs::read(path) {
799 Ok(bytes) => compute_md5(&String::from_utf8_lossy(&bytes)) != cached_hash,
801 Err(_) => true,
802 }
803}
804
805fn compute_md5(content: &str) -> String {
806 let mut hasher = Md5::new();
807 hasher.update(content.as_bytes());
808 format!("{:x}", hasher.finalize())
809}
810
811#[cfg(test)]
812mod tests {
813 use super::*;
814 use std::time::Duration;
815
816 #[test]
817 fn cache_stores_and_retrieves() {
818 let mut cache = SessionCache::new();
819 let result = cache.store("/test/file.rs", "fn main() {}");
820 assert!(!result.was_hit);
821 assert_eq!(result.line_count, 1);
822 assert!(cache.get("/test/file.rs").is_some());
823 }
824
825 #[test]
826 fn cache_hit_on_same_content() {
827 let mut cache = SessionCache::new();
828 cache.store("/test/file.rs", "content");
829 let result = cache.store("/test/file.rs", "content");
830 assert!(result.was_hit, "same content should be a cache hit");
831 }
832
833 #[test]
834 fn cache_miss_on_changed_content() {
835 let mut cache = SessionCache::new();
836 cache.store("/test/file.rs", "old content");
837 let result = cache.store("/test/file.rs", "new content");
838 assert!(!result.was_hit, "changed content should not be a cache hit");
839 }
840
841 #[test]
842 fn file_refs_are_sequential() {
843 let mut cache = SessionCache::new();
844 assert_eq!(cache.get_file_ref("/a.rs"), "F1");
845 assert_eq!(cache.get_file_ref("/b.rs"), "F2");
846 assert_eq!(cache.get_file_ref("/a.rs"), "F1"); }
848
849 #[test]
850 fn cache_clear_resets_everything() {
851 let mut cache = SessionCache::new();
852 cache.store("/a.rs", "a");
853 cache.store("/b.rs", "b");
854 let count = cache.clear();
855 assert_eq!(count, 2);
856 assert!(cache.get("/a.rs").is_none());
857 assert_eq!(cache.get_file_ref("/c.rs"), "F1"); }
859
860 #[test]
861 fn cache_invalidate_removes_entry() {
862 let mut cache = SessionCache::new();
863 cache.store("/test.rs", "test");
864 assert!(cache.invalidate("/test.rs"));
865 assert!(!cache.invalidate("/nonexistent.rs"));
866 }
867
868 #[test]
869 fn cache_stats_track_correctly() {
870 let mut cache = SessionCache::new();
871 cache.store("/a.rs", "hello");
872 cache.store("/a.rs", "hello"); let stats = cache.get_stats();
874 assert_eq!(stats.total_reads(), 2);
875 assert_eq!(stats.cache_hits(), 1);
876 assert!(stats.hit_rate() > 0.0);
877 }
878
879 #[test]
880 fn record_cache_hit_works_through_shared_ref() {
881 let mut cache = SessionCache::new();
882 cache.store("/x.rs", "hello world");
883 let shared: &SessionCache = &cache;
885 assert!(shared.record_cache_hit("/x.rs").is_some());
886 assert!(shared.record_cache_hit("/x.rs").is_some());
887 assert_eq!(cache.get("/x.rs").unwrap().read_count(), 3);
889 assert_eq!(cache.get_stats().cache_hits(), 2);
890 }
891
892 #[test]
893 fn concurrent_cache_hits_are_lossless() {
894 use std::sync::Arc;
895 let mut cache = SessionCache::new();
896 cache.store("/a.rs", "a");
897 cache.store("/b.rs", "b");
898 let cache = Arc::new(cache);
901 let threads = 8;
902 let iters = 1_000;
903 let handles: Vec<_> = (0..threads)
904 .map(|_| {
905 let c = Arc::clone(&cache);
906 std::thread::spawn(move || {
907 for _ in 0..iters {
908 c.record_cache_hit("/a.rs");
909 c.record_cache_hit("/b.rs");
910 }
911 })
912 })
913 .collect();
914 for h in handles {
915 h.join().unwrap();
916 }
917 let total = (threads * iters) as u64;
918 assert_eq!(cache.get_stats().cache_hits(), total * 2);
919 assert_eq!(cache.get("/a.rs").unwrap().read_count(), 1 + total as u32);
920 assert_eq!(cache.get("/b.rs").unwrap().read_count(), 1 + total as u32);
921 }
922
923 #[test]
924 fn md5_is_deterministic() {
925 let h1 = compute_md5("test content");
926 let h2 = compute_md5("test content");
927 assert_eq!(h1, h2);
928 assert_ne!(h1, compute_md5("different"));
929 }
930
931 #[test]
932 fn rrf_eviction_prefers_recent() {
933 let key_a = "a.rs".to_string();
934 let key_b = "b.rs".to_string();
935 let recent = CacheEntry::new("a", "h1".to_string(), 1, 10, "/a.rs".to_string(), None);
938 let old = CacheEntry::new("b", "h2".to_string(), 1, 10, "/b.rs".to_string(), None);
939 let t_old = Instant::now();
940 std::thread::sleep(std::time::Duration::from_millis(10));
941 let t_recent = Instant::now();
942 old.set_last_access(t_old);
943 recent.set_last_access(t_recent);
944 let now = Instant::now();
945 let entries: Vec<(&String, &CacheEntry)> = vec![(&key_a, &recent), (&key_b, &old)];
946 let scores = eviction_scores_rrf(&entries, now);
947 let score_a = scores.iter().find(|(p, _)| p == "a.rs").unwrap().1;
948 let score_b = scores.iter().find(|(p, _)| p == "b.rs").unwrap().1;
949 assert!(
950 score_a > score_b,
951 "recently accessed entries should score higher via RRF"
952 );
953 }
954
955 #[test]
956 fn rrf_eviction_prefers_frequent() {
957 let now = Instant::now();
958 let key_a = "a.rs".to_string();
959 let key_b = "b.rs".to_string();
960 let frequent = {
961 let e = CacheEntry::new("a", "h1".to_string(), 1, 10, "/a.rs".to_string(), None);
962 e.set_read_count(20);
963 e
964 };
965 let rare = CacheEntry::new("b", "h2".to_string(), 1, 10, "/b.rs".to_string(), None);
966 let entries: Vec<(&String, &CacheEntry)> = vec![(&key_a, &frequent), (&key_b, &rare)];
967 let scores = eviction_scores_rrf(&entries, now);
968 let score_a = scores.iter().find(|(p, _)| p == "a.rs").unwrap().1;
969 let score_b = scores.iter().find(|(p, _)| p == "b.rs").unwrap().1;
970 assert!(
971 score_a > score_b,
972 "frequently accessed entries should score higher via RRF"
973 );
974 }
975
976 #[test]
977 fn evict_if_needed_removes_lowest_score() {
978 std::env::set_var("LEAN_CTX_CACHE_MAX_TOKENS", "50");
979 let mut cache = SessionCache::new();
980 let big_content = "a]".repeat(30); cache.store("/old.rs", &big_content);
982 let new_content = "b ".repeat(30); cache.store("/new.rs", &new_content);
986 assert!(
991 cache.total_cached_tokens() <= 60,
992 "eviction should have kicked in"
993 );
994 std::env::remove_var("LEAN_CTX_CACHE_MAX_TOKENS");
995 }
996
997 #[test]
998 fn stale_detection_flags_newer_file() {
999 let dir = tempfile::tempdir().unwrap();
1000 let path = dir.path().join("stale.txt");
1001 let p = path.to_string_lossy().to_string();
1002
1003 std::fs::write(&path, "one").unwrap();
1004 let mut cache = SessionCache::new();
1005 cache.store(&p, "one");
1006
1007 let entry = cache.get(&p).unwrap();
1008 assert!(!is_cache_entry_stale(&p, entry.stored_mtime));
1009
1010 std::thread::sleep(Duration::from_secs(1));
1012 std::fs::write(&path, "two").unwrap();
1013
1014 let entry = cache.get(&p).unwrap();
1015 assert!(is_cache_entry_stale(&p, entry.stored_mtime));
1016 }
1017
1018 #[test]
1020 fn stale_detection_flags_backward_mtime() {
1021 let dir = tempfile::tempdir().unwrap();
1022 let path = dir.path().join("backward.txt");
1023 let p = path.to_string_lossy().to_string();
1024
1025 std::fs::write(&path, "one").unwrap();
1026 let mut cache = SessionCache::new();
1027 cache.store(&p, "one");
1028 let entry_mtime = cache.get(&p).unwrap().stored_mtime;
1029 assert!(!is_cache_entry_stale(&p, entry_mtime));
1030
1031 std::fs::write(&path, "zero").unwrap();
1033 let f = std::fs::OpenOptions::new().write(true).open(&path).unwrap();
1034 f.set_modified(SystemTime::now() - Duration::from_hours(1))
1035 .unwrap();
1036 drop(f);
1037
1038 assert!(
1039 is_cache_entry_stale(&p, entry_mtime),
1040 "older mtime must read as stale"
1041 );
1042 }
1043
1044 #[test]
1047 fn verified_staleness_catches_same_mtime_content_change() {
1048 let dir = tempfile::tempdir().unwrap();
1049 let path = dir.path().join("sneaky.txt");
1050 let p = path.to_string_lossy().to_string();
1051
1052 std::fs::write(&path, "one").unwrap();
1053 let original_mtime = std::fs::metadata(&path).unwrap().modified().unwrap();
1054 let mut cache = SessionCache::new();
1055 cache.store(&p, "one");
1056 let (mtime, hash) = {
1057 let e = cache.get(&p).unwrap();
1058 (e.stored_mtime, e.hash.clone())
1059 };
1060
1061 assert!(!is_cache_entry_stale_verified(&p, mtime, &hash));
1063
1064 std::fs::write(&path, "two").unwrap();
1066 let f = std::fs::OpenOptions::new().write(true).open(&path).unwrap();
1067 f.set_modified(original_mtime).unwrap();
1068 drop(f);
1069
1070 assert!(
1071 !is_cache_entry_stale(&p, mtime),
1072 "test premise: the mtime check alone is fooled"
1073 );
1074 assert!(
1075 is_cache_entry_stale_verified(&p, mtime, &hash),
1076 "hash verification must catch the change"
1077 );
1078 }
1079
1080 #[test]
1081 fn verified_staleness_flags_unreadable_file() {
1082 let mut cache = SessionCache::new();
1083 cache.store("/nonexistent/file.rs", "content");
1084 let (mtime, hash) = {
1085 let e = cache.get("/nonexistent/file.rs").unwrap();
1086 (e.stored_mtime, e.hash.clone())
1087 };
1088 assert!(is_cache_entry_stale_verified(
1089 "/nonexistent/file.rs",
1090 mtime,
1091 &hash
1092 ));
1093 }
1094
1095 #[test]
1096 fn compressed_outputs_cached_and_retrieved() {
1097 let mut cache = SessionCache::new();
1098 cache.store("/test.rs", "fn main() {}");
1099 cache.set_compressed("/test.rs", "map", "compressed map output".to_string());
1100 assert_eq!(
1101 cache.get_compressed("/test.rs", "map"),
1102 Some(&"compressed map output".to_string())
1103 );
1104 assert_eq!(cache.get_compressed("/test.rs", "signatures"), None);
1105 }
1106
1107 #[test]
1108 fn compressed_outputs_cleared_on_content_change() {
1109 let mut cache = SessionCache::new();
1110 cache.store("/test.rs", "old content");
1111 cache.set_compressed("/test.rs", "map", "old map".to_string());
1112 assert!(cache.get_compressed("/test.rs", "map").is_some());
1113
1114 cache.store("/test.rs", "new content");
1115 assert_eq!(cache.get_compressed("/test.rs", "map"), None);
1116 }
1117
1118 #[test]
1119 fn compressed_outputs_survive_same_content_store() {
1120 let mut cache = SessionCache::new();
1121 cache.store("/test.rs", "content");
1122 cache.set_compressed("/test.rs", "map", "cached map".to_string());
1123
1124 let result = cache.store("/test.rs", "content");
1125 assert!(result.was_hit);
1126 assert_eq!(
1127 cache.get_compressed("/test.rs", "map"),
1128 Some(&"cached map".to_string())
1129 );
1130 }
1131
1132 #[test]
1133 fn compressed_outputs_cleared_on_invalidate() {
1134 let mut cache = SessionCache::new();
1135 cache.store("/test.rs", "content");
1136 cache.set_compressed("/test.rs", "signatures", "cached sigs".to_string());
1137 cache.invalidate("/test.rs");
1138 assert_eq!(cache.get_compressed("/test.rs", "signatures"), None);
1139 }
1140
1141 #[test]
1142 fn compressed_outputs_cleared_on_clear() {
1143 let mut cache = SessionCache::new();
1144 cache.store("/a.rs", "a");
1145 cache.set_compressed("/a.rs", "map", "map_a".to_string());
1146 cache.clear();
1147 assert_eq!(cache.get_compressed("/a.rs", "map"), None);
1148 }
1149}