388 lines
17 KiB
Python
388 lines
17 KiB
Python
"""Client-side response caching primitives (SEP-2549, protocol revision 2026-07-28)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import time
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from typing import Any, Final, Literal, Protocol
|
|
|
|
import anyio
|
|
import anyio.lowlevel
|
|
from mcp_types import (
|
|
CacheableResult,
|
|
PromptListChangedNotification,
|
|
ResourceListChangedNotification,
|
|
ResourceUpdatedNotification,
|
|
ServerNotification,
|
|
ToolListChangedNotification,
|
|
)
|
|
from mcp_types.version import MODERN_PROTOCOL_VERSIONS
|
|
|
|
__all__ = [
|
|
"MAX_TTL_MS",
|
|
"CacheConfig",
|
|
"CacheEntry",
|
|
"CacheKey",
|
|
"CacheMode",
|
|
"InMemoryResponseCacheStore",
|
|
"ResponseCacheStore",
|
|
]
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
CacheMode = Literal["use", "refresh", "bypass"]
|
|
"""Per-call cache behavior: `"use"` serves and stores, `"refresh"` stores
|
|
without serving, `"bypass"` skips the cache entirely."""
|
|
|
|
MAX_TTL_MS: Final[int] = 24 * 60 * 60 * 1000
|
|
"""Cap on any entry's time-to-live (24 hours, in milliseconds); larger `ttlMs` values are clamped down."""
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class CacheKey:
|
|
"""Identity of one cached response; compare as the field tuple, never a flattened string (collision hazard)."""
|
|
|
|
method: str
|
|
|
|
params_key: str = ""
|
|
"""Result-affecting params discriminator: the uri for `resources/read`, `""` for the list methods."""
|
|
|
|
partition: str = ""
|
|
"""Coordinator-computed arm identifier; opaque to stores."""
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class CacheEntry:
|
|
"""One cached response with its freshness and sharing metadata."""
|
|
|
|
value: Any
|
|
"""The cached result; the SDK deep-copies on write and on serve, so a store may hold it as-is."""
|
|
|
|
scope: Literal["public", "private"]
|
|
"""Server-asserted `cacheScope`: only `"public"` entries may be shared across authorization contexts."""
|
|
|
|
expires_at: float | None
|
|
"""Epoch seconds after which the entry is stale; `None` is never fresh."""
|
|
|
|
|
|
class ResponseCacheStore(Protocol):
|
|
"""Storage contract for the client response cache.
|
|
|
|
Each `Client` calls its store from a single event loop; per-operation
|
|
atomicity is the implementation's responsibility. Operations may raise -
|
|
the SDK degrades to a miss rather than failing the call. A serializing
|
|
store must round-trip `value` back to the result model object (a
|
|
wrong-shape entry is a miss, never an error). A lookup may issue two
|
|
sequential `get` calls (private arm, then public).
|
|
"""
|
|
|
|
async def get(self, key: CacheKey) -> CacheEntry | None: ...
|
|
|
|
async def set(self, key: CacheKey, entry: CacheEntry) -> None: ...
|
|
|
|
async def delete(self, key: CacheKey) -> None: ...
|
|
|
|
async def clear(self) -> None: ...
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class CacheConfig:
|
|
"""Configuration for a `Client`'s response cache.
|
|
|
|
Raises:
|
|
ValueError: On a custom `store` without `partition`, an empty `target_id`, or a negative `default_ttl_ms`.
|
|
"""
|
|
|
|
store: ResponseCacheStore | None = None
|
|
"""Backing store; `None` means a per-client `InMemoryResponseCacheStore`.
|
|
A custom store requires an explicit `partition`."""
|
|
|
|
partition: str = ""
|
|
"""Authorization-context identifier isolating `"private"`-scoped entries
|
|
within a shared store. Derive it from a verified credential - never from
|
|
request-supplied data or the server URL. Fixed for the `Client`'s
|
|
lifetime: construct a new `Client` when the principal changes."""
|
|
|
|
target_id: str | None = None
|
|
"""Server-identity override for custom transports and proxies where the
|
|
SDK cannot derive one from a URL; must be non-empty when provided."""
|
|
|
|
default_ttl_ms: int = 0
|
|
"""TTL in milliseconds for results carrying no `ttlMs` hint; the default `0` leaves them uncached."""
|
|
|
|
clock: Callable[[], float] = time.time
|
|
"""Wall-clock source returning epoch seconds; injectable for expiry tests."""
|
|
|
|
share_public: bool = False
|
|
"""Serve server-marked `"public"` entries across every partition in the store.
|
|
|
|
WARNING: this trusts the server's `"public"` classification for every
|
|
principal sharing the store - a mislabeled response leaks across tenants.
|
|
Constructor-level only: the per-call `cache_mode` can never widen sharing."""
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.store is not None and not self.partition:
|
|
raise ValueError("a custom store requires an explicit partition")
|
|
if self.target_id == "":
|
|
raise ValueError("target_id must be a non-empty string or omitted")
|
|
if self.default_ttl_ms < 0:
|
|
raise ValueError(f"default_ttl_ms must be >= 0, got {self.default_ttl_ms}")
|
|
|
|
|
|
class InMemoryResponseCacheStore:
|
|
"""Default in-process `ResponseCacheStore`.
|
|
|
|
Method bodies are synchronous, so concurrent tasks never observe a torn
|
|
write. `max_entries` caps the whole store, evicting least-recently-used
|
|
at the cap (`0` disables it); `get` and `set` both refresh recency, so a
|
|
hot entry survives churn from other keys.
|
|
|
|
Raises:
|
|
ValueError: If `max_entries` is negative.
|
|
"""
|
|
|
|
def __init__(self, *, max_entries: int = 1024) -> None:
|
|
if max_entries < 0:
|
|
raise ValueError(f"max_entries must be >= 0, got {max_entries}")
|
|
self._max_entries = max_entries
|
|
self._entries: dict[CacheKey, CacheEntry] = {}
|
|
|
|
async def get(self, key: CacheKey) -> CacheEntry | None:
|
|
entry = self._entries.get(key)
|
|
if entry is not None:
|
|
# Pop-and-reinsert moves the key to the back: the dict's insertion order is the LRU ledger.
|
|
self._entries[key] = self._entries.pop(key)
|
|
return entry
|
|
|
|
async def set(self, key: CacheKey, entry: CacheEntry) -> None:
|
|
self._entries.pop(key, None)
|
|
self._entries[key] = entry
|
|
if self._max_entries and len(self._entries) > self._max_entries:
|
|
del self._entries[next(iter(self._entries))]
|
|
|
|
async def delete(self, key: CacheKey) -> None:
|
|
self._entries.pop(key, None)
|
|
|
|
async def clear(self) -> None:
|
|
self._entries.clear()
|
|
|
|
|
|
_GENERATION_MAP_CAP: Final[int] = 4096
|
|
"""Cap on the generation map; at the cap the oldest key's eviction-race guard is dropped (FIFO)."""
|
|
|
|
_STORE_CLEANUP_TIMEOUT: Final[float] = 5
|
|
"""Bound for must-complete store cleanup deletes (mirrors the dispatcher's final-write bound);
|
|
a wedged store delete must not hold client teardown uncancellably."""
|
|
|
|
|
|
class ClientResponseCache:
|
|
"""Coordinates the `Client` caching verbs with a `ResponseCacheStore`: keys, era gate, TTL/scope, eviction."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
store: ResponseCacheStore,
|
|
partition: str,
|
|
arm_id: str,
|
|
default_ttl_ms: int,
|
|
clock: Callable[[], float],
|
|
share_public: bool,
|
|
negotiated_version: Callable[[], str | None],
|
|
generation_map_cap: int = _GENERATION_MAP_CAP,
|
|
store_cleanup_timeout: float = _STORE_CLEANUP_TIMEOUT,
|
|
) -> None:
|
|
self._store = store
|
|
self._partition = partition
|
|
self._arm_id = arm_id
|
|
self._share_public = share_public
|
|
self._default_ttl_ms = default_ttl_ms
|
|
self._clock = clock
|
|
self._negotiated_version = negotiated_version
|
|
# A key is eviction-race-guarded iff registered here.
|
|
self._generations: dict[tuple[str, str], int] = {}
|
|
self._generation_map_cap = generation_map_cap
|
|
self._store_cleanup_timeout = store_cleanup_timeout
|
|
self._warned_store_ops: set[str] = set()
|
|
|
|
def _arm(self, scope: Literal["public", "private"]) -> str:
|
|
# JSON arrays so crafted arm_id/partition values cannot collide across field boundaries.
|
|
# The negotiated version era-scopes every arm: a session never serves an entry written
|
|
# under a different protocol era (its content differs - sieve-stripped fields, header
|
|
# filtering). Every caller runs post-connect; were that ever untrue, the supplier's
|
|
# None still partitions harmlessly.
|
|
fields: list[str | None] = [scope, self._negotiated_version(), self._arm_id]
|
|
if scope == "private" or not self._share_public:
|
|
fields.append(self._partition)
|
|
return json.dumps(fields)
|
|
|
|
async def read(self, method: str, params_key: str) -> CacheableResult | None:
|
|
"""Serve a fresh entry for the key, or `None`; the served result is a deep copy."""
|
|
# A hit completes without any other yielding await, so checkpoint here: a poll
|
|
# loop over a fresh entry must not starve spawned tasks (eviction dispatch).
|
|
await anyio.lowlevel.checkpoint()
|
|
# A wrong-shape entry raises as late as the copy, so the boundary wraps the whole read path.
|
|
try:
|
|
entry = await self._get_fresh(CacheKey(method, params_key, self._arm("private")))
|
|
if entry is None:
|
|
# After a scope flip, a stale private entry must not shadow a fresh public one.
|
|
entry = await self._get_fresh(CacheKey(method, params_key, self._arm("public")))
|
|
if entry is not None and entry.scope != "public":
|
|
# Never serve an entry the server scoped "private" out of the shared arm.
|
|
entry = None
|
|
copied: CacheableResult | None = None if entry is None else entry.value.model_copy(deep=True)
|
|
except Exception: # boundary around user store code: any read-path failure is a miss, never a failed call
|
|
self._warn_store_failure("get")
|
|
return None
|
|
self._warned_store_ops.discard("get")
|
|
return copied
|
|
|
|
async def _get_fresh(self, key: CacheKey) -> CacheEntry | None:
|
|
entry = await self._store.get(key)
|
|
if entry is None or entry.expires_at is None or entry.expires_at <= self._clock():
|
|
return None
|
|
return entry
|
|
|
|
def capture(self, method: str, params_key: str) -> int:
|
|
"""Register the key for eviction-race detection before the fetch; `write` takes the returned generation."""
|
|
gen_key = (method, params_key)
|
|
if gen_key not in self._generations:
|
|
if len(self._generations) >= self._generation_map_cap:
|
|
# FIFO overflow: the dropped key's race guard degrades to the accepted co-tenant class.
|
|
del self._generations[next(iter(self._generations))]
|
|
self._generations[gen_key] = 0
|
|
return self._generations[gen_key]
|
|
|
|
async def write(
|
|
self,
|
|
method: str,
|
|
params_key: str,
|
|
result: CacheableResult,
|
|
gen_at_capture: int,
|
|
mode: Literal["use", "refresh"],
|
|
) -> None:
|
|
"""Store a fetched result under the arm its resolved scope selects."""
|
|
gen_key = (method, params_key)
|
|
if self._generation_moved(gen_key, gen_at_capture):
|
|
return # the key was evicted while the fetch was in flight
|
|
ttl_ms, scope = self._resolve(result)
|
|
private_key = CacheKey(method, params_key, self._arm("private"))
|
|
public_key = CacheKey(method, params_key, self._arm("public"))
|
|
if ttl_ms <= 0:
|
|
if mode == "refresh":
|
|
# The refetch superseded the warm entry, which a cancellation must not leave serving.
|
|
await self._cleanup_delete(private_key, public_key)
|
|
return
|
|
own, opposite = (public_key, private_key) if scope == "public" else (private_key, public_key)
|
|
# Opposite arm first: a failed delete aborts before the set - never two arms answering for one key.
|
|
if not await self._delete(opposite):
|
|
# The own arm's entry is superseded too: best-effort delete, degrading to a full miss.
|
|
await self._cleanup_delete(own)
|
|
return
|
|
entry = CacheEntry(value=result.model_copy(deep=True), scope=scope, expires_at=self._clock() + ttl_ms / 1000)
|
|
try:
|
|
if not await self._set(own, entry):
|
|
# The fetch superseded any pre-existing own-arm entry, and the failed set
|
|
# left it in place: purge it (mirrors the opposite-arm-failure path).
|
|
await self._cleanup_delete(own)
|
|
finally:
|
|
# An eviction can land while the set commits - even when the await
|
|
# is cancelled - so re-check on every exit; the delete must complete
|
|
# so the pending cancellation cannot resurrect the evicted entry.
|
|
if self._generation_moved(gen_key, gen_at_capture):
|
|
await self._cleanup_delete(own)
|
|
|
|
async def evict_method(self, method: str) -> None:
|
|
"""Evict the method's cursor-less entry."""
|
|
await self.evict_key(method, "")
|
|
|
|
async def evict_key(self, method: str, params_key: str) -> None:
|
|
"""Evict one key from both arms.
|
|
|
|
Only the current era's arms are touched; other-era entries in a persistent store age out by TTL.
|
|
"""
|
|
gen_key = (method, params_key)
|
|
# Bump first so an in-flight fetch cannot write the evicted entry back.
|
|
# Unregistered keys skip the bump (uris must not grow the map) but not
|
|
# the deletes - a persistent store may hold uncaptured entries.
|
|
if gen_key in self._generations:
|
|
self._generations[gen_key] += 1
|
|
# Must complete: a cancellation between the deletes would leave one arm serving the evicted entry.
|
|
await self._cleanup_delete(
|
|
CacheKey(method, params_key, self._arm("private")),
|
|
CacheKey(method, params_key, self._arm("public")),
|
|
)
|
|
|
|
async def evict_for_notification(self, notification: ServerNotification) -> None:
|
|
"""Map a server notification to the entries it makes stale.
|
|
|
|
Eviction is eventual (spawned-task dispatch): the generation bump closes
|
|
the write-back race; a racing read may briefly serve the old entry.
|
|
"""
|
|
match notification:
|
|
case ToolListChangedNotification():
|
|
await self.evict_method("tools/list")
|
|
case PromptListChangedNotification():
|
|
await self.evict_method("prompts/list")
|
|
case ResourceListChangedNotification():
|
|
# Templates enumerate the same changed resource space.
|
|
await self.evict_method("resources/list")
|
|
await self.evict_method("resources/templates/list")
|
|
case ResourceUpdatedNotification():
|
|
await self.evict_key("resources/read", notification.params.uri)
|
|
case _:
|
|
pass
|
|
|
|
def _resolve(self, result: CacheableResult) -> tuple[int, Literal["public", "private"]]:
|
|
# A legacy peer can also put `ttlMs`/`cacheScope` keys on the wire, so
|
|
# wire presence is not a peer-era signal - hints count only when modern.
|
|
modern = self._negotiated_version() in MODERN_PROTOCOL_VERSIONS
|
|
if modern and "ttl_ms" in result.model_fields_set:
|
|
# An explicit `ttlMs: 0` stays 0, and negatives are unconstructible
|
|
# upstream (model ge=0, parse-seam floor) - only the cap applies.
|
|
ttl_ms = result.ttl_ms
|
|
else:
|
|
ttl_ms = self._default_ttl_ms
|
|
scope: Literal["public", "private"] = "public" if modern and result.cache_scope == "public" else "private"
|
|
return min(ttl_ms, MAX_TTL_MS), scope
|
|
|
|
def _generation_moved(self, gen_key: tuple[str, str], gen_at_capture: int) -> bool:
|
|
# A FIFO-dropped key fails open (the accepted co-tenant race) rather than discarding the fetch.
|
|
return self._generations.get(gen_key, gen_at_capture) != gen_at_capture
|
|
|
|
async def _set(self, key: CacheKey, entry: CacheEntry) -> bool:
|
|
try:
|
|
await self._store.set(key, entry)
|
|
except Exception: # boundary around user store code: nothing cached, the fetch already succeeded
|
|
self._warn_store_failure("set")
|
|
return False
|
|
self._warned_store_ops.discard("set")
|
|
return True
|
|
|
|
async def _cleanup_delete(self, *keys: CacheKey) -> None:
|
|
# Must-complete cleanup: shielded so a pending cancellation cannot skip the deletes,
|
|
# bounded so a wedged store delete cannot hold client teardown uncancellably.
|
|
with anyio.move_on_after(self._store_cleanup_timeout, shield=True) as scope:
|
|
for key in keys:
|
|
await self._delete(key)
|
|
if scope.cancelled_caught:
|
|
logger.warning("Response cache store delete timed out; the entry will age out by TTL")
|
|
|
|
async def _delete(self, key: CacheKey) -> bool:
|
|
try:
|
|
await self._store.delete(key)
|
|
except Exception: # boundary around user store code: callers decide whether a failed delete aborts
|
|
self._warn_store_failure("delete")
|
|
return False
|
|
self._warned_store_ops.discard("delete")
|
|
return True
|
|
|
|
def _warn_store_failure(self, kind: Literal["get", "set", "delete"]) -> None:
|
|
# One warning per failure burst, per op kind; re-armed only when that
|
|
# same kind succeeds, so a healthy delete cannot re-arm a broken set.
|
|
if kind not in self._warned_store_ops:
|
|
self._warned_store_ops.add(kind)
|
|
logger.warning("Response cache store operation failed; continuing without the cache", exc_info=True)
|