Add a client-side response cache honoring SEP-2549 caching hints (#3023)
This commit is contained in:
Max
GitHub
parent
67d7593df1
commit
b15b1d5f07
7 changed files with 3789 additions and 88 deletions
|
|
@ -13,7 +13,7 @@ from __future__ import annotations
|
|||
from collections.abc import Mapping
|
||||
from functools import cache
|
||||
from types import MappingProxyType, UnionType
|
||||
from typing import Any, Final, TypeVar
|
||||
from typing import Any, Final, Literal, TypeVar, get_args
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
|
|
@ -23,9 +23,11 @@ import mcp_types.v2026_07_28 as v2026
|
|||
from mcp_types.version import KNOWN_PROTOCOL_VERSIONS
|
||||
|
||||
__all__ = [
|
||||
"CACHEABLE_METHODS",
|
||||
"CLIENT_NOTIFICATIONS",
|
||||
"CLIENT_REQUESTS",
|
||||
"CLIENT_RESULTS",
|
||||
"CacheableMethod",
|
||||
"MONOLITH_NOTIFICATIONS",
|
||||
"MONOLITH_REQUESTS",
|
||||
"MONOLITH_RESULTS",
|
||||
|
|
@ -404,6 +406,24 @@ MONOLITH_RESULTS: Final[Mapping[str, type[types.Result] | UnionType]] = MappingP
|
|||
"""Monolith result model (or two-arm union) per request method."""
|
||||
|
||||
|
||||
CacheableMethod = Literal[
|
||||
"prompts/list",
|
||||
"resources/list",
|
||||
"resources/read",
|
||||
"resources/templates/list",
|
||||
"server/discover",
|
||||
"tools/list",
|
||||
]
|
||||
"""Methods whose results carry `ttlMs`/`cacheScope`; hand-written Literal, welded to `CACHEABLE_METHODS` by tests."""
|
||||
|
||||
CACHEABLE_METHODS: Final[frozenset[str]] = frozenset(
|
||||
method
|
||||
for method, row in MONOLITH_RESULTS.items()
|
||||
if any(issubclass(arm, types.CacheableResult) for arm in (get_args(row) if isinstance(row, UnionType) else (row,)))
|
||||
)
|
||||
"""Runtime mirror of `CacheableMethod`, derived from `MONOLITH_RESULTS`."""
|
||||
|
||||
|
||||
# --- Parse functions ---
|
||||
|
||||
# Envelope stubs merged into bodies for surface validation (surface classes are full frames).
|
||||
|
|
|
|||
|
|
@ -2,8 +2,28 @@
|
|||
|
||||
from mcp.client._input_required import InputRequiredRoundsExceededError
|
||||
from mcp.client._transport import Transport
|
||||
from mcp.client.caching import (
|
||||
CacheConfig,
|
||||
CacheEntry,
|
||||
CacheKey,
|
||||
CacheMode,
|
||||
InMemoryResponseCacheStore,
|
||||
ResponseCacheStore,
|
||||
)
|
||||
from mcp.client.client import Client
|
||||
from mcp.client.context import ClientRequestContext
|
||||
from mcp.client.session import ClientSession
|
||||
|
||||
__all__ = ["Client", "ClientRequestContext", "ClientSession", "InputRequiredRoundsExceededError", "Transport"]
|
||||
__all__ = [
|
||||
"CacheConfig",
|
||||
"CacheEntry",
|
||||
"CacheKey",
|
||||
"CacheMode",
|
||||
"Client",
|
||||
"ClientRequestContext",
|
||||
"ClientSession",
|
||||
"InMemoryResponseCacheStore",
|
||||
"InputRequiredRoundsExceededError",
|
||||
"ResponseCacheStore",
|
||||
"Transport",
|
||||
]
|
||||
|
|
|
|||
387
src/mcp/client/caching.py
Normal file
387
src/mcp/client/caching.py
Normal file
|
|
@ -0,0 +1,387 @@
|
|||
"""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)
|
||||
|
|
@ -2,14 +2,20 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from contextlib import AsyncExitStack
|
||||
from dataclasses import KW_ONLY, dataclass, field
|
||||
from typing import Any, Literal, TypeVar
|
||||
from typing import Any, Literal, TypeVar, cast
|
||||
|
||||
import anyio
|
||||
import anyio.lowlevel
|
||||
import mcp_types as types
|
||||
from mcp_types import (
|
||||
INVALID_PARAMS,
|
||||
CacheableResult,
|
||||
CallToolResult,
|
||||
CompleteResult,
|
||||
EmptyResult,
|
||||
|
|
@ -39,6 +45,7 @@ from mcp.client._input_required import DEFAULT_INPUT_REQUIRED_MAX_ROUNDS, run_in
|
|||
from mcp.client._memory import InMemoryTransport
|
||||
from mcp.client._probe import negotiate_auto
|
||||
from mcp.client._transport import Transport
|
||||
from mcp.client.caching import CacheConfig, CacheMode, ClientResponseCache, InMemoryResponseCacheStore
|
||||
from mcp.client.session import (
|
||||
ClientRequestContext,
|
||||
ClientSession,
|
||||
|
|
@ -54,8 +61,11 @@ from mcp.server.mcpserver import MCPServer
|
|||
from mcp.server.runner import modern_on_request
|
||||
from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair
|
||||
from mcp.shared.dispatcher import Dispatcher, ProgressFnT
|
||||
from mcp.shared.exceptions import MCPDeprecationWarning
|
||||
from mcp.shared.exceptions import MCPDeprecationWarning, MCPError
|
||||
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
|
||||
from mcp.shared.session import RequestResponder
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
ConnectMode = Literal["legacy", "auto"] | str
|
||||
"""``mode=`` value: ``"legacy"`` (initialize handshake), ``"auto"`` (discover, fall back to
|
||||
|
|
@ -64,6 +74,7 @@ forward-compat; ``Client.__post_init__`` rejects anything outside that set at co
|
|||
|
||||
_T = TypeVar("_T")
|
||||
_ResultT = TypeVar("_ResultT")
|
||||
_CacheableT = TypeVar("_CacheableT", bound=CacheableResult)
|
||||
|
||||
_Connector = Callable[[AsyncExitStack, ConnectMode, bool], Awaitable["Dispatcher[Any]"]]
|
||||
"""Resolved at ``__post_init__`` from the shape of ``server`` alone: enter whatever resources
|
||||
|
|
@ -115,6 +126,46 @@ def _connected(value: _T | None) -> _T:
|
|||
return value
|
||||
|
||||
|
||||
def _strip_userinfo(url: str) -> str:
|
||||
"""Drop any userinfo from the URL's authority component; byte-exact otherwise.
|
||||
|
||||
Credentials must not enter cache-key material; any further normalization could merge distinct servers.
|
||||
"""
|
||||
# Pure text, no urlsplit: it strips embedded tab/CR/LF before parsing, which would misalign slices.
|
||||
sep = url.find("//")
|
||||
if sep == -1:
|
||||
return url
|
||||
start = sep + 2
|
||||
end = len(url)
|
||||
for delimiter in "/?#":
|
||||
if (found := url.find(delimiter, start)) != -1:
|
||||
end = min(end, found)
|
||||
authority = url[start:end]
|
||||
if "@" not in authority:
|
||||
return url
|
||||
return url[:start] + authority.rpartition("@")[2] + url[end:]
|
||||
|
||||
|
||||
def _evicting_message_handler(cache: ClientResponseCache, user_handler: MessageHandlerFnT | None) -> MessageHandlerFnT:
|
||||
"""Wrap the session message handler with cache eviction on server notifications."""
|
||||
|
||||
async def handler(
|
||||
message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
|
||||
) -> None:
|
||||
if isinstance(message, types.ServerNotification):
|
||||
try:
|
||||
await cache.evict_for_notification(message)
|
||||
except Exception: # boundary: eviction reaches user store code; a cache fault must not block delivery
|
||||
logger.exception("Response cache eviction failed; the notification is still delivered")
|
||||
if user_handler is not None:
|
||||
await user_handler(message)
|
||||
else:
|
||||
# Mirrors ClientSession's default handler (session._default_message_handler).
|
||||
await anyio.lowlevel.checkpoint()
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
def _synthesize_discover(protocol_version: str) -> types.DiscoverResult:
|
||||
return types.DiscoverResult(
|
||||
supported_versions=[protocol_version],
|
||||
|
|
@ -221,10 +272,20 @@ class Client:
|
|||
"""SEP-2133 extension support to advertise under `ClientCapabilities.extensions`
|
||||
(identifier -> settings), e.g. `{"io.modelcontextprotocol/ui": {"mimeTypes": [...]}}`."""
|
||||
|
||||
cache: CacheConfig | Literal[False] | None = None
|
||||
"""Client-side response caching for the SEP-2549 cacheable methods (2026-07-28).
|
||||
|
||||
`None` (the default) honors server `ttlMs`/`cacheScope` hints with a per-client
|
||||
in-memory store; pass a `CacheConfig` to customize, or `False` to disable. The
|
||||
cacheable verbs take a per-call `cache_mode` (see `CacheMode`); calls carrying
|
||||
`meta` always reach the server. A `CacheConfig` with a custom `store` requires
|
||||
`target_id` when the server is not a URL (no identity can be derived)."""
|
||||
|
||||
_entered: bool = field(init=False, default=False)
|
||||
_session: ClientSession | None = field(init=False, default=None)
|
||||
_exit_stack: AsyncExitStack | None = field(init=False, default=None)
|
||||
_connect: _Connector = field(init=False, repr=False, compare=False)
|
||||
_response_cache: ClientResponseCache | None = field(init=False, default=None, repr=False, compare=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.mode not in ("legacy", "auto") and self.mode not in MODERN_PROTOCOL_VERSIONS:
|
||||
|
|
@ -247,16 +308,44 @@ class Client:
|
|||
else:
|
||||
self._connect = _connect_transport(srv)
|
||||
|
||||
if self.cache is not False:
|
||||
config = self.cache if self.cache is not None else CacheConfig()
|
||||
# Only the hash below leaves this scope - the raw identity may carry credentials; never log or store it.
|
||||
target_id = config.target_id
|
||||
if target_id is None and isinstance(self.server, str):
|
||||
target_id = _strip_userinfo(self.server)
|
||||
if target_id is None:
|
||||
if config.store is not None:
|
||||
raise ValueError(
|
||||
"a custom cache store requires CacheConfig.target_id when the server is not a URL: "
|
||||
"in-process servers and Transport instances get a random per-client identity, so "
|
||||
"their entries in a shared store could never be served to another client"
|
||||
)
|
||||
target_id = uuid.uuid4().hex
|
||||
self._response_cache = ClientResponseCache(
|
||||
store=config.store if config.store is not None else InMemoryResponseCacheStore(),
|
||||
partition=config.partition,
|
||||
arm_id=hashlib.sha256(target_id.encode()).hexdigest(),
|
||||
default_ttl_ms=config.default_ttl_ms,
|
||||
clock=config.clock,
|
||||
share_public=config.share_public,
|
||||
# Lazy: the negotiated version is unknown until __aenter__'s handshake.
|
||||
negotiated_version=lambda: self._session.protocol_version if self._session is not None else None,
|
||||
)
|
||||
|
||||
async def _build_session(self, exit_stack: AsyncExitStack) -> ClientSession:
|
||||
"""Enter the resolved connector and return an un-entered ClientSession."""
|
||||
dispatcher = await self._connect(exit_stack, self.mode, self.raise_exceptions)
|
||||
message_handler = self.message_handler
|
||||
if self._response_cache is not None:
|
||||
message_handler = _evicting_message_handler(self._response_cache, self.message_handler)
|
||||
return ClientSession(
|
||||
dispatcher=dispatcher,
|
||||
read_timeout_seconds=self.read_timeout_seconds,
|
||||
sampling_callback=self.sampling_callback,
|
||||
list_roots_callback=self.list_roots_callback,
|
||||
logging_callback=self.logging_callback,
|
||||
message_handler=self.message_handler,
|
||||
message_handler=message_handler,
|
||||
client_info=self.client_info,
|
||||
elicitation_callback=self.elicitation_callback,
|
||||
extensions=self.extensions,
|
||||
|
|
@ -361,23 +450,76 @@ class Client:
|
|||
"""Set the logging level on the server."""
|
||||
return await self.session.set_logging_level(level=level, meta=meta) # pyright: ignore[reportDeprecated]
|
||||
|
||||
async def _cached_fetch(
|
||||
self,
|
||||
method: str,
|
||||
*,
|
||||
cursor: str | None,
|
||||
meta: RequestParamsMeta | None,
|
||||
cache_mode: CacheMode,
|
||||
send: Callable[[], Awaitable[_CacheableT]],
|
||||
absorb: Callable[[_CacheableT], _CacheableT] | None = None,
|
||||
) -> _CacheableT:
|
||||
"""Serve one of the four list verbs through the response cache.
|
||||
|
||||
`absorb` (tools/list only) re-applies session-side derived state to a served cache hit.
|
||||
"""
|
||||
cache = self._response_cache
|
||||
if cache is None or cache_mode == "bypass":
|
||||
return await send()
|
||||
# A closed (or never-entered) client must raise, never serve cached entries.
|
||||
_ = self.session
|
||||
if meta is not None and cache_mode == "use":
|
||||
# meta (a progress token, tracing fields) expects a wire request; fetch and replace the entry.
|
||||
cache_mode = "refresh"
|
||||
if cursor is not None:
|
||||
# Continuation pages skip the cache, but an expired cursor means the listing changed (spec SHOULD evict).
|
||||
try:
|
||||
return await send()
|
||||
except MCPError as e:
|
||||
if e.code == INVALID_PARAMS:
|
||||
await cache.evict_method(method)
|
||||
raise
|
||||
if cache_mode == "use" and (hit := await cache.read(method, "")) is not None:
|
||||
# The hit is a private deep copy, so absorption may mutate it freely.
|
||||
served = cast(_CacheableT, hit)
|
||||
return served if absorb is None else absorb(served)
|
||||
gen = cache.capture(method, "")
|
||||
result = await send()
|
||||
await cache.write(method, "", result, gen, cache_mode)
|
||||
return result
|
||||
|
||||
async def list_resources(
|
||||
self,
|
||||
*,
|
||||
cursor: str | None = None,
|
||||
meta: RequestParamsMeta | None = None,
|
||||
cache_mode: CacheMode = "use",
|
||||
) -> ListResourcesResult:
|
||||
"""List available resources from the server."""
|
||||
return await self.session.list_resources(params=PaginatedRequestParams(cursor=cursor, _meta=meta))
|
||||
return await self._cached_fetch(
|
||||
"resources/list",
|
||||
cursor=cursor,
|
||||
meta=meta,
|
||||
cache_mode=cache_mode,
|
||||
send=lambda: self.session.list_resources(params=PaginatedRequestParams(cursor=cursor, _meta=meta)),
|
||||
)
|
||||
|
||||
async def list_resource_templates(
|
||||
self,
|
||||
*,
|
||||
cursor: str | None = None,
|
||||
meta: RequestParamsMeta | None = None,
|
||||
cache_mode: CacheMode = "use",
|
||||
) -> ListResourceTemplatesResult:
|
||||
"""List available resource templates from the server."""
|
||||
return await self.session.list_resource_templates(params=PaginatedRequestParams(cursor=cursor, _meta=meta))
|
||||
return await self._cached_fetch(
|
||||
"resources/templates/list",
|
||||
cursor=cursor,
|
||||
meta=meta,
|
||||
cache_mode=cache_mode,
|
||||
send=lambda: self.session.list_resource_templates(params=PaginatedRequestParams(cursor=cursor, _meta=meta)),
|
||||
)
|
||||
|
||||
async def read_resource(
|
||||
self,
|
||||
|
|
@ -386,6 +528,7 @@ class Client:
|
|||
input_responses: InputResponses | None = None,
|
||||
request_state: str | None = None,
|
||||
meta: RequestParamsMeta | None = None,
|
||||
cache_mode: CacheMode = "use",
|
||||
) -> ReadResourceResult:
|
||||
"""Read a resource from the server.
|
||||
|
||||
|
|
@ -400,6 +543,8 @@ class Client:
|
|||
resuming from a persisted `InputRequiredResult`).
|
||||
request_state: Opaque state to seed the first call with.
|
||||
meta: Additional metadata for the request.
|
||||
cache_mode: Cache behavior for this call (see `CacheMode`); seeded
|
||||
calls (`input_responses` or `request_state` set) ignore it.
|
||||
|
||||
Returns:
|
||||
The resource content.
|
||||
|
|
@ -414,7 +559,29 @@ class Client:
|
|||
uri, input_responses=r, request_state=s, meta=meta, allow_input_required=True
|
||||
)
|
||||
|
||||
return await self._drive_input_required(await retry(input_responses, request_state), retry)
|
||||
# Seeded calls resume a specific exchange and must never be cached (spec MUST).
|
||||
seeded = input_responses is not None or request_state is not None
|
||||
cache = None if seeded else self._response_cache
|
||||
if cache is None or cache_mode == "bypass":
|
||||
return await self._drive_input_required(await retry(input_responses, request_state), retry)
|
||||
# A closed (or never-entered) client must raise, never serve cached entries.
|
||||
_ = self.session
|
||||
if meta is not None and cache_mode == "use":
|
||||
# Calls carrying meta always reach the server (mirrors `_cached_fetch`).
|
||||
cache_mode = "refresh"
|
||||
if cache_mode == "use" and (hit := await cache.read("resources/read", uri)) is not None:
|
||||
# Only terminal first-round results are stored, so a hit legitimately skips the driver.
|
||||
return cast(ReadResourceResult, hit)
|
||||
gen = cache.capture("resources/read", uri)
|
||||
first = await retry(None, None)
|
||||
if not isinstance(first, InputRequiredResult):
|
||||
await cache.write("resources/read", uri, first, gen, cache_mode)
|
||||
elif cache_mode == "refresh":
|
||||
# The refresh superseded whatever was cached, but an input_required resolution
|
||||
# cannot be stored: purge the warm entry so it cannot be served again.
|
||||
await cache.evict_key("resources/read", uri)
|
||||
# Driver rounds carry inputResponses, so a terminal result reached through them is never cached (spec MUST).
|
||||
return await self._drive_input_required(first, retry)
|
||||
|
||||
async def subscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> EmptyResult:
|
||||
"""Subscribe to resource updates."""
|
||||
|
|
@ -481,9 +648,16 @@ class Client:
|
|||
*,
|
||||
cursor: str | None = None,
|
||||
meta: RequestParamsMeta | None = None,
|
||||
cache_mode: CacheMode = "use",
|
||||
) -> ListPromptsResult:
|
||||
"""List available prompts from the server."""
|
||||
return await self.session.list_prompts(params=PaginatedRequestParams(cursor=cursor, _meta=meta))
|
||||
return await self._cached_fetch(
|
||||
"prompts/list",
|
||||
cursor=cursor,
|
||||
meta=meta,
|
||||
cache_mode=cache_mode,
|
||||
send=lambda: self.session.list_prompts(params=PaginatedRequestParams(cursor=cursor, _meta=meta)),
|
||||
)
|
||||
|
||||
async def get_prompt(
|
||||
self,
|
||||
|
|
@ -565,9 +739,27 @@ class Client:
|
|||
"""
|
||||
return await self.session.complete(ref=ref, argument=argument, context_arguments=context_arguments)
|
||||
|
||||
async def list_tools(self, *, cursor: str | None = None, meta: RequestParamsMeta | None = None) -> ListToolsResult:
|
||||
async def list_tools(
|
||||
self,
|
||||
*,
|
||||
cursor: str | None = None,
|
||||
meta: RequestParamsMeta | None = None,
|
||||
cache_mode: CacheMode = "use",
|
||||
) -> ListToolsResult:
|
||||
"""List available tools from the server."""
|
||||
return await self.session.list_tools(params=PaginatedRequestParams(cursor=cursor, _meta=meta))
|
||||
return await self._cached_fetch(
|
||||
"tools/list",
|
||||
cursor=cursor,
|
||||
meta=meta,
|
||||
cache_mode=cache_mode,
|
||||
send=lambda: self.session.list_tools(params=PaginatedRequestParams(cursor=cursor, _meta=meta)),
|
||||
# A cache hit skips session.list_tools, so the session re-absorbs the served
|
||||
# listing to rebuild its derived per-tool state. Hits are cursorless, but a
|
||||
# cached page 1 can carry next_cursor - never prune on a partial listing.
|
||||
absorb=lambda hit: self.session._absorb_tool_listing( # pyright: ignore[reportPrivateUsage]
|
||||
hit, complete=hit.next_cursor is None
|
||||
),
|
||||
)
|
||||
|
||||
@deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning)
|
||||
async def send_roots_list_changed(self) -> None:
|
||||
|
|
|
|||
|
|
@ -55,6 +55,13 @@ DISCOVER_TIMEOUT_SECONDS = 10.0
|
|||
logger = logging.getLogger("client")
|
||||
|
||||
|
||||
def _clamp_inbound_ttl(raw: dict[str, Any]) -> None:
|
||||
"""Floor a negative inbound `ttlMs` to 0 before `ge=0` validation fails the call (2026-07-28 caching SHOULD)."""
|
||||
ttl = raw.get("ttlMs")
|
||||
if isinstance(ttl, int | float) and not isinstance(ttl, bool) and ttl < 0:
|
||||
raw["ttlMs"] = 0
|
||||
|
||||
|
||||
def _preconnect_stamp(data: dict[str, Any], opts: CallOptions) -> None:
|
||||
# initialize/discover forbid cancellation; other pre-handshake requests (lowlevel
|
||||
# ClientSession callers may skip the handshake entirely) keep the courtesy cancel.
|
||||
|
|
@ -331,6 +338,7 @@ class ClientSession:
|
|||
if metadata.on_resumption_token_update is not None:
|
||||
opts["on_resumption_token"] = metadata.on_resumption_token_update
|
||||
raw = await self._dispatcher.send_raw_request(method, data.get("params"), opts)
|
||||
_clamp_inbound_ttl(raw)
|
||||
# Literal fallback covers pre-handshake and stateless; matches runner.py.
|
||||
version = self._negotiated_version or "2025-11-25"
|
||||
try:
|
||||
|
|
@ -458,7 +466,10 @@ class ClientSession:
|
|||
"cancel_on_abandon": False,
|
||||
"headers": {MCP_PROTOCOL_VERSION_HEADER: version, MCP_METHOD_HEADER: data["method"]},
|
||||
}
|
||||
return await self._dispatcher.send_raw_request(data["method"], data.get("params"), opts)
|
||||
raw = await self._dispatcher.send_raw_request(data["method"], data.get("params"), opts)
|
||||
# Un-floored, a negative ttl fails the mode='auto' probe's validation and silently downgrades the handshake.
|
||||
_clamp_inbound_ttl(raw)
|
||||
return raw
|
||||
|
||||
async def discover(self) -> types.DiscoverResult:
|
||||
"""Probe `server/discover` and adopt the result.
|
||||
|
|
@ -895,7 +906,15 @@ class ClientSession:
|
|||
types.ListToolsRequest(params=params),
|
||||
types.ListToolsResult,
|
||||
)
|
||||
complete = (params is None or params.cursor is None) and result.next_cursor is None
|
||||
return self._absorb_tool_listing(result, complete=complete)
|
||||
|
||||
def _absorb_tool_listing(self, result: types.ListToolsResult, *, complete: bool) -> types.ListToolsResult:
|
||||
"""Filter the listing per the 2026 x-mcp-header MUST and rebuild derived per-tool state, in place.
|
||||
|
||||
Idempotent: cached values are already post-filter, so the response cache can re-absorb a served listing.
|
||||
`complete` (an uncursored single-page listing) prunes per-tool state down to the listing's tools.
|
||||
"""
|
||||
if self._negotiated_version in MODERN_PROTOCOL_VERSIONS:
|
||||
# 2026-07-28: clients MUST drop tools whose x-mcp-header annotations are invalid.
|
||||
kept: list[types.Tool] = []
|
||||
|
|
@ -911,11 +930,17 @@ class ClientSession:
|
|||
kept.append(tool)
|
||||
result.tools = kept
|
||||
|
||||
# Cache tool output schemas for future validation
|
||||
# Note: don't clear the cache, as we may be using a cursor
|
||||
# Cache tool output schemas for future validation; cursor pages only ever add.
|
||||
for tool in result.tools:
|
||||
self._tool_output_schemas[tool.name] = tool.output_schema
|
||||
|
||||
if complete:
|
||||
# The listing is the full tool universe, so state for unlisted tools is stale
|
||||
# (the server dropped them, or a shared-cache writer's filter did).
|
||||
names = {tool.name for tool in result.tools}
|
||||
self._x_mcp_header_maps = {k: v for k, v in self._x_mcp_header_maps.items() if k in names}
|
||||
self._tool_output_schemas = {k: v for k, v in self._tool_output_schemas.items() if k in names}
|
||||
|
||||
return result
|
||||
|
||||
@deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning)
|
||||
|
|
|
|||
|
|
@ -11,27 +11,13 @@ from __future__ import annotations
|
|||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Final, Literal, TypeVar, get_args
|
||||
from typing import Any, Literal, TypeVar
|
||||
|
||||
import mcp_types as types
|
||||
from mcp_types.methods import CACHEABLE_METHODS, CacheableMethod
|
||||
|
||||
__all__ = ["CACHEABLE_METHODS", "CacheHint", "CacheableMethod", "apply_cache_hint", "validate_cache_hints"]
|
||||
|
||||
CacheableMethod = Literal[
|
||||
"prompts/list",
|
||||
"resources/list",
|
||||
"resources/read",
|
||||
"resources/templates/list",
|
||||
"server/discover",
|
||||
"tools/list",
|
||||
]
|
||||
"""The methods whose results carry `ttlMs`/`cacheScope`. Closed set: the spec
|
||||
defines caching hints on exactly these six (tests pin it to which result models
|
||||
mix in `CacheableResult`)."""
|
||||
|
||||
CACHEABLE_METHODS: Final[frozenset[str]] = frozenset(get_args(CacheableMethod))
|
||||
"""Runtime mirror of `CacheableMethod`, for callers the type checker can't see."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CacheHint:
|
||||
|
|
@ -87,7 +73,8 @@ def validate_cache_hints(cache_hints: Mapping[Any, Any] | None) -> dict[str, Cac
|
|||
"""
|
||||
if cache_hints is None:
|
||||
return {}
|
||||
unknown = sorted(method for method in cache_hints if method not in CACHEABLE_METHODS)
|
||||
# repr-format keys so a non-string key raises this ValueError, not a TypeError from sorted/join.
|
||||
unknown = sorted(repr(method) for method in cache_hints if method not in CACHEABLE_METHODS)
|
||||
if unknown:
|
||||
raise ValueError(f"cache_hints keys must be cacheable methods (see CacheableMethod); got: {', '.join(unknown)}")
|
||||
validated: dict[str, CacheHint] = {}
|
||||
|
|
|
|||
|
|
@ -198,12 +198,15 @@ class ServerRunner(Generic[LifespanT]):
|
|||
if isinstance(result, ErrorData):
|
||||
# Raise inside the chain so middleware observes the failure.
|
||||
raise MCPError.from_error_data(result)
|
||||
# Fill cache hints on the typed result, before the serialize sieve
|
||||
# Fill cache hints on the handler result, before the serialize sieve
|
||||
# decides whether the negotiated version carries the fields at all.
|
||||
# `input_required` interim results are not `CacheableResult` models,
|
||||
# so the MRTR carve-out (no hints on them) holds by shape.
|
||||
if isinstance(result, CacheableResult) and (hint := self.server.cache_hints.get(method)) is not None:
|
||||
result = apply_cache_hint(result, hint)
|
||||
# MRTR carve-out: `input_required` interim results, typed or mapping, never get hints.
|
||||
if (hint := self.server.cache_hints.get(method)) is not None:
|
||||
if isinstance(result, CacheableResult):
|
||||
result = apply_cache_hint(result, hint)
|
||||
elif isinstance(result, Mapping) and result.get("resultType") != "input_required":
|
||||
# Hint keys first so wire keys the handler set win, matching `apply_cache_hint` precedence.
|
||||
result = {"ttlMs": hint.ttl_ms, "cacheScope": hint.scope, **result}
|
||||
# Dump and serialize inside the chain so the OpenTelemetry span (the
|
||||
# outermost middleware) records a failing handler return shape too.
|
||||
return self._serialize(method, version, result)
|
||||
|
|
|
|||
Loading…
Reference in a new issue