Skip to content

Commit 5f8a392

Browse files
committed
Add client response cache coordinator: scope arms, era gating, TTL resolution, eviction
1 parent d2d4a25 commit 5f8a392

2 files changed

Lines changed: 980 additions & 3 deletions

File tree

‎src/mcp/client/caching.py‎

Lines changed: 256 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,24 @@
99

1010
from __future__ import annotations
1111

12+
import json
13+
import logging
1214
import time
1315
from collections.abc import Callable
1416
from dataclasses import dataclass
1517
from typing import Any, Final, Literal, Protocol
1618

19+
import anyio
20+
from mcp_types import (
21+
CacheableResult,
22+
PromptListChangedNotification,
23+
ResourceListChangedNotification,
24+
ResourceUpdatedNotification,
25+
ServerNotification,
26+
ToolListChangedNotification,
27+
)
28+
from mcp_types.version import MODERN_PROTOCOL_VERSIONS
29+
1730
__all__ = [
1831
"MAX_TTL_MS",
1932
"CacheConfig",
@@ -24,6 +37,8 @@
2437
"ResponseCacheStore",
2538
]
2639

40+
logger = logging.getLogger(__name__)
41+
2742
CacheMode = Literal["use", "refresh", "bypass"]
2843
"""Per-call cache behavior: `"use"` serves fresh entries and stores fetches,
2944
`"refresh"` skips the read but stores the fetch, `"bypass"` touches the cache
@@ -81,6 +96,13 @@ class ResponseCacheStore(Protocol):
8196
atomicity are the implementation's responsibility. Operations may raise;
8297
the SDK degrades per its error discipline (a failing store never fails a
8398
successful fetch).
99+
100+
A store that serializes entries (any cross-process store must) is
101+
responsible for round-tripping them: `get` returns the entry as stored,
102+
with `value` still the result model object `set` received - the SDK has
103+
no rehydration hook to rebuild it from serialized data. An entry that
104+
comes back in the wrong shape (e.g. with a plain-dict value) degrades to
105+
a cache miss, never an error.
84106
"""
85107

86108
async def get(self, key: CacheKey) -> CacheEntry | None: ...
@@ -186,3 +208,237 @@ async def delete(self, key: CacheKey) -> None:
186208

187209
async def clear(self) -> None:
188210
self._entries.clear()
211+
212+
213+
_GENERATION_MAP_CAP: Final[int] = 4096
214+
"""Cap on the coordinator's eviction-race bookkeeping (the generation map).
215+
At the cap, registering a new key drops the oldest one, degrading the dropped
216+
key's race guard to the accepted co-tenant class."""
217+
218+
219+
class ClientResponseCache:
220+
"""Coordinator between the `Client` verbs and a `ResponseCacheStore`.
221+
222+
Owns key construction (the scope arms), the era gate, TTL/scope
223+
resolution, eviction, and the store error discipline. `Client` mints one
224+
per instance; the caching verbs and the notification wrap are the only
225+
callers.
226+
"""
227+
228+
def __init__(
229+
self,
230+
*,
231+
store: ResponseCacheStore,
232+
partition: str,
233+
arm_id: str,
234+
default_ttl_ms: int,
235+
clock: Callable[[], float],
236+
share_public: bool,
237+
negotiated_version: Callable[[], str | None],
238+
generation_map_cap: int = _GENERATION_MAP_CAP,
239+
) -> None:
240+
self._store = store
241+
self._default_ttl_ms = default_ttl_ms
242+
self._clock = clock
243+
self._negotiated_version = negotiated_version
244+
# Arms are JSON arrays so crafted arm_id/partition values cannot
245+
# collide across field boundaries. Private entries always carry the
246+
# partition; public entries do too unless the operator opted into
247+
# fleet-wide sharing of server-asserted-public results.
248+
self._private_arm = json.dumps(["private", arm_id, partition])
249+
self._public_arm = json.dumps(["public", arm_id] if share_public else ["public", arm_id, partition])
250+
# The generation map is the sole membership structure: a key is
251+
# race-guarded iff registered here.
252+
self._generations: dict[tuple[str, str], int] = {}
253+
self._generation_map_cap = generation_map_cap
254+
# Operation kinds ("get"/"set"/"delete") that warned and have not
255+
# succeeded since; membership suppresses repeat warnings for the kind.
256+
self._warned_store_ops: set[str] = set()
257+
258+
async def read(self, method: str, params_key: str) -> CacheableResult | None:
259+
"""Serve a fresh entry for the key, or `None`.
260+
261+
Called only under `cache_mode="use"`; returns a deep copy so a served
262+
result never aliases the stored one.
263+
"""
264+
# One boundary around the whole read path: a raising store `get` and
265+
# an entry rehydrated into the wrong shape (which raises only at the
266+
# freshness check or the copy) are the same "get" failure class -
267+
# warned once per burst, re-armed only by a fully successful read.
268+
try:
269+
entry = await self._get_fresh(CacheKey(method, params_key, self._private_arm))
270+
if entry is None:
271+
# Stale counts as a miss for fall-through too: after a server
272+
# scope flip (private -> public), a stale private leftover
273+
# must not shadow a fresh public entry.
274+
entry = await self._get_fresh(CacheKey(method, params_key, self._public_arm))
275+
if entry is not None and entry.scope != "public":
276+
# The arm routes, the scope verifies: never serve an entry the
277+
# server scoped "private" out of the shared arm, however it
278+
# got there.
279+
entry = None
280+
copied: CacheableResult | None = None if entry is None else entry.value.model_copy(deep=True)
281+
except Exception: # boundary around user store code: any read-path failure is a miss, never a failed call
282+
self._warn_store_failure("get")
283+
return None
284+
self._warned_store_ops.discard("get")
285+
return copied
286+
287+
async def _get_fresh(self, key: CacheKey) -> CacheEntry | None:
288+
entry = await self._store.get(key)
289+
if entry is None or entry.expires_at is None or entry.expires_at <= self._clock():
290+
return None
291+
return entry
292+
293+
def capture(self, method: str, params_key: str) -> int:
294+
"""Register the key for eviction-race detection, before the fetch is
295+
sent; the matching `write` passes the returned generation back."""
296+
gen_key = (method, params_key)
297+
if gen_key not in self._generations:
298+
if len(self._generations) >= self._generation_map_cap:
299+
# FIFO overflow: drop the oldest key, degrading its race guard
300+
# to the accepted co-tenant class (an eviction racing that
301+
# key's in-flight fetch is no longer detected at write time).
302+
del self._generations[next(iter(self._generations))]
303+
self._generations[gen_key] = 0
304+
return self._generations[gen_key]
305+
306+
async def write(
307+
self,
308+
method: str,
309+
params_key: str,
310+
result: CacheableResult,
311+
gen_at_capture: int,
312+
mode: Literal["use", "refresh"],
313+
) -> None:
314+
"""Store a fetched result under the arm its resolved scope selects."""
315+
gen_key = (method, params_key)
316+
if self._generation_moved(gen_key, gen_at_capture):
317+
return # the key was evicted while the fetch was in flight
318+
ttl_ms, scope = self._resolve(result)
319+
private_key = CacheKey(method, params_key, self._private_arm)
320+
public_key = CacheKey(method, params_key, self._public_arm)
321+
if ttl_ms <= 0:
322+
if mode == "refresh":
323+
# The refetch superseded whatever was cached; purge the warm
324+
# entry so it cannot be served again. Shielded: a cancellation
325+
# delivered between the two deletes would leave the opposite
326+
# arm warm for its full TTL.
327+
with anyio.CancelScope(shield=True):
328+
await self._delete(private_key)
329+
await self._delete(public_key)
330+
return
331+
own, opposite = (public_key, private_key) if scope == "public" else (private_key, public_key)
332+
# Opposite arm first: a failed (or cancelled) delete aborts before the
333+
# set, leaving a miss - never two arms answering for one key.
334+
if not await self._delete(opposite):
335+
return
336+
entry = CacheEntry(value=result.model_copy(deep=True), scope=scope, expires_at=self._clock() + ttl_ms / 1000)
337+
try:
338+
await self._set(own, entry)
339+
finally:
340+
# An eviction can land while an async store's set is committing,
341+
# and the set can commit even when its await is cancelled (the
342+
# request may already be on the wire) - so the re-check runs on
343+
# every exit, and the compensating delete is shielded so the
344+
# pending cancellation cannot abort it and resurrect the evicted
345+
# entry for its full TTL. (A delete after a set that raised is an
346+
# idempotent no-op.)
347+
if self._generation_moved(gen_key, gen_at_capture):
348+
with anyio.CancelScope(shield=True):
349+
await self._delete(own)
350+
351+
async def evict_method(self, method: str) -> None:
352+
"""Evict the method's cursor-less entry (notification- or
353+
cursor-expiry-driven)."""
354+
await self.evict_key(method, "")
355+
356+
async def evict_key(self, method: str, params_key: str) -> None:
357+
"""Evict one key from both arms."""
358+
gen_key = (method, params_key)
359+
# Bump before deleting so an in-flight fetch that captured earlier
360+
# cannot write the just-evicted entry back. Only registered keys bump
361+
# (arbitrary notification uris must not grow the map); the store
362+
# deletes always run - a persistent store may hold warm entries this
363+
# coordinator never captured.
364+
if gen_key in self._generations:
365+
self._generations[gen_key] += 1
366+
# Shielded: eviction runs in spawned notification tasks that die with
367+
# the session - a cancellation between the two deletes would leave one
368+
# arm serving the evicted entry until its TTL.
369+
with anyio.CancelScope(shield=True):
370+
await self._delete(CacheKey(method, params_key, self._private_arm))
371+
await self._delete(CacheKey(method, params_key, self._public_arm))
372+
373+
async def evict_for_notification(self, notification: ServerNotification) -> None:
374+
"""Map a server notification to the entries it makes stale.
375+
376+
Wire-path notifications are dispatched from spawned tasks, so eviction
377+
is eventual relative to in-flight responses: the generation bump
378+
closes the write-back race, while a read racing the notification may
379+
briefly serve the pre-eviction entry (accepted, latency-bounded).
380+
"""
381+
match notification:
382+
case ToolListChangedNotification():
383+
await self.evict_method("tools/list")
384+
case PromptListChangedNotification():
385+
await self.evict_method("prompts/list")
386+
case ResourceListChangedNotification():
387+
# Templates enumerate the same changed resource space.
388+
await self.evict_method("resources/list")
389+
await self.evict_method("resources/templates/list")
390+
case ResourceUpdatedNotification():
391+
await self.evict_key("resources/read", notification.params.uri)
392+
case _:
393+
pass
394+
395+
def _resolve(self, result: CacheableResult) -> tuple[int, Literal["public", "private"]]:
396+
# Hints count only on modern sessions: a legacy peer can also put
397+
# `ttlMs`/`cacheScope` keys on the wire (the 2025 surfaces validate
398+
# and discard unknown keys, so wire presence still reaches
399+
# `model_fields_set`) - wire presence is not a peer-era signal.
400+
modern = self._negotiated_version() in MODERN_PROTOCOL_VERSIONS
401+
if modern and "ttl_ms" in result.model_fields_set:
402+
# An explicit `ttlMs: 0` stays 0 (never overridden by the
403+
# default), and negatives are unconstructible here - the model
404+
# enforces ge=0 and the parse seam floors negative wire values -
405+
# so only the cap applies.
406+
ttl_ms = result.ttl_ms
407+
else:
408+
ttl_ms = self._default_ttl_ms
409+
scope: Literal["public", "private"] = "public" if modern and result.cache_scope == "public" else "private"
410+
return min(ttl_ms, MAX_TTL_MS), scope
411+
412+
def _generation_moved(self, gen_key: tuple[str, str], gen_at_capture: int) -> bool:
413+
# A key FIFO-dropped from the map can no longer be checked; the guard
414+
# fails open (the accepted co-tenant race class) rather than
415+
# discarding the fetch.
416+
return self._generations.get(gen_key, gen_at_capture) != gen_at_capture
417+
418+
async def _set(self, key: CacheKey, entry: CacheEntry) -> bool:
419+
try:
420+
await self._store.set(key, entry)
421+
except Exception: # boundary around user store code: nothing cached, the fetch already succeeded
422+
self._warn_store_failure("set")
423+
return False
424+
self._warned_store_ops.discard("set")
425+
return True
426+
427+
async def _delete(self, key: CacheKey) -> bool:
428+
try:
429+
await self._store.delete(key)
430+
except Exception: # boundary around user store code: callers decide whether a failed delete aborts
431+
self._warn_store_failure("delete")
432+
return False
433+
self._warned_store_ops.discard("delete")
434+
return True
435+
436+
def _warn_store_failure(self, kind: Literal["get", "set", "delete"]) -> None:
437+
# One warning per failure burst, tracked per operation kind: armed by
438+
# the kind's first failure, re-armed only when that same kind succeeds.
439+
# A dead store warns once, not once per request - and a store where
440+
# only `set` is broken warns once too, instead of its healthy deletes
441+
# re-arming the warning every write cycle.
442+
if kind not in self._warned_store_ops:
443+
self._warned_store_ops.add(kind)
444+
logger.warning("Response cache store operation failed; continuing without the cache", exc_info=True)

0 commit comments

Comments
 (0)