feat: add client-side TTL-honoring response cache (SEP-2549) (#1025)
* feat: add client-side TTL-honoring response cache (SEP-2549) Add a configurable client response cache in the service layer that honors SEP-2549 caching hints (ttlMs / cacheScope) for tools/list, prompts/list, resources/list, resources/templates/list, and resources/read. Closes #974 Co-Authored-by: John-Francis Nnadi <nnadifrancis23@gmail.com> * fix: address PR feedback --------- Co-authored-by: John-Francis Nnadi <nnadifrancis23@gmail.com>
This commit is contained in:
Alex Hancock
GitHub
parent
07abfdfd99
commit
151970775c
5 changed files with 1022 additions and 29 deletions
47
README.md
47
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).
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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<Self>,
|
||||
_notification: &Self::PeerNot,
|
||||
) -> impl Future<Output = ()> + MaybeSendFuture {
|
||||
async {}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn uses_legacy_lifecycle(
|
||||
|
|
@ -571,6 +584,8 @@ pub struct Peer<R: ServiceRole> {
|
|||
client_request_metadata: Arc<OnceLock<ClientRequestMetadata>>,
|
||||
request_metadata_required: Arc<std::sync::atomic::AtomicBool>,
|
||||
subscription_channels: Arc<std::sync::RwLock<SubscriptionChannelMap<R::PeerNot>>>,
|
||||
#[cfg(feature = "client")]
|
||||
response_cache: client::cache::PeerResponseCache<R>,
|
||||
}
|
||||
|
||||
impl<R: Clone + ServiceRole> Clone for Peer<R>
|
||||
|
|
@ -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<R: ServiceRole> Peer<R> {
|
|||
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();
|
||||
|
|
|
|||
|
|
@ -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<Self>, 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<RoleClient>;
|
||||
|
|
@ -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<PaginatedRequestParams>) -> 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<String> {
|
||||
// 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<PaginatedRequestParams>) -> 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<RoleClient> {
|
|||
/// 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<DiscoverResult, ServiceError> {
|
||||
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<String>,
|
||||
ttl_ms: Option<u64>,
|
||||
cache_scope: Option<CacheScope>,
|
||||
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<RoleClient> {
|
|||
&self,
|
||||
params: ReadResourceRequestParams,
|
||||
) -> Result<ReadResourceResponse, ServiceError> {
|
||||
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<RoleClient> {
|
|||
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<RoleClient> {
|
|||
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<PaginatedRequestParams>,
|
||||
) -> Result<ListPromptsResult, ServiceError> {
|
||||
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<PaginatedRequestParams>,
|
||||
) -> Result<ListResourcesResult, ServiceError> {
|
||||
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<PaginatedRequestParams>,
|
||||
) -> Result<ListResourceTemplatesResult, ServiceError> {
|
||||
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<ReadResourceResult, ServiceError> {
|
||||
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<PaginatedRequestParams>,
|
||||
) -> Result<ListToolsResult, ServiceError> {
|
||||
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<RoleClient> {
|
||||
let (peer, receiver) =
|
||||
Peer::<RoleClient>::new(Arc::new(AtomicU32RequestIdProvider::default()), None);
|
||||
drop(receiver);
|
||||
peer
|
||||
}
|
||||
|
||||
fn tools_result(ttl_ms: Option<u64>, cache_scope: Option<CacheScope>) -> 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::<PaginatedRequestParams>;
|
||||
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::<PaginatedRequestParams>;
|
||||
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::<PaginatedRequestParams>;
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
401
crates/rmcp/src/service/client/cache.rs
Normal file
401
crates/rmcp/src/service/client/cache.rs
Normal file
|
|
@ -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<String>,
|
||||
/// 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<String>) -> 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<str>),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
struct CacheKey {
|
||||
logical_key: String,
|
||||
partition: CachePartition,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct CachedPeerResponse<T> {
|
||||
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<R: ServiceRole> {
|
||||
entries: HashMap<CacheKey, CachedPeerResponse<R::PeerResp>>,
|
||||
config: ClientCacheConfig,
|
||||
generation: u64,
|
||||
}
|
||||
|
||||
impl<R: ServiceRole> Default for PeerResponseCacheState<R> {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
entries: HashMap::new(),
|
||||
config: ClientCacheConfig::default(),
|
||||
generation: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: ServiceRole> PeerResponseCacheState<R> {
|
||||
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<R> = Arc<tokio::sync::RwLock<PeerResponseCacheState<R>>>;
|
||||
|
||||
impl<R: ServiceRole> Peer<R> {
|
||||
fn private_partition(config: &ClientCacheConfig) -> Arc<str> {
|
||||
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<R::PeerResp> {
|
||||
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<R::PeerResp> {
|
||||
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<u64>,
|
||||
cache_scope: Option<CacheScope>,
|
||||
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<u64>,
|
||||
cache_scope: Option<CacheScope>,
|
||||
) {
|
||||
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<RoleClient> {
|
||||
/// 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();
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue