use crate::cache::{CacheKey as UnifiedCacheKey, DEFAULT_CACHE_TTL, EvictionPolicy, UnifiedCache};
use serde_json::Value;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub struct ToolCacheKey {
pub tool: String,
pub params_hash: u64,
pub target_path: String,
}
impl UnifiedCacheKey for ToolCacheKey {
fn to_cache_key(&self) -> String {
format!("{}:{}:{}", self.tool, self.params_hash, self.target_path)
}
}
impl ToolCacheKey {
#[inline]
pub fn new(tool: &str, params: &str, target_path: &str) -> Self {
let mut hasher = DefaultHasher::new();
params.hash(&mut hasher);
let params_hash = hasher.finish();
ToolCacheKey {
tool: tool.to_string(),
params_hash,
target_path: target_path.to_string(),
}
}
#[inline]
pub fn from_json(tool: &str, params: &Value, target_path: &str) -> Self {
let mut hasher = DefaultHasher::new();
if let Ok(bytes) = serde_json::to_vec(params) {
bytes.hash(&mut hasher);
} else {
params.to_string().hash(&mut hasher);
}
let params_hash = hasher.finish();
ToolCacheKey {
tool: tool.to_string(),
params_hash,
target_path: target_path.to_string(),
}
}
}
pub struct ToolResultCache {
inner: UnifiedCache<ToolCacheKey, String>,
}
impl ToolResultCache {
pub fn new(capacity: usize) -> Self {
Self {
inner: UnifiedCache::new(capacity, DEFAULT_CACHE_TTL, EvictionPolicy::Lru),
}
}
fn insert_owned(&mut self, key: ToolCacheKey, output: String) {
let size_bytes = size_of_val(&output) as u64;
self.inner.insert(key, output, size_bytes);
}
pub fn insert(&mut self, key: ToolCacheKey, output: String) {
self.insert_owned(key, output);
}
pub fn insert_arc(&mut self, key: ToolCacheKey, output: Arc<String>) {
self.insert_owned(key, (*output).clone());
}
pub fn get(&self, key: &ToolCacheKey) -> Option<Arc<String>> {
self.inner.get(key)
}
pub fn get_owned(&self, key: &ToolCacheKey) -> Option<String> {
self.inner.get_owned(key)
}
pub fn invalidate_for_path(&mut self, path: &str) {
self.inner
.remove_where(|key| key.target_path.starts_with(path));
}
pub fn invalidate_key(&mut self, key: &ToolCacheKey) {
self.inner.remove(key);
}
pub fn invalidate_for_paths<I, S>(&mut self, paths: I)
where
I: IntoIterator<Item = S>,
S: AsRef<str>,
{
let path_prefixes: Vec<String> = paths
.into_iter()
.map(|path| path.as_ref().trim().to_string())
.filter(|path| !path.is_empty())
.collect();
if path_prefixes.is_empty() {
return;
}
self.inner.remove_where(|key| {
path_prefixes
.iter()
.any(|prefix| key.target_path.starts_with(prefix))
});
}
pub fn clear(&mut self) {
self.inner.clear();
}
pub fn check_pressure_and_evict(&mut self) {
if self.inner.total_memory_bytes() > 50 * 1024 * 1024 {
self.inner.evict_under_pressure(30); }
}
pub fn stats(&self) -> crate::cache::CacheStats {
self.inner.stats()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::constants::tools;
#[test]
fn creates_cache_key() {
let key = ToolCacheKey::new(tools::UNIFIED_SEARCH, "pattern=test", "/workspace");
assert_eq!(key.tool, tools::UNIFIED_SEARCH);
assert_eq!(key.target_path, "/workspace");
}
#[test]
fn from_json_and_new_equivalence() {
let params = serde_json::json!({"a": 1, "b": [1,2,3]});
let params_str = serde_json::to_string(¶ms).unwrap();
let k1 = ToolCacheKey::new("tool", ¶ms_str, "/workspace");
let k2 = ToolCacheKey::from_json("tool", ¶ms, "/workspace");
assert_eq!(k1.tool, k2.tool);
assert_eq!(k1.target_path, k2.target_path);
assert_ne!(k1.params_hash, 0);
assert_ne!(k2.params_hash, 0);
}
#[test]
fn caches_and_retrieves_result() {
let mut cache = ToolResultCache::new(10);
let key = ToolCacheKey::new(tools::UNIFIED_SEARCH, "pattern=test", "/workspace");
let output = "line 1\nline 2".to_string();
cache.insert_arc(key.clone(), Arc::new(output.clone()));
assert_eq!(cache.get(&key).as_ref(), Some(&Arc::new(output)));
}
#[test]
fn returns_none_for_missing_key() {
let cache = ToolResultCache::new(10);
let key = ToolCacheKey::new(tools::UNIFIED_SEARCH, "pattern=test", "/workspace");
assert!(cache.get(&key).is_none());
}
#[test]
fn evicts_least_recently_used() {
let mut cache = ToolResultCache::new(3);
let key1 = ToolCacheKey::new("tool", "p1", "/a");
let key2 = ToolCacheKey::new("tool", "p2", "/b");
let key3 = ToolCacheKey::new("tool", "p3", "/c");
let key4 = ToolCacheKey::new("tool", "p4", "/d");
cache.insert(key1.clone(), "out1".to_string());
cache.insert(key2.clone(), "out2".to_string());
cache.insert(key3.clone(), "out3".to_string());
cache.insert(key4.clone(), "out4".to_string());
assert!(cache.get(&key1).is_none());
assert_eq!(cache.get(&key2).unwrap().as_ref(), "out2");
}
#[test]
fn invalidates_by_path() {
let mut cache = ToolResultCache::new(10);
let key1 = ToolCacheKey::new("tool", "p1", "/workspace/file1.rs");
let key2 = ToolCacheKey::new("tool", "p2", "/workspace/file2.rs");
let key3 = ToolCacheKey::new("tool", "p3", "/other/file3.rs");
cache.insert(key1.clone(), "out1".to_string());
cache.insert(key2.clone(), "out2".to_string());
cache.insert(key3.clone(), "out3".to_string());
cache.invalidate_for_path("/workspace/file1.rs");
assert!(cache.get(&key1).is_none());
assert_eq!(cache.get(&key2).unwrap().as_ref(), "out2");
assert_eq!(cache.get(&key3).unwrap().as_ref(), "out3");
}
#[test]
fn invalidates_exact_key_only() {
let mut cache = ToolResultCache::new(10);
let key1 = ToolCacheKey::new("tool", "p1", "/workspace/file.rs");
let key2 = ToolCacheKey::new("tool", "p2", "/workspace/file.rs");
cache.insert(key1.clone(), "out1".to_string());
cache.insert(key2.clone(), "out2".to_string());
cache.invalidate_key(&key1);
assert!(cache.get(&key1).is_none());
assert_eq!(cache.get(&key2).unwrap().as_ref(), "out2");
}
#[test]
fn invalidates_multiple_paths() {
let mut cache = ToolResultCache::new(10);
let key1 = ToolCacheKey::new("tool", "p1", "/workspace/file1.rs");
let key2 = ToolCacheKey::new("tool", "p2", "/workspace/file2.rs");
let key3 = ToolCacheKey::new("tool", "p3", "/workspace/file3.rs");
cache.insert(key1.clone(), "out1".to_string());
cache.insert(key2.clone(), "out2".to_string());
cache.insert(key3.clone(), "out3".to_string());
cache.invalidate_for_paths(["/workspace/file1.rs", "/workspace/file3.rs"]);
assert!(cache.get(&key1).is_none());
assert!(cache.get(&key3).is_none());
assert_eq!(cache.get(&key2).unwrap().as_ref(), "out2");
}
#[test]
fn tracks_access_count() {
let mut cache = ToolResultCache::new(10);
let key = ToolCacheKey::new("tool", "p1", "/a");
cache.insert(key.clone(), "output".to_string());
let initial_stats = cache.stats();
cache.get(&key);
cache.get(&key);
let final_stats = cache.stats();
assert!(final_stats.hits > initial_stats.hits);
}
#[test]
fn clears_cache() {
let mut cache = ToolResultCache::new(10);
let key = ToolCacheKey::new("tool", "p1", "/a");
cache.insert(key.clone(), "output".to_string());
assert_eq!(cache.stats().current_size, 1);
cache.clear();
assert_eq!(cache.stats().current_size, 0);
assert!(cache.get(&key).is_none());
}
#[test]
fn computes_stats() {
let mut cache = ToolResultCache::new(10);
let key1 = ToolCacheKey::new("tool", "p1", "/a");
let key2 = ToolCacheKey::new("tool", "p2", "/b");
cache.insert(key1.clone(), "out1".to_string());
cache.insert(key2.clone(), "out2".to_string());
cache.get(&key1);
cache.get(&key2);
cache.get(&key1);
let stats = cache.stats();
assert_eq!(stats.current_size, 2);
assert_eq!(stats.max_size, 10);
assert_eq!(stats.hits, 3);
assert_eq!(stats.misses, 0); }
#[test]
fn insert_arc_and_get_arc() {
let mut cache = ToolResultCache::new(10);
let key = ToolCacheKey::new("tool", "p1", "/a");
let arc = Arc::new("output".to_string());
cache.insert_arc(key.clone(), Arc::clone(&arc));
assert_eq!(cache.get(&key).unwrap(), arc);
}
#[test]
fn test_granular_cache_invalidation() {
let mut cache = ToolResultCache::new(100);
let key1 = ToolCacheKey::new("grep", "pattern=test", "/workspace/src/main.rs");
let key2 = ToolCacheKey::new("grep", "pattern=test", "/workspace/src/lib.rs");
let key3 = ToolCacheKey::new("list", "recursive=true", "/workspace/src/");
cache.insert(key1.clone(), "result1".to_string());
cache.insert(key2.clone(), "result2".to_string());
cache.insert(key3.clone(), "result3".to_string());
assert_eq!(cache.stats().current_size, 3);
cache.invalidate_for_path("/workspace/src/main.rs");
assert!(cache.get(&key1).is_none(), "Key1 should be removed");
assert!(
cache.get(&key2).is_some(),
"Key2 should still exist (different file)"
);
assert!(
cache.get(&key3).is_some(),
"Key3 should still exist (different tool)"
);
assert_eq!(cache.stats().current_size, 2);
}
#[test]
fn test_invalidate_prefix_removes_only_matched() {
let mut cache = ToolResultCache::new(100);
let key1 = ToolCacheKey::new("grep", "p1", "/workspace/a");
let key2 = ToolCacheKey::new("grep", "p2", "/workspace/b");
let key3 = ToolCacheKey::new("grep", "p3", "/other/c");
cache.insert(key1.clone(), "1".to_string());
cache.insert(key2.clone(), "2".to_string());
cache.insert(key3.clone(), "3".to_string());
cache.invalidate_for_path("/workspace");
assert!(cache.get(&key1).is_none());
assert!(cache.get(&key2).is_none());
assert!(cache.get(&key3).is_some());
}
#[test]
fn test_cache_hit_ratio_preserved_after_selective_invalidation() {
let mut cache = ToolResultCache::new(100);
for i in 0..10 {
let key = ToolCacheKey::new("tool", "params", &format!("/file_{}", i));
cache.insert(key, format!("result_{}", i));
}
let stats_before = cache.stats();
assert_eq!(stats_before.current_size, 10);
for i in 0..5 {
let key = ToolCacheKey::new("tool", "params", &format!("/file_{}", i));
let _ = cache.get(&key);
}
let stats_mid = cache.stats();
let hits_before_invalidation = stats_mid.hits;
cache.invalidate_for_path("/file_0");
for i in 1..5 {
let key = ToolCacheKey::new("tool", "params", &format!("/file_{}", i));
assert!(
cache.get(&key).is_some(),
"Cache for /file_{} should still be valid",
i
);
}
let stats_after = cache.stats();
assert_eq!(stats_after.current_size, 9);
assert!(stats_after.hits > hits_before_invalidation);
}
}