smcp_computer/mcp_clients/
resource_cache.rs1use serde_json::Value;
23use std::collections::HashMap;
24use std::sync::Arc;
25use std::time::{Duration, Instant};
26use tokio::sync::RwLock;
27
28#[derive(Debug, Clone)]
30pub struct CachedResource {
31 pub data: Value,
33 pub cached_at: Instant,
35 pub ttl: Duration,
37 pub version: u64,
39}
40
41impl CachedResource {
42 pub fn new(data: Value, ttl: Duration) -> Self {
44 Self {
45 data,
46 cached_at: Instant::now(),
47 ttl,
48 version: 1,
49 }
50 }
51
52 pub fn is_expired(&self) -> bool {
54 self.cached_at.elapsed() > self.ttl
55 }
56
57 pub fn remaining_ttl(&self) -> Duration {
59 if self.is_expired() {
60 Duration::ZERO
61 } else {
62 self.ttl.saturating_sub(self.cached_at.elapsed())
63 }
64 }
65
66 pub fn refresh(&mut self, new_data: Value) {
68 self.data = new_data;
69 self.cached_at = Instant::now();
70 self.version += 1;
71 }
72}
73
74#[derive(Debug, Clone)]
78pub struct ResourceCache {
79 cache: Arc<RwLock<HashMap<String, CachedResource>>>,
81 default_ttl: Duration,
83}
84
85impl ResourceCache {
86 pub fn new(default_ttl: Duration) -> Self {
91 Self {
92 cache: Arc::new(RwLock::new(HashMap::new())),
93 default_ttl,
94 }
95 }
96
97 pub async fn set(&self, uri: String, data: Value, ttl: Option<Duration>) {
104 let ttl = ttl.unwrap_or(self.default_ttl);
105 let mut cache = self.cache.write().await;
106 cache.insert(uri, CachedResource::new(data, ttl));
107 }
108
109 pub async fn get(&self, uri: &str) -> Option<Value> {
118 let mut cache = self.cache.write().await;
119
120 if let Some(cached) = cache.get(uri) {
121 if !cached.is_expired() {
122 return Some(cached.data.clone());
123 } else {
124 cache.remove(uri);
126 }
127 }
128 None
129 }
130
131 pub async fn refresh(&self, uri: &str, new_data: Value) -> Result<u64, String> {
141 let mut cache = self.cache.write().await;
142
143 if let Some(cached) = cache.get_mut(uri) {
144 cached.refresh(new_data);
145 Ok(cached.version)
146 } else {
147 Err(format!("Resource not cached: {}", uri))
148 }
149 }
150
151 pub async fn remove(&self, uri: &str) -> bool {
160 let mut cache = self.cache.write().await;
161 cache.remove(uri).is_some()
162 }
163
164 pub async fn clear(&self) {
166 let mut cache = self.cache.write().await;
167 cache.clear();
168 }
169
170 pub async fn contains(&self, uri: &str) -> bool {
179 let cache = self.cache.read().await;
180 if let Some(cached) = cache.get(uri) {
181 !cached.is_expired()
182 } else {
183 false
184 }
185 }
186
187 pub async fn get_entry(&self, uri: &str) -> Option<CachedResource> {
196 let cache = self.cache.read().await;
197 cache.get(uri).cloned()
198 }
199
200 pub async fn size(&self) -> usize {
202 let cache = self.cache.read().await;
203 cache.len()
204 }
205
206 pub async fn keys(&self) -> Vec<String> {
208 let cache = self.cache.read().await;
209 cache.keys().cloned().collect()
210 }
211
212 pub async fn cleanup_expired(&self) -> usize {
214 let mut cache = self.cache.write().await;
215 let initial_size = cache.len();
216
217 cache.retain(|_, cached| !cached.is_expired());
218
219 initial_size - cache.len()
220 }
221}
222
223impl Default for ResourceCache {
224 fn default() -> Self {
225 Self::new(Duration::from_secs(60)) }
227}
228
229#[cfg(test)]
230mod tests {
231 use super::*;
232
233 #[tokio::test]
234 async fn test_set_and_get_cache() {
235 let cache = ResourceCache::new(Duration::from_secs(60));
236
237 let data = Value::String("test data".to_string());
238 cache
239 .set("window://test".to_string(), data.clone(), None)
240 .await;
241
242 let retrieved = cache.get("window://test").await;
243 assert!(retrieved.is_some());
244 assert_eq!(retrieved.unwrap(), data);
245 }
246
247 #[tokio::test]
248 async fn test_cache_expiration() {
249 let cache = ResourceCache::new(Duration::from_millis(100)); let data = Value::String("test data".to_string());
252 cache.set("window://test".to_string(), data, None).await;
253
254 assert!(cache.get("window://test").await.is_some());
256
257 tokio::time::sleep(Duration::from_millis(150)).await;
259
260 assert!(cache.get("window://test").await.is_none());
262 }
263
264 #[tokio::test]
265 async fn test_refresh_cache() {
266 let cache = ResourceCache::new(Duration::from_secs(60));
267
268 let data1 = Value::String("version 1".to_string());
269 cache
270 .set("window://test".to_string(), data1.clone(), None)
271 .await;
272
273 let entry = cache.get_entry("window://test").await.unwrap();
274 assert_eq!(entry.version, 1);
275
276 let data2 = Value::String("version 2".to_string());
277 cache.refresh("window://test", data2.clone()).await.unwrap();
278
279 let entry = cache.get_entry("window://test").await.unwrap();
280 assert_eq!(entry.version, 2);
281 assert_eq!(entry.data, data2);
282 }
283
284 #[tokio::test]
285 async fn test_remove_cache() {
286 let cache = ResourceCache::new(Duration::from_secs(60));
287
288 cache
289 .set(
290 "window://test".to_string(),
291 Value::String("test".to_string()),
292 None,
293 )
294 .await;
295
296 assert!(cache.contains("window://test").await);
297
298 cache.remove("window://test").await;
299 assert!(!cache.contains("window://test").await);
300 }
301
302 #[tokio::test]
303 async fn test_cleanup_expired() {
304 let cache = ResourceCache::new(Duration::from_millis(100));
305
306 cache
307 .set(
308 "window://test1".to_string(),
309 Value::String("test1".to_string()),
310 None,
311 )
312 .await;
313 cache
314 .set(
315 "window://test2".to_string(),
316 Value::String("test2".to_string()),
317 Some(Duration::from_secs(60)), )
319 .await;
320
321 tokio::time::sleep(Duration::from_millis(150)).await;
322
323 let cleaned = cache.cleanup_expired().await;
324 assert_eq!(cleaned, 1);
325
326 assert!(!cache.contains("window://test1").await);
327 assert!(cache.contains("window://test2").await);
328 }
329}