1use std::collections::HashMap;
2use std::path::Path;
3
4use sha2::{Digest, Sha256};
5
6use crate::delta_encoder::DeltaEncoder;
7use crate::error::Result;
8use crate::pipeline::{CompressionPipeline, SessionContext};
9use crate::preset::Preset;
10use crate::session_store::SessionStore;
11use crate::types::CompressedContent;
12
13pub enum CacheResult {
20 Dedup {
22 inline_ref: String,
24 token_cost: u32,
26 },
27 Delta {
29 delta_text: String,
31 token_cost: u32,
33 similarity: f64,
35 },
36 Fresh { output: CompressedContent },
38}
39
40#[derive(Debug, Clone)]
42struct RefEntry {
43 last_sent_turn: u64,
45}
46
47pub struct CacheManager {
56 store: SessionStore,
57 max_size_bytes: u64,
58 delta_encoder: DeltaEncoder,
59 turn_counter: std::cell::Cell<u64>,
61 ref_tracker: std::cell::RefCell<HashMap<String, RefEntry>>,
63 max_ref_age_turns: u64,
67}
68
69impl CacheManager {
70 pub fn new(store: SessionStore, max_size_bytes: u64) -> Self {
76 Self {
77 store,
78 max_size_bytes,
79 delta_encoder: DeltaEncoder::new(),
80 turn_counter: std::cell::Cell::new(0),
81 ref_tracker: std::cell::RefCell::new(HashMap::new()),
82 max_ref_age_turns: 20,
83 }
84 }
85
86 pub fn with_ref_age(store: SessionStore, max_size_bytes: u64, max_ref_age_turns: u64) -> Self {
88 Self {
89 store,
90 max_size_bytes,
91 delta_encoder: DeltaEncoder::new(),
92 turn_counter: std::cell::Cell::new(0),
93 ref_tracker: std::cell::RefCell::new(HashMap::new()),
94 max_ref_age_turns,
95 }
96 }
97
98 fn sha256_hex(bytes: &[u8]) -> String {
100 let mut hasher = Sha256::new();
101 hasher.update(bytes);
102 format!("{:x}", hasher.finalize())
103 }
104
105 pub fn advance_turn(&self) {
107 self.turn_counter.set(self.turn_counter.get() + 1);
108 }
109
110 pub fn current_turn(&self) -> u64 {
112 self.turn_counter.get()
113 }
114
115 pub fn notify_compaction(&self) {
126 self.ref_tracker.borrow_mut().clear();
127 }
128
129 fn is_ref_fresh(&self, hash: &str) -> bool {
132 let tracker = self.ref_tracker.borrow();
133 if let Some(entry) = tracker.get(hash) {
134 let age = self.turn_counter.get().saturating_sub(entry.last_sent_turn);
135 age < self.max_ref_age_turns
136 } else {
137 false
138 }
139 }
140
141 fn record_ref_sent(&self, hash: &str) {
143 self.ref_tracker.borrow_mut().insert(
144 hash.to_string(),
145 RefEntry {
146 last_sent_turn: self.turn_counter.get(),
147 },
148 );
149 }
150
151 pub fn get_or_compress(
159 &self,
160 _path: &Path,
161 content: &[u8],
162 pipeline: &CompressionPipeline,
163 ) -> Result<CacheResult> {
164 let hash = Self::sha256_hex(content);
165
166 if self.store.get_cache_entry(&hash)?.is_some() {
168 if self.is_ref_fresh(&hash) {
169 let hash_prefix = &hash[..16];
171 let inline_ref = format!("§ref:{hash_prefix}§");
172 self.record_ref_sent(&hash);
174 return Ok(CacheResult::Dedup {
175 inline_ref,
176 token_cost: 13,
177 });
178 } else {
179 let text = String::from_utf8_lossy(content).into_owned();
182 let ctx = SessionContext {
183 session_id: "cache".to_string(),
184 };
185 let preset = Preset::default();
186 let compressed = pipeline.compress(&text, &ctx, &preset)?;
187 self.record_ref_sent(&hash);
189 return Ok(CacheResult::Fresh { output: compressed });
190 }
191 }
192
193 let text = String::from_utf8_lossy(content).into_owned();
195 if let Some(delta_result) = self.try_delta_encode(&text)? {
196 let ctx = SessionContext {
198 session_id: "cache".to_string(),
199 };
200 let preset = Preset::default();
201 let compressed = pipeline.compress(&text, &ctx, &preset)?;
202 self.store.save_cache_entry(&hash, &compressed)?;
203 self.record_ref_sent(&hash);
204
205 let token_cost = (delta_result.delta_text.len() / 4) as u32;
206 return Ok(CacheResult::Delta {
207 delta_text: delta_result.delta_text,
208 token_cost: token_cost.max(5),
209 similarity: delta_result.similarity,
210 });
211 }
212
213 let ctx = SessionContext {
214 session_id: "cache".to_string(),
215 };
216 let preset = Preset::default();
217 let compressed = pipeline.compress(&text, &ctx, &preset)?;
218 self.store.save_cache_entry(&hash, &compressed)?;
219 self.record_ref_sent(&hash);
221
222 Ok(CacheResult::Fresh { output: compressed })
223 }
224
225 fn try_delta_encode(
228 &self,
229 new_content: &str,
230 ) -> Result<Option<crate::delta_encoder::DeltaResult>> {
231 let entries = self.store.list_cache_entries_lru()?;
232
233 let check_count = entries.len().min(10);
235 for (hash, _) in entries.iter().rev().take(check_count) {
236 if let Some(cached) = self.store.get_cache_entry(hash)? {
237 let hash_prefix = &hash[..hash.len().min(16)];
238 if let Ok(Some(delta)) =
239 self.delta_encoder
240 .encode(&cached.data, new_content, hash_prefix)
241 {
242 if delta.delta_text.len() < new_content.len() {
244 return Ok(Some(delta));
245 }
246 }
247 }
248 }
249
250 Ok(None)
251 }
252
253 pub fn check_dedup(&self, content: &[u8]) -> Result<Option<String>> {
258 let hash = Self::sha256_hex(content);
259 if self.store.get_cache_entry(&hash)?.is_some() {
260 if self.is_ref_fresh(&hash) {
261 let hash_prefix = &hash[..16];
262 self.record_ref_sent(&hash);
263 Ok(Some(format!("§ref:{hash_prefix}§")))
264 } else {
265 Ok(None)
267 }
268 } else {
269 Ok(None)
270 }
271 }
272
273 pub fn store_compressed(
278 &self,
279 original_content: &[u8],
280 compressed: &CompressedContent,
281 ) -> Result<()> {
282 let hash = Self::sha256_hex(original_content);
283 self.store.save_cache_entry(&hash, compressed)?;
284 self.record_ref_sent(&hash);
285 Ok(())
286 }
287
288 pub fn invalidate(&self, path: &Path) -> Result<()> {
293 if !path.exists() {
294 return Ok(());
295 }
296 let bytes = std::fs::read(path)?;
297 let hash = Self::sha256_hex(&bytes);
298 self.store.delete_cache_entry(&hash)?;
299 Ok(())
300 }
301
302 pub fn evict_lru(&self) -> Result<u64> {
307 let entries = self.store.list_cache_entries_lru()?;
308
309 let total: u64 = entries.iter().map(|(_, sz)| sz).sum();
311 if total <= self.max_size_bytes {
312 return Ok(0);
313 }
314
315 let mut freed: u64 = 0;
316 let mut remaining = total;
317
318 for (hash, size) in &entries {
319 if remaining <= self.max_size_bytes {
320 break;
321 }
322 self.store.delete_cache_entry(hash)?;
323 freed += size;
324 remaining -= size;
325 }
326
327 Ok(freed)
328 }
329}
330
331#[cfg(test)]
334mod tests {
335 use super::*;
336 use crate::preset::{
337 BudgetConfig, CollapseArraysConfig, CompressionConfig, CondenseConfig,
338 CustomTransformsConfig, ModelConfig, PresetMeta, StripNullsConfig, TerseModeConfig,
339 ToolSelectionConfig, TruncateStringsConfig,
340 };
341 use crate::session_store::SessionStore;
342
343 fn in_memory_store() -> (SessionStore, tempfile::TempDir) {
344 let dir = tempfile::tempdir().unwrap();
345 let path = dir.path().join("test.db");
346 let store = SessionStore::open_or_create(&path).unwrap();
347 (store, dir)
348 }
349
350 fn test_preset() -> Preset {
351 Preset {
352 preset: PresetMeta {
353 name: "test".into(),
354 version: "1.0".into(),
355 description: String::new(),
356 },
357 compression: CompressionConfig {
358 stages: vec![],
359 keep_fields: None,
360 strip_fields: None,
361 condense: Some(CondenseConfig {
362 enabled: true,
363 max_repeated_lines: 3,
364 }),
365 git_diff_fold: None,
366 strip_nulls: Some(StripNullsConfig { enabled: true }),
367 flatten: None,
368 truncate_strings: Some(TruncateStringsConfig {
369 enabled: true,
370 max_length: 500,
371 }),
372 collapse_arrays: Some(CollapseArraysConfig {
373 enabled: true,
374 max_items: 5,
375 summary_template: "... and {remaining} more items".into(),
376 }),
377 custom_transforms: Some(CustomTransformsConfig { enabled: true }),
378 },
379 tool_selection: ToolSelectionConfig {
380 max_tools: 5,
381 similarity_threshold: 0.7,
382 default_tools: vec![],
383 },
384 budget: BudgetConfig {
385 warning_threshold: 0.70,
386 ceiling_threshold: 0.85,
387 default_window_size: 200_000,
388 agents: Default::default(),
389 },
390 terse_mode: TerseModeConfig {
391 enabled: false,
392 level: crate::preset::TerseLevel::Moderate,
393 },
394 model: ModelConfig {
395 family: "anthropic".into(),
396 primary: "claude-sonnet-4-20250514".into(),
397 local: String::new(),
398 complexity_threshold: 0.4,
399 pricing: None,
400 },
401 }
402 }
403
404 fn make_pipeline() -> CompressionPipeline {
405 CompressionPipeline::new(&test_preset())
406 }
407
408 #[test]
409 fn first_read_is_miss() {
410 let (store, _dir) = in_memory_store();
411 let cm = CacheManager::new(store, u64::MAX);
412 let pipeline = make_pipeline();
413 let content = b"hello world";
414 let result = cm
415 .get_or_compress(Path::new("file.txt"), content, &pipeline)
416 .unwrap();
417 assert!(matches!(result, CacheResult::Fresh { .. }));
418 }
419
420 #[test]
421 fn second_read_is_hit() {
422 let (store, _dir) = in_memory_store();
423 let cm = CacheManager::new(store, u64::MAX);
424 let pipeline = make_pipeline();
425 let content = b"hello world";
426 let path = Path::new("file.txt");
427
428 cm.get_or_compress(path, content, &pipeline).unwrap();
430
431 let result = cm.get_or_compress(path, content, &pipeline).unwrap();
433 match result {
434 CacheResult::Dedup {
435 inline_ref,
436 token_cost,
437 } => {
438 assert!(inline_ref.starts_with("§ref:"));
439 assert!(inline_ref.ends_with('§'));
440 assert_eq!(token_cost, 13);
441 }
442 CacheResult::Fresh { .. } | CacheResult::Delta { .. } => panic!("expected cache hit"),
443 }
444 }
445
446 #[test]
447 fn different_content_is_miss() {
448 let (store, _dir) = in_memory_store();
449 let cm = CacheManager::new(store, u64::MAX);
450 let pipeline = make_pipeline();
451 let path = Path::new("file.txt");
452
453 cm.get_or_compress(path, b"content v1", &pipeline).unwrap();
454 let result = cm
455 .get_or_compress(path, b"content v2", &pipeline)
456 .unwrap();
457 assert!(matches!(result, CacheResult::Fresh { .. } | CacheResult::Delta { .. }));
458 }
459
460 #[test]
461 fn evict_lru_frees_bytes_when_over_limit() {
462 let (store, _dir) = in_memory_store();
463 let cm = CacheManager::new(store, 1);
465 let pipeline = make_pipeline();
466 let path = Path::new("f.txt");
467
468 cm.get_or_compress(path, b"entry one", &pipeline).unwrap();
470 cm.get_or_compress(path, b"entry two", &pipeline).unwrap();
471 cm.get_or_compress(path, b"entry three", &pipeline).unwrap();
472
473 let freed = cm.evict_lru().unwrap();
474 assert!(freed > 0, "expected bytes to be freed");
475 }
476
477 #[test]
478 fn evict_lru_no_op_when_under_limit() {
479 let (store, _dir) = in_memory_store();
480 let cm = CacheManager::new(store, u64::MAX);
481 let pipeline = make_pipeline();
482
483 cm.get_or_compress(Path::new("f.txt"), b"data", &pipeline)
484 .unwrap();
485
486 let freed = cm.evict_lru().unwrap();
487 assert_eq!(freed, 0);
488 }
489
490 #[test]
491 fn invalidate_removes_entry() {
492 let dir = tempfile::tempdir().unwrap();
493 let file_path = dir.path().join("test.txt");
494 std::fs::write(&file_path, b"some content").unwrap();
495
496 let store_path = dir.path().join("store.db");
497 let store = SessionStore::open_or_create(&store_path).unwrap();
498 let cm = CacheManager::new(store, u64::MAX);
499 let pipeline = make_pipeline();
500
501 let content = std::fs::read(&file_path).unwrap();
503 cm.get_or_compress(&file_path, &content, &pipeline).unwrap();
504
505 let hit = cm
507 .get_or_compress(&file_path, &content, &pipeline)
508 .unwrap();
509 assert!(matches!(hit, CacheResult::Dedup { .. }));
510
511 cm.invalidate(&file_path).unwrap();
512
513 let miss = cm
514 .get_or_compress(&file_path, &content, &pipeline)
515 .unwrap();
516 assert!(matches!(miss, CacheResult::Fresh { .. }));
517 }
518
519 #[test]
520 fn invalidate_nonexistent_path_is_noop() {
521 let (store, _dir) = in_memory_store();
522 let cm = CacheManager::new(store, u64::MAX);
523 cm.invalidate(Path::new("/nonexistent/path/file.txt"))
525 .unwrap();
526 }
527
528 #[test]
531 fn stale_ref_returns_fresh_instead_of_dedup() {
532 let (store, _dir) = in_memory_store();
533 let cm = CacheManager::with_ref_age(store, u64::MAX, 3);
535 let pipeline = make_pipeline();
536 let content = b"hello world";
537 let path = Path::new("file.txt");
538
539 cm.get_or_compress(path, content, &pipeline).unwrap();
541
542 let result = cm.get_or_compress(path, content, &pipeline).unwrap();
544 assert!(matches!(result, CacheResult::Dedup { .. }), "ref should be fresh at turn 0");
545
546 for _ in 0..4 {
548 cm.advance_turn();
549 }
550
551 let result = cm.get_or_compress(path, content, &pipeline).unwrap();
553 assert!(matches!(result, CacheResult::Fresh { .. }),
554 "stale ref should return Fresh, not Dedup");
555 }
556
557 #[test]
558 fn notify_compaction_invalidates_all_refs() {
559 let (store, _dir) = in_memory_store();
560 let cm = CacheManager::new(store, u64::MAX);
561 let pipeline = make_pipeline();
562 let path = Path::new("file.txt");
563
564 cm.get_or_compress(path, b"content A", &pipeline).unwrap();
566 cm.get_or_compress(path, b"content B", &pipeline).unwrap();
567
568 assert!(matches!(
570 cm.get_or_compress(path, b"content A", &pipeline).unwrap(),
571 CacheResult::Dedup { .. }
572 ));
573 assert!(matches!(
574 cm.get_or_compress(path, b"content B", &pipeline).unwrap(),
575 CacheResult::Dedup { .. }
576 ));
577
578 cm.notify_compaction();
580
581 assert!(matches!(
583 cm.get_or_compress(path, b"content A", &pipeline).unwrap(),
584 CacheResult::Fresh { .. }
585 ));
586 assert!(matches!(
587 cm.get_or_compress(path, b"content B", &pipeline).unwrap(),
588 CacheResult::Fresh { .. }
589 ));
590 }
591
592 #[test]
593 fn ref_refreshed_after_resend() {
594 let (store, _dir) = in_memory_store();
595 let cm = CacheManager::with_ref_age(store, u64::MAX, 3);
596 let pipeline = make_pipeline();
597 let content = b"hello world";
598 let path = Path::new("file.txt");
599
600 cm.get_or_compress(path, content, &pipeline).unwrap();
602
603 for _ in 0..4 {
605 cm.advance_turn();
606 }
607
608 let result = cm.get_or_compress(path, content, &pipeline).unwrap();
610 assert!(matches!(result, CacheResult::Fresh { .. }));
611
612 let result = cm.get_or_compress(path, content, &pipeline).unwrap();
614 assert!(matches!(result, CacheResult::Dedup { .. }),
615 "ref should be fresh after re-send");
616 }
617
618 #[test]
619 fn check_dedup_returns_none_for_stale_ref() {
620 let (store, _dir) = in_memory_store();
621 let cm = CacheManager::with_ref_age(store, u64::MAX, 2);
622 let pipeline = make_pipeline();
623 let content = b"test content";
624 let path = Path::new("file.txt");
625
626 cm.get_or_compress(path, content, &pipeline).unwrap();
628
629 assert!(cm.check_dedup(content).unwrap().is_some());
631
632 for _ in 0..3 {
634 cm.advance_turn();
635 }
636
637 assert!(cm.check_dedup(content).unwrap().is_none(),
639 "stale ref should not be returned by check_dedup");
640 }
641
642 #[test]
643 fn advance_turn_increments_counter() {
644 let (store, _dir) = in_memory_store();
645 let cm = CacheManager::new(store, u64::MAX);
646 assert_eq!(cm.current_turn(), 0);
647 cm.advance_turn();
648 assert_eq!(cm.current_turn(), 1);
649 cm.advance_turn();
650 assert_eq!(cm.current_turn(), 2);
651 }
652
653 use proptest::prelude::*;
656
657 proptest! {
665 #[test]
670 fn prop_cache_deduplication(
671 content in proptest::collection::vec(any::<u8>(), 1..=1000usize),
672 ) {
673 let (store, _dir) = in_memory_store();
674 let cm = CacheManager::new(store, u64::MAX);
675 let pipeline = make_pipeline();
676 let path = Path::new("file.txt");
677
678 let first = cm.get_or_compress(path, &content, &pipeline).unwrap();
680 prop_assert!(
681 matches!(first, CacheResult::Fresh { .. }),
682 "first read should be a cache miss"
683 );
684
685 let second = cm.get_or_compress(path, &content, &pipeline).unwrap();
686 match second {
687 CacheResult::Dedup { inline_ref, token_cost } => {
688 prop_assert_eq!(
689 token_cost, 13,
690 "cache hit should report ~13 reference tokens"
691 );
692 prop_assert!(
693 inline_ref.starts_with("§ref:"),
694 "reference token should start with §ref:"
695 );
696 prop_assert!(
697 inline_ref.ends_with('§'),
698 "reference token should end with §"
699 );
700 }
701 CacheResult::Fresh { .. } | CacheResult::Delta { .. } => {
702 prop_assert!(false, "second read should be a cache hit, not a miss");
703 }
704 }
705 }
706 }
707
708 proptest! {
716 #[test]
721 fn prop_cache_invalidation_on_content_change(
722 content_a in proptest::collection::vec(any::<u8>(), 1..=500usize),
723 content_b in proptest::collection::vec(any::<u8>(), 1..=500usize),
724 ) {
725 prop_assume!(content_a != content_b);
727
728 let (store, _dir) = in_memory_store();
729 let cm = CacheManager::new(store, u64::MAX);
730 let pipeline = make_pipeline();
731 let path = Path::new("file.txt");
732
733 let r1 = cm.get_or_compress(path, &content_a, &pipeline).unwrap();
735 prop_assert!(
736 matches!(r1, CacheResult::Fresh { .. }),
737 "first read of content_a should be a miss"
738 );
739
740 let r2 = cm.get_or_compress(path, &content_a, &pipeline).unwrap();
741 prop_assert!(
742 matches!(r2, CacheResult::Dedup { .. }),
743 "second read of content_a should be a hit"
744 );
745
746 let r3 = cm.get_or_compress(path, &content_b, &pipeline).unwrap();
747 prop_assert!(
748 matches!(r3, CacheResult::Fresh { .. } | CacheResult::Delta { .. }),
749 "read with changed content should be a cache miss or delta"
750 );
751 }
752 }
753
754 proptest! {
762 #[test]
767 fn prop_cache_lru_eviction(
768 entries in proptest::collection::vec(
770 proptest::collection::vec(any::<u8>(), 10..=200usize),
771 2..=8usize,
772 ),
773 ) {
774 let mut unique_entries: Vec<Vec<u8>> = Vec::new();
776 for e in &entries {
777 if !unique_entries.contains(e) {
778 unique_entries.push(e.clone());
779 }
780 }
781 prop_assume!(unique_entries.len() >= 2);
782
783 let (store, _dir) = in_memory_store();
784 let cm = CacheManager::new(store, 1);
786 let pipeline = make_pipeline();
787 let path = Path::new("f.txt");
788
789 for entry in &unique_entries {
791 cm.get_or_compress(path, entry, &pipeline).unwrap();
792 }
793
794 let freed = cm.evict_lru().unwrap();
796
797 prop_assert!(freed > 0, "evict_lru should free bytes when over limit");
799
800 let freed_again = cm.evict_lru().unwrap();
803 prop_assert_eq!(
804 freed_again, 0,
805 "second evict_lru call should free 0 bytes (already at or below limit)"
806 );
807 }
808 }
809
810 proptest! {
819 #[test]
827 fn prop_cache_persistence_across_sessions(
828 content in proptest::collection::vec(any::<u8>(), 1..=500usize),
829 ) {
830 use crate::session_store::SessionStore;
831
832 let dir = tempfile::tempdir().unwrap();
833 let db_path = dir.path().join("cache.db");
834 let path = Path::new("file.txt");
835
836 {
838 let store = SessionStore::open_or_create(&db_path).unwrap();
839 let cm = CacheManager::new(store, u64::MAX);
840 let pipeline = make_pipeline();
841
842 let r = cm.get_or_compress(path, &content, &pipeline).unwrap();
843 prop_assert!(
844 matches!(r, CacheResult::Fresh { .. }),
845 "first read should be a miss"
846 );
847 }
848 {
852 let store = SessionStore::open_or_create(&db_path).unwrap();
853 let cm = CacheManager::new(store, u64::MAX);
854 let pipeline = make_pipeline();
855
856 let r1 = cm.get_or_compress(path, &content, &pipeline).unwrap();
860 prop_assert!(
861 matches!(r1, CacheResult::Fresh { .. }),
862 "first read after restart should re-send (ref tracker empty)"
863 );
864
865 let r2 = cm.get_or_compress(path, &content, &pipeline).unwrap();
868 match r2 {
869 CacheResult::Dedup { token_cost, .. } => {
870 prop_assert_eq!(
871 token_cost, 13,
872 "second read in same session should be dedup hit"
873 );
874 }
875 CacheResult::Fresh { .. } | CacheResult::Delta { .. } => {
876 prop_assert!(
877 false,
878 "second read should be a dedup hit after re-send"
879 );
880 }
881 }
882 }
883 }
884 }
885}