mcp_proxy/
cache.rs

1//! Response caching middleware for the proxy.
2//!
3//! Caches `ReadResource` and `CallTool` responses with per-backend TTL.
4//! Cache keys are derived from the request type, name/URI, arguments, and MCP
5//! continuation state.
6//!
7//! # Cache Backends
8//!
9//! The cache backend is configurable via `[cache]` in the proxy config:
10//!
11//! - `"memory"` (default): In-process moka cache. Fast, no external deps.
12//! - `"redis"`: External Redis cache. Shared across instances. Requires the
13//!   `redis-cache` feature flag.
14//! - `"sqlite"`: Local SQLite cache. Persistent across restarts. Requires the
15//!   `sqlite-cache` feature flag.
16//!
17//! # Per-Backend Configuration
18//!
19//! ```toml
20//! [[backends]]
21//! name = "slow-api"
22//! transport = "http"
23//! url = "http://localhost:8080"
24//!
25//! [backends.cache]
26//! resource_ttl_seconds = 300
27//! tool_ttl_seconds = 60
28//! max_entries = 1000
29//! ```
30
31use std::convert::Infallible;
32use std::future::Future;
33use std::pin::Pin;
34use std::sync::Arc;
35use std::sync::atomic::{AtomicU64, Ordering};
36use std::task::{Context, Poll};
37use std::time::Duration;
38
39use moka::future::Cache;
40use serde::Serialize;
41use tower::{Layer, Service};
42use tower_mcp::router::{RouterRequest, RouterResponse};
43use tower_mcp_types::protocol::{InputResponses, McpRequest};
44
45use crate::config::{BackendCacheConfig, CacheBackendConfig};
46
47/// Pluggable cache storage backend.
48///
49/// Each variant provides the same logical operations (get, insert, invalidate,
50/// count) but differs in where entries are stored:
51///
52/// - [`Memory`](CacheStore::Memory): in-process moka cache (default)
53/// - [`Redis`](CacheStore::Redis): external Redis server (requires `redis-cache` feature)
54/// - [`Sqlite`](CacheStore::Sqlite): local SQLite database (requires `sqlite-cache` feature)
55#[derive(Clone)]
56pub(crate) enum CacheStore {
57    /// In-process moka cache.
58    Memory(Cache<String, RouterResponse>),
59    /// Redis-backed cache.
60    #[cfg(feature = "redis-cache")]
61    Redis {
62        client: redis::Client,
63        prefix: String,
64        ttl: Duration,
65    },
66    /// SQLite-backed cache.
67    #[cfg(feature = "sqlite-cache")]
68    Sqlite {
69        conn: Arc<std::sync::Mutex<rusqlite::Connection>>,
70        ttl: Duration,
71    },
72}
73
74impl CacheStore {
75    /// Retrieve a cached response by key.
76    async fn get(&self, key: &str) -> Option<RouterResponse> {
77        match self {
78            CacheStore::Memory(cache) => cache.get(key).await,
79            #[cfg(feature = "redis-cache")]
80            CacheStore::Redis {
81                client,
82                prefix,
83                ttl: _,
84            } => {
85                let full_key = format!("{prefix}{key}");
86                let mut conn = client.get_multiplexed_async_connection().await.ok()?;
87                let data: Option<String> =
88                    redis::AsyncCommands::get(&mut conn, &full_key).await.ok()?;
89                data.and_then(|s| serde_json::from_str(&s).ok())
90            }
91            #[cfg(feature = "sqlite-cache")]
92            CacheStore::Sqlite { conn, ttl: _ } => {
93                let key = key.to_string();
94                let conn = conn.lock().ok()?;
95                let now = std::time::SystemTime::now()
96                    .duration_since(std::time::UNIX_EPOCH)
97                    .unwrap_or_default()
98                    .as_secs() as i64;
99                let result: Option<String> = conn
100                    .query_row(
101                        "SELECT value FROM cache_entries WHERE key = ?1 AND expires_at > ?2",
102                        rusqlite::params![key, now],
103                        |row| row.get(0),
104                    )
105                    .ok();
106                result.and_then(|s| serde_json::from_str(&s).ok())
107            }
108        }
109    }
110
111    /// Insert a response into the cache.
112    async fn insert(&self, key: String, value: RouterResponse) {
113        match self {
114            CacheStore::Memory(cache) => {
115                cache.insert(key, value).await;
116            }
117            #[cfg(feature = "redis-cache")]
118            CacheStore::Redis {
119                client,
120                prefix,
121                ttl,
122            } => {
123                let full_key = format!("{prefix}{key}");
124                if let Ok(json) = serde_json::to_string(&value)
125                    && let Ok(mut conn) = client.get_multiplexed_async_connection().await
126                {
127                    let _: Result<(), _> =
128                        redis::AsyncCommands::set_ex(&mut conn, &full_key, &json, ttl.as_secs())
129                            .await;
130                }
131            }
132            #[cfg(feature = "sqlite-cache")]
133            CacheStore::Sqlite { conn, ttl } => {
134                if let Ok(json) = serde_json::to_string(&value) {
135                    let expires_at = std::time::SystemTime::now()
136                        .duration_since(std::time::UNIX_EPOCH)
137                        .unwrap_or_default()
138                        .as_secs() as i64
139                        + ttl.as_secs() as i64;
140                    if let Ok(conn) = conn.lock() {
141                        let _ = conn.execute(
142                            "INSERT OR REPLACE INTO cache_entries (key, value, expires_at) VALUES (?1, ?2, ?3)",
143                            rusqlite::params![key, json, expires_at],
144                        );
145                    }
146                }
147            }
148        }
149    }
150
151    /// Remove all entries from the cache.
152    async fn invalidate_all(&self) {
153        match self {
154            CacheStore::Memory(cache) => {
155                cache.invalidate_all();
156            }
157            #[cfg(feature = "redis-cache")]
158            CacheStore::Redis {
159                client,
160                prefix,
161                ttl: _,
162            } => {
163                if let Ok(mut conn) = client.get_multiplexed_async_connection().await {
164                    let pattern = format!("{prefix}*");
165                    let keys: Vec<String> = redis::AsyncCommands::keys(&mut conn, &pattern)
166                        .await
167                        .unwrap_or_default();
168                    if !keys.is_empty() {
169                        let _: Result<(), _> = redis::AsyncCommands::del(&mut conn, &keys).await;
170                    }
171                }
172            }
173            #[cfg(feature = "sqlite-cache")]
174            CacheStore::Sqlite { conn, ttl: _ } => {
175                if let Ok(conn) = conn.lock() {
176                    let _ = conn.execute("DELETE FROM cache_entries", []);
177                }
178            }
179        }
180    }
181
182    /// Return the approximate number of entries in the cache.
183    async fn entry_count(&self) -> u64 {
184        match self {
185            CacheStore::Memory(cache) => cache.entry_count(),
186            #[cfg(feature = "redis-cache")]
187            CacheStore::Redis {
188                client,
189                prefix,
190                ttl: _,
191            } => {
192                if let Ok(mut conn) = client.get_multiplexed_async_connection().await {
193                    let pattern = format!("{prefix}*");
194                    let keys: Vec<String> = redis::AsyncCommands::keys(&mut conn, &pattern)
195                        .await
196                        .unwrap_or_default();
197                    keys.len() as u64
198                } else {
199                    0
200                }
201            }
202            #[cfg(feature = "sqlite-cache")]
203            CacheStore::Sqlite { conn, ttl: _ } => {
204                let now = std::time::SystemTime::now()
205                    .duration_since(std::time::UNIX_EPOCH)
206                    .unwrap_or_default()
207                    .as_secs() as i64;
208                if let Ok(conn) = conn.lock() {
209                    conn.query_row(
210                        "SELECT COUNT(*) FROM cache_entries WHERE expires_at > ?1",
211                        rusqlite::params![now],
212                        |row| row.get::<_, i64>(0),
213                    )
214                    .unwrap_or(0) as u64
215                } else {
216                    0
217                }
218            }
219        }
220    }
221}
222
223/// Build a [`CacheStore`] from the global cache backend configuration and
224/// a per-backend TTL.
225fn build_cache_store(
226    backend_config: &CacheBackendConfig,
227    ttl: Duration,
228    max_entries: u64,
229) -> CacheStore {
230    match backend_config.backend.as_str() {
231        #[cfg(feature = "redis-cache")]
232        "redis" => {
233            let url = backend_config.url.as_deref().unwrap_or("redis://127.0.0.1");
234            let client =
235                redis::Client::open(url).expect("invalid Redis URL in cache configuration");
236            CacheStore::Redis {
237                client,
238                prefix: backend_config.prefix.clone(),
239                ttl,
240            }
241        }
242        #[cfg(feature = "sqlite-cache")]
243        "sqlite" => {
244            let path = backend_config.url.as_deref().unwrap_or("cache.db");
245            let conn =
246                rusqlite::Connection::open(path).expect("failed to open SQLite cache database");
247            conn.execute_batch(
248                "CREATE TABLE IF NOT EXISTS cache_entries (
249                    key TEXT PRIMARY KEY,
250                    value TEXT NOT NULL,
251                    expires_at INTEGER NOT NULL
252                )",
253            )
254            .expect("failed to create SQLite cache table");
255            CacheStore::Sqlite {
256                conn: Arc::new(std::sync::Mutex::new(conn)),
257                ttl,
258            }
259        }
260        // Default: memory backend (also handles "memory" explicitly)
261        _ => CacheStore::Memory(
262            Cache::builder()
263                .max_capacity(max_entries)
264                .time_to_live(ttl)
265                .build(),
266        ),
267    }
268}
269
270/// Per-backend cache with separate resource and tool caches (different TTLs).
271#[derive(Clone)]
272struct BackendCache {
273    namespace: String,
274    resource_cache: Option<CacheStore>,
275    tool_cache: Option<CacheStore>,
276    stats: Arc<CacheStats>,
277}
278
279/// Atomic hit/miss counters for a backend cache.
280struct CacheStats {
281    hits: AtomicU64,
282    misses: AtomicU64,
283}
284
285impl CacheStats {
286    fn new() -> Self {
287        Self {
288            hits: AtomicU64::new(0),
289            misses: AtomicU64::new(0),
290        }
291    }
292}
293
294/// Snapshot of cache statistics for a single backend.
295///
296/// Returned by [`CacheHandle::stats()`] to report hit/miss rates
297/// and entry counts per cached namespace.
298#[derive(Serialize, Clone)]
299#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
300pub struct CacheStatsSnapshot {
301    /// Backend namespace this cache covers.
302    pub namespace: String,
303    /// Total cache hits.
304    pub hits: u64,
305    /// Total cache misses.
306    pub misses: u64,
307    /// Hit rate as a fraction (0.0-1.0).
308    pub hit_rate: f64,
309    /// Current number of cached entries.
310    pub entry_count: u64,
311}
312
313/// Shared handle for querying cache stats and clearing caches.
314#[derive(Clone)]
315pub struct CacheHandle {
316    caches: Arc<Vec<BackendCache>>,
317}
318
319impl CacheHandle {
320    /// Get a snapshot of cache statistics for all backends.
321    pub async fn stats(&self) -> Vec<CacheStatsSnapshot> {
322        let mut snapshots = Vec::with_capacity(self.caches.len());
323        for bc in self.caches.iter() {
324            let hits = bc.stats.hits.load(Ordering::Relaxed);
325            let misses = bc.stats.misses.load(Ordering::Relaxed);
326            let total = hits + misses;
327            let resource_count = match &bc.resource_cache {
328                Some(store) => store.entry_count().await,
329                None => 0,
330            };
331            let tool_count = match &bc.tool_cache {
332                Some(store) => store.entry_count().await,
333                None => 0,
334            };
335            snapshots.push(CacheStatsSnapshot {
336                namespace: bc.namespace.clone(),
337                hits,
338                misses,
339                hit_rate: if total > 0 {
340                    hits as f64 / total as f64
341                } else {
342                    0.0
343                },
344                entry_count: resource_count + tool_count,
345            });
346        }
347        snapshots
348    }
349
350    /// Clear all cache entries and reset stats.
351    pub async fn clear(&self) {
352        for bc in self.caches.iter() {
353            if let Some(store) = &bc.resource_cache {
354                store.invalidate_all().await;
355            }
356            if let Some(store) = &bc.tool_cache {
357                store.invalidate_all().await;
358            }
359            bc.stats.hits.store(0, Ordering::Relaxed);
360            bc.stats.misses.store(0, Ordering::Relaxed);
361        }
362    }
363}
364
365/// Build the shared cache state from per-backend configs.
366///
367/// Returns an `Arc<Vec<BackendCache>>` that can be shared between a
368/// [`CacheLayer`] (or [`CacheService`]) and its [`CacheHandle`].
369fn build_caches(
370    configs: Vec<(String, &BackendCacheConfig)>,
371    backend_config: &CacheBackendConfig,
372) -> Arc<Vec<BackendCache>> {
373    let caches: Vec<BackendCache> = configs
374        .into_iter()
375        .map(|(namespace, cfg)| {
376            let resource_cache = if cfg.resource_ttl_seconds > 0 {
377                Some(build_cache_store(
378                    backend_config,
379                    Duration::from_secs(cfg.resource_ttl_seconds),
380                    cfg.max_entries,
381                ))
382            } else {
383                None
384            };
385            let tool_cache = if cfg.tool_ttl_seconds > 0 {
386                Some(build_cache_store(
387                    backend_config,
388                    Duration::from_secs(cfg.tool_ttl_seconds),
389                    cfg.max_entries,
390                ))
391            } else {
392                None
393            };
394            BackendCache {
395                namespace,
396                resource_cache,
397                tool_cache,
398                stats: Arc::new(CacheStats::new()),
399            }
400        })
401        .collect();
402    Arc::new(caches)
403}
404
405/// Tower [`Layer`] that produces [`CacheService`] instances sharing the same
406/// cache state and [`CacheHandle`].
407///
408/// Because `CacheService::new()` returns a `(CacheService, CacheHandle)` tuple,
409/// a standard `Layer` cannot propagate the side-channel handle. `CacheLayer`
410/// solves this by creating the shared cache state up-front and handing out an
411/// `Arc`-cloned handle to the caller while cloning the same `Arc` into every
412/// service produced by [`Layer::layer`].
413///
414/// # Example
415///
416/// ```rust
417/// use mcp_proxy::cache::{CacheLayer, CacheHandle};
418/// use mcp_proxy::config::{BackendCacheConfig, CacheBackendConfig};
419///
420/// let cfg = BackendCacheConfig {
421///     resource_ttl_seconds: 300,
422///     tool_ttl_seconds: 60,
423///     max_entries: 1000,
424/// };
425/// let backend_cfg = CacheBackendConfig::default();
426///
427/// let (layer, handle) = CacheLayer::new(
428///     vec![("api/".to_string(), &cfg)],
429///     &backend_cfg,
430/// );
431///
432/// // `layer` implements `tower::Layer<S>` and can be used in a middleware stack.
433/// // `handle` can be used to query stats or clear the cache.
434/// ```
435#[derive(Clone)]
436pub struct CacheLayer {
437    caches: Arc<Vec<BackendCache>>,
438}
439
440impl CacheLayer {
441    /// Create a new cache layer and return it with a shareable [`CacheHandle`].
442    ///
443    /// The handle provides [`CacheHandle::stats()`] and [`CacheHandle::clear()`]
444    /// over the same underlying cache state used by every service the layer
445    /// produces.
446    pub fn new(
447        configs: Vec<(String, &BackendCacheConfig)>,
448        backend_config: &CacheBackendConfig,
449    ) -> (Self, CacheHandle) {
450        let caches = build_caches(configs, backend_config);
451        let handle = CacheHandle {
452            caches: Arc::clone(&caches),
453        };
454        (Self { caches }, handle)
455    }
456}
457
458impl<S> Layer<S> for CacheLayer {
459    type Service = CacheService<S>;
460
461    fn layer(&self, inner: S) -> Self::Service {
462        CacheService {
463            inner,
464            caches: Arc::clone(&self.caches),
465        }
466    }
467}
468
469/// Tower service that caches resource reads and tool call results.
470#[derive(Clone)]
471pub struct CacheService<S> {
472    inner: S,
473    caches: Arc<Vec<BackendCache>>,
474}
475
476impl<S> CacheService<S> {
477    /// Create a new cache service and return it with a shareable handle.
478    pub fn new(
479        inner: S,
480        configs: Vec<(String, &BackendCacheConfig)>,
481        backend_config: &CacheBackendConfig,
482    ) -> (Self, CacheHandle) {
483        let caches = build_caches(configs, backend_config);
484        let handle = CacheHandle {
485            caches: Arc::clone(&caches),
486        };
487        (Self { inner, caches }, handle)
488    }
489}
490
491/// Extract cache key and find the matching backend cache + stats.
492fn continuation_identity(
493    input_responses: &Option<InputResponses>,
494    request_state: &Option<String>,
495) -> String {
496    serde_json::to_string(&(input_responses, request_state)).unwrap_or_default()
497}
498
499fn resolve_cache<'a>(
500    caches: &'a [BackendCache],
501    req: &McpRequest,
502) -> Option<(&'a CacheStore, String, &'a Arc<CacheStats>)> {
503    match req {
504        McpRequest::ReadResource(params) => {
505            let continuation =
506                continuation_identity(&params.input_responses, &params.request_state);
507            let key = format!("res:{}:{continuation}", params.uri);
508            for bc in caches {
509                if params.uri.starts_with(&bc.namespace) {
510                    return bc.resource_cache.as_ref().map(|c| (c, key, &bc.stats));
511                }
512            }
513            None
514        }
515        McpRequest::CallTool(params) => {
516            let args = serde_json::to_string(&params.arguments).unwrap_or_default();
517            let continuation =
518                continuation_identity(&params.input_responses, &params.request_state);
519            let key = format!("tool:{}:{args}:{continuation}", params.name);
520            for bc in caches {
521                if params.name.starts_with(&bc.namespace) {
522                    return bc.tool_cache.as_ref().map(|c| (c, key, &bc.stats));
523                }
524            }
525            None
526        }
527        _ => None,
528    }
529}
530
531impl<S> Service<RouterRequest> for CacheService<S>
532where
533    S: Service<RouterRequest, Response = RouterResponse, Error = Infallible>
534        + Clone
535        + Send
536        + 'static,
537    S::Future: Send,
538{
539    type Response = RouterResponse;
540    type Error = Infallible;
541    type Future = Pin<Box<dyn Future<Output = Result<RouterResponse, Infallible>> + Send>>;
542
543    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
544        self.inner.poll_ready(cx)
545    }
546
547    fn call(&mut self, req: RouterRequest) -> Self::Future {
548        let caches = Arc::clone(&self.caches);
549
550        if let Some((store, key, stats)) = resolve_cache(&caches, &req.inner) {
551            let store = store.clone();
552            let stats = Arc::clone(stats);
553            let mut inner = self.inner.clone();
554
555            return Box::pin(async move {
556                // Cache hit -- return with current request ID
557                if let Some(cached) = store.get(&key).await {
558                    stats.hits.fetch_add(1, Ordering::Relaxed);
559                    return Ok(RouterResponse {
560                        id: req.id,
561                        inner: cached.inner,
562                    });
563                }
564
565                stats.misses.fetch_add(1, Ordering::Relaxed);
566                let result = inner.call(req).await;
567
568                // Only cache successful MCP responses
569                let Ok(ref resp) = result;
570                if resp.inner.is_ok() {
571                    store.insert(key, resp.clone()).await;
572                }
573
574                result
575            });
576        }
577
578        // No caching for this request type or backend
579        let fut = self.inner.call(req);
580        Box::pin(fut)
581    }
582}
583
584#[cfg(test)]
585mod tests {
586    use tower_mcp::protocol::{McpRequest, McpResponse};
587
588    use super::CacheService;
589    use crate::config::{BackendCacheConfig, CacheBackendConfig};
590    use crate::test_util::{MockService, call_service};
591
592    fn tool_call_with_state(name: &str, request_state: Option<&str>) -> McpRequest {
593        McpRequest::CallTool(tower_mcp::protocol::CallToolParams {
594            name: name.to_string(),
595            arguments: serde_json::json!({"key": "value"}),
596            input_responses: None,
597            request_state: request_state.map(str::to_owned),
598            meta: None,
599            task: None,
600        })
601    }
602
603    fn tool_call(name: &str) -> McpRequest {
604        tool_call_with_state(name, None)
605    }
606
607    fn default_backend_config() -> CacheBackendConfig {
608        CacheBackendConfig::default()
609    }
610
611    #[tokio::test]
612    async fn test_cache_hit_returns_same_result() {
613        let mock = MockService::with_tools(&["fs/read"]);
614        let cfg = BackendCacheConfig {
615            resource_ttl_seconds: 60,
616            tool_ttl_seconds: 60,
617            max_entries: 100,
618        };
619        let (mut svc, _handle) = CacheService::new(
620            mock,
621            vec![("fs/".to_string(), &cfg)],
622            &default_backend_config(),
623        );
624
625        let resp1 = call_service(&mut svc, tool_call("fs/read")).await;
626        let resp2 = call_service(&mut svc, tool_call("fs/read")).await;
627
628        // Both should succeed with same content
629        match (resp1.inner.unwrap(), resp2.inner.unwrap()) {
630            (McpResponse::CallTool(r1), McpResponse::CallTool(r2)) => {
631                assert_eq!(r1.all_text(), r2.all_text());
632            }
633            _ => panic!("expected CallTool responses"),
634        }
635    }
636
637    #[tokio::test]
638    async fn test_cache_disabled_passes_through() {
639        let mock = MockService::with_tools(&["fs/read"]);
640        let cfg = BackendCacheConfig {
641            resource_ttl_seconds: 0,
642            tool_ttl_seconds: 0,
643            max_entries: 100,
644        };
645        let (mut svc, _handle) = CacheService::new(
646            mock,
647            vec![("fs/".to_string(), &cfg)],
648            &default_backend_config(),
649        );
650
651        let resp = call_service(&mut svc, tool_call("fs/read")).await;
652        assert!(resp.inner.is_ok());
653    }
654
655    #[tokio::test]
656    async fn test_cache_non_matching_namespace_passes_through() {
657        let mock = MockService::with_tools(&["db/query"]);
658        let cfg = BackendCacheConfig {
659            resource_ttl_seconds: 60,
660            tool_ttl_seconds: 60,
661            max_entries: 100,
662        };
663        let (mut svc, _handle) = CacheService::new(
664            mock,
665            vec![("fs/".to_string(), &cfg)],
666            &default_backend_config(),
667        );
668
669        let resp = call_service(&mut svc, tool_call("db/query")).await;
670        assert!(resp.inner.is_ok());
671    }
672
673    #[tokio::test]
674    async fn test_cache_list_tools_not_cached() {
675        let mock = MockService::with_tools(&["fs/read"]);
676        let cfg = BackendCacheConfig {
677            resource_ttl_seconds: 60,
678            tool_ttl_seconds: 60,
679            max_entries: 100,
680        };
681        let (mut svc, _handle) = CacheService::new(
682            mock,
683            vec![("fs/".to_string(), &cfg)],
684            &default_backend_config(),
685        );
686
687        let resp = call_service(&mut svc, McpRequest::ListTools(Default::default())).await;
688        assert!(resp.inner.is_ok(), "list_tools should pass through");
689    }
690
691    #[tokio::test]
692    async fn test_cache_stats_tracks_hits_and_misses() {
693        let mock = MockService::with_tools(&["fs/read"]);
694        let cfg = BackendCacheConfig {
695            resource_ttl_seconds: 60,
696            tool_ttl_seconds: 60,
697            max_entries: 100,
698        };
699        let (mut svc, handle) = CacheService::new(
700            mock,
701            vec![("fs/".to_string(), &cfg)],
702            &default_backend_config(),
703        );
704
705        // First call = miss
706        let _ = call_service(&mut svc, tool_call("fs/read")).await;
707        let stats = handle.stats().await;
708        assert_eq!(stats.len(), 1);
709        assert_eq!(stats[0].hits, 0);
710        assert_eq!(stats[0].misses, 1);
711
712        // Second call = hit
713        let _ = call_service(&mut svc, tool_call("fs/read")).await;
714        let stats = handle.stats().await;
715        assert_eq!(stats[0].hits, 1);
716        assert_eq!(stats[0].misses, 1);
717        assert!((stats[0].hit_rate - 0.5).abs() < f64::EPSILON);
718    }
719
720    #[tokio::test]
721    async fn test_cache_separates_continuation_state() {
722        let mock = MockService::with_tools(&["fs/read"]);
723        let cfg = BackendCacheConfig {
724            resource_ttl_seconds: 60,
725            tool_ttl_seconds: 60,
726            max_entries: 100,
727        };
728        let (mut svc, handle) = CacheService::new(
729            mock,
730            vec![("fs/".to_string(), &cfg)],
731            &default_backend_config(),
732        );
733
734        let _ = call_service(&mut svc, tool_call_with_state("fs/read", Some("first"))).await;
735        let _ = call_service(&mut svc, tool_call_with_state("fs/read", Some("second"))).await;
736        let _ = call_service(&mut svc, tool_call_with_state("fs/read", Some("first"))).await;
737
738        let stats = handle.stats().await;
739        assert_eq!(stats[0].misses, 2);
740        assert_eq!(stats[0].hits, 1);
741    }
742
743    #[tokio::test]
744    async fn test_cache_clear_resets_stats() {
745        let mock = MockService::with_tools(&["fs/read"]);
746        let cfg = BackendCacheConfig {
747            resource_ttl_seconds: 60,
748            tool_ttl_seconds: 60,
749            max_entries: 100,
750        };
751        let (mut svc, handle) = CacheService::new(
752            mock,
753            vec![("fs/".to_string(), &cfg)],
754            &default_backend_config(),
755        );
756
757        let _ = call_service(&mut svc, tool_call("fs/read")).await;
758        let _ = call_service(&mut svc, tool_call("fs/read")).await;
759
760        handle.clear().await;
761        let stats = handle.stats().await;
762        assert_eq!(stats[0].hits, 0);
763        assert_eq!(stats[0].misses, 0);
764    }
765
766    #[tokio::test]
767    async fn test_cache_layer_produces_working_service() {
768        use super::CacheLayer;
769        use tower::Layer;
770
771        let cfg = BackendCacheConfig {
772            resource_ttl_seconds: 60,
773            tool_ttl_seconds: 60,
774            max_entries: 100,
775        };
776        let (layer, handle) =
777            CacheLayer::new(vec![("fs/".to_string(), &cfg)], &default_backend_config());
778
779        let mock = MockService::with_tools(&["fs/read"]);
780        let mut svc = layer.layer(mock);
781
782        // First call = miss
783        let _ = call_service(&mut svc, tool_call("fs/read")).await;
784        let stats = handle.stats().await;
785        assert_eq!(stats[0].misses, 1);
786        assert_eq!(stats[0].hits, 0);
787
788        // Second call = hit (cached)
789        let _ = call_service(&mut svc, tool_call("fs/read")).await;
790        let stats = handle.stats().await;
791        assert_eq!(stats[0].hits, 1);
792        assert_eq!(stats[0].misses, 1);
793    }
794
795    #[tokio::test]
796    async fn test_cache_layer_shares_state_across_services() {
797        use super::CacheLayer;
798        use tower::Layer;
799
800        let cfg = BackendCacheConfig {
801            resource_ttl_seconds: 60,
802            tool_ttl_seconds: 60,
803            max_entries: 100,
804        };
805        let (layer, handle) =
806            CacheLayer::new(vec![("fs/".to_string(), &cfg)], &default_backend_config());
807
808        // Create two services from the same layer
809        let mock1 = MockService::with_tools(&["fs/read"]);
810        let mut svc1 = layer.layer(mock1);
811
812        let mock2 = MockService::with_tools(&["fs/read"]);
813        let mut svc2 = layer.layer(mock2);
814
815        // Miss on svc1
816        let _ = call_service(&mut svc1, tool_call("fs/read")).await;
817        assert_eq!(handle.stats().await[0].misses, 1);
818
819        // Hit on svc2 (same underlying cache)
820        let _ = call_service(&mut svc2, tool_call("fs/read")).await;
821        assert_eq!(handle.stats().await[0].hits, 1);
822        assert_eq!(handle.stats().await[0].misses, 1);
823    }
824
825    #[tokio::test]
826    async fn test_cache_layer_handle_clear() {
827        use super::CacheLayer;
828        use tower::Layer;
829
830        let cfg = BackendCacheConfig {
831            resource_ttl_seconds: 60,
832            tool_ttl_seconds: 60,
833            max_entries: 100,
834        };
835        let (layer, handle) =
836            CacheLayer::new(vec![("fs/".to_string(), &cfg)], &default_backend_config());
837
838        let mock = MockService::with_tools(&["fs/read"]);
839        let mut svc = layer.layer(mock);
840
841        let _ = call_service(&mut svc, tool_call("fs/read")).await;
842        let _ = call_service(&mut svc, tool_call("fs/read")).await;
843        assert_eq!(handle.stats().await[0].hits, 1);
844
845        handle.clear().await;
846        let stats = handle.stats().await;
847        assert_eq!(stats[0].hits, 0);
848        assert_eq!(stats[0].misses, 0);
849    }
850
851    #[tokio::test]
852    async fn test_cache_store_memory_get_insert() {
853        use super::{CacheStore, build_cache_store};
854
855        let store = build_cache_store(&default_backend_config(), Duration::from_secs(60), 100);
856        assert!(matches!(store, CacheStore::Memory(_)));
857
858        // Initially empty
859        assert!(store.get("key1").await.is_none());
860        assert_eq!(store.entry_count().await, 0);
861    }
862
863    #[cfg(feature = "redis-cache")]
864    #[test]
865    fn test_cache_store_redis_construction() {
866        use super::build_cache_store;
867
868        let cfg = CacheBackendConfig {
869            backend: "redis".to_string(),
870            url: Some("redis://127.0.0.1:6379".to_string()),
871            prefix: "test:".to_string(),
872        };
873        let store = build_cache_store(&cfg, Duration::from_secs(60), 100);
874        assert!(matches!(store, super::CacheStore::Redis { .. }));
875    }
876
877    #[cfg(feature = "sqlite-cache")]
878    #[tokio::test]
879    async fn test_cache_store_sqlite_construction() {
880        use super::build_cache_store;
881
882        let dir = std::env::temp_dir().join(format!("mcp-proxy-test-{}", std::process::id()));
883        std::fs::create_dir_all(&dir).unwrap();
884        let db_path = dir.join("test_cache.db");
885
886        let cfg = CacheBackendConfig {
887            backend: "sqlite".to_string(),
888            url: Some(db_path.to_string_lossy().to_string()),
889            prefix: "test:".to_string(),
890        };
891        let store = build_cache_store(&cfg, Duration::from_secs(60), 100);
892        assert!(matches!(store, super::CacheStore::Sqlite { .. }));
893        assert_eq!(store.entry_count().await, 0);
894
895        // Clean up
896        let _ = std::fs::remove_dir_all(&dir);
897    }
898
899    use std::time::Duration;
900}
mcp_proxy/cache.rs

mcp_proxy/
cache.rs
