1use serde_json::Value;
7use std::collections::hash_map::DefaultHasher;
8use std::collections::HashMap;
9use std::hash::{Hash, Hasher};
10use std::sync::atomic::{AtomicU64, Ordering};
11use std::time::Duration;
12use tokio::sync::Mutex;
13use tokio::time::Instant;
14
15struct CacheEntry {
17 result: Value,
18 inserted_at: Instant,
19 ttl: Duration,
20}
21
22impl CacheEntry {
23 fn is_expired(&self) -> bool {
24 self.inserted_at.elapsed() > self.ttl
25 }
26}
27
28#[derive(Debug, Clone)]
30pub struct CacheStats {
31 pub hits: u64,
32 pub misses: u64,
33 pub entries: usize,
34}
35
36pub struct ResultCache {
38 entries: Mutex<HashMap<String, CacheEntry>>,
39 tool_ttls: Mutex<HashMap<String, Duration>>,
41 hits: AtomicU64,
42 misses: AtomicU64,
43}
44
45impl ResultCache {
46 pub fn new() -> Self {
48 Self {
49 entries: Mutex::new(HashMap::new()),
50 tool_ttls: Mutex::new(HashMap::new()),
51 hits: AtomicU64::new(0),
52 misses: AtomicU64::new(0),
53 }
54 }
55
56 pub async fn enable_caching(&self, tool: &str, ttl_secs: u64) {
58 let mut ttls = self.tool_ttls.lock().await;
59 ttls.insert(tool.to_string(), Duration::from_secs(ttl_secs));
60 }
61
62 pub async fn get(&self, tool: &str, params: &Value) -> Option<Value> {
64 let ttls = self.tool_ttls.lock().await;
65 if !ttls.contains_key(tool) {
66 return None;
67 }
68 drop(ttls);
69
70 let key = cache_key(tool, params);
71 let mut entries = self.entries.lock().await;
72
73 if let Some(entry) = entries.get(&key) {
74 if entry.is_expired() {
75 entries.remove(&key);
76 self.misses.fetch_add(1, Ordering::Relaxed);
77 None
78 } else {
79 self.hits.fetch_add(1, Ordering::Relaxed);
80 Some(entry.result.clone())
81 }
82 } else {
83 self.misses.fetch_add(1, Ordering::Relaxed);
84 None
85 }
86 }
87
88 pub async fn put(&self, tool: &str, params: &Value, result: Value) {
90 let ttls = self.tool_ttls.lock().await;
91 let ttl = match ttls.get(tool) {
92 Some(ttl) => *ttl,
93 None => return,
94 };
95 drop(ttls);
96
97 let key = cache_key(tool, params);
98 let mut entries = self.entries.lock().await;
99 entries.insert(
100 key,
101 CacheEntry {
102 result,
103 inserted_at: Instant::now(),
104 ttl,
105 },
106 );
107 }
108
109 pub async fn invalidate(&self, tool: &str) {
111 let prefix = format!("{}:", tool);
112 let mut entries = self.entries.lock().await;
113 entries.retain(|k, _| !k.starts_with(&prefix));
114 }
115
116 pub async fn invalidate_all(&self) {
118 let mut entries = self.entries.lock().await;
119 entries.clear();
120 }
121
122 pub async fn stats(&self) -> CacheStats {
124 let entries = self.entries.lock().await;
125 CacheStats {
126 hits: self.hits.load(Ordering::Relaxed),
127 misses: self.misses.load(Ordering::Relaxed),
128 entries: entries.len(),
129 }
130 }
131}
132
133impl Default for ResultCache {
134 fn default() -> Self {
135 Self::new()
136 }
137}
138
139fn cache_key(tool: &str, params: &Value) -> String {
141 let serialized = serde_json::to_string(params).unwrap_or_default();
142 let mut hasher = DefaultHasher::new();
143 serialized.hash(&mut hasher);
144 let hash = hasher.finish();
145 format!("{}:{:x}", tool, hash)
146}
147
148#[cfg(test)]
149mod tests {
150 use super::*;
151 use serde_json::json;
152
153 #[tokio::test]
154 async fn test_cache_hit_returns_stored_result() {
155 let cache = ResultCache::new();
156 cache.enable_caching("add", 60).await;
157
158 let params = json!({"a": 1, "b": 2});
159 let result = json!(3);
160
161 cache.put("add", ¶ms, result.clone()).await;
162
163 let cached = cache.get("add", ¶ms).await;
164 assert_eq!(cached, Some(result));
165
166 let stats = cache.stats().await;
167 assert_eq!(stats.hits, 1);
168 assert_eq!(stats.misses, 0);
169 assert_eq!(stats.entries, 1);
170 }
171
172 #[tokio::test(start_paused = true)]
173 async fn test_expired_entries_return_none() {
174 let cache = ResultCache::new();
175 cache.enable_caching("add", 1).await;
176
177 let params = json!({"a": 1, "b": 2});
178 cache.put("add", ¶ms, json!(3)).await;
179 assert_eq!(cache.get("add", ¶ms).await, Some(json!(3)));
180
181 tokio::time::advance(Duration::from_millis(1_001)).await;
182 assert_eq!(cache.get("add", ¶ms).await, None);
183
184 let stats = cache.stats().await;
185 assert_eq!(stats.hits, 1);
186 assert_eq!(stats.misses, 1);
187 }
188
189 #[tokio::test]
190 async fn test_different_params_produce_different_keys() {
191 let cache = ResultCache::new();
192 cache.enable_caching("add", 60).await;
193
194 let params_a = json!({"a": 1, "b": 2});
195 let params_b = json!({"a": 3, "b": 4});
196
197 cache.put("add", ¶ms_a, json!(3)).await;
198 cache.put("add", ¶ms_b, json!(7)).await;
199
200 assert_eq!(cache.get("add", ¶ms_a).await, Some(json!(3)));
201 assert_eq!(cache.get("add", ¶ms_b).await, Some(json!(7)));
202
203 let stats = cache.stats().await;
204 assert_eq!(stats.entries, 2);
205 }
206
207 #[tokio::test]
208 async fn test_invalidate_clears_tool_entries() {
209 let cache = ResultCache::new();
210 cache.enable_caching("add", 60).await;
211 cache.enable_caching("echo", 60).await;
212
213 cache.put("add", &json!({"a": 1}), json!(1)).await;
214 cache.put("echo", &json!({"msg": "hi"}), json!("hi")).await;
215
216 cache.invalidate("add").await;
217
218 assert_eq!(cache.get("add", &json!({"a": 1})).await, None);
219 assert_eq!(
220 cache.get("echo", &json!({"msg": "hi"})).await,
221 Some(json!("hi"))
222 );
223 }
224
225 #[tokio::test]
226 async fn test_invalidate_all_clears_everything() {
227 let cache = ResultCache::new();
228 cache.enable_caching("add", 60).await;
229 cache.enable_caching("echo", 60).await;
230
231 cache.put("add", &json!({"a": 1}), json!(1)).await;
232 cache.put("echo", &json!({"msg": "hi"}), json!("hi")).await;
233
234 cache.invalidate_all().await;
235
236 let stats = cache.stats().await;
237 assert_eq!(stats.entries, 0);
238 }
239
240 #[tokio::test]
241 async fn test_uncacheable_tool_returns_none() {
242 let cache = ResultCache::new();
243 cache.put("add", &json!({"a": 1}), json!(1)).await;
245 assert_eq!(cache.get("add", &json!({"a": 1})).await, None);
246 }
247}