From 94a6ae811b36fd4878ada5cd0611961d87b0e49a Mon Sep 17 00:00:00 2001 From: Dale Seo <5466341+DaleSeo@users.noreply.github.com> Date: Tue, 3 Nov 2026 11:33:38 -0500 Subject: [PATCH] feat(client): merge cache hints across list pages --- crates/rmcp/src/model.rs | 34 ++- crates/rmcp/src/service/client.rs | 372 ++++++++++++++++++------ crates/rmcp/src/service/client/cache.rs | 81 +++++- 3 files changed, 386 insertions(+), 101 deletions(-) diff --git a/crates/rmcp/src/model.rs b/crates/rmcp/src/model.rs index c5d6b72c1..952ad1907 100644 --- a/crates/rmcp/src/model.rs +++ b/crates/rmcp/src/model.rs @@ -1891,9 +1891,20 @@ where } } +/// One page of a paginated list result. +#[cfg(feature = "client")] +pub(crate) trait PaginatedResult { + type Item; + + fn ttl_ms(&self) -> Option; + fn cache_scope(&self) -> Option; + /// Splits the page into its items and the cursor for the next page. + fn into_page(self) -> (Vec, Option); +} + macro_rules! paginated_result { ($t:ident { - $i_item: ident: $t_item: ty + $i_item: ident: Vec<$t_item: ty> }) => { #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] #[serde(rename_all = "camelCase")] @@ -1934,7 +1945,7 @@ macro_rules! paginated_result { skip_serializing_if = "Option::is_none" )] pub cache_scope: Option, - pub $i_item: $t_item, + pub $i_item: Vec<$t_item>, } impl Default for $t { @@ -1943,8 +1954,25 @@ macro_rules! paginated_result { } } + #[cfg(feature = "client")] + impl PaginatedResult for $t { + type Item = $t_item; + + fn ttl_ms(&self) -> Option { + self.ttl_ms + } + + fn cache_scope(&self) -> Option { + self.cache_scope + } + + fn into_page(self) -> (Vec<$t_item>, Option) { + (self.$i_item, self.next_cursor) + } + } + impl $t { - pub fn with_all_items(items: $t_item) -> Self { + pub fn with_all_items(items: Vec<$t_item>) -> Self { Self { result_type: Some(ResultType::COMPLETE), meta: None, diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 21d68be89..dd9bb350d 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -2,7 +2,12 @@ #![expect(deprecated)] pub(super) mod cache; -use std::{borrow::Cow, num::NonZeroUsize, sync::Arc, time::Duration}; +use std::{ + borrow::Cow, + num::NonZeroUsize, + sync::Arc, + time::{Duration, Instant}, +}; use cache::CacheGeneration; pub use cache::{ClientCacheConfig, MAX_CLIENT_CACHE_TTL}; @@ -24,13 +29,14 @@ use crate::{ JsonRpcResponse, ListPromptsRequest, ListPromptsResult, ListResourceTemplatesRequest, ListResourceTemplatesResult, ListResourcesRequest, ListResourcesResult, ListSkillsRequest, ListSkillsRequestMethod, ListSkillsResult, ListToolsRequest, ListToolsResult, - NumberOrString, PaginatedRequestParams, ProgressNotification, ProgressNotificationParam, - ProtocolVersion, ReadResourceDirectoryRequest, ReadResourceDirectoryRequestMethod, - ReadResourceDirectoryRequestParams, ReadResourceDirectoryResult, ReadResourceRequest, - ReadResourceRequestParams, ReadResourceResponse, ReadResourceResult, Reference, RequestId, - RequestMetaObject, RootsListChangedNotification, ServerJsonRpcMessage, ServerNotification, - ServerPeerInfo, ServerRequest, ServerResult, SetLevelRequest, SetLevelRequestParams, - SubscribeRequest, SubscribeRequestParams, SubscriptionFilter, SubscriptionsListenRequest, + NumberOrString, PaginatedRequestParams, PaginatedResult, ProgressNotification, + ProgressNotificationParam, ProtocolVersion, ReadResourceDirectoryRequest, + ReadResourceDirectoryRequestMethod, ReadResourceDirectoryRequestParams, + ReadResourceDirectoryResult, ReadResourceRequest, ReadResourceRequestParams, + ReadResourceResponse, ReadResourceResult, Reference, RequestId, RequestMetaObject, + RootsListChangedNotification, ServerJsonRpcMessage, ServerNotification, ServerPeerInfo, + ServerRequest, ServerResult, SetLevelRequest, SetLevelRequestParams, SubscribeRequest, + SubscribeRequestParams, SubscriptionFilter, SubscriptionsListenRequest, SubscriptionsListenRequestParams, SubscriptionsListenResult, UnsubscribeRequest, UnsubscribeRequestParams, UpdateTaskParams, UpdateTaskRequest, }, @@ -1940,100 +1946,126 @@ impl Peer { method!(peer_not notify_roots_list_changed RootsListChangedNotification); } +/// Every item from a paginated list, with a cache hint covering all of its pages. +#[derive(Debug, Clone, PartialEq)] +#[non_exhaustive] +pub struct MergedList { + pub items: Vec, + /// Milliseconds until the earliest page expires, or `None` if any page omitted `ttlMs`. + pub ttl_ms: Option, + /// `Public` only if every page is public, or `None` if any page omitted `cacheScope`. + pub cache_scope: Option, +} + +/// Follows `nextCursor` from the first page to the last and merges the pages. +async fn list_all_pages( + mut list_page: impl FnMut(Option) -> Fut, +) -> Result, ServiceError> +where + P: PaginatedResult, + Fut: Future>, +{ + let params = |cursor| Some(PaginatedRequestParams { meta: None, cursor }); + let expiry = |ttl_ms: Option| { + ttl_ms.and_then(|ttl_ms| Instant::now().checked_add(Duration::from_millis(ttl_ms))) + }; + + let page = list_page(params(None)).await?; + let mut expires_at = expiry(page.ttl_ms()); + let mut cache_scope = page.cache_scope(); + let (mut items, mut cursor) = page.into_page(); + while cursor.is_some() { + let page = list_page(params(cursor)).await?; + expires_at = expires_at + .zip(expiry(page.ttl_ms())) + .map(|(merged, page)| merged.min(page)); + cache_scope = + cache_scope + .zip(page.cache_scope()) + .map(|(merged, page)| match (merged, page) { + (CacheScope::Public, CacheScope::Public) => CacheScope::Public, + _ => CacheScope::Private, + }); + let (page_items, next_cursor) = page.into_page(); + items.extend(page_items); + cursor = next_cursor; + } + + let ttl_ms = expires_at.map(|expires_at| { + let remaining = expires_at.saturating_duration_since(Instant::now()); + u64::try_from(remaining.as_millis()).unwrap_or(u64::MAX) + }); + Ok(MergedList { + items, + ttl_ms, + cache_scope, + }) +} + impl Peer { - /// A wrapper method for [`Peer::list_tools`]. - /// - /// This function will call [`Peer::list_tools`] multiple times until all tools are listed. + /// Calls [`Peer::list_tools`] until every page is listed. pub async fn list_all_tools(&self) -> Result, ServiceError> { - let mut tools = Vec::new(); - let mut cursor = None; - loop { - let result = self - .list_tools(Some(PaginatedRequestParams { meta: None, cursor })) - .await?; - tools.extend(result.tools); - cursor = result.next_cursor; - if cursor.is_none() { - break; - } - } - Ok(tools) + Ok(self.list_all_tools_with_cache_hint().await?.items) } - /// A wrapper method for [`Peer::list_prompts`]. - /// - /// This function will call [`Peer::list_prompts`] multiple times until all prompts are listed. + /// Like [`Peer::list_all_tools`], with a cache hint for the merged list. + pub async fn list_all_tools_with_cache_hint( + &self, + ) -> Result, ServiceError> { + list_all_pages(|params| self.list_tools(params)).await + } + + /// Calls [`Peer::list_prompts`] until every page is listed. pub async fn list_all_prompts(&self) -> Result, ServiceError> { - let mut prompts = Vec::new(); - let mut cursor = None; - loop { - let result = self - .list_prompts(Some(PaginatedRequestParams { meta: None, cursor })) - .await?; - prompts.extend(result.prompts); - cursor = result.next_cursor; - if cursor.is_none() { - break; - } - } - Ok(prompts) + Ok(self.list_all_prompts_with_cache_hint().await?.items) } - /// A wrapper method for [`Peer::list_resources`]. - /// - /// This function will call [`Peer::list_resources`] multiple times until all resources are listed. + /// Like [`Peer::list_all_prompts`], with a cache hint for the merged list. + pub async fn list_all_prompts_with_cache_hint( + &self, + ) -> Result, ServiceError> { + list_all_pages(|params| self.list_prompts(params)).await + } + + /// Calls [`Peer::list_resources`] until every page is listed. pub async fn list_all_resources(&self) -> Result, ServiceError> { - let mut resources = Vec::new(); - let mut cursor = None; - loop { - let result = self - .list_resources(Some(PaginatedRequestParams { meta: None, cursor })) - .await?; - resources.extend(result.resources); - cursor = result.next_cursor; - if cursor.is_none() { - break; - } - } - Ok(resources) + Ok(self.list_all_resources_with_cache_hint().await?.items) } - /// Call [`Peer::list_skills`] until every page is listed. + /// Like [`Peer::list_all_resources`], with a cache hint for the merged list. + pub async fn list_all_resources_with_cache_hint( + &self, + ) -> Result, ServiceError> { + list_all_pages(|params| self.list_resources(params)).await + } + + /// Calls [`Peer::list_skills`] until every page is listed. pub async fn list_all_skills(&self) -> Result, ServiceError> { - let mut skills = Vec::new(); - let mut cursor = None; - loop { - let result = self - .list_skills(Some(PaginatedRequestParams { meta: None, cursor })) - .await?; - skills.extend(result.skills); - cursor = result.next_cursor; - if cursor.is_none() { - break; - } - } - Ok(skills) + Ok(self.list_all_skills_with_cache_hint().await?.items) } - /// A wrapper method for [`Peer::list_resource_templates`]. - /// - /// This function will call [`Peer::list_resource_templates`] multiple times until all resource templates are listed. + /// Like [`Peer::list_all_skills`], with a cache hint for the merged list. + pub async fn list_all_skills_with_cache_hint( + &self, + ) -> Result, ServiceError> { + list_all_pages(|params| self.list_skills(params)).await + } + + /// Calls [`Peer::list_resource_templates`] until every page is listed. pub async fn list_all_resource_templates( &self, ) -> Result, ServiceError> { - let mut resource_templates = Vec::new(); - let mut cursor = None; - loop { - let result = self - .list_resource_templates(Some(PaginatedRequestParams { meta: None, cursor })) - .await?; - resource_templates.extend(result.resource_templates); - cursor = result.next_cursor; - if cursor.is_none() { - break; - } - } - Ok(resource_templates) + Ok(self + .list_all_resource_templates_with_cache_hint() + .await? + .items) + } + + /// Like [`Peer::list_all_resource_templates`], with a cache hint for the merged list. + pub async fn list_all_resource_templates_with_cache_hint( + &self, + ) -> Result, ServiceError> { + list_all_pages(|params| self.list_resource_templates(params)).await } /// Convenient method to get completion suggestions for a prompt argument @@ -2752,7 +2784,50 @@ mod tests { ) .await; - assert_eq!(peer.list_tools(params).await.unwrap(), expected); + let served = peer.list_tools(params).await.unwrap(); + assert_eq!( + ListToolsResult { + ttl_ms: expected.ttl_ms, + ..served + }, + expected + ); + } + + #[tokio::test] + async fn cached_page_reports_the_time_left_until_expiry() { + let peer = disconnected_peer(); + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, &None); + peer.cache_response( + key, + ServerResult::ListToolsResult(tools_result(Some(5_000), Some(CacheScope::Public))), + Some(5_000), + Some(CacheScope::Public), + ) + .await; + tokio::time::sleep(Duration::from_millis(20)).await; + + let ttl_ms = peer.list_tools(None).await.unwrap().ttl_ms.unwrap(); + assert!(ttl_ms <= 4_980, "ttl_ms = {ttl_ms}"); + } + + #[tokio::test] + async fn cached_page_without_ttl_stays_without_ttl() { + let peer = disconnected_peer(); + peer.set_response_cache_config( + ClientCacheConfig::default().with_default_ttl(Duration::from_secs(5)), + ) + .await; + let key = list_response_cache_key(TOOL_LIST_CACHE_PREFIX, &None); + peer.cache_response( + key, + ServerResult::ListToolsResult(tools_result(None, Some(CacheScope::Public))), + None, + Some(CacheScope::Public), + ) + .await; + + assert_eq!(peer.list_tools(None).await.unwrap().ttl_ms, None); } #[tokio::test] @@ -2847,7 +2922,13 @@ mod tests { .await; tokio::time::sleep(Duration::from_millis(5)).await; - assert_eq!(peer.list_tools(params).await.unwrap(), expected); + assert_eq!( + peer.list_tools(params).await.unwrap(), + ListToolsResult { + ttl_ms: Some(0), + ..expected + } + ); } #[tokio::test] @@ -2867,7 +2948,120 @@ mod tests { ) .await; - assert_eq!(peer.discover(meta).await.unwrap(), expected); + let served = peer.discover(meta).await.unwrap(); + assert_eq!( + DiscoverResult { + ttl_ms: expected.ttl_ms, + ..served + }, + expected + ); + } + + /// Caches a chain of tool pages, linked by cursors `page-1`, `page-2`, ... + async fn cache_tool_pages( + peer: &Peer, + pages: &[(Option, Option)], + ) { + for (index, &(ttl_ms, cache_scope)) in pages.iter().enumerate() { + let cursor = (index > 0).then(|| format!("page-{index}")); + let params = Some(PaginatedRequestParams { meta: None, cursor }); + let mut page = tools_result(ttl_ms, cache_scope); + page.tools = vec![crate::model::Tool::new( + format!("tool-{index}"), + "", + Arc::new(Default::default()), + )]; + page.next_cursor = (index + 1 < pages.len()).then(|| format!("page-{}", index + 1)); + peer.cache_response( + list_response_cache_key(TOOL_LIST_CACHE_PREFIX, ¶ms), + ServerResult::ListToolsResult(page), + ttl_ms.or(Some(5_000)), + cache_scope, + ) + .await; + } + } + + #[tokio::test] + async fn merged_list_collects_items_from_every_page() { + let peer = disconnected_peer(); + let hint = (Some(5_000), Some(CacheScope::Public)); + cache_tool_pages(&peer, &[hint, hint, hint]).await; + + let names: Vec<_> = peer + .list_all_tools_with_cache_hint() + .await + .unwrap() + .items + .into_iter() + .map(|tool| tool.name) + .collect(); + assert_eq!(names, ["tool-0", "tool-1", "tool-2"]); + } + + #[tokio::test] + async fn merged_ttl_is_the_earliest_page_expiry() { + let peer = disconnected_peer(); + let public = Some(CacheScope::Public); + cache_tool_pages(&peer, &[(Some(5_000), public), (Some(1_000), public)]).await; + + let ttl_ms = peer + .list_all_tools_with_cache_hint() + .await + .unwrap() + .ttl_ms + .unwrap(); + assert!((900..=1_000).contains(&ttl_ms), "ttl_ms = {ttl_ms}"); + } + + #[tokio::test] + async fn merged_ttl_is_absent_when_any_page_omits_it() { + let peer = disconnected_peer(); + let public = Some(CacheScope::Public); + cache_tool_pages(&peer, &[(Some(5_000), public), (None, public)]).await; + + let merged = peer.list_all_tools_with_cache_hint().await.unwrap(); + assert_eq!(merged.ttl_ms, None); + } + + #[tokio::test] + async fn merged_scope_is_public_when_every_page_is_public() { + let peer = disconnected_peer(); + let hint = (Some(5_000), Some(CacheScope::Public)); + cache_tool_pages(&peer, &[hint, hint]).await; + + let merged = peer.list_all_tools_with_cache_hint().await.unwrap(); + assert_eq!(merged.cache_scope, Some(CacheScope::Public)); + } + + #[tokio::test] + async fn merged_scope_is_private_when_any_page_is_private() { + let peer = disconnected_peer(); + cache_tool_pages( + &peer, + &[ + (Some(5_000), Some(CacheScope::Public)), + (Some(5_000), Some(CacheScope::Private)), + ], + ) + .await; + + let merged = peer.list_all_tools_with_cache_hint().await.unwrap(); + assert_eq!(merged.cache_scope, Some(CacheScope::Private)); + } + + #[tokio::test] + async fn merged_scope_is_absent_when_any_page_omits_it() { + let peer = disconnected_peer(); + cache_tool_pages( + &peer, + &[(Some(5_000), Some(CacheScope::Public)), (Some(5_000), None)], + ) + .await; + + let merged = peer.list_all_tools_with_cache_hint().await.unwrap(); + assert_eq!(merged.cache_scope, None); } } diff --git a/crates/rmcp/src/service/client/cache.rs b/crates/rmcp/src/service/client/cache.rs index f7137ccd6..4e85d1304 100644 --- a/crates/rmcp/src/service/client/cache.rs +++ b/crates/rmcp/src/service/client/cache.rs @@ -6,7 +6,7 @@ use std::{ use super::RoleClient; use crate::{ - model::CacheScope, + model::{CacheScope, ServerResult}, service::{Peer, ServiceRole}, }; @@ -104,6 +104,48 @@ impl ClientCacheConfig { } } +/// A cached response whose `ttlMs` can be rewritten when it is served. +pub(crate) trait RemainingTtl { + /// Replaces a server-provided `ttlMs` with `remaining`; leaves an absent one absent. + fn set_remaining_ttl(&mut self, remaining: Duration); +} + +impl RemainingTtl for ServerResult { + fn set_remaining_ttl(&mut self, remaining: Duration) { + let remaining_ms = u64::try_from(remaining.as_millis()).unwrap_or(u64::MAX); + match self { + ServerResult::DiscoverResult(crate::model::DiscoverResult { ttl_ms, .. }) + | ServerResult::ListToolsResult(crate::model::ListToolsResult { + ttl_ms: Some(ttl_ms), + .. + }) + | ServerResult::ListPromptsResult(crate::model::ListPromptsResult { + ttl_ms: Some(ttl_ms), + .. + }) + | ServerResult::ListResourcesResult(crate::model::ListResourcesResult { + ttl_ms: Some(ttl_ms), + .. + }) + | ServerResult::ListResourceTemplatesResult( + crate::model::ListResourceTemplatesResult { + ttl_ms: Some(ttl_ms), + .. + }, + ) + | ServerResult::ListSkillsResult(crate::model::ListSkillsResult { + ttl_ms: Some(ttl_ms), + .. + }) + | ServerResult::ReadResourceResult(crate::model::ReadResourceResult { + ttl_ms: Some(ttl_ms), + .. + }) => *ttl_ms = remaining_ms, + _ => {} + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Hash)] enum CachePartition { Public, @@ -124,6 +166,20 @@ struct CachedPeerResponse { scope: CacheScope, } +impl CachedPeerResponse { + fn fresh_value(&self, now: Instant) -> T { + let mut value = self.value.clone(); + value.set_remaining_ttl(self.expires_at.saturating_duration_since(now)); + value + } + + fn stale_value(&self) -> T { + let mut value = self.value.clone(); + value.set_remaining_ttl(Duration::ZERO); + value + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct CacheGeneration(u64); @@ -197,12 +253,15 @@ impl Peer { } /// Returns a fresh cached response, preferring the current private partition - /// before the public partition. + /// before the public partition. Its `ttlMs` is the time left until expiry. /// /// 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 { + pub(crate) async fn cached_response(&self, logical_key: &str) -> Option + where + R::PeerResp: RemainingTtl, + { let now = Instant::now(); let mut cache = self.response_cache.write().await; if !cache.config.enabled { @@ -216,7 +275,7 @@ impl Peer { ); let private_fresh = cache.entries.get(&private_key).and_then(|entry| { (entry.expires_at > now && entry.scope == CacheScope::Private) - .then(|| entry.value.clone()) + .then(|| entry.fresh_value(now)) }); if let Some(value) = private_fresh { return Some(value); @@ -228,7 +287,7 @@ impl Peer { 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()) + .then(|| entry.fresh_value(now)) }); if let Some(value) = public_fresh { return Some(value); @@ -240,13 +299,17 @@ impl Peer { } /// 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). + /// re-fetch fails (SEP-2549 permits serving stale responses on error). Its + /// `ttlMs` is zero. /// /// 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 { + pub(crate) async fn stale_cached_response(&self, logical_key: &str) -> Option + where + R::PeerResp: RemainingTtl, + { let cache = self.response_cache.read().await; if !cache.config.enabled || !cache.config.serve_stale_on_error { return None; @@ -259,14 +322,14 @@ impl Peer { if let Some(entry) = cache.entries.get(&private_key) && entry.scope == CacheScope::Private { - return Some(entry.value.clone()); + return Some(entry.stale_value()); } 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()); + return Some(entry.stale_value()); } None }