Skip to main content

fastmcp_server/
caching.rs

1//! Response caching middleware for MCP servers.
2//!
3//! This module provides a bounded, process-local response cache for MCP methods
4//! whose results are safe to share between otherwise unrelated requests.
5//!
6//! Cache keys include the complete verified [`fastmcp_core::AuthContext`] plus
7//! its hidden stable owner key. Session-backed requests additionally include a
8//! cryptographically opaque session-state identity and monotonic mutation
9//! revision. Stateless authenticated contexts are therefore isolated by both
10//! handler-visible facts and provider-scoped ownership. Requests with
11//! uncommitted authentication, or session state whose partition cannot be
12//! obtained, bypass lookup and storage. Direct anonymous contexts remain in a
13//! separate stateless domain for standalone middleware use and tests.
14//!
15//! # Cached Methods
16//!
17//! For a live context with a safe complete partition (or a standalone context
18//! with neither session nor auth), the default method policy permits caching:
19//! - `server/discover` - its final private `ttlMs` policy
20//! - `tools/list` - 5 minute TTL
21//! - `resources/list` - 5 minute TTL
22//! - `resources/templates/list` - 5 minute TTL
23//! - `prompts/list` - 5 minute TTL
24//! - `resources/read` - 1 hour TTL
25//! - `prompts/get` - 1 hour TTL
26//! - `tools/call` - disabled unless individual tool names are explicitly
27//!   allowlisted with [`ResponseCachingMiddleware::include_tools`]
28//!
29//! Exclusions always override the `tools/call` allowlist. Tool-call caching is
30//! intended only for tools whose results are deterministic, side-effect free,
31//! and independent of external mutable state. Even an allowlisted tool is
32//! bypassed when a complete safe cache partition cannot be derived.
33//!
34//! # Example
35//!
36//! ```ignore
37//! use fastmcp_rust::prelude::*;
38//! use fastmcp_rust::caching::ResponseCachingMiddleware;
39//!
40//! let caching = ResponseCachingMiddleware::new()
41//!     .list_ttl_secs(600)  // 10 minute TTL for list operations
42//!     .call_ttl_secs(3600) // 1 hour TTL for call/get/read operations
43//!     .include_tools(vec!["deterministic_lookup".to_string()]);
44//!
45//! Server::new("my-server", "1.0.0")
46//!     .middleware(caching)
47//!     .build()
48//!     .run_stdio();
49//! ```
50
51use std::collections::HashMap;
52use std::io::Write;
53use std::sync::atomic::{AtomicU64, Ordering};
54use std::sync::{Arc, Mutex};
55use std::time::{Duration, Instant};
56
57use fastmcp_core::{McpContext, McpError, McpResult, Sha256Digest, sha256_bounded};
58use fastmcp_protocol::protocol_policy::ProtocolEra;
59use fastmcp_protocol::{
60    CacheTtl, FINAL_PROTOCOL_VERSION, FINAL_PROTOCOL_VERSION_META_KEY, JsonRpcRequest,
61    SERVER_DISCOVER_METHOD,
62};
63
64use crate::{Middleware, MiddlewareDecision};
65
66/// Default TTL for list operations (5 minutes).
67pub const DEFAULT_LIST_TTL_SECS: u64 = 300;
68
69/// Default TTL for allowlisted call/get/read operations (1 hour).
70pub const DEFAULT_CALL_TTL_SECS: u64 = 3600;
71
72/// Maximum cache item size in bytes (1 MB).
73pub const DEFAULT_MAX_ITEM_SIZE: usize = 1024 * 1024;
74
75/// Maximum canonical input admitted while deriving one fixed-width cache key.
76const MAX_CACHE_KEY_INPUT_BYTES: usize = 10 * 1024 * 1024;
77
78/// Maximum JSON nesting and aggregate nodes admitted to cache serialization.
79const MAX_CACHE_JSON_DEPTH: usize = 128;
80const MAX_CACHE_JSON_NODES: usize = 100_000;
81
82/// Small writes share an initial allocation and subsequent growth doubles the
83/// current capacity. Every target is still capped by the caller's logical byte
84/// limit, so fragmented serializer output cannot trigger one allocation per
85/// fragment or reserve beyond the configured bound.
86const CACHE_BYTES_GROWTH_CHUNK: usize = 4 * 1024;
87
88/// Conservative accounting for the entry, duplicate map/order keys, hash-table
89/// bucket/control storage, the `Arc` allocation header, and allocator metadata.
90/// The encoded payload length is added separately.
91const CACHE_ENTRY_METADATA_BYTES: usize = 512;
92
93/// Domain separators for cache request and authorization/session partitions.
94const CACHE_REQUEST_KEY_DOMAIN: &[u8] = b"fastmcp-response-cache-request-v2\0";
95const CACHE_INVALIDATION_KEY_DOMAIN: &[u8] = b"fastmcp-response-cache-invalidation-v1\0";
96const CACHE_PARTITION_KEY_DOMAIN: &[u8] = b"fastmcp-response-cache-partition-v2\0";
97const CACHE_STATELESS_PARTITION_DOMAIN: &[u8] = b"fastmcp-response-cache-stateless-partition-v1\0";
98
99static NEXT_CACHE_INSTANCE_ID: AtomicU64 = AtomicU64::new(1);
100
101fn next_cache_instance_id() -> u64 {
102    NEXT_CACHE_INSTANCE_ID
103        .try_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
104            current.checked_add(1)
105        })
106        .unwrap_or(0)
107}
108
109/// A cached response with expiration time.
110#[derive(Clone)]
111struct CacheEntry {
112    encoded: Arc<[u8]>,
113    expires_at: Instant,
114    size_bytes: usize,
115}
116
117impl std::fmt::Debug for CacheEntry {
118    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
119        f.debug_struct("CacheEntry")
120            .field("payload_bytes", &self.encoded.len())
121            .field("expires_at", &self.expires_at)
122            .field("accounted_bytes", &self.size_bytes)
123            .finish()
124    }
125}
126
127impl CacheEntry {
128    fn new(value: serde_json::Value, ttl: Duration, max_size_bytes: usize) -> Option<Self> {
129        let encoded = encode_json_bounded(&value, max_size_bytes)?;
130        Self::new_encoded(encoded, ttl)
131    }
132
133    fn new_encoded(encoded: Arc<[u8]>, ttl: Duration) -> Option<Self> {
134        if ttl.is_zero() {
135            return None;
136        }
137        let expires_at = Instant::now().checked_add(ttl)?;
138        let size_bytes = encoded.len().checked_add(CACHE_ENTRY_METADATA_BYTES)?;
139        Some(Self {
140            encoded,
141            expires_at,
142            size_bytes,
143        })
144    }
145
146    fn is_expired(&self) -> bool {
147        Instant::now() >= self.expires_at
148    }
149}
150
151/// Cache key derived from method and parameters.
152#[derive(Debug, Clone, PartialEq, Eq, Hash)]
153struct CacheKey {
154    request_digest: Sha256Digest,
155    invalidation_digest: Sha256Digest,
156    partition_digest: Sha256Digest,
157    binding: CacheEntryBinding,
158}
159
160/// Extra identity attached only to final discovery cache entries.
161///
162/// Discovery is a final-only surface, but the cache can sit behind a dual-era
163/// transport. The selected era and the invalidation generation therefore
164/// participate in the key rather than being inferred from a previously cached
165/// payload.
166#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
167struct DiscoveryCacheBinding {
168    era: ProtocolEra,
169    generation: u64,
170}
171
172#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
173enum CacheEntryBinding {
174    Ordinary,
175    Discovery(DiscoveryCacheBinding),
176}
177
178impl CacheKey {
179    fn try_request_digest(
180        method: &str,
181        params: Option<&serde_json::Value>,
182    ) -> Option<Sha256Digest> {
183        Self::try_digest(CACHE_REQUEST_KEY_DOMAIN, method, params)
184    }
185
186    /// Derives the identity used to invalidate an entire paginated result set.
187    ///
188    /// The request key still includes the opaque cursor exactly, so distinct
189    /// pages cannot collide at lookup. Invalidation intentionally removes the
190    /// cursor from the semantic result-set identity so a catalog or resource
191    /// mutation cannot leave another page observable from the old generation.
192    fn try_invalidation_digest(
193        method: &str,
194        params: Option<&serde_json::Value>,
195    ) -> Option<Sha256Digest> {
196        let mut projection = params.cloned();
197        if let Some(serde_json::Value::Object(object)) = projection.as_mut() {
198            object.remove("cursor");
199            // A cursor-only request has an empty semantic parameter set; it
200            // must share the invalidation identity of the parameterless
201            // request, or `invalidate(method, None)` can never remove cached
202            // continuation pages after a catalog mutation.
203            if object.is_empty() {
204                projection = None;
205            }
206        }
207        Self::try_digest(CACHE_INVALIDATION_KEY_DOMAIN, method, projection.as_ref())
208    }
209
210    fn try_digest(
211        domain: &[u8],
212        method: &str,
213        params: Option<&serde_json::Value>,
214    ) -> Option<Sha256Digest> {
215        if params.is_some_and(|params| !cache_json_shape_is_bounded(params)) {
216            return None;
217        }
218        let mut canonical = BoundedCacheBytes::new(MAX_CACHE_KEY_INPUT_BYTES);
219        canonical.write_all(domain).ok()?;
220        let method_len = u64::try_from(method.len()).ok()?;
221        canonical.write_all(&method_len.to_be_bytes()).ok()?;
222        canonical.write_all(method.as_bytes()).ok()?;
223        match params {
224            None => canonical.write_all(&[0]).ok()?,
225            Some(params) => {
226                canonical.write_all(&[1]).ok()?;
227                serde_json::to_writer(&mut canonical, params).ok()?;
228            }
229        }
230        sha256_bounded(&canonical.bytes, MAX_CACHE_KEY_INPUT_BYTES).ok()
231    }
232
233    fn try_new_partitioned(
234        method: &str,
235        params: Option<&serde_json::Value>,
236        partition_digest: Sha256Digest,
237        binding: CacheEntryBinding,
238    ) -> Option<Self> {
239        Some(Self {
240            request_digest: Self::try_request_digest(method, params)?,
241            invalidation_digest: Self::try_invalidation_digest(method, params)?,
242            partition_digest,
243            binding,
244        })
245    }
246
247    #[cfg(test)]
248    fn try_new(method: &str, params: Option<&serde_json::Value>) -> Option<Self> {
249        let partition_digest = sha256_bounded(
250            CACHE_STATELESS_PARTITION_DOMAIN,
251            CACHE_STATELESS_PARTITION_DOMAIN.len(),
252        )
253        .ok()?;
254        Self::try_new_partitioned(
255            method,
256            params,
257            partition_digest,
258            CacheEntryBinding::Ordinary,
259        )
260    }
261
262    #[cfg(test)]
263    fn new(method: &str, params: Option<&serde_json::Value>) -> Self {
264        Self::try_new(method, params).expect("test cache key must fit the fixed input bound")
265    }
266}
267
268struct BoundedCacheBytes {
269    bytes: Vec<u8>,
270    max_bytes: usize,
271    #[cfg(test)]
272    growth_events: usize,
273}
274
275impl BoundedCacheBytes {
276    fn new(max_bytes: usize) -> Self {
277        Self {
278            bytes: Vec::new(),
279            max_bytes,
280            #[cfg(test)]
281            growth_events: 0,
282        }
283    }
284
285    fn ensure_capacity_for(&mut self, next_size: usize) -> std::io::Result<()> {
286        if next_size <= self.bytes.capacity() {
287            return Ok(());
288        }
289
290        let current_capacity = self.bytes.capacity();
291        let chunk_target = CACHE_BYTES_GROWTH_CHUNK.min(self.max_bytes);
292        let geometric_target = if current_capacity == 0 {
293            chunk_target
294        } else {
295            current_capacity
296                .checked_mul(2)
297                .unwrap_or(self.max_bytes)
298                .min(self.max_bytes)
299        };
300        let target_capacity = next_size.max(geometric_target).min(self.max_bytes);
301
302        // Allocate separately so even an allocator that reports more capacity
303        // than requested cannot leave this bounded writer above its logical
304        // limit. The existing buffer remains intact on every failure path.
305        let mut grown = Vec::new();
306        grown
307            .try_reserve_exact(target_capacity)
308            .map_err(|_| std::io::Error::other("cannot allocate bounded cache input"))?;
309        if grown.capacity() > self.max_bytes {
310            return Err(std::io::Error::other(
311                "cache input allocation exceeds configured limit",
312            ));
313        }
314        grown.extend_from_slice(&self.bytes);
315        self.bytes = grown;
316        #[cfg(test)]
317        {
318            self.growth_events = self.growth_events.saturating_add(1);
319        }
320        Ok(())
321    }
322}
323
324impl Write for BoundedCacheBytes {
325    fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
326        let next_size = self
327            .bytes
328            .len()
329            .checked_add(buffer.len())
330            .filter(|size| *size <= self.max_bytes)
331            .ok_or_else(|| std::io::Error::other("cache input exceeds configured limit"))?;
332        self.ensure_capacity_for(next_size)?;
333        self.bytes.extend_from_slice(buffer);
334        Ok(buffer.len())
335    }
336
337    fn flush(&mut self) -> std::io::Result<()> {
338        Ok(())
339    }
340}
341
342fn cache_json_shape_is_bounded(value: &serde_json::Value) -> bool {
343    let mut stack = Vec::new();
344    if stack.try_reserve_exact(1).is_err() {
345        return false;
346    }
347    stack.push((value, 0_usize));
348    let mut admitted_nodes = 1_usize;
349
350    while let Some((node, depth)) = stack.pop() {
351        let child_depth = match depth.checked_add(1) {
352            Some(depth) => depth,
353            None => return false,
354        };
355        match node {
356            serde_json::Value::Array(values) => {
357                if !values.is_empty() && child_depth > MAX_CACHE_JSON_DEPTH {
358                    return false;
359                }
360                admitted_nodes = match admitted_nodes
361                    .checked_add(values.len())
362                    .filter(|nodes| *nodes <= MAX_CACHE_JSON_NODES)
363                {
364                    Some(nodes) => nodes,
365                    None => return false,
366                };
367                if stack.try_reserve(values.len()).is_err() {
368                    return false;
369                }
370                stack.extend(values.iter().map(|value| (value, child_depth)));
371            }
372            serde_json::Value::Object(values) => {
373                if !values.is_empty() && child_depth > MAX_CACHE_JSON_DEPTH {
374                    return false;
375                }
376                admitted_nodes = match admitted_nodes
377                    .checked_add(values.len())
378                    .filter(|nodes| *nodes <= MAX_CACHE_JSON_NODES)
379                {
380                    Some(nodes) => nodes,
381                    None => return false,
382                };
383                if stack.try_reserve(values.len()).is_err() {
384                    return false;
385                }
386                stack.extend(values.values().map(|value| (value, child_depth)));
387            }
388            _ => {}
389        }
390    }
391
392    true
393}
394
395fn encode_json_bounded(value: &serde_json::Value, max_bytes: usize) -> Option<Arc<[u8]>> {
396    if !cache_json_shape_is_bounded(value) {
397        return None;
398    }
399    let mut encoded = BoundedCacheBytes::new(max_bytes);
400    serde_json::to_writer(&mut encoded, value).ok()?;
401    Some(Arc::from(encoded.bytes.into_boxed_slice()))
402}
403
404fn decode_cached_json(encoded: &[u8]) -> Option<serde_json::Value> {
405    serde_json::from_slice(encoded).ok()
406}
407
408#[derive(Clone, Copy)]
409enum CachePartitionPhase {
410    Request,
411    Response,
412}
413
414fn context_cache_partition(ctx: &McpContext, phase: CachePartitionPhase) -> Option<Sha256Digest> {
415    ctx.ensure_live().ok()?;
416    // Authentication admission is write-once. An uncommitted request must
417    // never consult or populate a cache merely because it has no session.
418    let auth_partition = ctx.cache_auth_partition()?;
419    let session_partition = match phase {
420        CachePartitionPhase::Request => ctx.begin_session_cache_partition(),
421        CachePartitionPhase::Response => ctx.complete_session_cache_partition(),
422    };
423    if let Some(auth) = auth_partition.as_ref() {
424        if auth.scopes.len() > MAX_CACHE_JSON_NODES
425            || auth
426                .claims
427                .as_ref()
428                .is_some_and(|claims| !cache_json_shape_is_bounded(claims))
429        {
430            return None;
431        }
432    }
433
434    let mut canonical = BoundedCacheBytes::new(MAX_CACHE_KEY_INPUT_BYTES);
435    canonical.write_all(CACHE_PARTITION_KEY_DOMAIN).ok()?;
436    if ctx.session_is_ephemeral() {
437        // Per-POST modern HTTP state exists so disable_*/enable_* can
438        // publish list_changed. It is not a durable cache identity.
439        canonical.write_all(CACHE_STATELESS_PARTITION_DOMAIN).ok()?;
440    } else {
441        match session_partition {
442            Some((opaque_session, state_revision)) => {
443                canonical.write_all(&[1]).ok()?;
444                canonical.write_all(&opaque_session).ok()?;
445                canonical.write_all(&state_revision.to_be_bytes()).ok()?;
446            }
447            None if ctx.has_session_state() => return None,
448            None => canonical.write_all(CACHE_STATELESS_PARTITION_DOMAIN).ok()?,
449        }
450    }
451    match auth_partition {
452        None => canonical.write_all(&[0]).ok()?,
453        Some(auth) => {
454            canonical.write_all(&[1]).ok()?;
455            match auth.session_owner() {
456                None => canonical.write_all(&[0]).ok()?,
457                Some(owner) => {
458                    canonical.write_all(&[1]).ok()?;
459                    canonical.write_all(owner.as_bytes()).ok()?;
460                }
461            }
462            serde_json::to_writer(&mut canonical, &auth).ok()?;
463        }
464    }
465    sha256_bounded(&canonical.bytes, MAX_CACHE_KEY_INPUT_BYTES).ok()
466}
467
468fn context_cache_commit_is_admissible(ctx: &McpContext) -> bool {
469    // Check the session partition before the final liveness read. In
470    // particular, `has_session_state()` intentionally reports false after a
471    // request lease closes; the final `ensure_live()` prevents that transition
472    // from being mistaken for a genuinely stateless request.
473    let session_partition_is_current =
474        !ctx.has_session_state() || ctx.complete_session_cache_partition().is_some();
475    session_partition_is_current && ctx.ensure_live().is_ok()
476}
477
478/// Returns whether request parameters carry state from a multi-round-trip
479/// continuation. Such requests are never deterministic cache lookups, even
480/// when their eventual result happens to be complete.
481fn request_carries_uncacheable_continuation(params: Option<&serde_json::Value>) -> bool {
482    let Some(serde_json::Value::Object(params)) = params else {
483        return false;
484    };
485    params.contains_key("inputResponses") || params.contains_key("requestState")
486}
487
488/// Returns whether a response can be stored by the internal memoization cache.
489///
490/// Modern responses must explicitly be `complete`; input-required and task
491/// branches never enter this cache. The absent discriminator remains accepted
492/// for the current exact-2024 compatibility surface, which did not carry
493/// `resultType`. Continuation-bearing payloads are rejected in either era.
494fn response_is_cacheable_complete(response: &serde_json::Value) -> bool {
495    let Some(response) = response.as_object() else {
496        return false;
497    };
498    if [
499        "inputResponses",
500        "requestState",
501        "task",
502        "taskId",
503        "taskStatus",
504        "requestScopedNotifications",
505        "notifications",
506    ]
507    .iter()
508    .any(|field| response.contains_key(*field))
509    {
510        return false;
511    }
512    match response.get("resultType") {
513        None => true,
514        Some(serde_json::Value::String(kind)) => kind == "complete",
515        Some(_) => false,
516    }
517}
518
519/// The cache policy attached to an exact final discovery result.
520#[derive(Debug, Clone, Copy, PartialEq, Eq)]
521enum FinalDiscoveryCachePolicy {
522    Private(Duration),
523    Public,
524}
525
526/// Local cache-admission decision for a complete final method result.
527///
528/// This is deliberately separate from the wire cache hints: a valid final
529/// `ttlMs` may exceed the bounded local duration domain and must then remain
530/// deliverable without creating a local cache entry.
531enum FinalCompleteCachePolicy {
532    Private(Duration),
533    Public,
534}
535
536/// Returns whether a result carries the exact final discovery shape relevant
537/// to caching. Full discovery validation remains owned by the protocol/server
538/// boundary; this narrower check prevents legacy lookalikes from selecting a
539/// final cache policy.
540fn is_final_discovery_result(response: &serde_json::Value) -> bool {
541    let Some(response) = response.as_object() else {
542        return false;
543    };
544    response
545        .get("capabilities")
546        .is_some_and(serde_json::Value::is_object)
547        && response
548            .get("supportedVersions")
549            .and_then(serde_json::Value::as_array)
550            .is_some_and(|versions| {
551                versions.len() == 1 && versions[0].as_str() == Some(FINAL_PROTOCOL_VERSION)
552            })
553}
554
555/// Reads the final discovery cache policy exactly as it appears on the wire.
556fn final_discovery_cache_policy(response: &serde_json::Value) -> Option<FinalDiscoveryCachePolicy> {
557    if !is_final_discovery_result(response) {
558        return None;
559    }
560    let response = response.as_object()?;
561    let ttl_ms = serde_json::from_value::<CacheTtl>(response.get("ttlMs")?.clone()).ok()?;
562    match response.get("cacheScope")?.as_str()? {
563        "private" => Some(FinalDiscoveryCachePolicy::Private(Duration::from_millis(
564            ttl_ms.try_as_millis().ok()?,
565        ))),
566        "public" => Some(FinalDiscoveryCachePolicy::Public),
567        _ => None,
568    }
569}
570
571fn final_complete_cache_policy(response: &serde_json::Value) -> Option<FinalCompleteCachePolicy> {
572    let response = response.as_object()?;
573    let ttl_ms = serde_json::from_value::<CacheTtl>(response.get("ttlMs")?.clone()).ok()?;
574    match response.get("cacheScope")?.as_str()? {
575        "private" => Some(FinalCompleteCachePolicy::Private(Duration::from_millis(
576            ttl_ms.try_as_millis().ok()?,
577        ))),
578        "public" => Some(FinalCompleteCachePolicy::Public),
579        _ => None,
580    }
581}
582
583/// Returns whether a final discovery result has schema-valid cache hints.
584///
585/// A wire-valid TTL can exceed the local runtime duration domain. That is not
586/// a reason to rewrite the peer response: callers must still observe its exact
587/// JSON-integer spelling, while this process simply declines to cache it.
588fn final_discovery_cache_hints_are_wire_valid(response: &serde_json::Value) -> bool {
589    let Some(response) = response.as_object() else {
590        return false;
591    };
592    let Some(ttl_ms) = response.get("ttlMs") else {
593        return false;
594    };
595    serde_json::from_value::<CacheTtl>(ttl_ms.clone()).is_ok()
596        && matches!(
597            response
598                .get("cacheScope")
599                .and_then(serde_json::Value::as_str),
600            Some("private" | "public")
601        )
602}
603
604/// Selects the request era used by the discovery cache.
605///
606/// Final transports remove the recognized version metadata before middleware
607/// runs, so missing metadata here represents an already-admitted final
608/// request. An explicitly supplied legacy or unsupported version is never
609/// allowed to reuse the final discovery cache.
610fn discovery_request_protocol_era(request: &JsonRpcRequest) -> Option<ProtocolEra> {
611    let version = request
612        .params
613        .as_ref()
614        .and_then(|params| params.get("_meta"))
615        .and_then(|metadata| metadata.get(FINAL_PROTOCOL_VERSION_META_KEY))
616        .and_then(serde_json::Value::as_str);
617    match version {
618        None | Some(FINAL_PROTOCOL_VERSION) => Some(ProtocolEra::Modern2026),
619        Some(version) if version == ProtocolEra::Legacy2024.version().as_str() => {
620            Some(ProtocolEra::Legacy2024)
621        }
622        Some(_) => None,
623    }
624}
625
626/// Returns whether a method requires `ttlMs` and `cacheScope` in a modern
627/// complete result. This list is protocol-facing and deliberately independent
628/// from the internal memoization allowlist.
629fn method_requires_protocol_cache_hints(method: &str) -> bool {
630    matches!(
631        method,
632        "server/discover"
633            | "tools/list"
634            | "prompts/list"
635            | "resources/list"
636            | "resources/read"
637            | "resources/templates/list"
638    )
639}
640
641/// Configuration for caching specific methods.
642#[derive(Debug, Clone)]
643pub struct MethodCacheConfig {
644    /// Whether caching is enabled for this method.
645    pub enabled: bool,
646    /// Time to live in seconds.
647    pub ttl_secs: u64,
648}
649
650impl Default for MethodCacheConfig {
651    fn default() -> Self {
652        Self {
653            enabled: true,
654            ttl_secs: DEFAULT_CALL_TTL_SECS,
655        }
656    }
657}
658
659/// Configuration for `tools/call` caching.
660///
661/// A tool is cacheable only when [`MethodCacheConfig::enabled`] is `true`, its
662/// name appears in [`Self::included_tools`], and its name does not appear in
663/// [`Self::excluded_tools`].
664#[derive(Debug, Clone, Default)]
665pub struct ToolCallCacheConfig {
666    /// Base configuration.
667    pub base: MethodCacheConfig,
668    /// Tools explicitly allowlisted for caching (empty disables tool caching).
669    pub included_tools: Vec<String>,
670    /// Tools to exclude (takes precedence over included).
671    pub excluded_tools: Vec<String>,
672}
673
674impl ToolCallCacheConfig {
675    /// Checks if a specific tool should be cached.
676    fn should_cache_tool(&self, tool_name: &str) -> bool {
677        if !self.base.enabled {
678            return false;
679        }
680
681        // Check exclusions first (takes precedence)
682        if self.excluded_tools.iter().any(|name| name == tool_name) {
683            return false;
684        }
685
686        // Tool calls are stateful by default. Only an explicit allowlist entry
687        // can opt a tool into caching.
688        self.included_tools.iter().any(|name| name == tool_name)
689    }
690}
691
692/// Simple LRU cache with TTL support.
693#[derive(Debug)]
694struct LruCache {
695    /// Map of keys to entries.
696    entries: HashMap<CacheKey, CacheEntry>,
697    /// Order of keys for LRU eviction (most recent at the end).
698    order: Vec<CacheKey>,
699    /// Maximum number of entries.
700    max_entries: usize,
701    /// Maximum total size in bytes.
702    max_size_bytes: usize,
703    /// Maximum size per item in bytes.
704    max_item_size: usize,
705    /// Current total size in bytes.
706    current_size_bytes: usize,
707}
708
709impl LruCache {
710    fn new(max_entries: usize, max_size_bytes: usize, max_item_size: usize) -> Self {
711        Self {
712            entries: HashMap::new(),
713            order: Vec::new(),
714            max_entries,
715            max_size_bytes,
716            max_item_size,
717            current_size_bytes: 0,
718        }
719    }
720
721    fn get_encoded(&mut self, key: &CacheKey) -> Option<Arc<[u8]>> {
722        // Check if entry exists and is not expired
723        if let Some(entry) = self.entries.get(key) {
724            if entry.is_expired() {
725                // Remove expired entry
726                self.remove(key);
727                return None;
728            }
729
730            // Move to end of order (most recently used)
731            if let Some(pos) = self.order.iter().position(|k| k == key) {
732                let k = self.order.remove(pos);
733                self.order.push(k);
734            }
735
736            return Some(Arc::clone(&entry.encoded));
737        }
738        None
739    }
740
741    #[cfg(test)]
742    fn get_value(&mut self, key: &CacheKey) -> Option<serde_json::Value> {
743        self.get_encoded(key)
744            .and_then(|encoded| decode_cached_json(&encoded))
745    }
746
747    fn insert(&mut self, key: CacheKey, value: serde_json::Value, ttl: Duration) {
748        let admission_limit = self.max_item_size.min(
749            self.max_size_bytes
750                .saturating_sub(CACHE_ENTRY_METADATA_BYTES),
751        );
752        let Some(entry) = CacheEntry::new(value, ttl, admission_limit) else {
753            // An unrepresentable expiration must not turn into a panic or an
754            // accidentally immortal entry.
755            return;
756        };
757        self.insert_entry(key, entry);
758    }
759
760    fn insert_encoded(&mut self, key: CacheKey, encoded: Arc<[u8]>, ttl: Duration) {
761        if encoded.len() > self.max_item_size {
762            return;
763        }
764        let Some(entry) = CacheEntry::new_encoded(encoded, ttl) else {
765            return;
766        };
767        self.insert_entry(key, entry);
768    }
769
770    fn insert_entry(&mut self, key: CacheKey, entry: CacheEntry) {
771        // Reject impossible configurations and entries that can never fit.
772        // These checks happen before replacing an existing value, so a rejected
773        // replacement cannot destroy a valid cached entry.
774        if self.max_entries == 0
775            || self.max_size_bytes == 0
776            || entry.encoded.len() > self.max_item_size
777            || entry.size_bytes > self.max_size_bytes
778        {
779            return;
780        }
781
782        // Expired entries should not force eviction of live entries.
783        self.evict_expired();
784
785        // Remove old entry if it exists
786        if self.entries.contains_key(&key) {
787            self.remove(&key);
788        }
789
790        // Evict entries if needed to make room
791        while self.entries.len() >= self.max_entries
792            || self
793                .current_size_bytes
794                .checked_add(entry.size_bytes)
795                .is_none_or(|size| size > self.max_size_bytes)
796        {
797            if self.order.is_empty() {
798                // An inconsistent accounting state must fail closed instead of
799                // admitting an entry beyond a configured bound.
800                return;
801            }
802            // Evict least recently used (first in order)
803            let oldest_key = self.order.remove(0);
804            if let Some(old_entry) = self.entries.remove(&oldest_key) {
805                self.current_size_bytes =
806                    self.current_size_bytes.saturating_sub(old_entry.size_bytes);
807            }
808        }
809
810        let Some(new_size) = self.current_size_bytes.checked_add(entry.size_bytes) else {
811            return;
812        };
813        if new_size > self.max_size_bytes || self.entries.len() >= self.max_entries {
814            return;
815        }
816
817        // Insert new entry only after all bounds have been rechecked.
818        self.current_size_bytes = new_size;
819        self.entries.insert(key.clone(), entry);
820        self.order.push(key);
821    }
822
823    fn remove(&mut self, key: &CacheKey) {
824        if let Some(entry) = self.entries.remove(key) {
825            self.current_size_bytes = self.current_size_bytes.saturating_sub(entry.size_bytes);
826            if let Some(pos) = self.order.iter().position(|k| k == key) {
827                self.order.remove(pos);
828            }
829        }
830    }
831
832    fn remove_invalidation_digest(&mut self, invalidation_digest: Sha256Digest) {
833        let mut retained_size = self.current_size_bytes;
834        self.entries.retain(|key, entry| {
835            if key.invalidation_digest == invalidation_digest {
836                retained_size = retained_size.saturating_sub(entry.size_bytes);
837                false
838            } else {
839                true
840            }
841        });
842        self.order
843            .retain(|key| key.invalidation_digest != invalidation_digest);
844        self.current_size_bytes = retained_size;
845    }
846
847    fn remove_discovery_entries(&mut self) {
848        let mut retained_size = self.current_size_bytes;
849        self.entries.retain(|key, entry| {
850            if matches!(key.binding, CacheEntryBinding::Discovery(_)) {
851                retained_size = retained_size.saturating_sub(entry.size_bytes);
852                false
853            } else {
854                true
855            }
856        });
857        self.order
858            .retain(|key| !matches!(key.binding, CacheEntryBinding::Discovery(_)));
859        self.current_size_bytes = retained_size;
860    }
861
862    fn evict_expired(&mut self) {
863        let mut retained_size = self.current_size_bytes;
864        self.entries.retain(|_, entry| {
865            if entry.is_expired() {
866                retained_size = retained_size.saturating_sub(entry.size_bytes);
867                false
868            } else {
869                true
870            }
871        });
872        let entries = &self.entries;
873        self.order.retain(|key| entries.contains_key(key));
874        self.current_size_bytes = retained_size;
875    }
876
877    fn clear(&mut self) {
878        self.entries.clear();
879        self.order.clear();
880        self.current_size_bytes = 0;
881    }
882
883    fn len(&self) -> usize {
884        self.entries.len()
885    }
886
887    #[allow(dead_code)]
888    fn is_empty(&self) -> bool {
889        self.entries.is_empty()
890    }
891}
892
893/// Cache statistics.
894#[derive(Debug, Clone, Default, PartialEq, Eq)]
895pub struct CacheStats {
896    /// Number of hits from cache-eligible partitioned or standalone lookups.
897    pub hits: u64,
898    /// Number of misses from cache-eligible partitioned or standalone lookups.
899    ///
900    /// Requests bypassed due to an incomplete partition or method policy are
901    /// not counted as misses.
902    pub misses: u64,
903    /// Number of entries currently in cache.
904    pub entries: usize,
905    /// Current cache size in bytes.
906    pub size_bytes: usize,
907}
908
909impl CacheStats {
910    /// Returns the hit rate as a percentage.
911    #[must_use]
912    pub fn hit_rate(&self) -> f64 {
913        let hits = self.hits as f64;
914        let total = hits + self.misses as f64;
915        if total == 0.0 {
916            0.0
917        } else {
918            (hits / total) * 100.0
919        }
920    }
921}
922
923/// Response caching middleware for MCP servers.
924///
925/// Caches eligible responses with configurable TTL and bounded LRU eviction.
926///
927/// Production contexts are isolated by opaque session-state identity, state
928/// mutation revision, and complete verified authentication facts. An
929/// incomplete partition fails closed. `tools/call` is additionally disabled by
930/// default and requires an explicit per-tool allowlist entry via
931/// [`Self::include_tools`].
932pub struct ResponseCachingMiddleware {
933    /// Process-local identity used only for per-request hit bookkeeping.
934    instance_id: u64,
935    /// Monotonic final-discovery invalidation generation.
936    ///
937    /// It is included in every final discovery key and is advanced under the
938    /// cache lock, preventing a response captured before invalidation from
939    /// becoming observable afterwards.
940    discovery_generation: AtomicU64,
941    /// Cache storage.
942    cache: Mutex<LruCache>,
943    /// TTL for list operations.
944    list_ttl: Duration,
945    /// TTL for allowlisted call/get/read operations.
946    call_ttl: Duration,
947    /// Configuration for tools/list caching.
948    tools_list_config: MethodCacheConfig,
949    /// Configuration for resources/list caching.
950    resources_list_config: MethodCacheConfig,
951    /// Configuration for prompts/list caching.
952    prompts_list_config: MethodCacheConfig,
953    /// Configuration for tools/call caching.
954    tools_call_config: ToolCallCacheConfig,
955    /// Configuration for resources/read caching.
956    resources_read_config: MethodCacheConfig,
957    /// Configuration for prompts/get caching.
958    prompts_get_config: MethodCacheConfig,
959    /// Statistics tracking.
960    stats: Mutex<CacheStats>,
961}
962
963impl std::fmt::Debug for ResponseCachingMiddleware {
964    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
965        f.debug_struct("ResponseCachingMiddleware")
966            .field("instance_available", &(self.instance_id != 0))
967            .field(
968                "discovery_generation",
969                &self.discovery_generation.load(Ordering::Acquire),
970            )
971            .field("list_ttl", &self.list_ttl)
972            .field("call_ttl", &self.call_ttl)
973            .finish_non_exhaustive()
974    }
975}
976
977impl Default for ResponseCachingMiddleware {
978    fn default() -> Self {
979        Self::new()
980    }
981}
982
983impl ResponseCachingMiddleware {
984    /// Creates response caching middleware with default bounds and TTLs.
985    ///
986    /// `tools/call` caching remains off until [`Self::include_tools`] is used.
987    #[must_use]
988    pub fn new() -> Self {
989        Self {
990            instance_id: next_cache_instance_id(),
991            discovery_generation: AtomicU64::new(1),
992            cache: Mutex::new(LruCache::new(
993                1000,
994                100 * 1024 * 1024,
995                DEFAULT_MAX_ITEM_SIZE,
996            )),
997            list_ttl: Duration::from_secs(DEFAULT_LIST_TTL_SECS),
998            call_ttl: Duration::from_secs(DEFAULT_CALL_TTL_SECS),
999            tools_list_config: MethodCacheConfig {
1000                enabled: true,
1001                ttl_secs: DEFAULT_LIST_TTL_SECS,
1002            },
1003            resources_list_config: MethodCacheConfig {
1004                enabled: true,
1005                ttl_secs: DEFAULT_LIST_TTL_SECS,
1006            },
1007            prompts_list_config: MethodCacheConfig {
1008                enabled: true,
1009                ttl_secs: DEFAULT_LIST_TTL_SECS,
1010            },
1011            tools_call_config: ToolCallCacheConfig::default(),
1012            resources_read_config: MethodCacheConfig {
1013                enabled: true,
1014                ttl_secs: DEFAULT_CALL_TTL_SECS,
1015            },
1016            prompts_get_config: MethodCacheConfig {
1017                enabled: true,
1018                ttl_secs: DEFAULT_CALL_TTL_SECS,
1019            },
1020            stats: Mutex::new(CacheStats::default()),
1021        }
1022    }
1023
1024    /// Sets the maximum number of cache entries (`0` disables storage).
1025    #[must_use]
1026    pub fn max_entries(self, max: usize) -> Self {
1027        let max_size = {
1028            let cache = self
1029                .cache
1030                .lock()
1031                .unwrap_or_else(std::sync::PoisonError::into_inner);
1032            cache.max_size_bytes
1033        };
1034        let max_item_size = {
1035            let cache = self
1036                .cache
1037                .lock()
1038                .unwrap_or_else(std::sync::PoisonError::into_inner);
1039            cache.max_item_size
1040        };
1041        Self {
1042            cache: Mutex::new(LruCache::new(max, max_size, max_item_size)),
1043            ..self
1044        }
1045    }
1046
1047    /// Sets the maximum cache size in bytes (`0` disables storage).
1048    #[must_use]
1049    pub fn max_size_bytes(self, max: usize) -> Self {
1050        let max_entries = {
1051            let cache = self
1052                .cache
1053                .lock()
1054                .unwrap_or_else(std::sync::PoisonError::into_inner);
1055            cache.max_entries
1056        };
1057        let max_item_size = {
1058            let cache = self
1059                .cache
1060                .lock()
1061                .unwrap_or_else(std::sync::PoisonError::into_inner);
1062            cache.max_item_size
1063        };
1064        Self {
1065            cache: Mutex::new(LruCache::new(max_entries, max, max_item_size)),
1066            ..self
1067        }
1068    }
1069
1070    /// Sets the maximum size per cache item in bytes (`0` disables storage).
1071    #[must_use]
1072    pub fn max_item_size(self, max: usize) -> Self {
1073        let max_entries = {
1074            let cache = self
1075                .cache
1076                .lock()
1077                .unwrap_or_else(std::sync::PoisonError::into_inner);
1078            cache.max_entries
1079        };
1080        let max_size = {
1081            let cache = self
1082                .cache
1083                .lock()
1084                .unwrap_or_else(std::sync::PoisonError::into_inner);
1085            cache.max_size_bytes
1086        };
1087        Self {
1088            cache: Mutex::new(LruCache::new(max_entries, max_size, max)),
1089            ..self
1090        }
1091    }
1092
1093    /// Sets the TTL for list operations (tools/list, resources/list, prompts/list).
1094    #[must_use]
1095    pub fn list_ttl_secs(mut self, secs: u64) -> Self {
1096        self.list_ttl = Duration::from_secs(secs);
1097        self.tools_list_config.ttl_secs = secs;
1098        self.resources_list_config.ttl_secs = secs;
1099        self.prompts_list_config.ttl_secs = secs;
1100        self
1101    }
1102
1103    /// Sets the TTL for read/get operations and explicitly allowlisted calls.
1104    #[must_use]
1105    pub fn call_ttl_secs(mut self, secs: u64) -> Self {
1106        self.call_ttl = Duration::from_secs(secs);
1107        self.tools_call_config.base.ttl_secs = secs;
1108        self.resources_read_config.ttl_secs = secs;
1109        self.prompts_get_config.ttl_secs = secs;
1110        self
1111    }
1112
1113    /// Disables caching for tools/list.
1114    #[must_use]
1115    pub fn disable_tools_list(mut self) -> Self {
1116        self.tools_list_config.enabled = false;
1117        self
1118    }
1119
1120    /// Disables caching for resources/list.
1121    #[must_use]
1122    pub fn disable_resources_list(mut self) -> Self {
1123        self.resources_list_config.enabled = false;
1124        self
1125    }
1126
1127    /// Disables caching for prompts/list.
1128    #[must_use]
1129    pub fn disable_prompts_list(mut self) -> Self {
1130        self.prompts_list_config.enabled = false;
1131        self
1132    }
1133
1134    /// Disables caching for tools/call.
1135    #[must_use]
1136    pub fn disable_tools_call(mut self) -> Self {
1137        self.tools_call_config.base.enabled = false;
1138        self
1139    }
1140
1141    /// Disables caching for resources/read.
1142    #[must_use]
1143    pub fn disable_resources_read(mut self) -> Self {
1144        self.resources_read_config.enabled = false;
1145        self
1146    }
1147
1148    /// Disables caching for prompts/get.
1149    #[must_use]
1150    pub fn disable_prompts_get(mut self) -> Self {
1151        self.prompts_get_config.enabled = false;
1152        self
1153    }
1154
1155    /// Explicitly allowlists tools for `tools/call` caching.
1156    ///
1157    /// An empty list disables `tools/call` caching. Exclusions configured with
1158    /// [`Self::exclude_tools`] take precedence over this allowlist.
1159    #[must_use]
1160    pub fn include_tools(mut self, tools: Vec<String>) -> Self {
1161        self.tools_call_config.included_tools = tools;
1162        self
1163    }
1164
1165    /// Excludes tools from `tools/call` caching, overriding the allowlist.
1166    #[must_use]
1167    pub fn exclude_tools(mut self, tools: Vec<String>) -> Self {
1168        self.tools_call_config.excluded_tools = tools;
1169        self
1170    }
1171
1172    /// Returns current cache statistics.
1173    #[must_use]
1174    pub fn stats(&self) -> CacheStats {
1175        let cache = self
1176            .cache
1177            .lock()
1178            .unwrap_or_else(std::sync::PoisonError::into_inner);
1179        let mut stats = self
1180            .stats
1181            .lock()
1182            .unwrap_or_else(std::sync::PoisonError::into_inner)
1183            .clone();
1184        stats.entries = cache.len();
1185        stats.size_bytes = cache.current_size_bytes;
1186        stats
1187    }
1188
1189    fn cache_entry_binding(&self, request: &JsonRpcRequest) -> Option<CacheEntryBinding> {
1190        if request.method != SERVER_DISCOVER_METHOD {
1191            return Some(CacheEntryBinding::Ordinary);
1192        }
1193
1194        let era = discovery_request_protocol_era(request)?;
1195        if era != ProtocolEra::Modern2026 {
1196            return None;
1197        }
1198        let generation = self.discovery_generation.load(Ordering::Acquire);
1199        (generation != 0).then_some(CacheEntryBinding::Discovery(DiscoveryCacheBinding {
1200            era,
1201            generation,
1202        }))
1203    }
1204
1205    fn cache_entry_binding_is_current(&self, binding: CacheEntryBinding) -> bool {
1206        match binding {
1207            CacheEntryBinding::Ordinary => true,
1208            CacheEntryBinding::Discovery(binding) => {
1209                binding.era == ProtocolEra::Modern2026
1210                    && binding.generation == self.discovery_generation.load(Ordering::Acquire)
1211            }
1212        }
1213    }
1214
1215    /// Advances the discovery generation, permanently disabling discovery
1216    /// caching if the counter can no longer advance without wrapping.
1217    fn advance_discovery_generation(&self) {
1218        if self
1219            .discovery_generation
1220            .try_update(Ordering::AcqRel, Ordering::Acquire, |generation| {
1221                (generation != 0)
1222                    .then(|| generation.checked_add(1))
1223                    .flatten()
1224            })
1225            .is_err()
1226        {
1227            self.discovery_generation.store(0, Ordering::Release);
1228        }
1229    }
1230
1231    /// Clears the entire cache.
1232    pub fn clear(&self) {
1233        let mut cache = self
1234            .cache
1235            .lock()
1236            .unwrap_or_else(std::sync::PoisonError::into_inner);
1237        self.advance_discovery_generation();
1238        cache.clear();
1239    }
1240
1241    /// Invalidates every final discovery response and advances its cache
1242    /// generation. A response or lookup holding an older generation cannot
1243    /// reuse or repopulate the invalidated discovery state.
1244    pub fn invalidate_discovery(&self) {
1245        let mut cache = self
1246            .cache
1247            .lock()
1248            .unwrap_or_else(std::sync::PoisonError::into_inner);
1249        self.advance_discovery_generation();
1250        cache.remove_discovery_entries();
1251    }
1252
1253    /// Invalidates every session/auth partition and every cursor page for a
1254    /// method and semantic-parameter set.
1255    pub fn invalidate(&self, method: &str, params: Option<&serde_json::Value>) {
1256        if method == SERVER_DISCOVER_METHOD {
1257            self.invalidate_discovery();
1258            return;
1259        }
1260        let Some(invalidation_digest) = CacheKey::try_invalidation_digest(method, params) else {
1261            return;
1262        };
1263        let mut cache = self
1264            .cache
1265            .lock()
1266            .unwrap_or_else(std::sync::PoisonError::into_inner);
1267        cache.remove_invalidation_digest(invalidation_digest);
1268    }
1269
1270    /// Checks if a method should be cached.
1271    fn should_cache_method(&self, method: &str, params: Option<&serde_json::Value>) -> bool {
1272        match method {
1273            "server/discover" | "tools/list" => self.tools_list_config.enabled,
1274            "resources/list" | "resources/templates/list" => self.resources_list_config.enabled,
1275            "prompts/list" => self.prompts_list_config.enabled,
1276            "resources/read" => self.resources_read_config.enabled,
1277            "prompts/get" => self.prompts_get_config.enabled,
1278            "tools/call" => {
1279                if !self.tools_call_config.base.enabled {
1280                    return false;
1281                }
1282                // Extract tool name from params
1283                if let Some(params) = params {
1284                    if let Some(tool_name) = params.get("name").and_then(|v| v.as_str()) {
1285                        return self.tools_call_config.should_cache_tool(tool_name);
1286                    }
1287                }
1288                false
1289            }
1290            _ => false,
1291        }
1292    }
1293
1294    /// Gets the TTL for a specific method.
1295    fn get_ttl(&self, method: &str) -> Duration {
1296        match method {
1297            "server/discover" | "tools/list" => {
1298                Duration::from_secs(self.tools_list_config.ttl_secs)
1299            }
1300            "resources/list" | "resources/templates/list" => {
1301                Duration::from_secs(self.resources_list_config.ttl_secs)
1302            }
1303            "prompts/list" => Duration::from_secs(self.prompts_list_config.ttl_secs),
1304            "tools/call" => Duration::from_secs(self.tools_call_config.base.ttl_secs),
1305            "resources/read" => Duration::from_secs(self.resources_read_config.ttl_secs),
1306            "prompts/get" => Duration::from_secs(self.prompts_get_config.ttl_secs),
1307            _ => self.call_ttl,
1308        }
1309    }
1310
1311    fn protocol_cache_ttl(&self, method: &str) -> CacheTtl {
1312        CacheTtl::milliseconds(u64::try_from(self.get_ttl(method).as_millis()).unwrap_or(u64::MAX))
1313    }
1314
1315    /// Normalizes modern protocol cache hints at the server boundary.
1316    ///
1317    /// A wire-valid final cache policy is an upstream result field, not a
1318    /// bounded local expiry. Preserve it byte-for-byte through serialization;
1319    /// the cache admission path can independently decline a value outside its
1320    /// runtime duration domain.
1321    fn apply_protocol_cache_hints(&self, method: &str, response: &mut serde_json::Value) {
1322        if method == SERVER_DISCOVER_METHOD {
1323            if is_final_discovery_result(response) {
1324                if !final_discovery_cache_hints_are_wire_valid(response)
1325                    && let Some(response) = response.as_object_mut()
1326                {
1327                    response.remove("ttlMs");
1328                    response.remove("cacheScope");
1329                }
1330                return;
1331            }
1332        }
1333        let Some(response) = response.as_object_mut() else {
1334            return;
1335        };
1336        if !method_requires_protocol_cache_hints(method)
1337            || response
1338                .get("resultType")
1339                .and_then(serde_json::Value::as_str)
1340                != Some("complete")
1341        {
1342            // Cache hints are valid only on the explicitly cacheable modern
1343            // complete-result branches. Do not preserve a handler- or peer-
1344            // supplied lookalike on input-required, task, legacy, or
1345            // non-cacheable method results.
1346            response.remove("ttlMs");
1347            response.remove("cacheScope");
1348            return;
1349        }
1350
1351        let has_wire_valid_cache_hints = response
1352            .get("ttlMs")
1353            .cloned()
1354            .and_then(|ttl| serde_json::from_value::<CacheTtl>(ttl).ok())
1355            .is_some()
1356            && matches!(
1357                response
1358                    .get("cacheScope")
1359                    .and_then(serde_json::Value::as_str),
1360                Some("private" | "public")
1361            );
1362        if has_wire_valid_cache_hints {
1363            return;
1364        }
1365
1366        let ttl = self.protocol_cache_ttl(method);
1367        response.insert(
1368            "ttlMs".to_owned(),
1369            serde_json::to_value(ttl).expect("cache TTL serializes to JSON"),
1370        );
1371        response.insert(
1372            "cacheScope".to_owned(),
1373            serde_json::Value::String("private".to_owned()),
1374        );
1375    }
1376
1377    fn record_hit(&self) {
1378        let mut stats = self
1379            .stats
1380            .lock()
1381            .unwrap_or_else(std::sync::PoisonError::into_inner);
1382        stats.hits = stats.hits.saturating_add(1);
1383    }
1384
1385    fn record_miss(&self) {
1386        let mut stats = self
1387            .stats
1388            .lock()
1389            .unwrap_or_else(std::sync::PoisonError::into_inner);
1390        stats.misses = stats.misses.saturating_add(1);
1391    }
1392}
1393
1394impl Middleware for ResponseCachingMiddleware {
1395    fn on_request(
1396        &self,
1397        ctx: &McpContext,
1398        request: &JsonRpcRequest,
1399    ) -> McpResult<MiddlewareDecision> {
1400        if self.instance_id == 0 {
1401            return Ok(MiddlewareDecision::Continue);
1402        }
1403        // Check if this method should be cached
1404        if !self.should_cache_method(&request.method, request.params.as_ref()) {
1405            return Ok(MiddlewareDecision::Continue);
1406        }
1407        if request_carries_uncacheable_continuation(request.params.as_ref()) {
1408            return Ok(MiddlewareDecision::Continue);
1409        }
1410        let Some(binding) = self.cache_entry_binding(request) else {
1411            return Ok(MiddlewareDecision::Continue);
1412        };
1413
1414        let Some(partition_digest) = context_cache_partition(ctx, CachePartitionPhase::Request)
1415        else {
1416            return Ok(MiddlewareDecision::Continue);
1417        };
1418
1419        // Try to get cached response
1420        let Some(key) = CacheKey::try_new_partitioned(
1421            &request.method,
1422            request.params.as_ref(),
1423            partition_digest,
1424            binding,
1425        ) else {
1426            return Ok(MiddlewareDecision::Continue);
1427        };
1428        let encoded = {
1429            let mut cache = self
1430                .cache
1431                .lock()
1432                .unwrap_or_else(std::sync::PoisonError::into_inner);
1433            cache.get_encoded(&key)
1434        };
1435
1436        if let Some(encoded) = encoded {
1437            if let Some(value) = decode_cached_json(&encoded) {
1438                if !self.cache_entry_binding_is_current(binding) {
1439                    self.record_miss();
1440                    return Ok(MiddlewareDecision::Continue);
1441                }
1442                // Session state can change while this request waits for the
1443                // cache mutex or decodes a cached payload. Revalidate after
1444                // both operations so a hit linearizes against the admitted
1445                // revision instead of serving an entry made stale before the
1446                // lookup completed. The final liveness check applies the same
1447                // completion rule to cancellation and request-lease closure.
1448                if !context_cache_commit_is_admissible(ctx) {
1449                    return Ok(MiddlewareDecision::Continue);
1450                }
1451                if !ctx.mark_response_cache_hit(self.instance_id) {
1452                    self.record_miss();
1453                    return Ok(MiddlewareDecision::Continue);
1454                }
1455                self.record_hit();
1456                return Ok(MiddlewareDecision::Respond(value));
1457            }
1458            let mut cache = self
1459                .cache
1460                .lock()
1461                .unwrap_or_else(std::sync::PoisonError::into_inner);
1462            cache.remove(&key);
1463        }
1464
1465        self.record_miss();
1466        Ok(MiddlewareDecision::Continue)
1467    }
1468
1469    fn on_response(
1470        &self,
1471        ctx: &McpContext,
1472        request: &JsonRpcRequest,
1473        mut response: serde_json::Value,
1474    ) -> McpResult<serde_json::Value> {
1475        if self.instance_id == 0 {
1476            return Ok(response);
1477        }
1478        self.apply_protocol_cache_hints(&request.method, &mut response);
1479        let final_discovery_policy = (request.method == SERVER_DISCOVER_METHOD
1480            && is_final_discovery_result(&response))
1481        .then(|| final_discovery_cache_policy(&response))
1482        .flatten();
1483        if request.method == SERVER_DISCOVER_METHOD
1484            && is_final_discovery_result(&response)
1485            && final_discovery_policy.is_none()
1486        {
1487            return Ok(response);
1488        }
1489        let is_final_complete = method_requires_protocol_cache_hints(&request.method)
1490            && response
1491                .get("resultType")
1492                .and_then(serde_json::Value::as_str)
1493                == Some("complete");
1494        let final_complete_policy = is_final_complete
1495            .then(|| final_complete_cache_policy(&response))
1496            .flatten();
1497        // A wire-valid arbitrary-width TTL is still delivered, but cannot be
1498        // represented by this process's bounded expiry clock. Never fall back
1499        // to a configured local TTL: that would create a cache entry whose
1500        // expiry has no relationship to the final result's wire policy.
1501        if is_final_complete && final_complete_policy.is_none() {
1502            return Ok(response);
1503        }
1504        // Only cache if this method is cacheable
1505        if !self.should_cache_method(&request.method, request.params.as_ref()) {
1506            return Ok(response);
1507        }
1508        if request_carries_uncacheable_continuation(request.params.as_ref())
1509            || !response_is_cacheable_complete(&response)
1510        {
1511            return Ok(response);
1512        }
1513        if ctx.response_was_cache_hit(self.instance_id) {
1514            return Ok(response);
1515        }
1516        let Some(binding) = self.cache_entry_binding(request) else {
1517            return Ok(response);
1518        };
1519
1520        let Some(partition_digest) = context_cache_partition(ctx, CachePartitionPhase::Response)
1521        else {
1522            return Ok(response);
1523        };
1524
1525        // Store in cache
1526        let Some(key) = CacheKey::try_new_partitioned(
1527            &request.method,
1528            request.params.as_ref(),
1529            partition_digest,
1530            binding,
1531        ) else {
1532            return Ok(response);
1533        };
1534        let ttl = match final_discovery_policy {
1535            Some(FinalDiscoveryCachePolicy::Private(ttl)) => ttl,
1536            Some(FinalDiscoveryCachePolicy::Public) => return Ok(response),
1537            None => match final_complete_policy {
1538                Some(FinalCompleteCachePolicy::Private(ttl)) => ttl,
1539                Some(FinalCompleteCachePolicy::Public) => return Ok(response),
1540                None => self.get_ttl(&request.method),
1541            },
1542        };
1543
1544        let admission_limit = {
1545            let cache = self
1546                .cache
1547                .lock()
1548                .unwrap_or_else(std::sync::PoisonError::into_inner);
1549            cache.max_item_size.min(
1550                cache
1551                    .max_size_bytes
1552                    .saturating_sub(CACHE_ENTRY_METADATA_BYTES),
1553            )
1554        };
1555        let Some(encoded) = encode_json_bounded(&response, admission_limit) else {
1556            return Ok(response);
1557        };
1558        let mut cache = self
1559            .cache
1560            .lock()
1561            .unwrap_or_else(std::sync::PoisonError::into_inner);
1562        // Encoding and waiting for the cache mutex may both be non-trivial for
1563        // a response near the configured limits. Revalidate at the commit
1564        // boundary so cancellation, lease closure, or session mutation cannot
1565        // populate the cache after winning either race.
1566        if !context_cache_commit_is_admissible(ctx) {
1567            return Ok(response);
1568        }
1569        if !self.cache_entry_binding_is_current(binding) {
1570            return Ok(response);
1571        }
1572        cache.insert_encoded(key, encoded, ttl);
1573
1574        Ok(response)
1575    }
1576
1577    fn on_error(&self, _ctx: &McpContext, _request: &JsonRpcRequest, error: McpError) -> McpError {
1578        // Don't cache errors, just pass them through
1579        error
1580    }
1581}
1582
1583#[cfg(test)]
1584mod tests {
1585    use super::*;
1586    use asupersync::Cx;
1587    use fastmcp_core::{AuthContext, SessionState};
1588
1589    fn maximum_geometric_growth_events(max_bytes: usize) -> usize {
1590        if max_bytes == 0 {
1591            return 0;
1592        }
1593
1594        let mut capacity = CACHE_BYTES_GROWTH_CHUNK.min(max_bytes);
1595        let mut events = 1_usize;
1596        while capacity < max_bytes {
1597            capacity = capacity.checked_mul(2).unwrap_or(max_bytes).min(max_bytes);
1598            events = events.saturating_add(1);
1599        }
1600        events
1601    }
1602
1603    fn test_context() -> McpContext {
1604        let cx = Cx::for_testing();
1605        let ctx = McpContext::new(cx, 1);
1606        assert!(ctx.commit_anonymous_auth());
1607        ctx
1608    }
1609
1610    fn partitioned_context(state: &SessionState, request_id: u64, auth: AuthContext) -> McpContext {
1611        McpContext::with_state(Cx::for_testing(), request_id, state.clone()).with_auth(auth)
1612    }
1613
1614    fn anonymous_partitioned_context(state: &SessionState, request_id: u64) -> McpContext {
1615        let ctx = McpContext::with_state(Cx::for_testing(), request_id, state.clone());
1616        assert!(ctx.commit_anonymous_auth());
1617        assert!(ctx.auth().is_none());
1618        ctx
1619    }
1620
1621    fn test_request(method: &str, params: Option<serde_json::Value>) -> JsonRpcRequest {
1622        JsonRpcRequest {
1623            jsonrpc: std::borrow::Cow::Borrowed(fastmcp_protocol::JSONRPC_VERSION),
1624            method: method.to_string(),
1625            params,
1626            id: Some(fastmcp_protocol::RequestId::Number(1)),
1627        }
1628    }
1629
1630    fn final_discovery_request(protocol_version: &str) -> JsonRpcRequest {
1631        test_request(
1632            SERVER_DISCOVER_METHOD,
1633            Some(serde_json::json!({
1634                "_meta": {
1635                    FINAL_PROTOCOL_VERSION_META_KEY: protocol_version,
1636                },
1637            })),
1638        )
1639    }
1640
1641    fn final_discovery_response(ttl_ms: u64, cache_scope: &str) -> serde_json::Value {
1642        serde_json::json!({
1643            "supportedVersions": [FINAL_PROTOCOL_VERSION],
1644            "capabilities": {},
1645            "ttlMs": ttl_ms,
1646            "cacheScope": cache_scope,
1647        })
1648    }
1649
1650    // ========================================
1651    // LruCache tests
1652    // ========================================
1653
1654    #[test]
1655    fn test_lru_cache_basic_operations() {
1656        let mut cache = LruCache::new(10, 1024 * 1024, 1024);
1657
1658        let key = CacheKey::new("test", None);
1659        let value = serde_json::json!({"result": "cached"});
1660
1661        // Insert and retrieve
1662        cache.insert(key.clone(), value.clone(), Duration::from_secs(60));
1663        let retrieved = cache.get_value(&key);
1664        assert_eq!(retrieved, Some(value));
1665    }
1666
1667    #[test]
1668    fn test_lru_cache_expiration() {
1669        let mut cache = LruCache::new(10, 1024 * 1024, 1024);
1670
1671        let key = CacheKey::new("test", None);
1672        let value = serde_json::json!({"result": "cached"});
1673
1674        // Insert with very short TTL
1675        cache.insert(key.clone(), value, Duration::from_millis(1));
1676
1677        // Wait for expiration
1678        std::thread::sleep(std::time::Duration::from_millis(10));
1679
1680        // Should be expired
1681        assert!(cache.get_value(&key).is_none());
1682    }
1683
1684    #[test]
1685    fn test_lru_cache_eviction() {
1686        let mut cache = LruCache::new(2, 1024 * 1024, 1024);
1687
1688        let key1 = CacheKey::new("test1", None);
1689        let key2 = CacheKey::new("test2", None);
1690        let key3 = CacheKey::new("test3", None);
1691
1692        cache.insert(
1693            key1.clone(),
1694            serde_json::json!("v1"),
1695            Duration::from_secs(60),
1696        );
1697        cache.insert(
1698            key2.clone(),
1699            serde_json::json!("v2"),
1700            Duration::from_secs(60),
1701        );
1702
1703        // Should evict key1 (LRU)
1704        cache.insert(
1705            key3.clone(),
1706            serde_json::json!("v3"),
1707            Duration::from_secs(60),
1708        );
1709
1710        assert!(cache.get_value(&key1).is_none());
1711        assert!(cache.get_value(&key2).is_some());
1712        assert!(cache.get_value(&key3).is_some());
1713    }
1714
1715    #[test]
1716    fn test_lru_cache_size_limit() {
1717        let mut cache = LruCache::new(100, CACHE_ENTRY_METADATA_BYTES + 16, 1024);
1718
1719        let key1 = CacheKey::new("test1", None);
1720        let key2 = CacheKey::new("test2", None);
1721
1722        // First entry should fit
1723        cache.insert(
1724            key1.clone(),
1725            serde_json::json!("short"),
1726            Duration::from_secs(60),
1727        );
1728        assert_eq!(cache.len(), 1);
1729
1730        // Second entry should cause eviction
1731        cache.insert(
1732            key2.clone(),
1733            serde_json::json!("another"),
1734            Duration::from_secs(60),
1735        );
1736        assert!(cache.get_value(&key1).is_none());
1737        assert_eq!(cache.get_value(&key2), Some(serde_json::json!("another")));
1738    }
1739
1740    #[test]
1741    fn test_lru_cache_oversized_item_rejected() {
1742        let mut cache = LruCache::new(10, 1024 * 1024, 10); // max 10 bytes per item
1743
1744        let key = CacheKey::new("test", None);
1745        let large_value = serde_json::json!({"data": "this is much longer than 10 bytes"});
1746
1747        cache.insert(key.clone(), large_value, Duration::from_secs(60));
1748
1749        // Should not be stored
1750        assert!(cache.get_value(&key).is_none());
1751    }
1752
1753    #[test]
1754    fn lru_cache_zero_entry_limit_rejects_every_insert() {
1755        let mut cache = LruCache::new(0, 1024, 1024);
1756        let key = CacheKey::new("test", None);
1757
1758        cache.insert(
1759            key.clone(),
1760            serde_json::json!("value"),
1761            Duration::from_secs(60),
1762        );
1763
1764        assert!(cache.get_value(&key).is_none());
1765        assert_eq!(cache.len(), 0);
1766        assert_eq!(cache.current_size_bytes, 0);
1767    }
1768
1769    #[test]
1770    fn lru_cache_zero_total_size_rejects_every_insert() {
1771        let mut cache = LruCache::new(10, 0, 1024);
1772        let key = CacheKey::new("test", None);
1773
1774        cache.insert(
1775            key.clone(),
1776            serde_json::json!("value"),
1777            Duration::from_secs(60),
1778        );
1779
1780        assert!(cache.get_value(&key).is_none());
1781        assert_eq!(cache.len(), 0);
1782        assert_eq!(cache.current_size_bytes, 0);
1783    }
1784
1785    #[test]
1786    fn lru_cache_item_larger_than_total_capacity_is_rejected() {
1787        let value = serde_json::json!("larger than capacity");
1788        let value_size = value.to_string().len();
1789        let mut cache = LruCache::new(10, value_size - 1, value_size + 100);
1790        let key = CacheKey::new("test", None);
1791
1792        cache.insert(key.clone(), value, Duration::from_secs(60));
1793
1794        assert!(cache.get_value(&key).is_none());
1795        assert_eq!(cache.current_size_bytes, 0);
1796    }
1797
1798    #[test]
1799    fn lru_cache_rejected_replacement_preserves_existing_entry_and_accounting() {
1800        let mut cache = LruCache::new(10, 1024, 12);
1801        let key = CacheKey::new("test", None);
1802        let original = serde_json::json!("small");
1803
1804        cache.insert(key.clone(), original.clone(), Duration::from_secs(60));
1805        let original_size = cache.current_size_bytes;
1806
1807        cache.insert(
1808            key.clone(),
1809            serde_json::json!("this replacement is too large"),
1810            Duration::from_secs(60),
1811        );
1812
1813        assert_eq!(cache.get_value(&key), Some(original));
1814        assert_eq!(cache.len(), 1);
1815        assert_eq!(cache.current_size_bytes, original_size);
1816    }
1817
1818    #[test]
1819    fn lru_cache_unrepresentable_ttl_is_rejected_without_mutation() {
1820        let mut cache = LruCache::new(10, 1024, 1024);
1821        let key = CacheKey::new("test", None);
1822
1823        cache.insert(key.clone(), serde_json::json!("value"), Duration::MAX);
1824
1825        assert!(cache.get_value(&key).is_none());
1826        assert_eq!(cache.len(), 0);
1827        assert_eq!(cache.current_size_bytes, 0);
1828    }
1829
1830    #[test]
1831    fn lru_cache_zero_ttl_is_never_observable() {
1832        let mut cache = LruCache::new(10, 1024, 1024);
1833        let key = CacheKey::new("test", None);
1834
1835        cache.insert(key.clone(), serde_json::json!("value"), Duration::ZERO);
1836
1837        assert!(cache.get_value(&key).is_none());
1838        assert_eq!(cache.len(), 0);
1839        assert_eq!(cache.current_size_bytes, 0);
1840    }
1841
1842    // ========================================
1843    // ResponseCachingMiddleware tests
1844    // ========================================
1845
1846    #[test]
1847    fn test_caching_middleware_caches_tools_list() {
1848        let middleware = ResponseCachingMiddleware::new();
1849        let ctx = test_context();
1850        let request = test_request("tools/list", None);
1851
1852        // First request: miss, continue
1853        let decision = middleware.on_request(&ctx, &request).unwrap();
1854        assert!(matches!(decision, MiddlewareDecision::Continue));
1855
1856        // Simulate response
1857        let response = serde_json::json!({"tools": []});
1858        middleware
1859            .on_response(&ctx, &request, response.clone())
1860            .unwrap();
1861
1862        // Second request: hit, respond from cache
1863        let decision = middleware.on_request(&ctx, &request).unwrap();
1864        assert!(
1865            matches!(decision, MiddlewareDecision::Respond(_)),
1866            "Expected cache hit"
1867        );
1868        let MiddlewareDecision::Respond(cached) = decision else {
1869            return;
1870        };
1871        assert_eq!(cached, response);
1872
1873        // Check stats
1874        let stats = middleware.stats();
1875        assert_eq!(stats.hits, 1);
1876        assert_eq!(stats.misses, 1);
1877    }
1878
1879    #[test]
1880    fn cache_hit_response_does_not_refresh_absolute_expiration() {
1881        let middleware = ResponseCachingMiddleware::new().list_ttl_secs(60);
1882        let ctx = test_context();
1883        let request = test_request("tools/list", None);
1884        let response = serde_json::json!({"tools": []});
1885
1886        assert!(matches!(
1887            middleware.on_request(&ctx, &request).unwrap(),
1888            MiddlewareDecision::Continue
1889        ));
1890        middleware
1891            .on_response(&ctx, &request, response.clone())
1892            .unwrap();
1893        let expires_before_hit = middleware
1894            .cache
1895            .lock()
1896            .unwrap_or_else(std::sync::PoisonError::into_inner)
1897            .entries
1898            .values()
1899            .next()
1900            .expect("cached entry")
1901            .expires_at;
1902
1903        assert!(matches!(
1904            middleware.on_request(&ctx, &request).unwrap(),
1905            MiddlewareDecision::Respond(_)
1906        ));
1907        middleware.on_response(&ctx, &request, response).unwrap();
1908
1909        let expires_after_hit = middleware
1910            .cache
1911            .lock()
1912            .unwrap_or_else(std::sync::PoisonError::into_inner)
1913            .entries
1914            .values()
1915            .next()
1916            .expect("cache hit must retain the original entry")
1917            .expires_at;
1918        assert_eq!(expires_after_hit, expires_before_hit);
1919    }
1920
1921    #[test]
1922    fn downstream_cache_hit_does_not_prevent_upstream_cache_warming() {
1923        let upstream = ResponseCachingMiddleware::new();
1924        let downstream = ResponseCachingMiddleware::new();
1925        let request = test_request("tools/list", None);
1926        let response = serde_json::json!({"tools": ["warm"]});
1927
1928        let prewarm = test_context();
1929        assert!(matches!(
1930            downstream.on_request(&prewarm, &request).unwrap(),
1931            MiddlewareDecision::Continue
1932        ));
1933        downstream
1934            .on_response(&prewarm, &request, response.clone())
1935            .unwrap();
1936
1937        let shared_request = test_context();
1938        assert!(matches!(
1939            upstream.on_request(&shared_request, &request).unwrap(),
1940            MiddlewareDecision::Continue
1941        ));
1942        let MiddlewareDecision::Respond(cached) =
1943            downstream.on_request(&shared_request, &request).unwrap()
1944        else {
1945            panic!("downstream cache was not prewarmed");
1946        };
1947        downstream
1948            .on_response(&shared_request, &request, cached.clone())
1949            .unwrap();
1950        upstream
1951            .on_response(&shared_request, &request, cached)
1952            .unwrap();
1953
1954        assert!(matches!(
1955            upstream.on_request(&test_context(), &request).unwrap(),
1956            MiddlewareDecision::Respond(value) if value == response
1957        ));
1958    }
1959
1960    #[test]
1961    fn test_caching_middleware_skips_non_cacheable_methods() {
1962        let middleware = ResponseCachingMiddleware::new();
1963        let ctx = test_context();
1964        let request = test_request("initialize", None);
1965
1966        // Should continue (not cached)
1967        let decision = middleware.on_request(&ctx, &request).unwrap();
1968        assert!(matches!(decision, MiddlewareDecision::Continue));
1969
1970        // Even after response, next request should not hit cache
1971        middleware
1972            .on_response(&ctx, &request, serde_json::json!({}))
1973            .unwrap();
1974
1975        let decision = middleware.on_request(&ctx, &request).unwrap();
1976        assert!(matches!(decision, MiddlewareDecision::Continue));
1977    }
1978
1979    #[test]
1980    fn test_caching_middleware_different_params_different_keys() {
1981        let middleware = ResponseCachingMiddleware::new()
1982            .include_tools(vec!["tool_a".to_string(), "tool_b".to_string()]);
1983        let ctx = test_context();
1984
1985        let request1 = test_request(
1986            "tools/call",
1987            Some(serde_json::json!({"name": "tool_a", "arguments": {}})),
1988        );
1989        let request2 = test_request(
1990            "tools/call",
1991            Some(serde_json::json!({"name": "tool_b", "arguments": {}})),
1992        );
1993
1994        // Cache response for request1
1995        middleware.on_request(&ctx, &request1).unwrap();
1996        let response1 = serde_json::json!({"result": "a"});
1997        middleware
1998            .on_response(&ctx, &request1, response1.clone())
1999            .unwrap();
2000
2001        // Request2 should not hit cache
2002        let decision = middleware.on_request(&ctx, &request2).unwrap();
2003        assert!(matches!(decision, MiddlewareDecision::Continue));
2004
2005        // Request1 should hit cache
2006        let decision = middleware.on_request(&ctx, &request1).unwrap();
2007        assert!(
2008            matches!(decision, MiddlewareDecision::Respond(_)),
2009            "Expected cache hit"
2010        );
2011        let MiddlewareDecision::Respond(cached) = decision else {
2012            return;
2013        };
2014        assert_eq!(cached, response1);
2015    }
2016
2017    #[test]
2018    fn test_caching_middleware_tool_exclusion() {
2019        let middleware = ResponseCachingMiddleware::new()
2020            .include_tools(vec![
2021                "excluded_tool".to_string(),
2022                "included_tool".to_string(),
2023            ])
2024            .exclude_tools(vec!["excluded_tool".to_string()]);
2025        let ctx = test_context();
2026
2027        let excluded_request = test_request(
2028            "tools/call",
2029            Some(serde_json::json!({"name": "excluded_tool", "arguments": {}})),
2030        );
2031        let included_request = test_request(
2032            "tools/call",
2033            Some(serde_json::json!({"name": "included_tool", "arguments": {}})),
2034        );
2035
2036        // Excluded tool should not be cached
2037        middleware.on_request(&ctx, &excluded_request).unwrap();
2038        middleware
2039            .on_response(&ctx, &excluded_request, serde_json::json!({}))
2040            .unwrap();
2041
2042        let decision = middleware.on_request(&ctx, &excluded_request).unwrap();
2043        assert!(matches!(decision, MiddlewareDecision::Continue));
2044
2045        // Included tool should be cached
2046        middleware.on_request(&ctx, &included_request).unwrap();
2047        let response = serde_json::json!({"result": "included"});
2048        middleware
2049            .on_response(&ctx, &included_request, response.clone())
2050            .unwrap();
2051
2052        let decision = middleware.on_request(&ctx, &included_request).unwrap();
2053        assert!(
2054            matches!(decision, MiddlewareDecision::Respond(_)),
2055            "Expected cache hit for included tool"
2056        );
2057        let MiddlewareDecision::Respond(cached) = decision else {
2058            return;
2059        };
2060        assert_eq!(cached, response);
2061    }
2062
2063    #[test]
2064    fn test_caching_middleware_disable_method() {
2065        let middleware = ResponseCachingMiddleware::new().disable_tools_list();
2066        let ctx = test_context();
2067        let request = test_request("tools/list", None);
2068
2069        // Should not cache
2070        middleware.on_request(&ctx, &request).unwrap();
2071        middleware
2072            .on_response(&ctx, &request, serde_json::json!({}))
2073            .unwrap();
2074
2075        let decision = middleware.on_request(&ctx, &request).unwrap();
2076        assert!(matches!(decision, MiddlewareDecision::Continue));
2077    }
2078
2079    #[test]
2080    fn tools_call_is_not_cached_without_an_explicit_allowlist() {
2081        let middleware = ResponseCachingMiddleware::new();
2082        let ctx = test_context();
2083        let request = test_request(
2084            "tools/call",
2085            Some(serde_json::json!({"name": "stateful_tool", "arguments": {}})),
2086        );
2087
2088        assert!(matches!(
2089            middleware.on_request(&ctx, &request).unwrap(),
2090            MiddlewareDecision::Continue
2091        ));
2092        middleware
2093            .on_response(
2094                &ctx,
2095                &request,
2096                serde_json::json!({"result": "must not be stored"}),
2097            )
2098            .unwrap();
2099
2100        assert!(matches!(
2101            middleware.on_request(&ctx, &request).unwrap(),
2102            MiddlewareDecision::Continue
2103        ));
2104        assert_eq!(middleware.stats().entries, 0);
2105    }
2106
2107    #[test]
2108    fn explicitly_allowlisted_tool_can_cache_in_unpartitioned_context() {
2109        let middleware =
2110            ResponseCachingMiddleware::new().include_tools(vec!["pure_tool".to_string()]);
2111        let ctx = test_context();
2112        let request = test_request(
2113            "tools/call",
2114            Some(serde_json::json!({"name": "pure_tool", "arguments": {"x": 1}})),
2115        );
2116        let response = serde_json::json!({"result": 2});
2117
2118        assert!(matches!(
2119            middleware.on_request(&ctx, &request).unwrap(),
2120            MiddlewareDecision::Continue
2121        ));
2122        middleware
2123            .on_response(&ctx, &request, response.clone())
2124            .unwrap();
2125
2126        let MiddlewareDecision::Respond(cached) = middleware.on_request(&ctx, &request).unwrap()
2127        else {
2128            panic!("explicitly allowlisted tool did not produce a cache hit");
2129        };
2130        assert_eq!(cached, response);
2131    }
2132
2133    #[test]
2134    fn stateless_cache_partitions_anonymous_and_authenticated_requests() {
2135        let middleware = ResponseCachingMiddleware::new();
2136        let anonymous_ctx = test_context();
2137        let authenticated_ctx = McpContext::new(Cx::for_testing(), 2)
2138            .with_auth(AuthContext::with_subject("principal-a"));
2139        let request = test_request("tools/list", None);
2140        let public_response = serde_json::json!({"tools": ["public"]});
2141        let private_response = serde_json::json!({"tools": ["private"]});
2142        assert!(authenticated_ctx.auth().is_some());
2143
2144        middleware
2145            .on_response(&anonymous_ctx, &request, public_response.clone())
2146            .unwrap();
2147
2148        // An authenticated request must not read an anonymous entry.
2149        assert!(matches!(
2150            middleware.on_request(&authenticated_ctx, &request).unwrap(),
2151            MiddlewareDecision::Continue
2152        ));
2153
2154        // Its response receives a separate stateless authorization partition.
2155        middleware
2156            .on_response(&authenticated_ctx, &request, private_response.clone())
2157            .unwrap();
2158        let MiddlewareDecision::Respond(cached) =
2159            middleware.on_request(&anonymous_ctx, &request).unwrap()
2160        else {
2161            panic!("authenticated response unexpectedly replaced the public entry");
2162        };
2163        assert_eq!(cached, public_response);
2164
2165        let authenticated_retry = McpContext::new(Cx::for_testing(), 3)
2166            .with_auth(AuthContext::with_subject("principal-a"));
2167        let MiddlewareDecision::Respond(cached) = middleware
2168            .on_request(&authenticated_retry, &request)
2169            .unwrap()
2170        else {
2171            panic!("same authenticated stateless partition did not hit");
2172        };
2173        assert_eq!(cached, private_response);
2174        assert_eq!(middleware.stats().entries, 2);
2175    }
2176
2177    #[test]
2178    fn stateless_cache_partitions_identical_auth_facts_by_session_owner() {
2179        let middleware = ResponseCachingMiddleware::new();
2180        let first_auth = AuthContext::with_subject("same-display")
2181            .with_session_owner(Sha256Digest::from_bytes([1; 32]));
2182        let second_auth = AuthContext::with_subject("same-display")
2183            .with_session_owner(Sha256Digest::from_bytes([2; 32]));
2184        let first_ctx = McpContext::new(Cx::for_testing(), 1).with_auth(first_auth.clone());
2185        let second_ctx = McpContext::new(Cx::for_testing(), 2).with_auth(second_auth);
2186        let request = test_request("tools/list", None);
2187        let first_response = serde_json::json!({"tools": ["owner-one"]});
2188
2189        middleware
2190            .on_response(&first_ctx, &request, first_response.clone())
2191            .unwrap();
2192        assert!(matches!(
2193            middleware.on_request(&second_ctx, &request).unwrap(),
2194            MiddlewareDecision::Continue
2195        ));
2196
2197        let same_owner_retry = McpContext::new(Cx::for_testing(), 3).with_auth(first_auth);
2198        let MiddlewareDecision::Respond(cached) =
2199            middleware.on_request(&same_owner_retry, &request).unwrap()
2200        else {
2201            panic!("the same stateless owner partition did not hit");
2202        };
2203        assert_eq!(cached, first_response);
2204    }
2205
2206    #[test]
2207    fn stateless_cache_frames_absent_owner_separately_from_zero_owner() {
2208        let middleware = ResponseCachingMiddleware::new();
2209        let ownerless_auth = AuthContext::with_subject("same-display");
2210        let zero_owner_auth = ownerless_auth
2211            .clone()
2212            .with_session_owner(Sha256Digest::from_bytes([0; 32]));
2213        let ownerless_ctx = McpContext::new(Cx::for_testing(), 1).with_auth(ownerless_auth);
2214        let zero_owner_ctx = McpContext::new(Cx::for_testing(), 2).with_auth(zero_owner_auth);
2215        let request = test_request("tools/list", None);
2216        let ownerless_response = serde_json::json!({"tools": ["ownerless"]});
2217
2218        middleware
2219            .on_response(&ownerless_ctx, &request, ownerless_response.clone())
2220            .unwrap();
2221
2222        assert!(matches!(
2223            middleware.on_request(&zero_owner_ctx, &request).unwrap(),
2224            MiddlewareDecision::Continue
2225        ));
2226        let MiddlewareDecision::Respond(cached) =
2227            middleware.on_request(&ownerless_ctx, &request).unwrap()
2228        else {
2229            panic!("the ownerless cache partition was no longer retrievable");
2230        };
2231        assert_eq!(cached, ownerless_response);
2232        assert_eq!(middleware.stats().entries, 1);
2233    }
2234
2235    #[test]
2236    fn session_without_committed_auth_bypasses_lookup_and_storage() {
2237        let middleware = ResponseCachingMiddleware::new();
2238        let anonymous_ctx = test_context();
2239        let session_ctx = McpContext::with_state(Cx::for_testing(), 2, SessionState::new());
2240        let request = test_request("resources/list", None);
2241        let public_response = serde_json::json!({"resources": ["public"]});
2242        assert!(session_ctx.has_session_state());
2243
2244        middleware
2245            .on_response(&anonymous_ctx, &request, public_response.clone())
2246            .unwrap();
2247
2248        // A session-backed request must not read an unpartitioned entry.
2249        assert!(matches!(
2250            middleware.on_request(&session_ctx, &request).unwrap(),
2251            MiddlewareDecision::Continue
2252        ));
2253
2254        // Its response must not overwrite the unpartitioned entry either.
2255        middleware
2256            .on_response(
2257                &session_ctx,
2258                &request,
2259                serde_json::json!({"resources": ["session-private"]}),
2260            )
2261            .unwrap();
2262        let MiddlewareDecision::Respond(cached) =
2263            middleware.on_request(&anonymous_ctx, &request).unwrap()
2264        else {
2265            panic!("session response unexpectedly replaced the public entry");
2266        };
2267        assert_eq!(cached, public_response);
2268        assert_eq!(middleware.stats().entries, 1);
2269    }
2270
2271    #[test]
2272    fn allowlisted_tool_still_bypasses_uncommitted_context_partitions() {
2273        let middleware =
2274            ResponseCachingMiddleware::new().include_tools(vec!["pure_tool".to_string()]);
2275        let stateless_ctx = McpContext::new(Cx::for_testing(), 1);
2276        let session_ctx = McpContext::with_state(Cx::for_testing(), 2, SessionState::new());
2277        let request = test_request(
2278            "tools/call",
2279            Some(serde_json::json!({"name": "pure_tool", "arguments": {}})),
2280        );
2281
2282        for ctx in [&stateless_ctx, &session_ctx] {
2283            assert!(matches!(
2284                middleware.on_request(ctx, &request).unwrap(),
2285                MiddlewareDecision::Continue
2286            ));
2287            middleware
2288                .on_response(ctx, &request, serde_json::json!({"result": "private"}))
2289                .unwrap();
2290        }
2291
2292        assert_eq!(middleware.stats().entries, 0);
2293        assert!(matches!(
2294            middleware.on_request(&test_context(), &request).unwrap(),
2295            MiddlewareDecision::Continue
2296        ));
2297    }
2298
2299    #[test]
2300    fn production_context_caches_within_one_session_and_auth_partition() {
2301        let middleware = ResponseCachingMiddleware::new();
2302        let state = SessionState::new();
2303        let first = anonymous_partitioned_context(&state, 10);
2304        let second = anonymous_partitioned_context(&state, 11);
2305        let request = test_request("tools/list", None);
2306        let response = serde_json::json!({"tools": ["session-tool"]});
2307
2308        assert!(matches!(
2309            middleware.on_request(&first, &request).unwrap(),
2310            MiddlewareDecision::Continue
2311        ));
2312        middleware
2313            .on_response(&first, &request, response.clone())
2314            .unwrap();
2315
2316        let MiddlewareDecision::Respond(cached) = middleware.on_request(&second, &request).unwrap()
2317        else {
2318            panic!("same session/auth partition did not produce a cache hit");
2319        };
2320        assert_eq!(cached, response);
2321    }
2322
2323    #[test]
2324    fn request_local_sessions_share_the_stateless_cache_partition() {
2325        let middleware =
2326            ResponseCachingMiddleware::new().include_tools(vec!["pure_tool".to_string()]);
2327        let first = anonymous_partitioned_context(&SessionState::ephemeral(), 30);
2328        let second = anonymous_partitioned_context(&SessionState::ephemeral(), 31);
2329        let request = test_request(
2330            "tools/call",
2331            Some(serde_json::json!({"name": "pure_tool", "arguments": {"n": 1}})),
2332        );
2333        let response = serde_json::json!({"resultType": "complete", "content": [{"type": "text", "text": "1"}]});
2334
2335        assert!(first.session_is_ephemeral());
2336        assert!(second.session_is_ephemeral());
2337        assert!(matches!(
2338            middleware.on_request(&first, &request).unwrap(),
2339            MiddlewareDecision::Continue
2340        ));
2341        middleware
2342            .on_response(&first, &request, response.clone())
2343            .unwrap();
2344
2345        let MiddlewareDecision::Respond(cached) = middleware.on_request(&second, &request).unwrap()
2346        else {
2347            panic!("a second request-local modern HTTP session must hit the first complete result");
2348        };
2349        assert_eq!(cached, response);
2350        assert!(second.response_was_served_from_cache());
2351    }
2352
2353    #[test]
2354    fn durable_sessions_do_not_share_request_local_cache_entries() {
2355        let middleware = ResponseCachingMiddleware::new();
2356        let ephemeral = anonymous_partitioned_context(&SessionState::ephemeral(), 40);
2357        let durable = anonymous_partitioned_context(&SessionState::new(), 41);
2358        let request = test_request("tools/list", None);
2359        let response = serde_json::json!({"tools": ["request-local"]});
2360
2361        middleware
2362            .on_response(&ephemeral, &request, response)
2363            .unwrap();
2364        assert!(matches!(
2365            middleware.on_request(&durable, &request).unwrap(),
2366            MiddlewareDecision::Continue
2367        ));
2368    }
2369
2370    #[test]
2371    fn production_cache_isolates_sessions_and_complete_auth_facts() {
2372        let middleware = ResponseCachingMiddleware::new();
2373        let first_state = SessionState::new();
2374        let second_state = SessionState::new();
2375        let mut alice_auth = AuthContext::with_subject("alice");
2376        alice_auth.scopes = vec!["read".to_string()];
2377        alice_auth.claims = Some(serde_json::json!({"tenant": "one"}));
2378        let mut changed_claims = alice_auth.clone();
2379        changed_claims.claims = Some(serde_json::json!({"tenant": "two"}));
2380        let alice = partitioned_context(&first_state, 20, alice_auth.clone());
2381        let other_session = partitioned_context(&second_state, 21, alice_auth);
2382        let other_claims = partitioned_context(&first_state, 22, changed_claims);
2383        let request = test_request("resources/list", None);
2384
2385        assert!(matches!(
2386            middleware.on_request(&alice, &request).unwrap(),
2387            MiddlewareDecision::Continue
2388        ));
2389        middleware
2390            .on_response(
2391                &alice,
2392                &request,
2393                serde_json::json!({"resources": ["alice-only"]}),
2394            )
2395            .unwrap();
2396
2397        for isolated in [&other_session, &other_claims] {
2398            assert!(matches!(
2399                middleware.on_request(isolated, &request).unwrap(),
2400                MiddlewareDecision::Continue
2401            ));
2402        }
2403        assert_eq!(middleware.stats().entries, 1);
2404    }
2405
2406    #[test]
2407    fn session_state_mutation_invalidates_prior_revision() {
2408        let middleware = ResponseCachingMiddleware::new();
2409        let state = SessionState::new();
2410        let before = anonymous_partitioned_context(&state, 30);
2411        let request = test_request("prompts/list", None);
2412
2413        assert!(matches!(
2414            middleware.on_request(&before, &request).unwrap(),
2415            MiddlewareDecision::Continue
2416        ));
2417        middleware
2418            .on_response(
2419                &before,
2420                &request,
2421                serde_json::json!({"prompts": ["before"]}),
2422            )
2423            .unwrap();
2424        assert!(state.set("feature", "changed"));
2425        let after = anonymous_partitioned_context(&state, 31);
2426
2427        assert!(matches!(
2428            middleware.on_request(&after, &request).unwrap(),
2429            MiddlewareDecision::Continue
2430        ));
2431    }
2432
2433    #[test]
2434    fn response_is_not_cached_when_state_changes_during_dispatch() {
2435        let middleware = ResponseCachingMiddleware::new();
2436        let state = SessionState::new();
2437        let mutating_request = anonymous_partitioned_context(&state, 35);
2438        let request = test_request("resources/read", Some(serde_json::json!({"uri": "x"})));
2439
2440        assert!(matches!(
2441            middleware.on_request(&mutating_request, &request).unwrap(),
2442            MiddlewareDecision::Continue
2443        ));
2444        assert!(state.set("handler-mutation", true));
2445        middleware
2446            .on_response(
2447                &mutating_request,
2448                &request,
2449                serde_json::json!({"contents": ["computed-before-or-during-mutation"]}),
2450            )
2451            .unwrap();
2452
2453        assert_eq!(middleware.stats().entries, 0);
2454        let next = anonymous_partitioned_context(&state, 36);
2455        assert!(matches!(
2456            middleware.on_request(&next, &request).unwrap(),
2457            MiddlewareDecision::Continue
2458        ));
2459    }
2460
2461    #[test]
2462    fn cache_hit_is_rejected_if_state_changes_while_lookup_waits() {
2463        let middleware = Arc::new(ResponseCachingMiddleware::new());
2464        let state = SessionState::new();
2465        let populate = anonymous_partitioned_context(&state, 37);
2466        let request = test_request("resources/list", None);
2467        let response = serde_json::json!({"resources": ["before-mutation"]});
2468
2469        assert!(matches!(
2470            middleware.on_request(&populate, &request).unwrap(),
2471            MiddlewareDecision::Continue
2472        ));
2473        middleware
2474            .on_response(&populate, &request, response)
2475            .unwrap();
2476
2477        let lookup_ctx = anonymous_partitioned_context(&state, 38);
2478        let admission_observer = lookup_ctx.clone();
2479        let lookup_middleware = Arc::clone(&middleware);
2480        let lookup_request = request.clone();
2481        let cache_guard = middleware
2482            .cache
2483            .lock()
2484            .unwrap_or_else(std::sync::PoisonError::into_inner);
2485        let lookup = std::thread::spawn(move || {
2486            lookup_middleware
2487                .on_request(&lookup_ctx, &lookup_request)
2488                .expect("cache lookup should not fail")
2489        });
2490
2491        let deadline = Instant::now() + Duration::from_secs(5);
2492        while admission_observer
2493            .complete_session_cache_partition()
2494            .is_none()
2495        {
2496            assert!(
2497                Instant::now() < deadline,
2498                "lookup did not capture its session partition"
2499            );
2500            std::thread::yield_now();
2501        }
2502
2503        assert!(state.set("changed-while-cache-locked", true));
2504        drop(cache_guard);
2505
2506        let decision = lookup.join().expect("cache lookup thread");
2507        assert!(matches!(decision, MiddlewareDecision::Continue));
2508        assert_eq!(middleware.stats().hits, 0);
2509    }
2510
2511    #[test]
2512    fn invalidate_removes_every_partition_for_request_identity() {
2513        let middleware = ResponseCachingMiddleware::new();
2514        let first_state = SessionState::new();
2515        let second_state = SessionState::new();
2516        let first = anonymous_partitioned_context(&first_state, 40);
2517        let second = anonymous_partitioned_context(&second_state, 41);
2518        let request = test_request("tools/list", Some(serde_json::json!({"cursor": "same"})));
2519
2520        for (ctx, response) in [
2521            (&first, serde_json::json!({"tools": ["first"]})),
2522            (&second, serde_json::json!({"tools": ["second"]})),
2523        ] {
2524            assert!(matches!(
2525                middleware.on_request(ctx, &request).unwrap(),
2526                MiddlewareDecision::Continue
2527            ));
2528            middleware.on_response(ctx, &request, response).unwrap();
2529        }
2530        assert_eq!(middleware.stats().entries, 2);
2531
2532        middleware.invalidate("tools/list", request.params.as_ref());
2533
2534        assert_eq!(middleware.stats().entries, 0);
2535        for ctx in [&first, &second] {
2536            assert!(matches!(
2537                middleware.on_request(ctx, &request).unwrap(),
2538                MiddlewareDecision::Continue
2539            ));
2540        }
2541    }
2542
2543    #[test]
2544    fn test_caching_middleware_clear() {
2545        let middleware = ResponseCachingMiddleware::new();
2546        let ctx = test_context();
2547        let request = test_request("tools/list", None);
2548
2549        // Cache a response
2550        middleware.on_request(&ctx, &request).unwrap();
2551        middleware
2552            .on_response(&ctx, &request, serde_json::json!({}))
2553            .unwrap();
2554
2555        // Verify cached
2556        let decision = middleware.on_request(&ctx, &request).unwrap();
2557        assert!(matches!(decision, MiddlewareDecision::Respond(_)));
2558
2559        // Clear cache
2560        middleware.clear();
2561
2562        // Should miss now
2563        let decision = middleware.on_request(&ctx, &request).unwrap();
2564        assert!(matches!(decision, MiddlewareDecision::Continue));
2565    }
2566
2567    #[test]
2568    fn test_caching_middleware_invalidate() {
2569        let middleware = ResponseCachingMiddleware::new();
2570        let ctx = test_context();
2571        let request = test_request("tools/list", None);
2572
2573        // Cache a response
2574        middleware.on_request(&ctx, &request).unwrap();
2575        middleware
2576            .on_response(&ctx, &request, serde_json::json!({}))
2577            .unwrap();
2578
2579        // Invalidate specific entry
2580        let semantic_page_set = serde_json::json!({});
2581        middleware.invalidate("tools/list", Some(&semantic_page_set));
2582
2583        // Should miss now
2584        let decision = middleware.on_request(&ctx, &request).unwrap();
2585        assert!(matches!(decision, MiddlewareDecision::Continue));
2586    }
2587
2588    #[test]
2589    fn test_cache_stats_hit_rate() {
2590        let stats = CacheStats {
2591            hits: 75,
2592            misses: 25,
2593            entries: 10,
2594            size_bytes: 1000,
2595        };
2596
2597        assert!((stats.hit_rate() - 75.0).abs() < 0.001);
2598    }
2599
2600    // ── CacheStats edge cases ──────────────────────────────────────────
2601
2602    #[test]
2603    fn cache_stats_hit_rate_zero_total() {
2604        let stats = CacheStats::default();
2605        assert!(stats.hit_rate().abs() < f64::EPSILON);
2606    }
2607
2608    #[test]
2609    fn cache_stats_hit_rate_does_not_overflow_saturated_counters() {
2610        let stats = CacheStats {
2611            hits: u64::MAX,
2612            misses: u64::MAX,
2613            entries: 0,
2614            size_bytes: 0,
2615        };
2616        assert!((stats.hit_rate() - 50.0).abs() < f64::EPSILON);
2617    }
2618
2619    #[test]
2620    fn cache_stats_debug() {
2621        let stats = CacheStats::default();
2622        let debug = format!("{:?}", stats);
2623        assert!(debug.contains("CacheStats"));
2624    }
2625
2626    // ── CacheKey ───────────────────────────────────────────────────────
2627
2628    #[test]
2629    fn cache_key_same_method_same_params_are_equal() {
2630        let k1 = CacheKey::new("tools/list", Some(&serde_json::json!({"a": 1})));
2631        let k2 = CacheKey::new("tools/list", Some(&serde_json::json!({"a": 1})));
2632        assert_eq!(k1, k2);
2633    }
2634
2635    #[test]
2636    fn cache_key_different_params_differ() {
2637        let k1 = CacheKey::new("tools/list", Some(&serde_json::json!({"a": 1})));
2638        let k2 = CacheKey::new("tools/list", Some(&serde_json::json!({"a": 2})));
2639        assert_ne!(k1, k2);
2640    }
2641
2642    #[test]
2643    fn cache_key_method_and_param_presence_are_domain_separated() {
2644        let no_params = CacheKey::new("test", None);
2645        let null_params = CacheKey::new("test", Some(&serde_json::Value::Null));
2646        let other_method = CacheKey::new("other", None);
2647        assert_ne!(no_params, null_params);
2648        assert_ne!(no_params, other_method);
2649    }
2650
2651    #[test]
2652    fn cache_key_debug_and_clone() {
2653        let k = CacheKey::new("test", None);
2654        let debug = format!("{:?}", k);
2655        assert!(debug.contains("CacheKey"));
2656        assert!(!debug.contains("test"));
2657        let cloned = k.clone();
2658        assert_eq!(k, cloned);
2659    }
2660
2661    // ── bounded cache-key derivation ───────────────────────────────────
2662
2663    #[test]
2664    fn cache_key_derivation_is_deterministic() {
2665        let v = serde_json::json!({"key": "value", "num": 42});
2666        let h1 = CacheKey::new("tools/list", Some(&v));
2667        let h2 = CacheKey::new("tools/list", Some(&v));
2668        assert_eq!(h1, h2);
2669    }
2670
2671    #[test]
2672    fn cache_key_derivation_distinguishes_values() {
2673        let h1 = CacheKey::new("tools/list", Some(&serde_json::json!(1)));
2674        let h2 = CacheKey::new("tools/list", Some(&serde_json::json!(2)));
2675        assert_ne!(h1, h2);
2676    }
2677
2678    #[test]
2679    fn cache_key_derivation_rejects_oversized_input_before_retention() {
2680        let oversized_method = "x".repeat(MAX_CACHE_KEY_INPUT_BYTES + 1);
2681        assert!(CacheKey::try_new(&oversized_method, None).is_none());
2682
2683        let mut exact = BoundedCacheBytes::new(4);
2684        exact.write_all(b"1234").expect("exact boundary fits");
2685        assert!(exact.write_all(b"5").is_err());
2686        assert_eq!(exact.bytes, b"1234");
2687        assert!(exact.bytes.capacity() <= exact.max_bytes);
2688    }
2689
2690    #[test]
2691    fn bounded_cache_bytes_fragmented_writes_grow_geometrically_within_limit() {
2692        const LIMIT: usize = 128 * 1024 + 37;
2693        let mut encoded = BoundedCacheBytes::new(LIMIT);
2694
2695        for _ in 0..LIMIT {
2696            encoded
2697                .write_all(b"x")
2698                .expect("each byte remains inside the logical limit");
2699        }
2700
2701        assert_eq!(encoded.bytes.len(), LIMIT);
2702        assert!(encoded.bytes.capacity() <= LIMIT);
2703        assert!(
2704            encoded.growth_events <= maximum_geometric_growth_events(LIMIT),
2705            "{} growth events exceeded the logarithmic bound",
2706            encoded.growth_events
2707        );
2708        assert_eq!(encoded.bytes[0], b'x');
2709        assert_eq!(encoded.bytes[LIMIT - 1], b'x');
2710
2711        let length_before_rejection = encoded.bytes.len();
2712        let capacity_before_rejection = encoded.bytes.capacity();
2713        let growth_before_rejection = encoded.growth_events;
2714        assert!(encoded.write_all(b"x").is_err());
2715        assert_eq!(encoded.bytes.len(), length_before_rejection);
2716        assert_eq!(encoded.bytes.capacity(), capacity_before_rejection);
2717        assert_eq!(encoded.growth_events, growth_before_rejection);
2718    }
2719
2720    #[test]
2721    fn bounded_cache_bytes_large_flat_json_has_bounded_growth() {
2722        let value = serde_json::json!({"data": "x".repeat(768 * 1024)});
2723        let expected = serde_json::to_vec(&value).expect("test JSON serializes");
2724        let logical_limit = expected.len();
2725        let mut encoded = BoundedCacheBytes::new(logical_limit);
2726
2727        serde_json::to_writer(&mut encoded, &value).expect("flat JSON fits exact limit");
2728
2729        assert_eq!(encoded.bytes, expected);
2730        assert!(encoded.bytes.capacity() <= logical_limit);
2731        assert!(
2732            encoded.growth_events <= maximum_geometric_growth_events(logical_limit),
2733            "{} growth events exceeded the logarithmic bound",
2734            encoded.growth_events
2735        );
2736    }
2737
2738    #[test]
2739    fn cache_entry_size_measurement_stops_at_item_limit() {
2740        let value = serde_json::json!("0123456789");
2741        assert!(encode_json_bounded(&value, value.to_string().len()).is_some());
2742        assert!(encode_json_bounded(&value, value.to_string().len() - 1).is_none());
2743    }
2744
2745    #[test]
2746    fn cache_serialization_rejects_excessive_json_depth() {
2747        let mut value = serde_json::Value::Null;
2748        for _ in 0..=MAX_CACHE_JSON_DEPTH {
2749            value = serde_json::Value::Array(vec![value]);
2750        }
2751
2752        assert!(encode_json_bounded(&value, DEFAULT_MAX_ITEM_SIZE).is_none());
2753    }
2754
2755    // ── LruCache additional tests ──────────────────────────────────────
2756
2757    #[test]
2758    fn lru_cache_clear() {
2759        let mut cache = LruCache::new(10, 1024 * 1024, 1024);
2760        cache.insert(
2761            CacheKey::new("a", None),
2762            serde_json::json!(1),
2763            Duration::from_secs(60),
2764        );
2765        cache.insert(
2766            CacheKey::new("b", None),
2767            serde_json::json!(2),
2768            Duration::from_secs(60),
2769        );
2770        assert_eq!(cache.len(), 2);
2771        assert!(!cache.is_empty());
2772
2773        cache.clear();
2774        assert_eq!(cache.len(), 0);
2775        assert!(cache.is_empty());
2776        assert_eq!(cache.current_size_bytes, 0);
2777    }
2778
2779    #[test]
2780    fn lru_cache_remove_nonexistent() {
2781        let mut cache = LruCache::new(10, 1024 * 1024, 1024);
2782        let key = CacheKey::new("nonexistent", None);
2783        cache.remove(&key); // should not panic
2784        assert_eq!(cache.len(), 0);
2785    }
2786
2787    #[test]
2788    fn lru_cache_insert_duplicate_replaces() {
2789        let mut cache = LruCache::new(10, 1024 * 1024, 1024);
2790        let key = CacheKey::new("test", None);
2791        cache.insert(
2792            key.clone(),
2793            serde_json::json!("v1"),
2794            Duration::from_secs(60),
2795        );
2796        cache.insert(
2797            key.clone(),
2798            serde_json::json!("v2"),
2799            Duration::from_secs(60),
2800        );
2801        assert_eq!(cache.len(), 1);
2802        assert_eq!(cache.get_value(&key), Some(serde_json::json!("v2")));
2803    }
2804
2805    #[test]
2806    fn lru_cache_get_miss_returns_none() {
2807        let mut cache = LruCache::new(10, 1024 * 1024, 1024);
2808        assert!(cache.get_value(&CacheKey::new("missing", None)).is_none());
2809    }
2810
2811    #[test]
2812    fn lru_cache_lru_order_updated_on_access() {
2813        let mut cache = LruCache::new(2, 1024 * 1024, 1024);
2814        let k1 = CacheKey::new("a", None);
2815        let k2 = CacheKey::new("b", None);
2816        cache.insert(k1.clone(), serde_json::json!(1), Duration::from_secs(60));
2817        cache.insert(k2.clone(), serde_json::json!(2), Duration::from_secs(60));
2818
2819        // Access k1, making k2 the LRU
2820        let _ = cache.get_value(&k1);
2821
2822        // Insert k3, should evict k2 (LRU)
2823        let k3 = CacheKey::new("c", None);
2824        cache.insert(k3.clone(), serde_json::json!(3), Duration::from_secs(60));
2825        assert!(cache.get_value(&k1).is_some()); // k1 was accessed recently
2826        assert!(cache.get_value(&k2).is_none()); // k2 was evicted
2827        assert!(cache.get_value(&k3).is_some());
2828    }
2829
2830    // ── ToolCallCacheConfig ────────────────────────────────────────────
2831
2832    #[test]
2833    fn should_cache_tool_disabled_returns_false() {
2834        let config = ToolCallCacheConfig {
2835            base: MethodCacheConfig {
2836                enabled: false,
2837                ttl_secs: 60,
2838            },
2839            ..ToolCallCacheConfig::default()
2840        };
2841        assert!(!config.should_cache_tool("any_tool"));
2842    }
2843
2844    #[test]
2845    fn should_cache_tool_excluded_returns_false() {
2846        let config = ToolCallCacheConfig {
2847            base: MethodCacheConfig {
2848                enabled: true,
2849                ttl_secs: 60,
2850            },
2851            excluded_tools: vec!["excluded".to_string()],
2852            included_tools: vec!["excluded".to_string(), "other".to_string()],
2853        };
2854        assert!(!config.should_cache_tool("excluded"));
2855        assert!(config.should_cache_tool("other"));
2856    }
2857
2858    #[test]
2859    fn should_cache_tool_include_list_filters() {
2860        let config = ToolCallCacheConfig {
2861            base: MethodCacheConfig {
2862                enabled: true,
2863                ttl_secs: 60,
2864            },
2865            included_tools: vec!["allowed".to_string()],
2866            excluded_tools: vec![],
2867        };
2868        assert!(config.should_cache_tool("allowed"));
2869        assert!(!config.should_cache_tool("not_allowed"));
2870    }
2871
2872    #[test]
2873    fn should_cache_tool_exclude_takes_precedence_over_include() {
2874        let config = ToolCallCacheConfig {
2875            base: MethodCacheConfig {
2876                enabled: true,
2877                ttl_secs: 60,
2878            },
2879            included_tools: vec!["tool".to_string()],
2880            excluded_tools: vec!["tool".to_string()],
2881        };
2882        assert!(!config.should_cache_tool("tool"));
2883    }
2884
2885    // ── MethodCacheConfig ──────────────────────────────────────────────
2886
2887    #[test]
2888    fn method_cache_config_default() {
2889        let config = MethodCacheConfig::default();
2890        assert!(config.enabled);
2891        assert_eq!(config.ttl_secs, DEFAULT_CALL_TTL_SECS);
2892    }
2893
2894    #[test]
2895    fn method_cache_config_debug() {
2896        let config = MethodCacheConfig::default();
2897        let debug = format!("{:?}", config);
2898        assert!(debug.contains("MethodCacheConfig"));
2899    }
2900
2901    // ── ResponseCachingMiddleware construction ──────────────────────────
2902
2903    #[test]
2904    fn default_equals_new() {
2905        let d = ResponseCachingMiddleware::default();
2906        let n = ResponseCachingMiddleware::new();
2907        assert_eq!(d.list_ttl, n.list_ttl);
2908        assert_eq!(d.call_ttl, n.call_ttl);
2909    }
2910
2911    #[test]
2912    fn debug_output() {
2913        let m = ResponseCachingMiddleware::new();
2914        let debug = format!("{:?}", m);
2915        assert!(debug.contains("ResponseCachingMiddleware"));
2916        assert!(debug.contains("list_ttl"));
2917        assert!(debug.contains("call_ttl"));
2918    }
2919
2920    // ── Fluent setters ─────────────────────────────────────────────────
2921
2922    #[test]
2923    fn list_ttl_secs_updates_all_list_configs() {
2924        let m = ResponseCachingMiddleware::new().list_ttl_secs(600);
2925        assert_eq!(m.list_ttl, Duration::from_secs(600));
2926        assert_eq!(m.tools_list_config.ttl_secs, 600);
2927        assert_eq!(m.resources_list_config.ttl_secs, 600);
2928        assert_eq!(m.prompts_list_config.ttl_secs, 600);
2929    }
2930
2931    #[test]
2932    fn call_ttl_secs_updates_all_call_configs() {
2933        let m = ResponseCachingMiddleware::new().call_ttl_secs(7200);
2934        assert_eq!(m.call_ttl, Duration::from_secs(7200));
2935        assert_eq!(m.tools_call_config.base.ttl_secs, 7200);
2936        assert_eq!(m.resources_read_config.ttl_secs, 7200);
2937        assert_eq!(m.prompts_get_config.ttl_secs, 7200);
2938    }
2939
2940    #[test]
2941    fn max_entries_setter() {
2942        let m = ResponseCachingMiddleware::new().max_entries(50);
2943        let cache = m
2944            .cache
2945            .lock()
2946            .unwrap_or_else(std::sync::PoisonError::into_inner);
2947        assert_eq!(cache.max_entries, 50);
2948    }
2949
2950    #[test]
2951    fn max_size_bytes_setter() {
2952        let m = ResponseCachingMiddleware::new().max_size_bytes(2048);
2953        let cache = m
2954            .cache
2955            .lock()
2956            .unwrap_or_else(std::sync::PoisonError::into_inner);
2957        assert_eq!(cache.max_size_bytes, 2048);
2958    }
2959
2960    #[test]
2961    fn max_item_size_setter() {
2962        let m = ResponseCachingMiddleware::new().max_item_size(512);
2963        let cache = m
2964            .cache
2965            .lock()
2966            .unwrap_or_else(std::sync::PoisonError::into_inner);
2967        assert_eq!(cache.max_item_size, 512);
2968    }
2969
2970    // ── Disable method variants ────────────────────────────────────────
2971
2972    #[test]
2973    fn disable_resources_list() {
2974        let m = ResponseCachingMiddleware::new().disable_resources_list();
2975        assert!(!m.resources_list_config.enabled);
2976        assert!(m.tools_list_config.enabled); // others unchanged
2977    }
2978
2979    #[test]
2980    fn disable_prompts_list() {
2981        let m = ResponseCachingMiddleware::new().disable_prompts_list();
2982        assert!(!m.prompts_list_config.enabled);
2983    }
2984
2985    #[test]
2986    fn disable_tools_call() {
2987        let m = ResponseCachingMiddleware::new().disable_tools_call();
2988        assert!(!m.tools_call_config.base.enabled);
2989    }
2990
2991    #[test]
2992    fn disable_resources_read() {
2993        let m = ResponseCachingMiddleware::new().disable_resources_read();
2994        assert!(!m.resources_read_config.enabled);
2995    }
2996
2997    #[test]
2998    fn disable_prompts_get() {
2999        let m = ResponseCachingMiddleware::new().disable_prompts_get();
3000        assert!(!m.prompts_get_config.enabled);
3001    }
3002
3003    // ── include_tools / exclude_tools ──────────────────────────────────
3004
3005    #[test]
3006    fn include_tools_restricts_caching() {
3007        let m = ResponseCachingMiddleware::new().include_tools(vec!["allowed_tool".to_string()]);
3008        let _ctx = test_context();
3009
3010        // allowed_tool should be cached
3011        let req = test_request(
3012            "tools/call",
3013            Some(serde_json::json!({"name": "allowed_tool"})),
3014        );
3015        assert!(m.should_cache_method(&req.method, req.params.as_ref()));
3016
3017        // other_tool should not be cached
3018        let req2 = test_request(
3019            "tools/call",
3020            Some(serde_json::json!({"name": "other_tool"})),
3021        );
3022        assert!(!m.should_cache_method(&req2.method, req2.params.as_ref()));
3023
3024        // non-tool methods still work
3025        let req3 = test_request("tools/list", None);
3026        assert!(m.should_cache_method(&req3.method, req3.params.as_ref()));
3027    }
3028
3029    // ── should_cache_method edge cases ─────────────────────────────────
3030
3031    #[test]
3032    fn should_cache_tools_call_without_name_returns_false() {
3033        let m = ResponseCachingMiddleware::new();
3034        // tools/call with params but no "name" field
3035        assert!(!m.should_cache_method("tools/call", Some(&serde_json::json!({"arguments": {}}))));
3036    }
3037
3038    #[test]
3039    fn should_cache_tools_call_with_no_params_returns_false() {
3040        let m = ResponseCachingMiddleware::new();
3041        assert!(!m.should_cache_method("tools/call", None));
3042    }
3043
3044    #[test]
3045    fn should_cache_unknown_method_returns_false() {
3046        let m = ResponseCachingMiddleware::new();
3047        assert!(!m.should_cache_method("unknown/method", None));
3048    }
3049
3050    #[test]
3051    fn should_cache_all_known_cacheable_methods() {
3052        let m = ResponseCachingMiddleware::new();
3053        assert!(m.should_cache_method("tools/list", None));
3054        assert!(m.should_cache_method("resources/list", None));
3055        assert!(m.should_cache_method("prompts/list", None));
3056        assert!(m.should_cache_method("resources/read", None));
3057        assert!(m.should_cache_method("prompts/get", None));
3058    }
3059
3060    // ── get_ttl ────────────────────────────────────────────────────────
3061
3062    #[test]
3063    fn get_ttl_list_methods() {
3064        let m = ResponseCachingMiddleware::new().list_ttl_secs(120);
3065        assert_eq!(m.get_ttl("tools/list"), Duration::from_secs(120));
3066        assert_eq!(m.get_ttl("resources/list"), Duration::from_secs(120));
3067        assert_eq!(m.get_ttl("prompts/list"), Duration::from_secs(120));
3068    }
3069
3070    #[test]
3071    fn get_ttl_call_methods() {
3072        let m = ResponseCachingMiddleware::new().call_ttl_secs(900);
3073        assert_eq!(m.get_ttl("tools/call"), Duration::from_mins(15));
3074        assert_eq!(m.get_ttl("resources/read"), Duration::from_mins(15));
3075        assert_eq!(m.get_ttl("prompts/get"), Duration::from_mins(15));
3076    }
3077
3078    #[test]
3079    fn get_ttl_unknown_method_uses_call_ttl() {
3080        let m = ResponseCachingMiddleware::new().call_ttl_secs(999);
3081        assert_eq!(m.get_ttl("unknown/method"), Duration::from_secs(999));
3082    }
3083
3084    // ── on_error passes through ────────────────────────────────────────
3085
3086    #[test]
3087    fn on_error_passes_through() {
3088        let m = ResponseCachingMiddleware::new();
3089        let ctx = test_context();
3090        let req = test_request("tools/list", None);
3091        let err = McpError::internal_error("test error");
3092        let result = m.on_error(&ctx, &req, err);
3093        assert!(result.message.contains("test error"));
3094    }
3095
3096    // ── stats tracks entries and size ──────────────────────────────────
3097
3098    #[test]
3099    fn stats_tracks_entries_and_size() {
3100        let m = ResponseCachingMiddleware::new();
3101        let ctx = test_context();
3102
3103        let stats = m.stats();
3104        assert_eq!(stats.entries, 0);
3105        assert_eq!(stats.size_bytes, 0);
3106
3107        let req = test_request("tools/list", None);
3108        m.on_request(&ctx, &req).unwrap();
3109        m.on_response(&ctx, &req, serde_json::json!({"tools": []}))
3110            .unwrap();
3111
3112        let stats = m.stats();
3113        assert_eq!(stats.entries, 1);
3114        assert!(stats.size_bytes > 0);
3115        assert_eq!(stats.misses, 1);
3116    }
3117
3118    // ── Middleware caches resources/list and prompts/list ───────────────
3119
3120    #[test]
3121    fn caches_resources_list() {
3122        let m = ResponseCachingMiddleware::new();
3123        let ctx = test_context();
3124        let req = test_request("resources/list", None);
3125
3126        m.on_request(&ctx, &req).unwrap();
3127        m.on_response(&ctx, &req, serde_json::json!({"resources": []}))
3128            .unwrap();
3129
3130        let decision = m.on_request(&ctx, &req).unwrap();
3131        assert!(matches!(decision, MiddlewareDecision::Respond(_)));
3132    }
3133
3134    #[test]
3135    fn caches_prompts_list() {
3136        let m = ResponseCachingMiddleware::new();
3137        let ctx = test_context();
3138        let req = test_request("prompts/list", None);
3139
3140        m.on_request(&ctx, &req).unwrap();
3141        m.on_response(&ctx, &req, serde_json::json!({"prompts": []}))
3142            .unwrap();
3143
3144        let decision = m.on_request(&ctx, &req).unwrap();
3145        assert!(matches!(decision, MiddlewareDecision::Respond(_)));
3146    }
3147
3148    // ── CacheEntry debug/clone ─────────────────────────────────────────
3149
3150    #[test]
3151    fn cache_entry_debug_and_clone() {
3152        let value = serde_json::json!("CACHE-SECRET-CANARY");
3153        let entry = CacheEntry::new(
3154            value.clone(),
3155            Duration::from_secs(60),
3156            DEFAULT_MAX_ITEM_SIZE,
3157        )
3158        .expect("short test TTL must be representable");
3159        let debug = format!("{:?}", entry);
3160        assert!(debug.contains("CacheEntry"));
3161        assert!(
3162            !debug.contains("CACHE-SECRET-CANARY"),
3163            "cached payloads must stay out of Debug"
3164        );
3165        let cloned = entry.clone();
3166        assert_eq!(decode_cached_json(&cloned.encoded), Some(value));
3167        assert_eq!(
3168            cloned.size_bytes,
3169            cloned.encoded.len() + CACHE_ENTRY_METADATA_BYTES
3170        );
3171    }
3172
3173    #[test]
3174    fn cache_entry_not_expired_initially() {
3175        let entry = CacheEntry::new(
3176            serde_json::json!(1),
3177            Duration::from_secs(60),
3178            DEFAULT_MAX_ITEM_SIZE,
3179        )
3180        .expect("short test TTL must be representable");
3181        assert!(!entry.is_expired());
3182    }
3183
3184    #[test]
3185    fn caches_resources_read() {
3186        let m = ResponseCachingMiddleware::new();
3187        let ctx = test_context();
3188        let req = test_request(
3189            "resources/read",
3190            Some(serde_json::json!({"uri": "file:///a.txt"})),
3191        );
3192
3193        m.on_request(&ctx, &req).unwrap();
3194        m.on_response(&ctx, &req, serde_json::json!({"contents": []}))
3195            .unwrap();
3196
3197        let decision = m.on_request(&ctx, &req).unwrap();
3198        assert!(matches!(decision, MiddlewareDecision::Respond(_)));
3199    }
3200
3201    #[test]
3202    fn caches_prompts_get() {
3203        let m = ResponseCachingMiddleware::new();
3204        let ctx = test_context();
3205        let req = test_request("prompts/get", Some(serde_json::json!({"name": "greeting"})));
3206
3207        m.on_request(&ctx, &req).unwrap();
3208        m.on_response(&ctx, &req, serde_json::json!({"messages": []}))
3209            .unwrap();
3210
3211        let decision = m.on_request(&ctx, &req).unwrap();
3212        assert!(matches!(decision, MiddlewareDecision::Respond(_)));
3213    }
3214
3215    #[test]
3216    fn lru_cache_evict_expired_frees_entries() {
3217        let mut cache = LruCache::new(10, 1024 * 1024, 1024);
3218        // Insert two entries with tiny TTL
3219        cache.insert(
3220            CacheKey::new("a", None),
3221            serde_json::json!(1),
3222            Duration::from_millis(1),
3223        );
3224        cache.insert(
3225            CacheKey::new("b", None),
3226            serde_json::json!(2),
3227            Duration::from_millis(1),
3228        );
3229        assert_eq!(cache.len(), 2);
3230
3231        std::thread::sleep(std::time::Duration::from_millis(10));
3232        cache.evict_expired();
3233
3234        assert_eq!(cache.len(), 0);
3235        assert_eq!(cache.current_size_bytes, 0);
3236    }
3237
3238    #[test]
3239    fn lru_cache_insert_replaces_updates_size() {
3240        let mut cache = LruCache::new(10, 1024 * 1024, 1024);
3241        let key = CacheKey::new("k", None);
3242        cache.insert(
3243            key.clone(),
3244            serde_json::json!("short"),
3245            Duration::from_secs(60),
3246        );
3247        let size_after_first = cache.current_size_bytes;
3248
3249        cache.insert(
3250            key.clone(),
3251            serde_json::json!("much longer value here"),
3252            Duration::from_secs(60),
3253        );
3254        let size_after_second = cache.current_size_bytes;
3255
3256        // Size should reflect only the new entry (old was removed first)
3257        assert_ne!(size_after_first, size_after_second);
3258        assert_eq!(cache.len(), 1);
3259    }
3260
3261    #[test]
3262    fn tool_call_cache_config_debug_and_clone() {
3263        let config = ToolCallCacheConfig {
3264            base: MethodCacheConfig {
3265                enabled: true,
3266                ttl_secs: 120,
3267            },
3268            included_tools: vec!["t1".to_string()],
3269            excluded_tools: vec!["t2".to_string()],
3270        };
3271        let debug = format!("{:?}", config);
3272        assert!(debug.contains("ToolCallCacheConfig"));
3273        let cloned = config.clone();
3274        assert_eq!(cloned.included_tools, vec!["t1".to_string()]);
3275        assert_eq!(cloned.excluded_tools, vec!["t2".to_string()]);
3276    }
3277
3278    #[test]
3279    fn cache_stats_clone() {
3280        let stats = CacheStats {
3281            hits: 10,
3282            misses: 5,
3283            entries: 3,
3284            size_bytes: 100,
3285        };
3286        let cloned = stats.clone();
3287        assert_eq!(cloned.hits, 10);
3288        assert_eq!(cloned.misses, 5);
3289        assert_eq!(cloned.entries, 3);
3290        assert_eq!(cloned.size_bytes, 100);
3291    }
3292
3293    #[test]
3294    fn should_cache_tool_empty_allowlist_disables_all() {
3295        let config = ToolCallCacheConfig {
3296            base: MethodCacheConfig {
3297                enabled: true,
3298                ttl_secs: 60,
3299            },
3300            included_tools: vec![],
3301            excluded_tools: vec![],
3302        };
3303        assert!(!config.should_cache_tool("any_tool"));
3304        assert!(!config.should_cache_tool("another_tool"));
3305    }
3306
3307    #[test]
3308    fn cache_01_a_positive() {
3309        let middleware = ResponseCachingMiddleware::new()
3310            .list_ttl_secs(120)
3311            .call_ttl_secs(900);
3312        let ctx = test_context();
3313        let methods = [
3314            ("server/discover", None, 120_000_u64),
3315            ("tools/list", None, 120_000),
3316            ("prompts/list", None, 120_000),
3317            ("resources/list", None, 120_000),
3318            (
3319                "resources/read",
3320                Some(serde_json::json!({"uri": "file:///catalog"})),
3321                900_000,
3322            ),
3323            ("resources/templates/list", None, 120_000),
3324        ];
3325
3326        for (method, params, expected_ttl_ms) in methods {
3327            let request = test_request(method, params);
3328            let response = middleware
3329                .on_response(
3330                    &ctx,
3331                    &request,
3332                    serde_json::json!({"resultType": "complete", "items": []}),
3333                )
3334                .expect("the middleware must preserve a complete result");
3335
3336            assert_eq!(
3337                response.get("ttlMs"),
3338                Some(&serde_json::json!(expected_ttl_ms))
3339            );
3340            assert_eq!(
3341                response.get("cacheScope"),
3342                Some(&serde_json::json!("private"))
3343            );
3344        }
3345    }
3346
3347    #[test]
3348    fn cache_01_final_wire_valid_hints_remain_lossless_outside_runtime_expiry() {
3349        let middleware = ResponseCachingMiddleware::new().list_ttl_secs(120);
3350        let ctx = test_context();
3351        let request = test_request("tools/list", None);
3352        let huge_ttl: serde_json::Value =
3353            serde_json::from_str("922337203685477580812345678901234567890")
3354                .expect("arbitrary-width JSON integer fixture");
3355        let response = serde_json::json!({
3356            "resultType": "complete",
3357            "tools": [],
3358            "ttlMs": huge_ttl,
3359            "cacheScope": "private",
3360        });
3361
3362        let delivered = middleware
3363            .on_response(&ctx, &request, response.clone())
3364            .expect("a wire-valid but uncacheable final TTL remains deliverable");
3365
3366        assert_eq!(
3367            delivered, response,
3368            "runtime expiry bounds must not rewrite upstream final cache hints"
3369        );
3370        assert_eq!(
3371            delivered["ttlMs"].to_string(),
3372            "922337203685477580812345678901234567890"
3373        );
3374        assert_eq!(
3375            middleware.stats().entries,
3376            0,
3377            "an unrepresentable private final TTL must not create a local-expiry entry"
3378        );
3379
3380        let fractional = serde_json::json!({
3381            "resultType": "complete",
3382            "tools": [],
3383            "ttlMs": 120_000.5,
3384            "cacheScope": "private",
3385        });
3386        let normalized = middleware
3387            .on_response(&ctx, &request, fractional)
3388            .expect("invalid cache hints are replaced by the local policy");
3389        assert_eq!(normalized["ttlMs"], serde_json::json!(120_000));
3390        assert_eq!(normalized["cacheScope"], serde_json::json!("private"));
3391        assert_eq!(
3392            middleware.stats().entries,
3393            1,
3394            "the paired representable local policy remains cacheable"
3395        );
3396    }
3397
3398    #[test]
3399    fn every_final_result_with_an_unrepresentable_private_ttl_skips_local_cache_state() {
3400        let middleware = ResponseCachingMiddleware::new();
3401        let ctx = test_context();
3402        let huge_ttl: serde_json::Value = serde_json::from_str("18446744073709551616000")
3403            .expect("arbitrary-width JSON integer fixture");
3404        let methods = [
3405            ("tools/list", None),
3406            ("resources/list", None),
3407            ("resources/templates/list", None),
3408            ("prompts/list", None),
3409            (
3410                "resources/read",
3411                Some(serde_json::json!({"uri": "file:///huge-ttl"})),
3412            ),
3413        ];
3414
3415        for (method, params) in methods {
3416            let request = test_request(method, params);
3417            let response = serde_json::json!({
3418                "resultType": "complete",
3419                "items": [],
3420                "ttlMs": huge_ttl.clone(),
3421                "cacheScope": "private",
3422            });
3423            let delivered = middleware
3424                .on_response(&ctx, &request, response.clone())
3425                .expect("wire-valid final result remains deliverable");
3426
3427            assert_eq!(delivered, response);
3428            assert_eq!(
3429                middleware.stats().entries,
3430                0,
3431                "{method} must not create a local entry for an unrepresentable TTL"
3432            );
3433        }
3434
3435        let discovery = serde_json::json!({
3436            "supportedVersions": [FINAL_PROTOCOL_VERSION],
3437            "capabilities": {},
3438            "ttlMs": huge_ttl.clone(),
3439            "cacheScope": "private",
3440        });
3441        let delivered = middleware
3442            .on_response(
3443                &ctx,
3444                &final_discovery_request(FINAL_PROTOCOL_VERSION),
3445                discovery.clone(),
3446            )
3447            .expect("wire-valid discovery result remains deliverable");
3448        assert_eq!(delivered, discovery);
3449        assert_eq!(
3450            middleware.stats().entries,
3451            0,
3452            "server/discover must also skip local state for an unrepresentable TTL"
3453        );
3454    }
3455
3456    #[test]
3457    fn cache_01_a_planted_negative() {
3458        let middleware = ResponseCachingMiddleware::new().list_ttl_secs(120);
3459        let ctx = test_context();
3460        let request = test_request("tools/list", None);
3461
3462        // The sole forbidden dimension differs from the positive case: this
3463        // is an input-required result, not a cacheable complete result.
3464        let response = middleware
3465            .on_response(
3466                &ctx,
3467                &request,
3468                serde_json::json!({
3469                    "resultType": "inputRequired",
3470                    "ttlMs": 120_000,
3471                    "cacheScope": "public",
3472                    "items": []
3473                }),
3474            )
3475            .expect("input-required results remain ordinary middleware output");
3476
3477        assert_eq!(
3478            response.get("resultType"),
3479            Some(&serde_json::json!("inputRequired"))
3480        );
3481        assert!(response.get("ttlMs").is_none());
3482        assert!(response.get("cacheScope").is_none());
3483        assert_eq!(middleware.stats().entries, 0);
3484        assert!(matches!(
3485            middleware
3486                .on_request(&ctx, &request)
3487                .expect("lookup is safe"),
3488            MiddlewareDecision::Continue
3489        ));
3490    }
3491
3492    #[test]
3493    fn cache_01_b_positive() {
3494        let middleware = ResponseCachingMiddleware::new();
3495        let ctx = test_context();
3496        let first_page = test_request("tools/list", Some(serde_json::json!({"cursor": "a"})));
3497        let second_page = test_request("tools/list", Some(serde_json::json!({"cursor": "b"})));
3498
3499        for (request, tool_name) in [(&first_page, "first"), (&second_page, "second")] {
3500            middleware
3501                .on_response(
3502                    &ctx,
3503                    request,
3504                    serde_json::json!({
3505                        "resultType": "complete",
3506                        "tools": [{"name": tool_name}]
3507                    }),
3508                )
3509                .expect("each page is individually cacheable");
3510        }
3511
3512        assert!(matches!(
3513            middleware
3514                .on_request(&ctx, &first_page)
3515                .expect("first lookup is safe"),
3516            MiddlewareDecision::Respond(_)
3517        ));
3518        assert!(matches!(
3519            middleware
3520                .on_request(&ctx, &second_page)
3521                .expect("second lookup is safe"),
3522            MiddlewareDecision::Respond(_)
3523        ));
3524
3525        middleware.invalidate("tools/list", None);
3526
3527        assert_eq!(middleware.stats().entries, 0);
3528        assert!(matches!(
3529            middleware
3530                .on_request(&ctx, &first_page)
3531                .expect("first post-invalidation lookup is safe"),
3532            MiddlewareDecision::Continue
3533        ));
3534        assert!(matches!(
3535            middleware
3536                .on_request(&ctx, &second_page)
3537                .expect("second post-invalidation lookup is safe"),
3538            MiddlewareDecision::Continue
3539        ));
3540    }
3541
3542    #[test]
3543    fn cache_01_b_planted_negative() {
3544        let middleware = ResponseCachingMiddleware::new();
3545        let ctx = test_context();
3546        let request = test_request(
3547            "tools/list",
3548            Some(serde_json::json!({
3549                "cursor": "a",
3550                "requestState": {"opaque": "continuation"}
3551            })),
3552        );
3553
3554        // The only semantic change from the cacheable page is continuation
3555        // state. It must neither read nor populate an internal cache entry.
3556        assert!(matches!(
3557            middleware
3558                .on_request(&ctx, &request)
3559                .expect("continuation lookup is safe"),
3560            MiddlewareDecision::Continue
3561        ));
3562        let response = middleware
3563            .on_response(
3564                &ctx,
3565                &request,
3566                serde_json::json!({"resultType": "complete", "tools": []}),
3567            )
3568            .expect("the continuation result remains deliverable");
3569
3570        assert_eq!(
3571            response.get("cacheScope"),
3572            Some(&serde_json::json!("private"))
3573        );
3574        assert_eq!(middleware.stats().entries, 0);
3575        assert!(matches!(
3576            middleware
3577                .on_request(&ctx, &request)
3578                .expect("repeat continuation lookup is safe"),
3579            MiddlewareDecision::Continue
3580        ));
3581    }
3582
3583    #[test]
3584    fn cache_discovery_final_policy_positive() {
3585        let middleware = ResponseCachingMiddleware::new().list_ttl_secs(300);
3586        let ctx = test_context();
3587        let request = final_discovery_request(FINAL_PROTOCOL_VERSION);
3588        let ttl = Duration::from_secs(30);
3589        let response = final_discovery_response(30_000, "private");
3590
3591        let delivered = middleware
3592            .on_response(&ctx, &request, response.clone())
3593            .expect("final discovery response remains deliverable");
3594
3595        assert_eq!(delivered, response, "final hints remain public-observable");
3596        let cache = middleware
3597            .cache
3598            .lock()
3599            .unwrap_or_else(std::sync::PoisonError::into_inner);
3600        let entry = cache
3601            .entries
3602            .values()
3603            .next()
3604            .expect("private final discovery response is cached");
3605        assert!(
3606            entry.expires_at <= Instant::now().checked_add(ttl).expect("short TTL is valid"),
3607            "the stored expiry must use final ttlMs rather than the list default"
3608        );
3609        drop(cache);
3610
3611        assert!(matches!(
3612            middleware
3613                .on_request(&ctx, &request)
3614                .expect("final discovery lookup is safe"),
3615            MiddlewareDecision::Respond(value) if value == response
3616        ));
3617    }
3618
3619    #[test]
3620    fn cache_discovery_public_scope_planted_negative() {
3621        let middleware = ResponseCachingMiddleware::new();
3622        let ctx = test_context();
3623        let request = final_discovery_request(FINAL_PROTOCOL_VERSION);
3624        let response = final_discovery_response(30_000, "public");
3625
3626        // This differs from the cacheable final response only in cacheScope.
3627        // A public wire claim is delivered faithfully but cannot authorize this
3628        // private middleware cache.
3629        let delivered = middleware
3630            .on_response(&ctx, &request, response.clone())
3631            .expect("public discovery response remains deliverable");
3632
3633        assert_eq!(delivered, response);
3634        let before_lookup = middleware.stats();
3635        assert_eq!(before_lookup.entries, 0);
3636        assert_eq!(before_lookup.size_bytes, 0);
3637        assert!(matches!(
3638            middleware
3639                .on_request(&ctx, &request)
3640                .expect("public-scope lookup is safe"),
3641            MiddlewareDecision::Continue
3642        ));
3643        let after_lookup = middleware.stats();
3644        assert_eq!(after_lookup.entries, before_lookup.entries);
3645        assert_eq!(after_lookup.size_bytes, before_lookup.size_bytes);
3646    }
3647
3648    #[test]
3649    fn cache_discovery_huge_and_fractional_ttls_leave_existing_entry_unchanged() {
3650        let middleware = ResponseCachingMiddleware::new();
3651        let ctx = test_context();
3652        let request = final_discovery_request(FINAL_PROTOCOL_VERSION);
3653        let cached = final_discovery_response(30_000, "private");
3654
3655        middleware
3656            .on_response(&ctx, &request, cached.clone())
3657            .expect("baseline final discovery response is cacheable");
3658        let before = middleware.stats();
3659        assert_eq!(before.entries, 1);
3660
3661        let mut huge = cached.clone();
3662        huge["ttlMs"] =
3663            serde_json::from_str("18446744073709551616").expect("unbounded JSON integer fixture");
3664        let delivered = middleware
3665            .on_response(&ctx, &request, huge.clone())
3666            .expect("huge final discovery TTL remains deliverable");
3667
3668        assert_eq!(
3669            delivered, huge,
3670            "a wire-valid TTL outside the local runtime domain stays lossless"
3671        );
3672        assert_eq!(
3673            delivered["ttlMs"].to_string(),
3674            "18446744073709551616",
3675            "the huge TTL retains its exact peer spelling"
3676        );
3677        assert_eq!(
3678            middleware.stats(),
3679            before,
3680            "an uncacheable huge TTL does not replace the cached response"
3681        );
3682
3683        let mut fractional = huge;
3684        fractional["ttlMs"] = serde_json::from_str("18446744073709551616.5")
3685            .expect("fractional paired-negative fixture");
3686        let delivered = middleware
3687            .on_response(&ctx, &request, fractional)
3688            .expect("invalid final discovery hints remain deliverable");
3689
3690        assert!(delivered.get("ttlMs").is_none());
3691        assert!(delivered.get("cacheScope").is_none());
3692        assert_eq!(
3693            middleware.stats(),
3694            before,
3695            "changing only the TTL to a fractional value leaves cached state unchanged"
3696        );
3697        assert!(matches!(
3698            middleware
3699                .on_request(&ctx, &request)
3700                .expect("baseline cache lookup is safe"),
3701            MiddlewareDecision::Respond(value) if value == cached
3702        ));
3703    }
3704
3705    #[test]
3706    fn cache_discovery_stale_generation_planted_negative() {
3707        let middleware = ResponseCachingMiddleware::new();
3708        let ctx = test_context();
3709        let request = final_discovery_request(FINAL_PROTOCOL_VERSION);
3710        let response = final_discovery_response(30_000, "private");
3711
3712        middleware
3713            .on_response(&ctx, &request, response)
3714            .expect("initial final discovery response is cacheable");
3715        let binding_before_invalidation = middleware
3716            .cache_entry_binding(&request)
3717            .expect("final discovery has a cache binding");
3718        assert_eq!(middleware.stats().entries, 1);
3719
3720        // The request and result are unchanged; advancing only the discovery
3721        // generation makes the previous entry stale.
3722        middleware.invalidate(SERVER_DISCOVER_METHOD, request.params.as_ref());
3723        let binding_after_invalidation = middleware
3724            .cache_entry_binding(&request)
3725            .expect("fresh generation has a cache binding");
3726        assert_ne!(binding_after_invalidation, binding_before_invalidation);
3727
3728        let before_lookup = middleware.stats();
3729        assert_eq!(before_lookup.entries, 0);
3730        assert_eq!(before_lookup.size_bytes, 0);
3731        assert!(matches!(
3732            middleware
3733                .on_request(&ctx, &request)
3734                .expect("stale discovery lookup is safe"),
3735            MiddlewareDecision::Continue
3736        ));
3737        let after_lookup = middleware.stats();
3738        assert_eq!(after_lookup.entries, before_lookup.entries);
3739        assert_eq!(after_lookup.size_bytes, before_lookup.size_bytes);
3740    }
3741
3742    #[test]
3743    fn cache_discovery_cross_era_planted_negative() {
3744        let middleware = ResponseCachingMiddleware::new();
3745        let final_ctx = test_context();
3746        let legacy_ctx = test_context();
3747        let final_request = final_discovery_request(FINAL_PROTOCOL_VERSION);
3748        let mut legacy_request = final_request.clone();
3749        legacy_request
3750            .params
3751            .as_mut()
3752            .expect("test request has metadata")["_meta"][FINAL_PROTOCOL_VERSION_META_KEY] =
3753            serde_json::json!(ProtocolEra::Legacy2024.version().as_str());
3754        let response = final_discovery_response(30_000, "private");
3755
3756        middleware
3757            .on_response(&final_ctx, &final_request, response.clone())
3758            .expect("final discovery response is cacheable");
3759        let before_legacy = middleware.stats();
3760        assert_eq!(before_legacy.entries, 1);
3761
3762        // The only changed input is the exact protocol era. It must neither
3763        // reuse the final entry nor add an entry in the final generation.
3764        let delivered = middleware
3765            .on_response(&legacy_ctx, &legacy_request, response.clone())
3766            .expect("legacy response remains deliverable without caching");
3767        assert_eq!(delivered, response);
3768        assert!(matches!(
3769            middleware
3770                .on_request(&legacy_ctx, &legacy_request)
3771                .expect("legacy lookup is safe"),
3772            MiddlewareDecision::Continue
3773        ));
3774        let after_legacy = middleware.stats();
3775        assert_eq!(after_legacy.entries, before_legacy.entries);
3776        assert_eq!(after_legacy.size_bytes, before_legacy.size_bytes);
3777        assert!(matches!(
3778            middleware
3779                .on_request(&final_ctx, &final_request)
3780                .expect("final lookup remains safe"),
3781            MiddlewareDecision::Respond(value) if value == response
3782        ));
3783    }
3784}