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
29pub(crate) const DEFAULT_CACHE_MAX_TOKENS: usize = 500_000;
31
32fn resolve_cache_max_tokens(env: Option<&str>, configured: usize) -> usize {
39 if let Some(raw) = env
40 && let Ok(n) = raw.trim().parse::<usize>()
41 && n > 0
42 {
43 return n;
44 }
45 if configured > 0 {
46 configured
47 } else {
48 DEFAULT_CACHE_MAX_TOKENS
49 }
50}
51
52pub(crate) fn max_cache_tokens() -> usize {
57 resolve_cache_max_tokens(
58 std::env::var("LEAN_CTX_CACHE_MAX_TOKENS").ok().as_deref(),
59 crate::core::config::Config::load().cache_max_tokens,
60 )
61}
62
63#[derive(Debug)]
69pub struct CacheEntry {
70 compressed_content: Vec<u8>,
71 pub hash: String,
72 pub line_count: usize,
73 pub original_tokens: usize,
74 read_count: AtomicU32,
75 pub path: String,
76 last_access: AtomicU64,
77 pub stored_mtime: Option<SystemTime>,
78 pub compressed_outputs: HashMap<String, String>,
80 pub full_content_delivered: bool,
83 pub delivered_conversation: Option<String>,
88 pub last_mode: String,
90}
91
92const ZSTD_LEVEL: i32 = 3;
93
94fn zstd_compress(data: &str) -> Vec<u8> {
95 zstd::encode_all(data.as_bytes(), ZSTD_LEVEL).unwrap_or_else(|_| data.as_bytes().to_vec())
96}
97
98fn zstd_decompress(data: &[u8]) -> Option<String> {
99 zstd::decode_all(data)
100 .ok()
101 .and_then(|v| String::from_utf8(v).ok())
102}
103
104impl CacheEntry {
105 pub fn new(
107 content: &str,
108 hash: String,
109 line_count: usize,
110 original_tokens: usize,
111 path: String,
112 stored_mtime: Option<SystemTime>,
113 ) -> Self {
114 let compressed_content = zstd_compress(content);
115 Self {
116 compressed_content,
117 hash,
118 line_count,
119 original_tokens,
120 read_count: AtomicU32::new(1),
121 path,
122 last_access: AtomicU64::new(encode_instant(Instant::now())),
123 stored_mtime,
124 compressed_outputs: HashMap::new(),
125 full_content_delivered: false,
126 delivered_conversation: None,
127 last_mode: String::new(),
128 }
129 }
130
131 pub fn read_count(&self) -> u32 {
133 self.read_count.load(Ordering::Relaxed)
134 }
135
136 pub fn bump_read_count(&self) -> u32 {
138 self.read_count.fetch_add(1, Ordering::Relaxed) + 1
139 }
140
141 pub fn set_read_count(&self, n: u32) {
143 self.read_count.store(n, Ordering::Relaxed);
144 }
145
146 pub fn last_access(&self) -> Instant {
148 decode_instant(self.last_access.load(Ordering::Relaxed))
149 }
150
151 pub fn touch(&self) {
153 self.last_access
154 .store(encode_instant(Instant::now()), Ordering::Relaxed);
155 }
156
157 pub fn set_last_access(&self, when: Instant) {
159 self.last_access
160 .store(encode_instant(when), Ordering::Relaxed);
161 }
162
163 pub fn content(&self) -> Option<String> {
165 zstd_decompress(&self.compressed_content)
166 }
167
168 pub fn set_content(&mut self, content: &str) {
170 self.compressed_content = zstd_compress(content);
171 }
172
173 pub fn compressed_size(&self) -> usize {
175 self.compressed_content.len()
176 }
177}
178
179#[derive(Debug, Clone)]
181pub struct StoreResult {
182 pub line_count: usize,
183 pub original_tokens: usize,
184 pub read_count: u32,
185 pub was_hit: bool,
186 pub full_content_delivered: bool,
188}
189
190impl CacheEntry {
191 pub fn eviction_score_legacy(&self, now: Instant) -> f64 {
193 let elapsed = now
194 .checked_duration_since(self.last_access())
195 .unwrap_or_default()
196 .as_secs_f64();
197 let recency = 1.0 / (1.0 + elapsed.sqrt());
198 let frequency = (self.read_count() as f64 + 1.0).ln();
199 let size_value = (self.original_tokens as f64 + 1.0).ln();
200 recency * 0.4 + frequency * 0.3 + size_value * 0.3
201 }
202
203 pub fn get_compressed(&self, mode_key: &str) -> Option<&String> {
204 self.compressed_outputs.get(mode_key)
205 }
206
207 pub fn set_compressed(&mut self, mode_key: &str, output: String) {
208 const MAX_COMPRESSED_VARIANTS: usize = 3;
209 if self.compressed_outputs.len() >= MAX_COMPRESSED_VARIANTS
210 && !self.compressed_outputs.contains_key(mode_key)
211 && let Some(oldest_key) = self.compressed_outputs.keys().next().cloned()
212 {
213 self.compressed_outputs.remove(&oldest_key);
214 }
215 self.compressed_outputs.insert(mode_key.to_string(), output);
216 }
217
218 pub fn mark_full_delivered(&mut self, conversation: Option<String>) {
219 self.full_content_delivered = true;
220 self.delivered_conversation = conversation;
221 }
222}
223
224const RRF_K: f64 = 60.0;
225
226const HEBBIAN_PROTECT_WEIGHT: f64 = 0.05;
231const HEBBIAN_ACTIVE_SET: usize = 8;
234
235pub fn eviction_scores_rrf(entries: &[(&String, &CacheEntry)], now: Instant) -> Vec<(String, f64)> {
240 if entries.is_empty() {
241 return Vec::new();
242 }
243
244 let n = entries.len();
245
246 let mut recency_order: Vec<usize> = (0..n).collect();
247 recency_order.sort_by(|&a, &b| {
248 let elapsed_a = now
249 .checked_duration_since(entries[a].1.last_access())
250 .unwrap_or_default()
251 .as_secs_f64();
252 let elapsed_b = now
253 .checked_duration_since(entries[b].1.last_access())
254 .unwrap_or_default()
255 .as_secs_f64();
256 elapsed_a
257 .partial_cmp(&elapsed_b)
258 .unwrap_or(std::cmp::Ordering::Equal)
259 });
260
261 let mut frequency_order: Vec<usize> = (0..n).collect();
262 frequency_order.sort_by(|&a, &b| entries[b].1.read_count().cmp(&entries[a].1.read_count()));
263
264 let mut size_order: Vec<usize> = (0..n).collect();
265 size_order.sort_by(|&a, &b| {
266 entries[b]
267 .1
268 .original_tokens
269 .cmp(&entries[a].1.original_tokens)
270 });
271
272 let mut recency_ranks = vec![0usize; n];
273 let mut frequency_ranks = vec![0usize; n];
274 let mut size_ranks = vec![0usize; n];
275
276 for (rank, &idx) in recency_order.iter().enumerate() {
277 recency_ranks[idx] = rank;
278 }
279 for (rank, &idx) in frequency_order.iter().enumerate() {
280 frequency_ranks[idx] = rank;
281 }
282 for (rank, &idx) in size_order.iter().enumerate() {
283 size_ranks[idx] = rank;
284 }
285
286 entries
287 .iter()
288 .enumerate()
289 .map(|(i, (path, _))| {
290 let score = 1.0 / (RRF_K + recency_ranks[i] as f64)
291 + 1.0 / (RRF_K + frequency_ranks[i] as f64)
292 + 1.0 / (RRF_K + size_ranks[i] as f64);
293 ((*path).clone(), score)
294 })
295 .collect()
296}
297
298fn apply_hebbian_bonus(scores: &mut [(String, f64)], bonus: &HashMap<String, f64>) {
301 if bonus.is_empty() {
302 return;
303 }
304 for s in scores.iter_mut() {
305 if let Some(b) = bonus.get(&s.0) {
306 s.1 += *b;
307 }
308 }
309}
310
311#[derive(Debug, Default)]
316pub struct CacheStats {
317 total_reads: AtomicU64,
318 cache_hits: AtomicU64,
319 total_original_tokens: AtomicU64,
320 total_sent_tokens: AtomicU64,
321 files_tracked: AtomicU64,
322}
323
324impl CacheStats {
325 pub fn total_reads(&self) -> u64 {
327 self.total_reads.load(Ordering::Relaxed)
328 }
329
330 pub fn cache_hits(&self) -> u64 {
332 self.cache_hits.load(Ordering::Relaxed)
333 }
334
335 pub fn total_original_tokens(&self) -> u64 {
337 self.total_original_tokens.load(Ordering::Relaxed)
338 }
339
340 pub fn total_sent_tokens(&self) -> u64 {
342 self.total_sent_tokens.load(Ordering::Relaxed)
343 }
344
345 pub fn files_tracked(&self) -> u64 {
347 self.files_tracked.load(Ordering::Relaxed)
348 }
349
350 pub fn hit_rate(&self) -> f64 {
352 let total = self.total_reads();
353 if total == 0 {
354 return 0.0;
355 }
356 (self.cache_hits() as f64 / total as f64) * 100.0
357 }
358
359 pub fn tokens_saved(&self) -> u64 {
361 self.total_original_tokens()
362 .saturating_sub(self.total_sent_tokens())
363 }
364
365 pub fn savings_percent(&self) -> f64 {
367 let original = self.total_original_tokens();
368 if original == 0 {
369 return 0.0;
370 }
371 (self.tokens_saved() as f64 / original as f64) * 100.0
372 }
373}
374
375#[derive(Clone, Debug)]
377pub struct SharedBlock {
378 pub canonical_path: String,
379 pub canonical_ref: String,
380 pub start_line: usize,
381 pub end_line: usize,
382 pub content: String,
383}
384
385pub struct SessionCache {
388 entries: HashMap<String, CacheEntry>,
389 file_refs: HashMap<String, String>,
390 next_ref: usize,
391 stats: CacheStats,
392 shared_blocks: Vec<SharedBlock>,
393 co_access: crate::core::hebbian_cache::CoAccessMatrix,
397}
398
399impl Default for SessionCache {
400 fn default() -> Self {
401 Self::new()
402 }
403}
404
405impl SessionCache {
406 pub fn new() -> Self {
408 Self {
409 entries: HashMap::new(),
410 file_refs: HashMap::new(),
411 next_ref: 1,
412 shared_blocks: Vec::new(),
413 stats: CacheStats::default(),
414 co_access: crate::core::hebbian_cache::CoAccessMatrix::new(),
415 }
416 }
417
418 pub fn record_co_access(&mut self, path: &str) {
422 let key = normalize_key(path);
423 self.co_access
424 .record_access(crate::core::hebbian_cache::path_hash(&key));
425 }
426
427 pub fn flush_co_access(&mut self) {
430 self.co_access.end_burst();
431 }
432
433 pub(crate) fn hebbian_eviction_bonus(&self) -> HashMap<String, f64> {
439 use crate::core::hebbian_cache::path_hash;
440 if self.entries.is_empty() {
441 return HashMap::new();
442 }
443 let mut by_recency: Vec<(&String, Instant)> = self
444 .entries
445 .iter()
446 .map(|(k, e)| (k, e.last_access()))
447 .collect();
448 by_recency.sort_by_key(|(_, t)| std::cmp::Reverse(*t));
449 let active: Vec<u64> = by_recency
450 .iter()
451 .take(HEBBIAN_ACTIVE_SET)
452 .map(|(k, _)| path_hash(k))
453 .collect();
454
455 let mut out = HashMap::new();
456 for k in self.entries.keys() {
457 let h = path_hash(k);
458 let peers: Vec<u64> = active.iter().copied().filter(|&a| a != h).collect();
460 let strength = self.co_access.association_strength(h, &peers);
461 if strength > 0.0 {
462 out.insert(k.clone(), f64::from(strength) * HEBBIAN_PROTECT_WEIGHT);
463 }
464 }
465 if !out.is_empty() {
466 crate::core::introspect::tick("hebbian_cache");
467 }
468 out
469 }
470
471 pub fn get_file_ref(&mut self, path: &str) -> String {
473 let key = normalize_key(path);
474 if let Some(r) = self.file_refs.get(&key) {
475 return r.clone();
476 }
477 let r = format!("F{}", self.next_ref);
478 self.next_ref += 1;
479 self.file_refs.insert(key, r.clone());
480 r
481 }
482
483 pub fn get_file_ref_readonly(&self, path: &str) -> Option<String> {
485 self.file_refs.get(&normalize_key(path)).cloned()
486 }
487
488 pub fn get(&self, path: &str) -> Option<&CacheEntry> {
490 self.entries.get(&normalize_key(path))
491 }
492
493 pub fn get_mut(&mut self, path: &str) -> Option<&mut CacheEntry> {
495 self.entries.get_mut(&normalize_key(path))
496 }
497
498 pub fn get_full_content(&self, path: &str) -> Option<String> {
501 self.entries
502 .get(&normalize_key(path))
503 .and_then(CacheEntry::content)
504 }
505
506 pub fn current_full_content(&self, path: &str) -> Option<(String, usize)> {
519 let entry = self.entries.get(&normalize_key(path))?;
520 if is_cache_entry_stale_verified(&entry.path, entry.stored_mtime, &entry.hash)
521 && let Ok(fresh) = crate::core::io_boundary::read_file_lossy(&entry.path)
522 {
523 let tokens = count_tokens(&fresh);
528 return Some((fresh, tokens));
529 }
530 Some((entry.content()?, entry.original_tokens))
531 }
532
533 pub fn record_cache_hit(&self, path: &str) -> Option<&CacheEntry> {
539 let key = normalize_key(path);
540 let ref_label = self
541 .file_refs
542 .get(&key)
543 .cloned()
544 .unwrap_or_else(|| "F?".to_string());
545 let entry = self.entries.get(&key)?;
546 let new_count = entry.bump_read_count();
547 entry.touch();
548 self.stats.total_reads.fetch_add(1, Ordering::Relaxed);
549 self.stats.cache_hits.fetch_add(1, Ordering::Relaxed);
550 self.stats
551 .total_original_tokens
552 .fetch_add(entry.original_tokens as u64, Ordering::Relaxed);
553 let hit_msg = format!("{ref_label} cached {new_count}t {}L", entry.line_count);
554 self.stats
555 .total_sent_tokens
556 .fetch_add(count_tokens(&hit_msg) as u64, Ordering::Relaxed);
557 crate::core::events::emit_cache_hit(path, entry.original_tokens as u64);
558 Some(entry)
559 }
560
561 pub fn store(&mut self, path: &str, content: &str) -> StoreResult {
563 let key = normalize_key(path);
564 self.co_access
567 .record_access(crate::core::hebbian_cache::path_hash(&key));
568 let hash = compute_md5(content);
569 let line_count = content.lines().count();
570 let original_tokens = count_tokens(content);
571 let stored_mtime = std::fs::metadata(path).and_then(|m| m.modified()).ok();
572 let now = Instant::now();
573
574 self.stats.total_reads.fetch_add(1, Ordering::Relaxed);
575 self.stats
576 .total_original_tokens
577 .fetch_add(original_tokens as u64, Ordering::Relaxed);
578
579 if let Some(existing) = self.entries.get_mut(&key) {
580 existing.set_last_access(now);
581 if stored_mtime.is_some() {
582 existing.stored_mtime = stored_mtime;
583 }
584 if existing.hash == hash {
585 let new_count = existing.bump_read_count();
586 self.stats.cache_hits.fetch_add(1, Ordering::Relaxed);
587 let hit_msg = format!(
588 "{} cached {new_count}t {}L",
589 self.file_refs.get(&key).unwrap_or(&"F?".to_string()),
590 existing.line_count,
591 );
592 self.stats
593 .total_sent_tokens
594 .fetch_add(count_tokens(&hit_msg) as u64, Ordering::Relaxed);
595 return StoreResult {
596 line_count: existing.line_count,
597 original_tokens: existing.original_tokens,
598 read_count: new_count,
599 was_hit: true,
600 full_content_delivered: existing.full_content_delivered,
601 };
602 }
603 existing.compressed_outputs.clear();
604 existing.set_content(content);
605 existing.hash = hash;
606 existing.line_count = line_count;
607 existing.original_tokens = original_tokens;
608 let new_count = existing.bump_read_count();
609 existing.full_content_delivered = false;
610 existing.delivered_conversation = None;
611 if stored_mtime.is_some() {
612 existing.stored_mtime = stored_mtime;
613 }
614 self.stats
615 .total_sent_tokens
616 .fetch_add(original_tokens as u64, Ordering::Relaxed);
617 return StoreResult {
618 line_count,
619 original_tokens,
620 read_count: new_count,
621 was_hit: false,
622 full_content_delivered: false,
623 };
624 }
625
626 self.evict_if_needed(original_tokens);
627 self.get_file_ref(&key);
628
629 let entry = CacheEntry::new(
630 content,
631 hash,
632 line_count,
633 original_tokens,
634 key.clone(),
635 stored_mtime,
636 );
637
638 self.entries.insert(key, entry);
639 self.stats.files_tracked.fetch_add(1, Ordering::Relaxed);
640 self.stats
641 .total_sent_tokens
642 .fetch_add(original_tokens as u64, Ordering::Relaxed);
643 StoreResult {
644 line_count,
645 original_tokens,
646 read_count: 1,
647 was_hit: false,
648 full_content_delivered: false,
649 }
650 }
651
652 pub fn total_cached_tokens(&self) -> usize {
654 self.entries.values().map(|e| e.original_tokens).sum()
655 }
656
657 pub fn evict_if_needed(&mut self, incoming_tokens: usize) {
660 let max_tokens = max_cache_tokens();
661 let current = self.total_cached_tokens();
662 if current + incoming_tokens <= max_tokens {
663 return;
664 }
665
666 let now = Instant::now();
667 let all: Vec<(&String, &CacheEntry)> = self.entries.iter().collect();
668 let mut scores = eviction_scores_rrf(&all, now);
669 apply_hebbian_bonus(&mut scores, &self.hebbian_eviction_bonus());
670 scores.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
672
673 let mut freed = 0usize;
674 let mut redelivered = 0u64;
675 let target = (current + incoming_tokens).saturating_sub(max_tokens);
676
677 for (path, _score) in &scores {
678 if freed >= target {
679 break;
680 }
681 if let Some(entry) = self.entries.remove(path) {
682 freed += entry.original_tokens;
683 if entry.full_content_delivered {
684 redelivered += 1;
685 }
686 self.file_refs.remove(path);
687 }
688 }
689 crate::core::cache_telemetry::record_eviction(redelivered);
690 }
691
692 pub fn get_all_entries(&self) -> Vec<(&String, &CacheEntry)> {
694 self.entries.iter().collect()
695 }
696
697 pub fn get_stats(&self) -> &CacheStats {
699 &self.stats
700 }
701
702 pub fn file_ref_map(&self) -> &HashMap<String, String> {
704 &self.file_refs
705 }
706
707 pub fn set_shared_blocks(&mut self, blocks: Vec<SharedBlock>) {
709 self.shared_blocks = blocks;
710 }
711
712 pub fn get_shared_blocks(&self) -> &[SharedBlock] {
714 &self.shared_blocks
715 }
716
717 pub fn apply_dedup(&self, path: &str, content: &str) -> Option<String> {
719 if self.shared_blocks.is_empty() {
720 return None;
721 }
722 let refs: Vec<&SharedBlock> = self
723 .shared_blocks
724 .iter()
725 .filter(|b| b.canonical_path != path && content.contains(&b.content))
726 .collect();
727 if refs.is_empty() {
728 return None;
729 }
730 let mut result = content.to_string();
731 for block in refs {
732 result = result.replacen(
733 &block.content,
734 &format!(
735 "[= {}:{}-{}]",
736 block.canonical_ref, block.start_line, block.end_line
737 ),
738 1,
739 );
740 }
741 Some(result)
742 }
743
744 pub fn invalidate(&mut self, path: &str) -> bool {
746 self.entries.remove(&normalize_key(path)).is_some()
747 }
748
749 pub fn get_compressed(&self, path: &str, mode_key: &str) -> Option<&String> {
751 self.entries
752 .get(&normalize_key(path))?
753 .get_compressed(mode_key)
754 }
755
756 pub fn mark_full_delivered(&mut self, path: &str) {
761 let conversation = crate::core::conversation::current_conversation_id();
762 let key = normalize_key(path);
763 let file_ref = self.file_refs.get(&key).cloned();
764 if let Some(entry) = self.entries.get_mut(&key) {
765 entry.mark_full_delivered(conversation.clone());
766 crate::core::read_stub_index::record(crate::core::read_stub_index::StubRecord::new(
770 key.clone(),
771 entry.hash.clone(),
772 entry.stored_mtime,
773 entry.line_count,
774 file_ref.unwrap_or_default(),
775 conversation,
776 ));
777 }
778 }
779
780 pub fn set_compressed(&mut self, path: &str, mode_key: &str, output: String) {
782 if let Some(entry) = self.entries.get_mut(&normalize_key(path)) {
783 entry.set_compressed(mode_key, output);
784 }
785 }
786
787 pub fn reset_delivery_flags(&mut self) -> usize {
791 let mut count = 0;
792 for entry in self.entries.values_mut() {
793 if entry.full_content_delivered {
794 entry.full_content_delivered = false;
795 count += 1;
796 }
797 }
798 count
799 }
800
801 pub fn is_full_delivered(&self, path: &str) -> bool {
803 self.entries
804 .get(&normalize_key(path))
805 .is_some_and(|e| e.full_content_delivered)
806 }
807
808 pub fn count_full_delivered(&self) -> usize {
812 self.entries
813 .values()
814 .filter(|e| e.full_content_delivered)
815 .count()
816 }
817
818 pub fn trim_compressed_outputs(&mut self) -> usize {
821 let mut trimmed = 0;
822 for entry in self.entries.values_mut() {
823 if !entry.compressed_outputs.is_empty() {
824 entry.compressed_outputs.clear();
825 trimmed += 1;
826 }
827 }
828 trimmed
829 }
830
831 pub fn evict_probationary(&mut self) -> usize {
834 let to_remove: Vec<String> = self
835 .entries
836 .iter()
837 .filter(|(_, e)| e.read_count() <= 1)
838 .map(|(k, _)| k.clone())
839 .collect();
840 let count = to_remove.len();
841 let mut redelivered = 0u64;
842 for key in &to_remove {
843 if self
844 .entries
845 .remove(key)
846 .is_some_and(|e| e.full_content_delivered)
847 {
848 redelivered += 1;
849 }
850 self.file_refs.remove(key);
851 }
852 crate::core::cache_telemetry::record_eviction(redelivered);
853 count
854 }
855
856 pub fn evict_to_budget(&mut self, target_tokens: usize) {
858 let current = self.total_cached_tokens();
859 if current <= target_tokens {
860 return;
861 }
862 let now = Instant::now();
863 let all: Vec<(&String, &CacheEntry)> = self.entries.iter().collect();
864 let mut scores = eviction_scores_rrf(&all, now);
865 apply_hebbian_bonus(&mut scores, &self.hebbian_eviction_bonus());
866 scores.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
867
868 let mut freed = 0usize;
869 let mut redelivered = 0u64;
870 let target_free = current.saturating_sub(target_tokens);
871 for (path, _score) in &scores {
872 if freed >= target_free {
873 break;
874 }
875 if let Some(entry) = self.entries.remove(path) {
876 freed += entry.original_tokens;
877 if entry.full_content_delivered {
878 redelivered += 1;
879 }
880 self.file_refs.remove(path);
881 }
882 }
883 crate::core::cache_telemetry::record_eviction(redelivered);
884 }
885
886 pub fn approximate_bytes(&self) -> usize {
888 let entries_bytes: usize = self
889 .entries
890 .values()
891 .map(|e| {
892 e.compressed_content.len()
893 + e.hash.len()
894 + e.path.len()
895 + e.compressed_outputs
896 .iter()
897 .map(|(k, v)| k.len() + v.len())
898 .sum::<usize>()
899 + 128 })
901 .sum();
902 let refs_bytes: usize = self.file_refs.iter().map(|(k, v)| k.len() + v.len()).sum();
903 let blocks_bytes: usize = self
904 .shared_blocks
905 .iter()
906 .map(|b| b.canonical_path.len() + b.canonical_ref.len() + b.content.len() + 32)
907 .sum();
908 entries_bytes + refs_bytes + blocks_bytes
909 }
910
911 const MAX_SHARED_BLOCKS: usize = 100;
912
913 pub fn trim_shared_blocks(&mut self) {
915 if self.shared_blocks.len() > Self::MAX_SHARED_BLOCKS {
916 let excess = self.shared_blocks.len() - Self::MAX_SHARED_BLOCKS;
917 self.shared_blocks.drain(..excess);
918 }
919 }
920
921 pub fn clear(&mut self) -> usize {
923 let count = self.entries.len();
924 self.entries.clear();
925 self.file_refs.clear();
926 self.shared_blocks.clear();
927 self.next_ref = 1;
928 self.stats = CacheStats::default();
929 count
930 }
931}
932
933pub fn file_mtime(path: &str) -> Option<SystemTime> {
934 std::fs::metadata(path).and_then(|m| m.modified()).ok()
935}
936
937pub fn is_cache_entry_stale(path: &str, cached_mtime: Option<SystemTime>) -> bool {
938 let current = file_mtime(path);
939 match (cached_mtime, current) {
940 (None, None) => false,
942 (Some(_), None) | (None, Some(_)) => true,
944 (Some(cached), Some(current)) => current != cached,
947 }
948}
949
950const VERIFY_HASH_CAP_BYTES: u64 = 8 * 1024 * 1024;
953
954fn cache_verify_enabled() -> bool {
955 std::env::var("LEAN_CTX_CACHE_VERIFY").map_or(true, |v| v != "0")
956}
957
958pub fn is_cache_entry_stale_verified(
973 path: &str,
974 cached_mtime: Option<SystemTime>,
975 cached_hash: &str,
976) -> bool {
977 if is_cache_entry_stale(path, cached_mtime) {
978 return true;
979 }
980 if cached_hash.is_empty() || !cache_verify_enabled() {
981 return false;
982 }
983 let Ok(meta) = std::fs::metadata(path) else {
984 return true;
986 };
987 if meta.len() > VERIFY_HASH_CAP_BYTES {
988 return false;
989 }
990 match std::fs::read(path) {
991 Ok(bytes) => compute_md5(&String::from_utf8_lossy(&bytes)) != cached_hash,
993 Err(_) => true,
994 }
995}
996
997fn compute_md5(content: &str) -> String {
998 let mut hasher = Md5::new();
999 hasher.update(content.as_bytes());
1000 crate::core::agent_identity::hex_encode(&hasher.finalize())
1001}
1002
1003#[cfg(test)]
1004mod tests {
1005 use super::*;
1006 use std::time::Duration;
1007
1008 #[test]
1009 fn cache_stores_and_retrieves() {
1010 let mut cache = SessionCache::new();
1011 let result = cache.store("/test/file.rs", "fn main() {}");
1012 assert!(!result.was_hit);
1013 assert_eq!(result.line_count, 1);
1014 assert!(cache.get("/test/file.rs").is_some());
1015 }
1016
1017 #[test]
1018 fn cache_hit_on_same_content() {
1019 let mut cache = SessionCache::new();
1020 cache.store("/test/file.rs", "content");
1021 let result = cache.store("/test/file.rs", "content");
1022 assert!(result.was_hit, "same content should be a cache hit");
1023 }
1024
1025 #[test]
1026 fn cache_miss_on_changed_content() {
1027 let mut cache = SessionCache::new();
1028 cache.store("/test/file.rs", "old content");
1029 let result = cache.store("/test/file.rs", "new content");
1030 assert!(!result.was_hit, "changed content should not be a cache hit");
1031 }
1032
1033 #[test]
1034 fn file_refs_are_sequential() {
1035 let mut cache = SessionCache::new();
1036 assert_eq!(cache.get_file_ref("/a.rs"), "F1");
1037 assert_eq!(cache.get_file_ref("/b.rs"), "F2");
1038 assert_eq!(cache.get_file_ref("/a.rs"), "F1"); }
1040
1041 #[test]
1042 fn cache_clear_resets_everything() {
1043 let mut cache = SessionCache::new();
1044 cache.store("/a.rs", "a");
1045 cache.store("/b.rs", "b");
1046 let count = cache.clear();
1047 assert_eq!(count, 2);
1048 assert!(cache.get("/a.rs").is_none());
1049 assert_eq!(cache.get_file_ref("/c.rs"), "F1"); }
1051
1052 #[test]
1053 fn cache_invalidate_removes_entry() {
1054 let mut cache = SessionCache::new();
1055 cache.store("/test.rs", "test");
1056 assert!(cache.invalidate("/test.rs"));
1057 assert!(!cache.invalidate("/nonexistent.rs"));
1058 }
1059
1060 #[test]
1061 fn cache_stats_track_correctly() {
1062 let mut cache = SessionCache::new();
1063 cache.store("/a.rs", "hello");
1064 cache.store("/a.rs", "hello"); let stats = cache.get_stats();
1066 assert_eq!(stats.total_reads(), 2);
1067 assert_eq!(stats.cache_hits(), 1);
1068 assert!(stats.hit_rate() > 0.0);
1069 }
1070
1071 #[test]
1072 fn current_full_content_serves_cached_when_fresh() {
1073 let dir = tempfile::tempdir().unwrap();
1074 let file = dir.path().join("handover.md");
1075 std::fs::write(&file, "HANDOVER V1\n").unwrap();
1076 let path = file.to_str().unwrap();
1077
1078 let mut cache = SessionCache::new();
1079 cache.store(path, "HANDOVER V1\n");
1080
1081 let (content, tokens) = cache.current_full_content(path).unwrap();
1082 assert_eq!(content, "HANDOVER V1\n");
1083 assert!(tokens > 0);
1084 }
1085
1086 #[test]
1087 fn current_full_content_rereads_when_file_changed() {
1088 let dir = tempfile::tempdir().unwrap();
1091 let file = dir.path().join("handover.md");
1092 std::fs::write(&file, "HANDOVER V1\n").unwrap();
1093 let path = file.to_str().unwrap();
1094
1095 let mut cache = SessionCache::new();
1096 cache.store(path, "HANDOVER V1\n");
1097
1098 std::thread::sleep(std::time::Duration::from_millis(10));
1100 std::fs::write(&file, "HANDOVER V2 CHANGED\n").unwrap();
1101
1102 let (content, _) = cache.current_full_content(path).unwrap();
1103 assert_eq!(
1104 content, "HANDOVER V2 CHANGED\n",
1105 "stale cached copy must be re-read from disk, not served as-is"
1106 );
1107 }
1108
1109 #[test]
1110 fn current_full_content_none_without_entry() {
1111 let cache = SessionCache::new();
1112 assert!(cache.current_full_content("/no/such/file.rs").is_none());
1113 }
1114
1115 #[test]
1116 fn current_full_content_falls_back_to_cache_when_file_unreadable() {
1117 let dir = tempfile::tempdir().unwrap();
1122 let canon = dir.path().canonicalize().unwrap();
1123 let file = canon.join("gone.md");
1124 std::fs::write(&file, "ORIGINAL\n").unwrap();
1125 let path = file.to_str().unwrap().to_string();
1126
1127 let mut cache = SessionCache::new();
1128 cache.store(&path, "ORIGINAL\n");
1129 std::fs::remove_file(&file).unwrap();
1130
1131 let (content, _) = cache.current_full_content(&path).unwrap();
1132 assert_eq!(
1133 content, "ORIGINAL\n",
1134 "unreadable file must fall back to last-known cached content"
1135 );
1136 }
1137
1138 #[test]
1139 fn record_cache_hit_works_through_shared_ref() {
1140 let mut cache = SessionCache::new();
1141 cache.store("/x.rs", "hello world");
1142 let shared: &SessionCache = &cache;
1144 assert!(shared.record_cache_hit("/x.rs").is_some());
1145 assert!(shared.record_cache_hit("/x.rs").is_some());
1146 assert_eq!(cache.get("/x.rs").unwrap().read_count(), 3);
1148 assert_eq!(cache.get_stats().cache_hits(), 2);
1149 }
1150
1151 #[test]
1152 fn concurrent_cache_hits_are_lossless() {
1153 use std::sync::Arc;
1154 let mut cache = SessionCache::new();
1155 cache.store("/a.rs", "a");
1156 cache.store("/b.rs", "b");
1157 let cache = Arc::new(cache);
1160 let threads = 8;
1161 let iters = 1_000;
1162 let handles: Vec<_> = (0..threads)
1163 .map(|_| {
1164 let c = Arc::clone(&cache);
1165 std::thread::spawn(move || {
1166 for _ in 0..iters {
1167 c.record_cache_hit("/a.rs");
1168 c.record_cache_hit("/b.rs");
1169 }
1170 })
1171 })
1172 .collect();
1173 for h in handles {
1174 h.join().unwrap();
1175 }
1176 let total = (threads * iters) as u64;
1177 assert_eq!(cache.get_stats().cache_hits(), total * 2);
1178 assert_eq!(cache.get("/a.rs").unwrap().read_count(), 1 + total as u32);
1179 assert_eq!(cache.get("/b.rs").unwrap().read_count(), 1 + total as u32);
1180 }
1181
1182 #[test]
1183 fn hebbian_eviction_bonus_is_wired() {
1184 let mut cache = SessionCache::new();
1187 cache.store("/a.rs", "fn a() {}");
1188 cache.store("/b.rs", "fn b() {}");
1189 cache.flush_co_access(); let bonus = cache.hebbian_eviction_bonus();
1191 assert!(
1192 !bonus.is_empty(),
1193 "co-accessed reads must yield a Hebbian eviction bonus (#3 wired)"
1194 );
1195 }
1196
1197 #[test]
1198 fn md5_is_deterministic() {
1199 let h1 = compute_md5("test content");
1200 let h2 = compute_md5("test content");
1201 assert_eq!(h1, h2);
1202 assert_ne!(h1, compute_md5("different"));
1203 }
1204
1205 #[test]
1206 fn rrf_eviction_prefers_recent() {
1207 let key_a = "a.rs".to_string();
1208 let key_b = "b.rs".to_string();
1209 let recent = CacheEntry::new("a", "h1".to_string(), 1, 10, "/a.rs".to_string(), None);
1212 let old = CacheEntry::new("b", "h2".to_string(), 1, 10, "/b.rs".to_string(), None);
1213 let t_old = Instant::now();
1214 std::thread::sleep(std::time::Duration::from_millis(10));
1215 let t_recent = Instant::now();
1216 old.set_last_access(t_old);
1217 recent.set_last_access(t_recent);
1218 let now = Instant::now();
1219 let entries: Vec<(&String, &CacheEntry)> = vec![(&key_a, &recent), (&key_b, &old)];
1220 let scores = eviction_scores_rrf(&entries, now);
1221 let score_a = scores.iter().find(|(p, _)| p == "a.rs").unwrap().1;
1222 let score_b = scores.iter().find(|(p, _)| p == "b.rs").unwrap().1;
1223 assert!(
1224 score_a > score_b,
1225 "recently accessed entries should score higher via RRF"
1226 );
1227 }
1228
1229 #[test]
1230 fn rrf_eviction_prefers_frequent() {
1231 let now = Instant::now();
1232 let key_a = "a.rs".to_string();
1233 let key_b = "b.rs".to_string();
1234 let frequent = {
1235 let e = CacheEntry::new("a", "h1".to_string(), 1, 10, "/a.rs".to_string(), None);
1236 e.set_read_count(20);
1237 e
1238 };
1239 let rare = CacheEntry::new("b", "h2".to_string(), 1, 10, "/b.rs".to_string(), None);
1240 let entries: Vec<(&String, &CacheEntry)> = vec![(&key_a, &frequent), (&key_b, &rare)];
1241 let scores = eviction_scores_rrf(&entries, now);
1242 let score_a = scores.iter().find(|(p, _)| p == "a.rs").unwrap().1;
1243 let score_b = scores.iter().find(|(p, _)| p == "b.rs").unwrap().1;
1244 assert!(
1245 score_a > score_b,
1246 "frequently accessed entries should score higher via RRF"
1247 );
1248 }
1249
1250 #[test]
1251 fn cache_budget_resolver_precedence() {
1252 assert_eq!(resolve_cache_max_tokens(Some("250000"), 999), 250_000);
1254 assert_eq!(resolve_cache_max_tokens(Some(" 80000 "), 0), 80_000);
1255 assert_eq!(resolve_cache_max_tokens(Some("0"), 123_456), 123_456);
1257 assert_eq!(resolve_cache_max_tokens(Some(""), 123_456), 123_456);
1258 assert_eq!(resolve_cache_max_tokens(Some("lots"), 123_456), 123_456);
1259 assert_eq!(resolve_cache_max_tokens(None, 42_000), 42_000);
1261 assert_eq!(resolve_cache_max_tokens(None, 0), DEFAULT_CACHE_MAX_TOKENS);
1263 assert_eq!(
1264 resolve_cache_max_tokens(Some("0"), 0),
1265 DEFAULT_CACHE_MAX_TOKENS
1266 );
1267 }
1268
1269 #[test]
1270 fn evict_if_needed_removes_lowest_score() {
1271 crate::test_env::set_var("LEAN_CTX_CACHE_MAX_TOKENS", "50");
1272 let mut cache = SessionCache::new();
1273 let big_content = "a]".repeat(30); cache.store("/old.rs", &big_content);
1275 let new_content = "b ".repeat(30); cache.store("/new.rs", &new_content);
1279 assert!(
1284 cache.total_cached_tokens() <= 60,
1285 "eviction should have kicked in"
1286 );
1287 crate::test_env::remove_var("LEAN_CTX_CACHE_MAX_TOKENS");
1288 }
1289
1290 #[test]
1291 fn stale_detection_flags_newer_file() {
1292 let dir = tempfile::tempdir().unwrap();
1293 let path = dir.path().join("stale.txt");
1294 let p = path.to_string_lossy().to_string();
1295
1296 std::fs::write(&path, "one").unwrap();
1297 let mut cache = SessionCache::new();
1298 cache.store(&p, "one");
1299
1300 let entry = cache.get(&p).unwrap();
1301 assert!(!is_cache_entry_stale(&p, entry.stored_mtime));
1302
1303 std::thread::sleep(Duration::from_secs(1));
1305 std::fs::write(&path, "two").unwrap();
1306
1307 let entry = cache.get(&p).unwrap();
1308 assert!(is_cache_entry_stale(&p, entry.stored_mtime));
1309 }
1310
1311 #[test]
1313 fn stale_detection_flags_backward_mtime() {
1314 let dir = tempfile::tempdir().unwrap();
1315 let path = dir.path().join("backward.txt");
1316 let p = path.to_string_lossy().to_string();
1317
1318 std::fs::write(&path, "one").unwrap();
1319 let mut cache = SessionCache::new();
1320 cache.store(&p, "one");
1321 let entry_mtime = cache.get(&p).unwrap().stored_mtime;
1322 assert!(!is_cache_entry_stale(&p, entry_mtime));
1323
1324 std::fs::write(&path, "zero").unwrap();
1326 let f = std::fs::OpenOptions::new().write(true).open(&path).unwrap();
1327 f.set_modified(SystemTime::now() - Duration::from_hours(1))
1328 .unwrap();
1329 drop(f);
1330
1331 assert!(
1332 is_cache_entry_stale(&p, entry_mtime),
1333 "older mtime must read as stale"
1334 );
1335 }
1336
1337 #[test]
1340 fn verified_staleness_catches_same_mtime_content_change() {
1341 let dir = tempfile::tempdir().unwrap();
1342 let path = dir.path().join("sneaky.txt");
1343 let p = path.to_string_lossy().to_string();
1344
1345 std::fs::write(&path, "one").unwrap();
1346 let original_mtime = std::fs::metadata(&path).unwrap().modified().unwrap();
1347 let mut cache = SessionCache::new();
1348 cache.store(&p, "one");
1349 let (mtime, hash) = {
1350 let e = cache.get(&p).unwrap();
1351 (e.stored_mtime, e.hash.clone())
1352 };
1353
1354 assert!(!is_cache_entry_stale_verified(&p, mtime, &hash));
1356
1357 std::fs::write(&path, "two").unwrap();
1359 let f = std::fs::OpenOptions::new().write(true).open(&path).unwrap();
1360 f.set_modified(original_mtime).unwrap();
1361 drop(f);
1362
1363 assert!(
1364 !is_cache_entry_stale(&p, mtime),
1365 "test premise: the mtime check alone is fooled"
1366 );
1367 assert!(
1368 is_cache_entry_stale_verified(&p, mtime, &hash),
1369 "hash verification must catch the change"
1370 );
1371 }
1372
1373 #[test]
1374 fn verified_staleness_flags_unreadable_file() {
1375 let mut cache = SessionCache::new();
1376 cache.store("/nonexistent/file.rs", "content");
1377 let (mtime, hash) = {
1378 let e = cache.get("/nonexistent/file.rs").unwrap();
1379 (e.stored_mtime, e.hash.clone())
1380 };
1381 assert!(is_cache_entry_stale_verified(
1382 "/nonexistent/file.rs",
1383 mtime,
1384 &hash
1385 ));
1386 }
1387
1388 #[test]
1389 fn compressed_outputs_cached_and_retrieved() {
1390 let mut cache = SessionCache::new();
1391 cache.store("/test.rs", "fn main() {}");
1392 cache.set_compressed("/test.rs", "map", "compressed map output".to_string());
1393 assert_eq!(
1394 cache.get_compressed("/test.rs", "map"),
1395 Some(&"compressed map output".to_string())
1396 );
1397 assert_eq!(cache.get_compressed("/test.rs", "signatures"), None);
1398 }
1399
1400 #[test]
1401 fn compressed_outputs_cleared_on_content_change() {
1402 let mut cache = SessionCache::new();
1403 cache.store("/test.rs", "old content");
1404 cache.set_compressed("/test.rs", "map", "old map".to_string());
1405 assert!(cache.get_compressed("/test.rs", "map").is_some());
1406
1407 cache.store("/test.rs", "new content");
1408 assert_eq!(cache.get_compressed("/test.rs", "map"), None);
1409 }
1410
1411 #[test]
1412 fn compressed_outputs_survive_same_content_store() {
1413 let mut cache = SessionCache::new();
1414 cache.store("/test.rs", "content");
1415 cache.set_compressed("/test.rs", "map", "cached map".to_string());
1416
1417 let result = cache.store("/test.rs", "content");
1418 assert!(result.was_hit);
1419 assert_eq!(
1420 cache.get_compressed("/test.rs", "map"),
1421 Some(&"cached map".to_string())
1422 );
1423 }
1424
1425 #[test]
1426 fn compressed_outputs_cleared_on_invalidate() {
1427 let mut cache = SessionCache::new();
1428 cache.store("/test.rs", "content");
1429 cache.set_compressed("/test.rs", "signatures", "cached sigs".to_string());
1430 cache.invalidate("/test.rs");
1431 assert_eq!(cache.get_compressed("/test.rs", "signatures"), None);
1432 }
1433
1434 #[test]
1435 fn compressed_outputs_cleared_on_clear() {
1436 let mut cache = SessionCache::new();
1437 cache.store("/a.rs", "a");
1438 cache.set_compressed("/a.rs", "map", "map_a".to_string());
1439 cache.clear();
1440 assert_eq!(cache.get_compressed("/a.rs", "map"), None);
1441 }
1442}