|
9 | 9 |
|
10 | 10 | from __future__ import annotations |
11 | 11 |
|
| 12 | +import json |
| 13 | +import logging |
12 | 14 | import time |
13 | 15 | from collections.abc import Callable |
14 | 16 | from dataclasses import dataclass |
15 | 17 | from typing import Any, Final, Literal, Protocol |
16 | 18 |
|
| 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 | + |
17 | 30 | __all__ = [ |
18 | 31 | "MAX_TTL_MS", |
19 | 32 | "CacheConfig", |
|
24 | 37 | "ResponseCacheStore", |
25 | 38 | ] |
26 | 39 |
|
| 40 | +logger = logging.getLogger(__name__) |
| 41 | + |
27 | 42 | CacheMode = Literal["use", "refresh", "bypass"] |
28 | 43 | """Per-call cache behavior: `"use"` serves fresh entries and stores fetches, |
29 | 44 | `"refresh"` skips the read but stores the fetch, `"bypass"` touches the cache |
@@ -81,6 +96,13 @@ class ResponseCacheStore(Protocol): |
81 | 96 | atomicity are the implementation's responsibility. Operations may raise; |
82 | 97 | the SDK degrades per its error discipline (a failing store never fails a |
83 | 98 | 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. |
84 | 106 | """ |
85 | 107 |
|
86 | 108 | async def get(self, key: CacheKey) -> CacheEntry | None: ... |
@@ -186,3 +208,237 @@ async def delete(self, key: CacheKey) -> None: |
186 | 208 |
|
187 | 209 | async def clear(self) -> None: |
188 | 210 | 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