diff --git a/README.md b/README.md index 13219440..57f69c5b 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,7 @@ For the full MCP specification, see [modelcontextprotocol.io](https://modelconte - [Notifications](#notifications) - [Subscriptions](#subscriptions) - [Tasks](#tasks-long-running-tool-invocations) +- [Caching](#caching) - [Examples](#examples) - [OAuth Support](#oauth-support) - [Related Resources](#related-resources) @@ -1003,6 +1004,52 @@ async fn call_tool(&self, request: CallToolRequestParams, context: RequestContex See [`servers_task_stdio`](examples/servers/src/task_stdio.rs) and the matching [`clients_task_stdio`](examples/clients/src/task_stdio.rs) for a runnable end-to-end example. +## Caching + +`rmcp` clients transparently cache responses that carry the +[SEP-2549](https://modelcontextprotocol.io/specification/draft/server/utilities/caching) +caching hints (`ttlMs` / `cacheScope`) for `server/discover`, `tools/list`, +`prompts/list`, `resources/list`, `resources/templates/list`, and `resources/read`. + +Caching is on by default but only stores a response when the server sends a +positive `ttlMs`, so servers that omit the hint behave exactly as before. Entries +expire after their TTL, are partitioned by cache scope, and are invalidated +automatically by the matching `list_changed` / `resource updated` notifications. + +No call-site changes are needed — existing calls benefit automatically: + +```rust, ignore +let tools = peer.list_tools(None).await?; // served from cache while fresh +let res = peer.read_resource(params).await?; // cached per-URI +``` + +Tune or disable it per connection via the `Peer`: + +```rust, ignore +use std::time::Duration; +use rmcp::ClientCacheConfig; + +// Customize behavior. +peer.set_response_cache_config( + ClientCacheConfig::default() + .with_default_ttl(Duration::from_secs(30)) // TTL for servers that omit ttlMs + .with_max_ttl(Duration::from_secs(3600)) // upper bound on any TTL + .with_max_entries(1024) + .with_private_partition(user_id) // separate private caches per principal + .with_serve_stale_on_error(false), // surface errors instead of stale data +).await; + +// Or turn it off entirely. +peer.set_response_cache_config(ClientCacheConfig::disabled()).await; + +// Manually flush. +peer.clear_response_cache().await; +``` + +> **Note:** with the default `serve_stale_on_error`, a failed re-fetch returns the +> last cached response (even if expired) as `Ok(..)` instead of an error. Set +> `with_serve_stale_on_error(false)` if callers must observe fetch failures. + ## Examples See [examples](examples/README.md). diff --git a/crates/rmcp/src/lib.rs b/crates/rmcp/src/lib.rs index ca195a6e..7c9b7b19 100644 --- a/crates/rmcp/src/lib.rs +++ b/crates/rmcp/src/lib.rs @@ -18,8 +18,8 @@ pub use handler::server::ServerHandler; pub use handler::server::wrapper::Json; #[cfg(feature = "client")] pub use service::{ - ClientLifecycleMode, ClientServiceExt, RoleClient, select_protocol_version, serve_client, - serve_client_with_lifecycle, + ClientCacheConfig, ClientLifecycleMode, ClientServiceExt, MAX_CLIENT_CACHE_TTL, RoleClient, + select_protocol_version, serve_client, serve_client_with_lifecycle, }; #[cfg(any(feature = "client", feature = "server"))] pub use service::{Peer, Service, ServiceError, ServiceExt}; diff --git a/crates/rmcp/src/service.rs b/crates/rmcp/src/service.rs index 4eeaeadf..75f34d30 100644 --- a/crates/rmcp/src/service.rs +++ b/crates/rmcp/src/service.rs @@ -136,6 +136,19 @@ pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone { fn peer_cancelled_params(_notification: &Self::PeerNot) -> Option<&CancelledNotificationParam> { None } + /// Invalidate any response cache affected by an inbound peer notification. + /// + /// The serve loop calls this for every notification *before* subscription + /// routing, so cache invalidation still runs when a notification is + /// delivered through a `listen` subscription channel rather than the + /// [`Service::handle_notification`] callbacks. + #[doc(hidden)] + fn invalidate_response_cache( + _peer: &Peer, + _notification: &Self::PeerNot, + ) -> impl Future + MaybeSendFuture { + async {} + } } pub(crate) fn uses_legacy_lifecycle( @@ -571,6 +584,8 @@ pub struct Peer { client_request_metadata: Arc>, request_metadata_required: Arc, subscription_channels: Arc>>, + #[cfg(feature = "client")] + response_cache: client::cache::PeerResponseCache, } impl Clone for Peer @@ -587,6 +602,8 @@ where client_request_metadata: self.client_request_metadata.clone(), request_metadata_required: self.request_metadata_required.clone(), subscription_channels: self.subscription_channels.clone(), + #[cfg(feature = "client")] + response_cache: self.response_cache.clone(), } } } @@ -661,6 +678,8 @@ impl Peer { client_request_metadata: Default::default(), request_metadata_required: Default::default(), subscription_channels: Default::default(), + #[cfg(feature = "client")] + response_cache: Default::default(), }, rx, ) @@ -1402,6 +1421,7 @@ where .. })) => { tracing::info!(?notification, "received notification"); + R::invalidate_response_cache(&peer, ¬ification).await; let cancellation_request_id = if let Some(cancelled) = R::peer_cancelled_params(¬ification) { let request_id = cancelled.request_id.clone(); diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 88336ed5..166207b7 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -1,30 +1,35 @@ // Sampling/Roots/Logging are SEP-2577-deprecated; internal references are expected. #![expect(deprecated)] +pub(super) mod cache; + use std::{borrow::Cow, num::NonZeroUsize, sync::Arc, time::Duration}; +use cache::CacheGeneration; +pub use cache::{ClientCacheConfig, MAX_CLIENT_CACHE_TTL}; use thiserror::Error; use super::*; use crate::{ model::{ - ArgumentInfo, CallToolRequest, CallToolRequestParams, CallToolResponse, CallToolResult, - CancelTaskParams, CancelTaskRequest, CancelledNotification, CancelledNotificationParam, - ClientInfo, ClientJsonRpcMessage, ClientNotification, ClientRequest, ClientResult, - CompleteRequest, CompleteRequestParams, CompleteResult, CompletionContext, CompletionInfo, - DEFAULT_MRTR_MAX_ROUNDS, DiscoverRequest, DiscoverRequestParams, DiscoverResult, ErrorData, - GetExtensions, GetMeta, GetPromptRequest, GetPromptRequestParams, GetPromptResponse, - GetPromptResult, GetTaskParams, GetTaskRequest, GetTaskResult, InitializeRequest, - InitializedNotification, InputRequest, InputRequiredResult, InputResponses, - JsonRpcResponse, ListPromptsRequest, ListPromptsResult, ListResourceTemplatesRequest, - ListResourceTemplatesResult, ListResourcesRequest, ListResourcesResult, ListToolsRequest, - ListToolsResult, NumberOrString, PaginatedRequestParams, ProgressNotification, - ProgressNotificationParam, ProtocolVersion, ReadResourceRequest, ReadResourceRequestParams, - ReadResourceResponse, ReadResourceResult, Reference, RequestId, RequestMetaObject, - RootsListChangedNotification, ServerInfo, ServerJsonRpcMessage, ServerNotification, - ServerRequest, ServerResult, SetLevelRequest, SetLevelRequestParams, SubscribeRequest, - SubscribeRequestParams, SubscriptionFilter, SubscriptionsListenRequest, - SubscriptionsListenRequestParams, SubscriptionsListenResult, UnsubscribeRequest, - UnsubscribeRequestParams, UpdateTaskParams, UpdateTaskRequest, + ArgumentInfo, CacheScope, CallToolRequest, CallToolRequestParams, CallToolResponse, + CallToolResult, CancelTaskParams, CancelTaskRequest, CancelledNotification, + CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage, ClientNotification, + ClientRequest, ClientResult, CompleteRequest, CompleteRequestParams, CompleteResult, + CompletionContext, CompletionInfo, DEFAULT_MRTR_MAX_ROUNDS, DiscoverRequest, + DiscoverRequestParams, DiscoverResult, ErrorData, GetExtensions, GetMeta, GetPromptRequest, + GetPromptRequestParams, GetPromptResponse, GetPromptResult, GetTaskParams, GetTaskRequest, + GetTaskResult, InitializeRequest, InitializedNotification, InputRequest, + InputRequiredResult, InputResponses, JsonRpcResponse, ListPromptsRequest, + ListPromptsResult, ListResourceTemplatesRequest, ListResourceTemplatesResult, + ListResourcesRequest, ListResourcesResult, ListToolsRequest, ListToolsResult, + NumberOrString, PaginatedRequestParams, ProgressNotification, ProgressNotificationParam, + ProtocolVersion, ReadResourceRequest, ReadResourceRequestParams, ReadResourceResponse, + ReadResourceResult, Reference, RequestId, RequestMetaObject, RootsListChangedNotification, + ServerInfo, ServerJsonRpcMessage, ServerNotification, ServerRequest, ServerResult, + SetLevelRequest, SetLevelRequestParams, SubscribeRequest, SubscribeRequestParams, + SubscriptionFilter, SubscriptionsListenRequest, SubscriptionsListenRequestParams, + SubscriptionsListenResult, UnsubscribeRequest, UnsubscribeRequestParams, UpdateTaskParams, + UpdateTaskRequest, }, transport::DynamicTransportError, }; @@ -207,6 +212,25 @@ impl ServiceRole for RoleClient { _ => None, } } + + async fn invalidate_response_cache(peer: &Peer, notification: &Self::PeerNot) { + match notification { + ServerNotification::ResourceUpdatedNotification(notification) => { + peer.invalidate_resource_read_cache(¬ification.params.uri) + .await; + } + ServerNotification::ResourceListChangedNotification(_) => { + peer.invalidate_resource_list_cache().await; + } + ServerNotification::ToolListChangedNotification(_) => { + peer.invalidate_tool_cache().await; + } + ServerNotification::PromptListChangedNotification(_) => { + peer.invalidate_prompt_cache().await; + } + _ => {} + } + } } pub type ServerSink = Peer; @@ -821,6 +845,52 @@ where } } +const DISCOVER_CACHE_PREFIX: &str = "server/discover:"; +const TOOL_LIST_CACHE_PREFIX: &str = "tools/list:"; +const PROMPT_LIST_CACHE_PREFIX: &str = "prompts/list:"; +const RESOURCE_LIST_CACHE_PREFIX: &str = "resources/list:"; +const RESOURCE_TEMPLATE_LIST_CACHE_PREFIX: &str = "resources/templates/list:"; +const RESOURCE_READ_CACHE_PREFIX: &str = "resources/read:"; + +// Cache keys are built only from the request method plus the parameters that +// affect the result (SEP-2549). Request `_meta` (progress tokens, trace +// context, etc.) does not affect the result, so it is deliberately excluded to +// avoid fragmenting the cache across otherwise-identical requests. +fn discover_cache_key() -> String { + // `server/discover` carries no result-affecting parameters. + DISCOVER_CACHE_PREFIX.to_string() +} + +fn list_response_cache_key(prefix: &str, params: &Option) -> String { + // Only the pagination cursor affects which page is returned. + let cursor = params.as_ref().and_then(|params| params.cursor.as_deref()); + let cursor = + serde_json::to_string(&cursor).expect("serializing a pagination cursor cannot fail"); + format!("{prefix}{cursor}") +} + +fn resource_read_cache_key(params: &ReadResourceRequestParams) -> Option { + // MRTR retries depend on inputs that are not part of the cache key and MUST + // NOT be cached. + if params.input_responses.is_some() || params.request_state.is_some() { + return None; + } + // Only the URI affects the result. + Some(resource_read_cache_prefix_for_uri(¶ms.uri)) +} + +fn resource_read_cache_prefix_for_uri(uri: &str) -> String { + let uri = serde_json::to_string(uri).expect("serializing a resource URI cannot fail"); + format!("{RESOURCE_READ_CACHE_PREFIX}{uri}:") +} + +fn request_uses_cursor(params: &Option) -> bool { + params + .as_ref() + .and_then(|params| params.cursor.as_ref()) + .is_some() +} + macro_rules! method { ($(#[$meta:meta])* peer_req $method:ident $Req:ident() => $Resp: ident ) => { $(#[$meta])* @@ -1016,17 +1086,80 @@ impl Peer { /// The high-level client currently exposes this peer only after initialization; /// pre-initialization probing is planned as follow-up work. pub async fn discover(&self, meta: RequestMetaObject) -> Result { + let cache_key = discover_cache_key(); + if let Some(ServerResult::DiscoverResult(result)) = self.cached_response(&cache_key).await { + return Ok(result); + } + let generation = self.capture_response_cache_generation().await; let mut request = DiscoverRequest::new(DiscoverRequestParams {}); request.extensions.insert(meta); let result = self .send_request(ClientRequest::DiscoverRequest(request)) - .await?; + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + if let Some(ServerResult::DiscoverResult(result)) = + self.stale_cached_response(&cache_key).await + { + return Ok(result); + } + return Err(error); + } + }; match result { - ServerResult::DiscoverResult(result) => Ok(result), + ServerResult::DiscoverResult(result) => { + self.cache_result( + Some(cache_key), + Some(result.ttl_ms), + Some(result.cache_scope), + generation, + ServerResult::DiscoverResult(result.clone()), + ) + .await; + Ok(result) + } _ => Err(ServiceError::UnexpectedResponse), } } + async fn cache_result( + &self, + cache_key: Option, + ttl_ms: Option, + cache_scope: Option, + generation: CacheGeneration, + result: ServerResult, + ) { + let Some(cache_key) = cache_key else { + return; + }; + self.cache_response_with_generation(cache_key, result, ttl_ms, cache_scope, generation) + .await; + } + + pub(crate) async fn invalidate_tool_cache(&self) { + self.invalidate_cached_responses(TOOL_LIST_CACHE_PREFIX) + .await; + } + + pub(crate) async fn invalidate_prompt_cache(&self) { + self.invalidate_cached_responses(PROMPT_LIST_CACHE_PREFIX) + .await; + } + + pub(crate) async fn invalidate_resource_list_cache(&self) { + self.invalidate_cached_responses(RESOURCE_LIST_CACHE_PREFIX) + .await; + self.invalidate_cached_responses(RESOURCE_TEMPLATE_LIST_CACHE_PREFIX) + .await; + } + + pub(crate) async fn invalidate_resource_read_cache(&self, uri: &str) { + self.invalidate_cached_responses(&resource_read_cache_prefix_for_uri(uri)) + .await; + } + /// Send one `tools/call` request and return either a final result or an MRTR /// `InputRequiredResult` without driving any follow-up rounds. pub async fn call_tool_once( @@ -1118,15 +1251,45 @@ impl Peer { &self, params: ReadResourceRequestParams, ) -> Result { + let cache_key = resource_read_cache_key(¶ms); + if let Some(key) = cache_key.as_deref() + && let Some(ServerResult::ReadResourceResult(result)) = self.cached_response(key).await + { + return Ok(ReadResourceResponse::Complete(result)); + } + + let generation = self.capture_response_cache_generation().await; let result = self .send_request(ClientRequest::ReadResourceRequest(ReadResourceRequest { method: Default::default(), params, extensions: Default::default(), })) - .await?; + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + if let Some(key) = cache_key.as_deref() + && let Some(ServerResult::ReadResourceResult(result)) = + self.stale_cached_response(key).await + { + return Ok(ReadResourceResponse::Complete(result)); + } + return Err(error); + } + }; match result { - ServerResult::ReadResourceResult(result) => Ok(ReadResourceResponse::Complete(result)), + ServerResult::ReadResourceResult(result) => { + self.cache_result( + cache_key, + result.ttl_ms, + result.cache_scope, + generation, + ServerResult::ReadResourceResult(result.clone()), + ) + .await; + Ok(ReadResourceResponse::Complete(result)) + } ServerResult::InputRequiredResult(result) => { Ok(ReadResourceResponse::InputRequired(result)) } @@ -1143,10 +1306,6 @@ impl Peer { peer_req set_level SetLevelRequest(SetLevelRequestParams) ); method!(peer_req get_prompt GetPromptRequest(GetPromptRequestParams) => GetPromptResult); - method!(peer_req list_prompts ListPromptsRequest(PaginatedRequestParams)? => ListPromptsResult); - method!(peer_req list_resources ListResourcesRequest(PaginatedRequestParams)? => ListResourcesResult); - method!(peer_req list_resource_templates ListResourceTemplatesRequest(PaginatedRequestParams)? => ListResourceTemplatesResult); - method!(peer_req read_resource ReadResourceRequest(ReadResourceRequestParams) => ReadResourceResult); method!( #[deprecated( note = "resources/subscribe is legacy-only; use Peer::listen for protocol version 2026-07-28" @@ -1160,7 +1319,219 @@ impl Peer { peer_req unsubscribe UnsubscribeRequest(UnsubscribeRequestParams) ); method!(peer_req call_tool CallToolRequest(CallToolRequestParams) => CallToolResult); - method!(peer_req list_tools ListToolsRequest(PaginatedRequestParams)? => ListToolsResult); + + pub async fn list_prompts( + &self, + params: Option, + ) -> Result { + let cache_key = list_response_cache_key(PROMPT_LIST_CACHE_PREFIX, ¶ms); + if let Some(ServerResult::ListPromptsResult(result)) = + self.cached_response(&cache_key).await + { + return Ok(result); + } + let generation = self.capture_response_cache_generation().await; + let uses_cursor = request_uses_cursor(¶ms); + let result = self + .send_request(ClientRequest::ListPromptsRequest(ListPromptsRequest { + method: Default::default(), + params, + extensions: Default::default(), + })) + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + if uses_cursor { + self.invalidate_prompt_cache().await; + return Err(error); + } + if let Some(ServerResult::ListPromptsResult(result)) = + self.stale_cached_response(&cache_key).await + { + return Ok(result); + } + return Err(error); + } + }; + match result { + ServerResult::ListPromptsResult(result) => { + self.cache_result( + Some(cache_key), + result.ttl_ms, + result.cache_scope, + generation, + ServerResult::ListPromptsResult(result.clone()), + ) + .await; + Ok(result) + } + _ => Err(ServiceError::UnexpectedResponse), + } + } + + pub async fn list_resources( + &self, + params: Option, + ) -> Result { + let cache_key = list_response_cache_key(RESOURCE_LIST_CACHE_PREFIX, ¶ms); + if let Some(ServerResult::ListResourcesResult(result)) = + self.cached_response(&cache_key).await + { + return Ok(result); + } + let generation = self.capture_response_cache_generation().await; + let uses_cursor = request_uses_cursor(¶ms); + let result = self + .send_request(ClientRequest::ListResourcesRequest(ListResourcesRequest { + method: Default::default(), + params, + extensions: Default::default(), + })) + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + if uses_cursor { + self.invalidate_cached_responses(RESOURCE_LIST_CACHE_PREFIX) + .await; + return Err(error); + } + if let Some(ServerResult::ListResourcesResult(result)) = + self.stale_cached_response(&cache_key).await + { + return Ok(result); + } + return Err(error); + } + }; + match result { + ServerResult::ListResourcesResult(result) => { + self.cache_result( + Some(cache_key), + result.ttl_ms, + result.cache_scope, + generation, + ServerResult::ListResourcesResult(result.clone()), + ) + .await; + Ok(result) + } + _ => Err(ServiceError::UnexpectedResponse), + } + } + + pub async fn list_resource_templates( + &self, + params: Option, + ) -> Result { + let cache_key = list_response_cache_key(RESOURCE_TEMPLATE_LIST_CACHE_PREFIX, ¶ms); + if let Some(ServerResult::ListResourceTemplatesResult(result)) = + self.cached_response(&cache_key).await + { + return Ok(result); + } + let generation = self.capture_response_cache_generation().await; + let uses_cursor = request_uses_cursor(¶ms); + let result = self + .send_request(ClientRequest::ListResourceTemplatesRequest( + ListResourceTemplatesRequest { + method: Default::default(), + params, + extensions: Default::default(), + }, + )) + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + if uses_cursor { + self.invalidate_cached_responses(RESOURCE_TEMPLATE_LIST_CACHE_PREFIX) + .await; + return Err(error); + } + if let Some(ServerResult::ListResourceTemplatesResult(result)) = + self.stale_cached_response(&cache_key).await + { + return Ok(result); + } + return Err(error); + } + }; + match result { + ServerResult::ListResourceTemplatesResult(result) => { + self.cache_result( + Some(cache_key), + result.ttl_ms, + result.cache_scope, + generation, + ServerResult::ListResourceTemplatesResult(result.clone()), + ) + .await; + Ok(result) + } + _ => Err(ServiceError::UnexpectedResponse), + } + } + + pub async fn read_resource( + &self, + params: ReadResourceRequestParams, + ) -> Result { + match self.read_resource_once(params).await? { + ReadResourceResponse::Complete(result) => Ok(result), + ReadResourceResponse::InputRequired(_) => Err(ServiceError::UnexpectedResponse), + } + } + + pub async fn list_tools( + &self, + params: Option, + ) -> Result { + let cache_key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms); + if let Some(ServerResult::ListToolsResult(result)) = self.cached_response(&cache_key).await + { + return Ok(result); + } + let generation = self.capture_response_cache_generation().await; + let uses_cursor = request_uses_cursor(¶ms); + let result = self + .send_request(ClientRequest::ListToolsRequest(ListToolsRequest { + method: Default::default(), + params, + extensions: Default::default(), + })) + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + if uses_cursor { + self.invalidate_tool_cache().await; + return Err(error); + } + if let Some(ServerResult::ListToolsResult(result)) = + self.stale_cached_response(&cache_key).await + { + return Ok(result); + } + return Err(error); + } + }; + match result { + ServerResult::ListToolsResult(result) => { + self.cache_result( + Some(cache_key), + result.ttl_ms, + result.cache_scope, + generation, + ServerResult::ListToolsResult(result.clone()), + ) + .await; + Ok(result) + } + _ => Err(ServiceError::UnexpectedResponse), + } + } method!(peer_not notify_cancelled CancelledNotification(CancelledNotificationParam)); method!(peer_not notify_progress ProgressNotification(ProgressNotificationParam)); @@ -1636,3 +2007,157 @@ where )) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn disconnected_peer() -> Peer { + let (peer, receiver) = + Peer::::new(Arc::new(AtomicU32RequestIdProvider::default()), None); + drop(receiver); + peer + } + + fn tools_result(ttl_ms: Option, cache_scope: Option) -> ListToolsResult { + let mut result = ListToolsResult::with_all_items(Vec::new()); + result.ttl_ms = ttl_ms; + result.cache_scope = cache_scope; + result + } + + #[tokio::test] + async fn fresh_cached_page_is_served_without_transport_io() { + let peer = disconnected_peer(); + let params = None::; + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms); + let expected = tools_result(Some(5_000), Some(CacheScope::Public)); + peer.cache_response( + key, + ServerResult::ListToolsResult(expected.clone()), + expected.ttl_ms, + expected.cache_scope, + ) + .await; + + assert_eq!(peer.list_tools(params).await.unwrap(), expected); + } + + #[tokio::test] + async fn expired_entry_falls_through_to_the_transport() { + let peer = disconnected_peer(); + peer.set_response_cache_config( + ClientCacheConfig::default().with_serve_stale_on_error(false), + ) + .await; + let params = None::; + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms); + peer.cache_response( + key, + ServerResult::ListToolsResult(tools_result(Some(1), Some(CacheScope::Public))), + Some(1), + Some(CacheScope::Public), + ) + .await; + tokio::time::sleep(Duration::from_millis(5)).await; + + assert!(matches!( + peer.list_tools(params).await, + Err(ServiceError::TransportClosed) + )); + } + + #[tokio::test] + async fn private_entries_are_isolated_between_authorization_partitions() { + let peer = disconnected_peer(); + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, &None); + + peer.set_response_cache_config( + ClientCacheConfig::default().with_private_partition("auth-a"), + ) + .await; + peer.cache_response( + key.clone(), + ServerResult::ListToolsResult(tools_result(Some(5_000), Some(CacheScope::Private))), + Some(5_000), + Some(CacheScope::Private), + ) + .await; + assert!(peer.cached_response(&key).await.is_some()); + + // Switching to a different authorization context must not expose the + // first partition's private entry. + peer.set_response_cache_config( + ClientCacheConfig::default().with_private_partition("auth-b"), + ) + .await; + assert!(peer.cached_response(&key).await.is_none()); + } + + #[tokio::test] + async fn list_change_notification_discards_every_cached_page() { + let peer = disconnected_peer(); + for cursor in [None, Some("page-a".into()), Some("page-b".into())] { + let params = + cursor.map(|cursor| PaginatedRequestParams::default().with_cursor(Some(cursor))); + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms); + peer.cache_response( + key, + ServerResult::ListToolsResult(tools_result(Some(5_000), Some(CacheScope::Public))), + Some(5_000), + Some(CacheScope::Public), + ) + .await; + } + + peer.invalidate_tool_cache().await; + + for cursor in [None, Some("page-a".into()), Some("page-b".into())] { + let params = + cursor.map(|cursor| PaginatedRequestParams::default().with_cursor(Some(cursor))); + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms); + assert!(peer.cached_response(&key).await.is_none()); + } + } + + #[tokio::test] + async fn expired_entry_is_served_when_refetch_fails() { + let peer = disconnected_peer(); + let params = None::; + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms); + let expected = tools_result(Some(1), Some(CacheScope::Public)); + peer.cache_response( + key, + ServerResult::ListToolsResult(expected.clone()), + Some(1), + Some(CacheScope::Public), + ) + .await; + tokio::time::sleep(Duration::from_millis(5)).await; + + assert_eq!(peer.list_tools(params).await.unwrap(), expected); + } + + #[tokio::test] + async fn discover_serves_a_fresh_cached_response_without_transport_io() { + let peer = disconnected_peer(); + let meta = RequestMetaObject::default(); + let key = discover_cache_key(); + let expected = DiscoverResult::new( + vec![ProtocolVersion::default()], + Default::default(), + crate::model::Implementation::from_build_env(), + ) + .with_ttl_ms(5_000) + .with_cache_scope(CacheScope::Public); + peer.cache_response( + key, + ServerResult::DiscoverResult(expected.clone()), + Some(5_000), + Some(CacheScope::Public), + ) + .await; + + assert_eq!(peer.discover(meta).await.unwrap(), expected); + } +} diff --git a/crates/rmcp/src/service/client/cache.rs b/crates/rmcp/src/service/client/cache.rs new file mode 100644 index 00000000..f7137ccd --- /dev/null +++ b/crates/rmcp/src/service/client/cache.rs @@ -0,0 +1,401 @@ +use std::{ + collections::HashMap, + sync::Arc, + time::{Duration, Instant}, +}; + +use super::RoleClient; +use crate::{ + model::CacheScope, + service::{Peer, ServiceRole}, +}; + +/// Maximum server-provided cache TTL honoured by the client response cache. +pub const MAX_CLIENT_CACHE_TTL: Duration = Duration::from_secs(24 * 60 * 60); + +/// Configuration for the built-in MCP client response cache. +/// +/// A cache is allocated per client [`Peer`]. Public responses may be reused +/// throughout that client connection. Private responses are additionally +/// partitioned by `private_partition`; changing the partition drops every +/// private entry while preserving public entries. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub struct ClientCacheConfig { + /// Enables cache reads and writes. + pub enabled: bool, + /// TTL used when a backwards-compatible server omits `ttlMs`. + /// + /// The default is zero, which leaves such responses immediately stale. + pub default_ttl: Duration, + /// Upper bound applied to both server-provided and default TTLs. + pub max_ttl: Duration, + /// Stable opaque identity for the current authorization context. + /// + /// A single-principal client may leave this unset because each client owns + /// its own in-memory store. Gateways or clients that change principals on an + /// existing connection should set this value and update it whenever the + /// authorization context changes. + pub private_partition: Option, + /// Maximum number of responses retained by the in-memory cache. + /// + /// A value of zero disables the size limit. + pub max_entries: usize, + /// Serves an expired cached response when a re-fetch fails. + /// + /// SEP-2549 permits clients to serve stale responses if errors occur while + /// re-fetching (for example, network issues or server downtime). When this + /// is enabled the client retains expired entries so it can fall back to the + /// last known response instead of surfacing the transport or server error. + /// A successful re-fetch always overwrites the stale entry. + pub serve_stale_on_error: bool, +} + +impl Default for ClientCacheConfig { + fn default() -> Self { + Self { + enabled: true, + default_ttl: Duration::ZERO, + max_ttl: MAX_CLIENT_CACHE_TTL, + private_partition: None, + max_entries: 512, + serve_stale_on_error: true, + } + } +} + +impl ClientCacheConfig { + /// Returns a configuration that disables all cache reads and writes. + pub fn disabled() -> Self { + Self { + enabled: false, + ..Self::default() + } + } + + /// Sets the TTL used when a response omits `ttlMs`. + pub fn with_default_ttl(mut self, default_ttl: Duration) -> Self { + self.default_ttl = default_ttl; + self + } + + /// Sets the maximum TTL the client will honour. + pub fn with_max_ttl(mut self, max_ttl: Duration) -> Self { + self.max_ttl = max_ttl; + self + } + + /// Sets the stable partition for private responses. + pub fn with_private_partition(mut self, partition: impl Into) -> Self { + self.private_partition = Some(partition.into()); + self + } + + /// Sets the maximum number of retained responses. + pub fn with_max_entries(mut self, max_entries: usize) -> Self { + self.max_entries = max_entries; + self + } + + /// Controls whether an expired response may be served when a re-fetch fails. + pub fn with_serve_stale_on_error(mut self, serve_stale_on_error: bool) -> Self { + self.serve_stale_on_error = serve_stale_on_error; + self + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +enum CachePartition { + Public, + Private(Arc), +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct CacheKey { + logical_key: String, + partition: CachePartition, +} + +#[derive(Debug, Clone)] +struct CachedPeerResponse { + value: T, + expires_at: Instant, + inserted_at: Instant, + scope: CacheScope, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct CacheGeneration(u64); + +#[derive(Debug)] +pub(crate) struct PeerResponseCacheState { + entries: HashMap>, + config: ClientCacheConfig, + generation: u64, +} + +impl Default for PeerResponseCacheState { + fn default() -> Self { + Self { + entries: HashMap::new(), + config: ClientCacheConfig::default(), + generation: 0, + } + } +} + +impl PeerResponseCacheState { + fn trim_to_limit(&mut self) { + while self.config.max_entries > 0 && self.entries.len() > self.config.max_entries { + let Some(oldest_key) = self + .entries + .iter() + .min_by_key(|(_, entry)| entry.inserted_at) + .map(|(key, _)| key.clone()) + else { + break; + }; + self.entries.remove(&oldest_key); + } + } +} + +pub(crate) type PeerResponseCache = Arc>>; + +impl Peer { + fn private_partition(config: &ClientCacheConfig) -> Arc { + Arc::from(config.private_partition.as_deref().unwrap_or("connection")) + } + + fn cache_key(logical_key: &str, partition: CachePartition) -> CacheKey { + CacheKey { + logical_key: logical_key.to_owned(), + partition, + } + } + + fn scoped_cache_key( + logical_key: &str, + scope: CacheScope, + config: &ClientCacheConfig, + ) -> CacheKey { + let partition = match scope { + CacheScope::Public => CachePartition::Public, + CacheScope::Private => CachePartition::Private(Self::private_partition(config)), + }; + Self::cache_key(logical_key, partition) + } + + /// Captures the cache generation before a request crosses the transport. + /// + /// Any configuration change, explicit clear, or notification invalidation + /// advances the generation. A response from an older generation is not + /// written back, preventing an in-flight stale response from undoing an + /// invalidation. + pub(crate) async fn capture_response_cache_generation(&self) -> CacheGeneration { + CacheGeneration(self.response_cache.read().await.generation) + } + + /// Returns a fresh cached response, preferring the current private partition + /// before the public partition. + /// + /// Expired entries are removed on access unless `serve_stale_on_error` is + /// enabled, in which case they are retained so a later re-fetch failure can + /// fall back to them via [`Peer::stale_cached_response`]. + pub(crate) async fn cached_response(&self, logical_key: &str) -> Option { + let now = Instant::now(); + let mut cache = self.response_cache.write().await; + if !cache.config.enabled { + return None; + } + let keep_stale = cache.config.serve_stale_on_error; + + let private_key = Self::cache_key( + logical_key, + CachePartition::Private(Self::private_partition(&cache.config)), + ); + let private_fresh = cache.entries.get(&private_key).and_then(|entry| { + (entry.expires_at > now && entry.scope == CacheScope::Private) + .then(|| entry.value.clone()) + }); + if let Some(value) = private_fresh { + return Some(value); + } + if !keep_stale { + cache.entries.remove(&private_key); + } + + let public_key = Self::cache_key(logical_key, CachePartition::Public); + let public_fresh = cache.entries.get(&public_key).and_then(|entry| { + (entry.expires_at > now && entry.scope == CacheScope::Public) + .then(|| entry.value.clone()) + }); + if let Some(value) = public_fresh { + return Some(value); + } + if !keep_stale { + cache.entries.remove(&public_key); + } + None + } + + /// Returns a cached response ignoring its TTL, for use as a fallback when a + /// re-fetch fails (SEP-2549 permits serving stale responses on error). + /// + /// Returns `None` when the cache is disabled or `serve_stale_on_error` is + /// turned off. The private partition is preferred over the public one. The + /// entry is left in place so repeated failures keep serving it until a + /// successful re-fetch overwrites it or a notification invalidates it. + pub(crate) async fn stale_cached_response(&self, logical_key: &str) -> Option { + let cache = self.response_cache.read().await; + if !cache.config.enabled || !cache.config.serve_stale_on_error { + return None; + } + + let private_key = Self::cache_key( + logical_key, + CachePartition::Private(Self::private_partition(&cache.config)), + ); + if let Some(entry) = cache.entries.get(&private_key) + && entry.scope == CacheScope::Private + { + return Some(entry.value.clone()); + } + + let public_key = Self::cache_key(logical_key, CachePartition::Public); + if let Some(entry) = cache.entries.get(&public_key) + && entry.scope == CacheScope::Public + { + return Some(entry.value.clone()); + } + None + } + + /// Stores a response when the configured effective TTL is positive. + /// + /// Missing `cacheScope` is treated as private. This is deliberately more + /// conservative than the model's backwards-compatible wire default and + /// prevents an older or malformed server response from becoming shareable. + pub(crate) async fn cache_response_with_generation( + &self, + logical_key: String, + value: R::PeerResp, + ttl_ms: Option, + cache_scope: Option, + generation: CacheGeneration, + ) { + let now = Instant::now(); + let mut cache = self.response_cache.write().await; + if !cache.config.enabled || generation.0 != cache.generation { + return; + } + + let requested_ttl = ttl_ms + .map(Duration::from_millis) + .unwrap_or(cache.config.default_ttl); + let ttl = requested_ttl.min(cache.config.max_ttl); + if ttl.is_zero() { + return; + } + let Some(expires_at) = now.checked_add(ttl) else { + return; + }; + let scope = cache_scope.unwrap_or(CacheScope::Private); + let target_key = Self::scoped_cache_key(&logical_key, scope, &cache.config); + let opposite_key = match scope { + CacheScope::Public => Self::cache_key( + &logical_key, + CachePartition::Private(Self::private_partition(&cache.config)), + ), + CacheScope::Private => Self::cache_key(&logical_key, CachePartition::Public), + }; + + if !cache.config.serve_stale_on_error { + cache.entries.retain(|_, entry| entry.expires_at > now); + } + cache.entries.remove(&opposite_key); + + if cache.config.max_entries > 0 + && !cache.entries.contains_key(&target_key) + && cache.entries.len() >= cache.config.max_entries + && let Some(oldest_key) = cache + .entries + .iter() + .min_by_key(|(_, entry)| entry.inserted_at) + .map(|(key, _)| key.clone()) + { + cache.entries.remove(&oldest_key); + } + + cache.entries.insert( + target_key, + CachedPeerResponse { + value, + expires_at, + inserted_at: now, + scope, + }, + ); + } + + #[cfg(test)] + pub(crate) async fn cache_response( + &self, + logical_key: String, + value: R::PeerResp, + ttl_ms: Option, + cache_scope: Option, + ) { + let generation = self.capture_response_cache_generation().await; + self.cache_response_with_generation(logical_key, value, ttl_ms, cache_scope, generation) + .await; + } + + pub(crate) async fn invalidate_cached_responses(&self, prefix: &str) { + let mut cache = self.response_cache.write().await; + cache.generation = cache.generation.wrapping_add(1); + cache + .entries + .retain(|key, _| !key.logical_key.starts_with(prefix)); + } +} + +impl Peer { + /// Replaces the response-cache configuration. + /// + /// Changing the private partition invalidates private entries from the old + /// authorization context. Disabling the cache clears every entry. Any + /// configuration change also suppresses writes from requests that were + /// already in flight under the previous configuration. + pub async fn set_response_cache_config(&self, config: ClientCacheConfig) { + let mut cache = self.response_cache.write().await; + let config_changed = cache.config != config; + let partition_changed = cache.config.private_partition != config.private_partition; + let ttl_policy_changed = cache.config.default_ttl != config.default_ttl + || cache.config.max_ttl != config.max_ttl; + cache.config = config; + if config_changed { + cache.generation = cache.generation.wrapping_add(1); + } + if !cache.config.enabled || ttl_policy_changed { + cache.entries.clear(); + } else if partition_changed { + cache + .entries + .retain(|_, entry| entry.scope == CacheScope::Public); + } + cache.trim_to_limit(); + } + + /// Returns a snapshot of the active response-cache configuration. + pub async fn response_cache_config(&self) -> ClientCacheConfig { + self.response_cache.read().await.config.clone() + } + + /// Clears every cached client response without changing the configuration. + pub async fn clear_response_cache(&self) { + let mut cache = self.response_cache.write().await; + cache.generation = cache.generation.wrapping_add(1); + cache.entries.clear(); + } +}