modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add a client extension API (#3034)

Browse source
This commit is contained in:
Max 2026-06-30 21:31:02 +01:00 • committed by GitHub
commit 4df609119f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 3410 additions and 180 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/mcp-types/mcp_types/__init__.py
View file

@ -8,6 +8,7 @@ https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema/dr
from mcp_types._types import (
CLIENT_CAPABILITIES_META_KEY,
CLIENT_INFO_META_KEY,
CORE_RESULT_TYPES,
DEFAULT_NEGOTIATED_VERSION,
LOG_LEVEL_META_KEY,
PROTOCOL_VERSION_META_KEY,
@ -231,6 +232,7 @@ __all__ = [
"CLIENT_CAPABILITIES_META_KEY",
"LOG_LEVEL_META_KEY",
# Type aliases and variables
"CORE_RESULT_TYPES",
"ContentBlock",
"ElicitRequestedSchema",
"ElicitRequestParams",

15
src/mcp-types/mcp_types/_types.py
View file

@ -8,7 +8,7 @@ the negotiated version. Per-field docstrings note version availability. The
from __future__ import annotations
from typing import Annotated, Any, Final, Generic, Literal, TypeAlias, TypeVar
from typing import Annotated, Any, ClassVar, Final, Generic, Literal, TypeAlias, TypeVar, get_args
from pydantic import (
BaseModel,
@ -128,6 +128,12 @@ class Request(MCPModel, Generic[RequestParamsT, MethodT]):
method: MethodT
params: RequestParamsT
name_param: ClassVar[str | None] = None
"""Wire-params key mirrored into the `Mcp-Name` header on sends; SEP-2663 requires it for tasks/*.
Subclasses override by bare assignment: re-annotating as `ClassVar` trips pyright's invariance check.
"""
class PaginatedRequest(Request[PaginatedRequestParams | None, MethodT], Generic[MethodT]):
"""Base class for paginated requests, matching the schema's PaginatedRequest interface."""
@ -144,7 +150,9 @@ class Notification(MCPModel, Generic[NotificationParamsT, MethodT]):
params: NotificationParamsT
ResultType = Literal["complete", "input_required"] | str
_CoreResultType = Literal["complete", "input_required"]
ResultType = _CoreResultType | str
"""Tags a `Result` so the client knows how to parse it (2026-07-28).
"complete" means the result is final; "input_required" means it is an
@ -152,6 +160,9 @@ ResultType = Literal["complete", "input_required"] | str
Absent `resultType` is equivalent to "complete".
"""
CORE_RESULT_TYPES: Final[frozenset[str]] = frozenset(get_args(_CoreResultType))
"""The `resultType` tags owned by the core protocol vocabulary; extension claims may not re-key them."""
class Result(MCPModel):
"""Base class for JSON-RPC results.

14
src/mcp/client/__init__.py
View file

@ -12,6 +12,14 @@ from mcp.client.caching import (
)
from mcp.client.client import Client
from mcp.client.context import ClientRequestContext
from mcp.client.extension import (
ClaimContext,
ClientExtension,
NotificationBinding,
ResultClaim,
UnexpectedClaimedResult,
advertise,
)
from mcp.client.session import ClientSession
__all__ = [
@ -19,11 +27,17 @@ __all__ = [
"CacheEntry",
"CacheKey",
"CacheMode",
"ClaimContext",
"Client",
"ClientExtension",
"ClientRequestContext",
"ClientSession",
"InMemoryResponseCacheStore",
"InputRequiredRoundsExceededError",
"NotificationBinding",
"ResponseCacheStore",
"ResultClaim",
"Transport",
"UnexpectedClaimedResult",
"advertise",
]

112
src/mcp/client/client.py
View file

@ -5,7 +5,7 @@ from __future__ import annotations
import hashlib
import logging
import uuid
from collections.abc import Awaitable, Callable, Mapping
from collections.abc import Awaitable, Callable, Mapping, Sequence
from contextlib import AsyncExitStack
from dataclasses import KW_ONLY, dataclass, field
from typing import Any, Literal, TypeVar, cast
@ -36,6 +36,7 @@ from mcp_types import (
ReadResourceResult,
RequestParamsMeta,
ResourceTemplateReference,
Result,
ServerCapabilities,
)
from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, MODERN_PROTOCOL_VERSIONS
@ -46,6 +47,7 @@ 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.extension import ClaimContext, ClientExtension, NotificationBinding, ResultClaim
from mcp.client.session import (
ClientRequestContext,
ClientSession,
@ -62,6 +64,7 @@ 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, MCPError
from mcp.shared.extension import validate_extension_identifier
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
from mcp.shared.session import RequestResponder
@ -188,6 +191,72 @@ async def _no_inbound_client_notifications(_dctx: Any, _method: str, _params: Ma
"""
@dataclass(frozen=True)
class _FoldedExtensions:
"""`Client.extensions` instances folded into the shapes `ClientSession` consumes."""
ad: dict[str, dict[str, Any]] | None
claims: dict[str, tuple[ResultClaim[Any], ...]] | None
bindings: tuple[NotificationBinding[Any], ...] | None
by_model: Mapping[type[Result], ResultClaim[Any]]
def _fold_extensions(extensions: Sequence[ClientExtension] | None) -> _FoldedExtensions:
"""Fold extension contributions at construction, naming both owners on duplicate tags or methods."""
if isinstance(extensions, Mapping):
raise TypeError(
"extensions= takes a sequence of ClientExtension instances. The mapping form was "
"replaced: use advertise(identifier, settings) for advertise-only entries"
)
if not extensions:
return _FoldedExtensions(ad=None, claims=None, bindings=None, by_model={})
ad: dict[str, dict[str, Any]] = {}
claims: dict[str, tuple[ResultClaim[Any], ...]] = {}
bindings: list[NotificationBinding[Any]] = []
by_model: dict[type[Result], ResultClaim[Any]] = {}
claim_owners: dict[str, str] = {}
binding_owners: dict[str, str] = {}
for extension in extensions:
identifier = getattr(extension, "identifier", None)
if identifier is None:
raise ValueError(
f"{type(extension).__name__} has no `identifier`; a ClientExtension must set the "
"`identifier` class attribute (or assign one in `__init__`) before it can be used"
)
validate_extension_identifier(identifier, owner=type(extension).__name__)
if identifier in ad:
raise ValueError(f"extension identifier {identifier!r} is passed more than once")
ad[identifier] = extension.settings()
extension_claims = tuple(extension.claims())
for claim in extension_claims:
tag = claim.result_type
if tag in claim_owners:
owner = claim_owners[tag]
both = (
f"extension {identifier!r} claims"
if owner == identifier
else (f"extensions {owner!r} and {identifier!r} both claim")
)
raise ValueError(f"{both} resultType {tag!r}; a wire tag can have only one resolver")
claim_owners[tag] = identifier
# Each model pins its result_type Literal to one tag, so this index cannot collide.
by_model[claim.model] = claim
if extension_claims:
claims[identifier] = extension_claims
for binding in extension.notifications():
if binding.method in binding_owners:
owner = binding_owners[binding.method]
both = (
f"extension {identifier!r} binds"
if owner == identifier
else (f"extensions {owner!r} and {identifier!r} both bind")
)
raise ValueError(f"{both} notification method {binding.method!r}; a method can have only one observer")
binding_owners[binding.method] = identifier
bindings.append(binding)
return _FoldedExtensions(ad=ad, claims=claims or None, bindings=tuple(bindings) or None, by_model=by_model)
@dataclass
class Client:
"""A high-level MCP client for connecting to MCP servers.
@ -268,9 +337,12 @@ class Client:
`read_resource` give up. Use `client.session.<method>(..., allow_input_required=True)`
to drive the loop manually instead."""
extensions: dict[str, dict[str, Any]] | None = None
"""SEP-2133 extension support to advertise under `ClientCapabilities.extensions`
(identifier -> settings), e.g. `{"io.modelcontextprotocol/ui": {"mimeTypes": [...]}}`."""
extensions: Sequence[ClientExtension] | None = None
"""Opt-in client extensions (SEP-2133).
Each instance contributes its capability ad, its result claims (resolved
transparently by `call_tool`), and its notification bindings. For an
ad-only entry use `mcp.client.advertise(identifier, settings)`."""
cache: CacheConfig | Literal[False] | None = None
"""Client-side response caching for the SEP-2549 cacheable methods (2026-07-28).
@ -286,6 +358,7 @@ class Client:
_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)
_folded_extensions: _FoldedExtensions = field(init=False, repr=False, compare=False)
def __post_init__(self) -> None:
if self.mode not in ("legacy", "auto") and self.mode not in MODERN_PROTOCOL_VERSIONS:
@ -298,6 +371,8 @@ class Client:
f"mode must be 'legacy', 'auto', or one of {list(MODERN_PROTOCOL_VERSIONS)}; got {self.mode!r}{hint}"
)
self._folded_extensions = _fold_extensions(self.extensions)
srv = self.server
if isinstance(srv, MCPServer):
srv = srv._lowlevel_server # pyright: ignore[reportPrivateUsage]
@ -348,7 +423,9 @@ class Client:
message_handler=message_handler,
client_info=self.client_info,
elicitation_callback=self.elicitation_callback,
extensions=self.extensions,
extensions=self._folded_extensions.ad,
result_claims=self._folded_extensions.claims,
notification_bindings=self._folded_extensions.bindings,
)
async def __aenter__(self) -> Client:
@ -613,6 +690,11 @@ class Client:
state is still subject to the server's TTL, request binding, and key
lifetime; a server on the default process-local key rejects it after a restart.
Result shapes claimed by this client's `extensions` are finished by the
owning claim's resolver, whose `CallToolResult` is returned; resolver
exceptions propagate as-is. To receive the claimed shape yourself, use
`client.session.call_tool(..., allow_claimed=True)`.
Args:
name: The name of the tool to call.
arguments: Arguments to pass to the tool.
@ -631,7 +713,7 @@ class Client:
MCPError: A callback returned `ErrorData` for an embedded input request.
"""
async def retry(r: InputResponses | None, s: str | None) -> CallToolResult | InputRequiredResult:
async def retry(r: InputResponses | None, s: str | None) -> CallToolResult | InputRequiredResult | Result:
return await self.session.call_tool(
name,
arguments,
@ -641,9 +723,23 @@ class Client:
request_state=s,
meta=meta,
allow_input_required=True,
# Input rounds resolve before a claimed result, so a claim may end any round.
allow_claimed=True,
)
return await self._drive_input_required(await retry(input_responses, request_state), retry)
result = await self._drive_input_required(await retry(input_responses, request_state), retry)
if isinstance(result, CallToolResult):
return result
# Only claimed shapes reach this point, so the lookup is total.
claim = self._folded_extensions.by_model[type(result)]
final = await claim.resolve(
result,
ClaimContext(session=self.session, tool_name=name, read_timeout_seconds=read_timeout_seconds),
)
if not final.is_error:
# Match the direct path: revalidate the output schema, but never for isError results.
await self.session.validate_tool_result(name, final)
return final
async def list_prompts(
self,
@ -717,7 +813,7 @@ class Client:
async def dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData:
ctx = ClientRequestContext(session=session, request_id=key, meta=req.params.meta if req.params else None)
return await session._dispatch_input_request(ctx, req) # pyright: ignore[reportPrivateUsage]
return await session.dispatch_input_request(ctx, req)
return await run_input_required_driver(
first, dispatch=dispatch, retry=retry, max_rounds=self.input_required_max_rounds

196
src/mcp/client/extension.py Normal file
View file

@ -0,0 +1,196 @@
"""Opt-in extension interface for MCP clients.
Subclass `ClientExtension`, set `identifier`, override the hooks you need, and
pass instances to `Client(extensions=[...])`. For an identifier-only
capability ad, use `advertise()`.
"""
from __future__ import annotations
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, Generic, Literal, TypeVar, get_args
from mcp_types import CORE_RESULT_TYPES, CallToolResult, InputRequiredResult, Result
from mcp_types.version import MODERN_PROTOCOL_VERSIONS
from pydantic import AliasChoices, AliasPath, BaseModel
from pydantic.fields import FieldInfo
from mcp.shared.extension import validate_extension_identifier
if TYPE_CHECKING:
from mcp.client.session import ClientSession
__all__ = [
"ClaimContext",
"ClientExtension",
"NotificationBinding",
"ResultClaim",
"UnexpectedClaimedResult",
"advertise",
]
_CLAIM_METHODS: Final[frozenset[str]] = frozenset({"tools/call"})
"""The closed set of verbs a claim may attach to; widen together with the `method` Literal."""
_RESERVED_WIRE_ALIASES: Final[frozenset[str]] = frozenset({"requestState", "inputRequests"})
"""Typed optional fields of the core result surface that pre-validates every inbound result."""
def _wire_keys(name: str, field: FieldInfo) -> frozenset[str]:
"""Every top-level wire key this field can read from or write to."""
keys = {field.alias or name}
if field.serialization_alias:
keys.add(field.serialization_alias)
validation_alias = field.validation_alias
choices = validation_alias.choices if isinstance(validation_alias, AliasChoices) else [validation_alias]
for choice in choices:
if isinstance(choice, AliasPath):
choice = choice.path[0]
if isinstance(choice, str):
keys.add(choice)
return frozenset(keys)
ClaimedT = TypeVar("ClaimedT", bound=Result)
NotifyParamsT = TypeVar("NotifyParamsT", bound=BaseModel)
@dataclass(frozen=True, kw_only=True)
class ClaimContext:
"""Host-injected context for one `ResultClaim.resolve` call."""
session: ClientSession
tool_name: str
read_timeout_seconds: float | None
@dataclass(frozen=True, kw_only=True)
class ResultClaim(Generic[ClaimedT]):
"""One extra result shape on one spec verb, keyed by the wire `resultType`.
Active only while the declaring extension is constructed into the client and
the negotiated protocol version admits it. `resolve` finishes a claimed
result, may send follow-ups through `ctx.session`, and must return the
verb's ordinary result. All field constraints are enforced at construction.
"""
result_type: str
model: type[ClaimedT]
resolve: Callable[[ClaimedT, ClaimContext], Awaitable[CallToolResult]]
method: Literal["tools/call"] = "tools/call"
protocol_versions: frozenset[str] | None = None
def __post_init__(self) -> None:
if self.method not in _CLAIM_METHODS:
raise ValueError(f"claims attach to {sorted(_CLAIM_METHODS)} only; got method {self.method!r}")
if self.result_type in CORE_RESULT_TYPES:
raise ValueError(f"resultType {self.result_type!r} is core protocol vocabulary")
if Result not in self.model.__mro__: # runtime guard; the ClaimedT bound only constrains checked callers
raise ValueError(f"{self.model.__name__} must subclass mcp_types.Result")
if issubclass(self.model, CallToolResult | InputRequiredResult):
raise ValueError("claim models must not subclass core result types")
for name, model_field in self.model.model_fields.items():
for clash in sorted(_wire_keys(name, model_field) & _RESERVED_WIRE_ALIASES):
raise ValueError(
f"{self.model.__name__}.{name} aliases {clash!r}, a typed field of the core "
"result surface; a colliding value would fail core validation before the "
"claim adapter runs"
)
field = self.model.model_fields.get("result_type")
if field is None or get_args(field.annotation) != (self.result_type,):
raise ValueError(f"{self.model.__name__}.result_type must be Literal[{self.result_type!r}]")
if self.protocol_versions is not None and not self.protocol_versions:
raise ValueError("empty protocol_versions could never activate; use None for all")
if self.protocol_versions is not None and not self.protocol_versions.issubset(MODERN_PROTOCOL_VERSIONS):
unrecognized = sorted(self.protocol_versions.difference(MODERN_PROTOCOL_VERSIONS))
raise ValueError(
f"protocol_versions {unrecognized} are not modern protocol revisions; claimed shapes "
"cannot be delivered on a legacy wire (None means every modern version)"
)
class UnexpectedClaimedResult(RuntimeError):
"""A claimed (extension) result arrived on a `call_tool` that did not opt in.
The parsed value is carried as `result`; the server may already hold state it
references. Opt in via `Client(extensions=[...])` or `allow_claimed=True`.
"""
def __init__(self, result: Result) -> None:
super().__init__(
f"Server returned a claimed result ({type(result).__name__}); pass the owning extension to "
"Client(extensions=[...]) for transparent resolution, or call with allow_claimed=True "
"and handle the shape. The carried result may reference server-side state needing cleanup."
)
self.result = result
@dataclass(frozen=True, kw_only=True)
class NotificationBinding(Generic[NotifyParamsT]):
"""Deliver server notifications for `method` (the bare wire name) to `handler`.
Observation-only: validated params arrive one at a time per binding, in
dispatch order, through a bounded queue that drops the oldest with a warning
on overflow. Stream transports dispatch each notification independently, so
near-simultaneous notifications may be dispatched out of wire order. Methods
the negotiated version's core tables handle are never delivered to bindings.
"""
method: str
params_type: type[NotifyParamsT]
handler: Callable[[NotifyParamsT], Awaitable[None]]
class ClientExtension:
"""Base class for an opt-in client extension; override only what you need.
The surface is declarative, fixed at construction, and never receives the client.
"""
#: Reverse-DNS extension identifier, advertised under `ClientCapabilities.extensions`.
identifier: str
def __init_subclass__(cls, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs)
# Per-instance identifiers (assigned in __init__) are validated at consumption instead.
if (identifier := cls.__dict__.get("identifier")) is not None:
validate_extension_identifier(identifier, owner=cls.__name__)
def settings(self) -> dict[str, Any]:
"""Per-extension settings advertised at `ClientCapabilities.extensions[identifier]`.
Read once at `Client` construction. A claim-bearing extension is
advertised only at protocol versions where at least one of its claims
is active.
"""
return {}
def claims(self) -> Sequence[ResultClaim[Any]]:
"""Extra result shapes this extension claims, with their resolvers."""
return ()
def notifications(self) -> Sequence[NotificationBinding[Any]]:
"""Server notifications this extension observes."""
return ()
class _AdvertiseOnly(ClientExtension):
"""Ad-only extension returned by `advertise()`."""
def __init__(self, identifier: str, settings: dict[str, Any]) -> None:
self.identifier = identifier
self._settings = settings
def settings(self) -> dict[str, Any]:
return self._settings
def advertise(identifier: str, settings: dict[str, Any] | None = None) -> ClientExtension:
"""Advertise an extension identifier (with optional settings) and nothing else.
Advertising an extension you do not implement asserts wire support you do
not have; for behavioral extensions construct the real extension instead.
"""
validate_extension_identifier(identifier, owner="advertise")
return _AdvertiseOnly(identifier, {} if settings is None else settings)

292
src/mcp/client/session.py
View file

@ -1,15 +1,18 @@
from __future__ import annotations
import logging
from collections.abc import Callable, Mapping
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from functools import reduce
from operator import or_
from types import TracebackType
from typing import Any, Literal, Protocol, cast, overload
from typing import Annotated, Any, Final, Literal, Protocol, cast, overload
import anyio
import anyio.abc
import anyio.lowlevel
import mcp_types as types
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp_types import (
CLIENT_CAPABILITIES_META_KEY,
CLIENT_INFO_META_KEY,
@ -27,10 +30,11 @@ from mcp_types.version import (
LATEST_MODERN_VERSION,
MODERN_PROTOCOL_VERSIONS,
)
from pydantic import BaseModel, TypeAdapter, ValidationError
from pydantic import BaseModel, Discriminator, Tag, TypeAdapter, ValidationError
from typing_extensions import Self, TypeVar, deprecated
from mcp.client._transport import ReadStream, WriteStream
from mcp.client.extension import NotificationBinding, ResultClaim, UnexpectedClaimedResult
from mcp.shared._compat import resync_tracer
from mcp.shared.dispatcher import CallOptions, DispatchContext, Dispatcher, ProgressFnT
from mcp.shared.exceptions import MCPDeprecationWarning, MCPError
@ -51,6 +55,7 @@ from mcp.shared.transport_context import TransportContext
DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0")
DISCOVER_TIMEOUT_SECONDS = 10.0
_NOTIFICATION_QUEUE_SIZE: Final = 256
logger = logging.getLogger("client")
@ -189,7 +194,8 @@ async def _default_logging_callback(
ClientResponse: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(types.ClientResult | types.ErrorData)
_CallToolResultAdapter: TypeAdapter[types.CallToolResult | types.InputRequiredResult] = TypeAdapter(
# Typed against the wide parse union so adopt-built claim adapters share this attribute type.
_CallToolResultAdapter: TypeAdapter[types.CallToolResult | types.InputRequiredResult | types.Result] = TypeAdapter(
types.CallToolResult | types.InputRequiredResult
)
_GetPromptResultAdapter: TypeAdapter[types.GetPromptResult | types.InputRequiredResult] = TypeAdapter(
@ -200,6 +206,89 @@ _ReadResourceResultAdapter: TypeAdapter[types.ReadResourceResult | types.InputRe
)
def _claim_active(claim: ResultClaim[Any], version: str) -> bool:
"""A claim is active at modern versions only, narrowed by its optional version subset."""
return version in MODERN_PROTOCOL_VERSIONS and (
claim.protocol_versions is None or version in claim.protocol_versions
)
def _active_claims_at(
claims_by_extension: Mapping[str, tuple[ResultClaim[Any], ...]], version: str
) -> dict[str, ResultClaim[Any]]:
"""Claims active at `version`, keyed by wire tag; empty at any legacy version."""
return {
claim.result_type: claim
for claims in claims_by_extension.values()
for claim in claims
if _claim_active(claim, version)
}
def _build_call_tool_adapter(
active: Mapping[str, ResultClaim[Any]],
) -> TypeAdapter[types.CallToolResult | types.InputRequiredResult | types.Result]:
"""Build a discriminated tools/call adapter: a core arm plus one arm per active claim."""
if not active:
return _CallToolResultAdapter
tags = frozenset(active)
core_arm = "core"
while core_arm in tags: # the routing sentinel must never collide with a claimed tag
core_arm += "-"
def _route(value: Any) -> str:
# pydantic hands the discriminator either the raw dict or an already-built model.
# Unknown or non-string tags route to the core arm and fail core validation there.
if isinstance(value, dict):
tag = cast("dict[str, Any]", value).get("resultType")
else:
tag = getattr(value, "result_type", None)
return tag if isinstance(tag, str) and tag in tags else core_arm
arms: list[Any] = [Annotated[types.CallToolResult | types.InputRequiredResult, Tag(core_arm)]]
arms += [Annotated[claim.model, Tag(tag)] for tag, claim in active.items()]
# reduce(or_) rather than Union star-unpack, which needs py3.11+.
return TypeAdapter(Annotated[reduce(or_, arms), Discriminator(_route)])
def _index_claims(
result_claims: Mapping[str, Sequence[ResultClaim[Any]]] | None,
extensions: dict[str, dict[str, Any]] | None,
) -> dict[str, tuple[ResultClaim[Any], ...]]:
"""Validate and copy the claims-by-extension mapping."""
indexed: dict[str, tuple[ResultClaim[Any], ...]] = {}
seen: set[str] = set()
for identifier, claims in (result_claims or {}).items():
if extensions is None or identifier not in extensions:
raise ValueError(
f"result_claims key {identifier!r} has no extensions entry; a claim is only "
"advertised through its extension's capability ad"
)
if not claims:
raise ValueError(
f"result_claims[{identifier!r}] is empty and would drop the extension from "
"the capability ad at every version. Omit the key instead"
)
for claim in claims:
if claim.result_type in seen:
raise ValueError(f"duplicate result claim for resultType {claim.result_type!r}")
seen.add(claim.result_type)
indexed[identifier] = tuple(claims)
return indexed
def _index_bindings(
notification_bindings: Sequence[NotificationBinding[Any]] | None,
) -> dict[str, NotificationBinding[Any]]:
"""Index bindings by wire method, rejecting duplicates."""
indexed: dict[str, NotificationBinding[Any]] = {}
for binding in notification_bindings or ():
if binding.method in indexed:
raise ValueError(f"duplicate notification binding for method {binding.method!r}")
indexed[binding.method] = binding
return indexed
def _input_required_unexpected(method: str) -> RuntimeError:
return RuntimeError(
"Server returned InputRequiredResult; pass allow_input_required=True to receive it "
@ -216,6 +305,9 @@ class ClientSession:
correlation; this class owns the typed MCP layer and the constructor
callbacks. Transport `Exception` items reach `message_handler` only when
the session builds its own dispatcher from a stream pair.
Extension `result_claims` fold into tools/call parsing at `adopt()`;
`notification_bindings` observe vendor notifications via bounded FIFOs.
"""
def __init__(
@ -232,13 +324,22 @@ class ClientSession:
*,
sampling_capabilities: types.SamplingCapability | None = None,
extensions: dict[str, dict[str, Any]] | None = None,
result_claims: Mapping[str, Sequence[ResultClaim[Any]]] | None = None,
notification_bindings: Sequence[NotificationBinding[Any]] | None = None,
dispatcher: Dispatcher[Any] | None = None,
) -> None:
self._session_read_timeout_seconds = read_timeout_seconds
self._client_info = client_info or DEFAULT_CLIENT_INFO
self._sampling_callback = sampling_callback or _default_sampling_callback
self._sampling_capabilities = sampling_capabilities
self._extensions = extensions
self._extensions = dict(extensions) if extensions is not None else None
self._result_claims = _index_claims(result_claims, extensions)
self._notification_bindings = _index_bindings(notification_bindings)
self._active_claims: dict[str, ResultClaim[Any]] = {}
self._call_tool_adapter = _CallToolResultAdapter
self._binding_queues: dict[
str, tuple[MemoryObjectSendStream[BaseModel], MemoryObjectReceiveStream[BaseModel]]
] = {}
self._elicitation_callback = elicitation_callback or _default_elicitation_callback
self._list_roots_callback = list_roots_callback or _default_list_roots_callback
self._logging_callback = logging_callback or _default_logging_callback
@ -274,7 +375,14 @@ class ClientSession:
self._task_group = anyio.create_task_group()
await self._task_group.__aenter__()
try:
# Queues must exist before the dispatcher starts: _on_notify enqueues into this dict.
for binding in self._notification_bindings.values():
send, receive = anyio.create_memory_object_stream[BaseModel](_NOTIFICATION_QUEUE_SIZE)
self._binding_queues[binding.method] = (send, receive)
await self._task_group.start(self._dispatcher.run, self._on_request, self._on_notify)
for binding in self._notification_bindings.values():
_, receive = self._binding_queues[binding.method]
self._task_group.start_soon(self._deliver_bound_notifications, binding, receive)
except BaseException:
# Unwind the entered task group before propagating: a cancellation
# landing here (e.g. `move_on_after` around connect) would abandon
@ -285,7 +393,10 @@ class ClientSession:
# Shield the group's own scope (a new one would break LIFO exit)
# so a pending outer cancellation cannot re-fire inside __aexit__.
task_group.cancel_scope.shield = True
await task_group.__aexit__(None, None, None)
try:
await task_group.__aexit__(None, None, None)
finally:
self._close_binding_queues()
raise
return self
@ -295,16 +406,38 @@ class ClientSession:
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> bool | None:
# Exit must not block: cancel the dispatcher and in-flight callbacks.
# Exit must not block: cancel the dispatcher, binding consumers, and in-flight callbacks.
assert self._task_group is not None
self._task_group.cancel_scope.cancel()
result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
try:
result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
finally:
self._close_binding_queues()
await resync_tracer()
return result
def _close_binding_queues(self) -> None:
# Unclosed memory object streams warn at garbage collection; close is idempotent.
for send, receive in self._binding_queues.values():
send.close()
receive.close()
self._binding_queues.clear()
async def _deliver_bound_notifications(
self, binding: NotificationBinding[Any], receive: MemoryObjectReceiveStream[BaseModel]
) -> None:
"""Consume one binding's FIFO, decoupled from the dispatcher so handlers can do session I/O."""
while True:
params = await receive.receive()
try:
await binding.handler(params)
except Exception:
# A raising handler costs only that delivery, as in _on_notify.
logger.exception("notification binding handler for %r raised", binding.method)
async def send_request(
self,
request: types.ClientRequest,
request: types.ClientRequest | types.Request[Any, Any],
result_type: type[ReceiveResultT] | TypeAdapter[ReceiveResultT],
request_read_timeout_seconds: float | None = None,
metadata: ClientMessageMetadata | None = None,
@ -318,11 +451,20 @@ class ClientSession:
Raises:
MCPError: Error response, read timeout, or connection closed.
RuntimeError: Called before entering the context manager.
ValueError: The request declares `name_param` but its params carry no string name.
"""
data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
method: str = data["method"]
opts: CallOptions = {}
self._stamp(data, opts)
# The stamp runs first, so its NAME_BEARING_METHODS rows win; a missing name fails loud.
headers = opts.setdefault("headers", {})
if (key := type(request).name_param) is not None and MCP_NAME_HEADER not in headers:
params_data: dict[str, Any] = data.get("params") or {}
name = params_data.get(key)
if not isinstance(name, str):
raise ValueError(f"{method} requires params[{key!r}] for Mcp-Name")
headers[MCP_NAME_HEADER] = encode_header_value(name)
timeout = (
request_read_timeout_seconds
if request_read_timeout_seconds is not None
@ -360,7 +502,21 @@ class ClientSession:
self._stamp(data, opts)
await self._dispatcher.notify(data["method"], data.get("params"), opts)
def _build_capabilities(self) -> types.ClientCapabilities:
def _build_capabilities(self, version: str) -> types.ClientCapabilities:
"""Build the capability ad for a wire speaking `version`.
Claim-bearing identifiers whose claims are all inactive at `version` drop, so
the client never advertises result shapes it would reject; claim-less
identifiers always advertise.
"""
extensions = self._extensions
if extensions is not None and self._result_claims:
extensions = {
identifier: settings
for identifier, settings in extensions.items()
if identifier not in self._result_claims
or any(_claim_active(claim, version) for claim in self._result_claims[identifier])
} or None
sampling = (
(self._sampling_capabilities or types.SamplingCapability())
if self._sampling_callback is not _default_sampling_callback
@ -380,7 +536,7 @@ class ClientSession:
else None
)
return types.ClientCapabilities(
sampling=sampling, elicitation=elicitation, experimental=None, extensions=self._extensions, roots=roots
sampling=sampling, elicitation=elicitation, experimental=None, extensions=extensions, roots=roots
)
async def initialize(self) -> types.InitializeResult:
@ -390,7 +546,8 @@ class ClientSession:
types.InitializeRequest(
params=types.InitializeRequestParams(
protocol_version=LATEST_HANDSHAKE_VERSION,
capabilities=self._build_capabilities(),
# The handshake negotiates only legacy versions, where no claim is active.
capabilities=self._build_capabilities(LATEST_HANDSHAKE_VERSION),
client_info=self._client_info,
),
),
@ -424,17 +581,30 @@ class ClientSession:
f"No mutually supported modern protocol version "
f"(server: {result.supported_versions}, client: {list(MODERN_PROTOCOL_VERSIONS)})"
)
version = mutual[-1]
client_info = self._client_info.model_dump(by_alias=True, mode="json", exclude_none=True)
capabilities = self._build_capabilities().model_dump(by_alias=True, mode="json", exclude_none=True)
self._stamp = _make_modern_stamp(mutual[-1], client_info, capabilities, self._resolve_param_headers)
capabilities = self._build_capabilities(version).model_dump(by_alias=True, mode="json", exclude_none=True)
self._stamp = _make_modern_stamp(version, client_info, capabilities, self._resolve_param_headers)
self._discover_result = result
self._initialize_result = None
self._negotiated_version = mutual[-1]
else:
self._stamp = _make_handshake_stamp(result.protocol_version)
version = result.protocol_version
self._stamp = _make_handshake_stamp(version)
self._initialize_result = result
self._discover_result = None
self._negotiated_version = result.protocol_version
self._negotiated_version = version
# Both arms reach here, so re-adoption resets cleanly; legacy versions activate no claims.
# Core-vocabulary tags are unconstructible (ResultClaim.__post_init__), so no exclusion needed.
self._active_claims = _active_claims_at(self._result_claims, version)
self._call_tool_adapter = _build_call_tool_adapter(self._active_claims)
for method in self._notification_bindings:
# Bindings are consulted only for methods core does not know, so this one can never fire.
if (method, version) in _methods.SERVER_NOTIFICATIONS:
logger.warning(
"notification binding for %r will never fire at %s: the core protocol defines this method",
method,
version,
)
async def send_discover(self, version: str) -> dict[str, Any]:
"""Send a single ``server/discover`` at ``version`` and return the raw result dict.
@ -450,7 +620,7 @@ class ClientSession:
synthesized into a JSON-RPC error by the transport).
"""
client_info = self._client_info.model_dump(by_alias=True, mode="json", exclude_none=True)
capabilities = self._build_capabilities().model_dump(by_alias=True, mode="json", exclude_none=True)
capabilities = self._build_capabilities(version).model_dump(by_alias=True, mode="json", exclude_none=True)
request = types.DiscoverRequest(
params=types.RequestParams(
_meta={
@ -704,6 +874,7 @@ class ClientSession:
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: Literal[False] = False,
allow_claimed: Literal[False] = False,
) -> types.CallToolResult: ...
@overload
@ -718,8 +889,39 @@ class ClientSession:
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: bool,
allow_claimed: Literal[False] = False,
) -> types.CallToolResult | types.InputRequiredResult: ...
@overload
async def call_tool(
self,
name: str,
arguments: dict[str, Any] | None = None,
read_timeout_seconds: float | None = None,
progress_callback: ProgressFnT | None = None,
*,
input_responses: types.InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: Literal[False] = False,
allow_claimed: bool,
) -> types.CallToolResult | types.Result: ...
@overload
async def call_tool(
self,
name: str,
arguments: dict[str, Any] | None = None,
read_timeout_seconds: float | None = None,
progress_callback: ProgressFnT | None = None,
*,
input_responses: types.InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: bool,
allow_claimed: bool,
) -> types.CallToolResult | types.InputRequiredResult | types.Result: ...
async def call_tool(
self,
name: str,
@ -731,7 +933,8 @@ class ClientSession:
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: bool = False,
) -> types.CallToolResult | types.InputRequiredResult:
allow_claimed: bool = False,
) -> types.CallToolResult | types.InputRequiredResult | types.Result:
"""Send a tools/call request with optional progress callback support.
On a modern (2026-07-28) connection, arguments annotated with `x-mcp-header`
@ -745,10 +948,13 @@ class ClientSession:
allow_input_required: When ``False`` (default), an `InputRequiredResult`
from the server raises `RuntimeError`; when ``True``, it is returned
so the caller can resolve the requests and retry.
allow_claimed: When `False` (default), a claimed extension result raises
`UnexpectedClaimedResult`; when `True`, the parsed claim model is returned.
Raises:
RuntimeError: If the server returns an `InputRequiredResult` and
``allow_input_required`` is ``False``.
UnexpectedClaimedResult: Claimed result with `allow_claimed` False; carries the parsed value.
"""
result = await self.send_request(
types.CallToolRequest(
@ -760,16 +966,19 @@ class ClientSession:
_meta=meta,
),
),
_CallToolResultAdapter,
self._call_tool_adapter,
request_read_timeout_seconds=read_timeout_seconds,
progress_callback=progress_callback,
)
if isinstance(result, types.CallToolResult) and not result.is_error:
await self._validate_tool_result(name, result)
await self.validate_tool_result(name, result)
# The input_required arm stays first; a claimed shape is terminal for the multi-round-trip driver.
if isinstance(result, types.InputRequiredResult) and not allow_input_required:
raise _input_required_unexpected("call_tool")
if not isinstance(result, types.CallToolResult | types.InputRequiredResult) and not allow_claimed:
raise UnexpectedClaimedResult(result)
return result
def _resolve_param_headers(self, name: str, arguments: Mapping[str, Any]) -> dict[str, str]:
@ -779,8 +988,12 @@ class ClientSession:
return {}
return mcp_param_headers(header_map, arguments)
async def _validate_tool_result(self, name: str, result: types.CallToolResult) -> None:
"""Validate the structured content of a tool result against its output schema."""
async def validate_tool_result(self, name: str, result: types.CallToolResult) -> None:
"""Revalidate a `CallToolResult` against the tool's declared output schema.
Raises:
RuntimeError: Structured content is missing or does not conform to the schema.
"""
if name not in self._tool_output_schemas:
# refresh output schema cache
await self.list_tools()
@ -970,7 +1183,7 @@ class ClientSession:
ctx = ClientRequestContext(
session=self, request_id=dctx.request_id, meta=request.params.meta if request.params else None
)
response = await self._dispatch_input_request(ctx, request)
response = await self.dispatch_input_request(ctx, request)
client_response = ClientResponse.validate_python(response)
if isinstance(client_response, types.ErrorData):
raise MCPError.from_error_data(client_response)
@ -982,16 +1195,18 @@ class ClientSession:
raise MCPError(code=INTERNAL_ERROR, message="Client callback returned an invalid result") from None
return dumped
async def _dispatch_input_request(
self, ctx: ClientRequestContext, req: types.InputRequest
async def dispatch_input_request(
self, ctx: ClientRequestContext, request: types.InputRequest
) -> types.InputResponse | types.ErrorData:
"""Route a server-initiated input request to the matching constructor callback.
"""Route an input request through the client's callback table.
Shared by the legacy server→client RPC path (`_on_request`) and the
2026-07-28 multi-round-trip driver, which dispatches the embedded
`InputRequiredResult.input_requests` through the same callbacks.
Returns the callback's `InputResponse`, or `ErrorData` when the callback declines.
"""
match req:
match request:
case types.CreateMessageRequest(params=p):
return await self._sampling_callback(ctx, p)
case types.ElicitRequest(params=p):
@ -1008,7 +1223,26 @@ class ClientSession:
try:
notification = cast(types.ServerNotification, _methods.parse_server_notification(method, version, params))
except KeyError:
logger.debug("dropped %r: not defined at %s", method, version)
# Only methods unknown to the negotiated version's core tables reach the bindings.
binding = self._notification_bindings.get(method)
if binding is None:
logger.debug("dropped %r: not defined at %s", method, version)
return
try:
bound_params = binding.params_type.model_validate(params or {})
except ValidationError:
logger.warning("Failed to validate notification: %s", method, exc_info=True)
return
send, receive = self._binding_queues[method]
try:
# Must not await: DirectDispatcher calls _on_notify inline; blocking deadlocks in-process servers.
send.send_nowait(bound_params)
except anyio.WouldBlock:
# Evict the oldest event; no checkpoint since the failed send,
# so the buffer is still full and the retry cannot block.
receive.receive_nowait()
logger.warning("notification queue for %r is full; dropped the oldest event", method)
send.send_nowait(bound_params)
return
except ValidationError:
logger.warning("Failed to validate notification: %s", method, exc_info=True)

24
src/mcp/server/extension.py
View file

@ -19,7 +19,6 @@ extensions remain importable without constructing an `MCPServer`.
from __future__ import annotations
import re
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
@ -30,31 +29,14 @@ from pydantic import BaseModel
from mcp.server.context import CallNext, HandlerResult, ServerMiddleware, ServerRequestContext
# Re-exported from `mcp.shared.extension` (shared with the client surface) for existing importers.
from mcp.shared.extension import validate_extension_identifier as validate_extension_identifier
if TYPE_CHECKING:
from mcp.server.mcpserver.resources import Resource
RequestHandler = Callable[[ServerRequestContext[Any, Any], Any], Awaitable[HandlerResult]]
# Extension identifiers follow the `_meta` key grammar with a mandatory prefix
# (SEP-2133 / basic/index.mdx): dot-separated labels, each starting with a
# letter and ending with a letter or digit (hyphens interior), then `/`, then a
# name that starts and ends alphanumeric (`.`/`_`/`-` interior).
_LABEL = r"[A-Za-z](?:[A-Za-z0-9-]*[A-Za-z0-9])?"
_NAME = r"[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?"
_IDENTIFIER_RE = re.compile(rf"{_LABEL}(?:\.{_LABEL})*/{_NAME}")
def validate_extension_identifier(identifier: Any, *, owner: str) -> None:
"""Raise `TypeError` unless `identifier` is a `vendor-prefix/name` string.
SEP-2133 requires extension identifiers to carry a reverse-DNS prefix.
"""
if not isinstance(identifier, str) or not _IDENTIFIER_RE.fullmatch(identifier):
raise TypeError(
f"{owner}.identifier must be a `vendor-prefix/name` string "
f"(reverse-DNS prefix required), got {identifier!r}"
)
@dataclass(frozen=True)
class ToolBinding:

28
src/mcp/shared/extension.py Normal file
View file

@ -0,0 +1,28 @@
"""Extension-identifier grammar shared by the server and client extension surfaces."""
from __future__ import annotations
import re
from typing import Any
__all__ = ["validate_extension_identifier"]
# Extension identifiers follow the `_meta` key grammar with a mandatory prefix
# (SEP-2133 / basic/index.mdx): dot-separated labels, each starting with a
# letter and ending with a letter or digit (hyphens interior), then `/`, then a
# name that starts and ends alphanumeric (`.`/`_`/`-` interior).
_LABEL = r"[A-Za-z](?:[A-Za-z0-9-]*[A-Za-z0-9])?"
_NAME = r"[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?"
_IDENTIFIER_RE = re.compile(rf"{_LABEL}(?:\.{_LABEL})*/{_NAME}")
def validate_extension_identifier(identifier: Any, *, owner: str) -> None:
"""Raise `TypeError` unless `identifier` is a `vendor-prefix/name` string.
SEP-2133 requires extension identifiers to carry a reverse-DNS prefix.
"""
if not isinstance(identifier, str) or not _IDENTIFIER_RE.fullmatch(identifier):
raise TypeError(
f"{owner}.identifier must be a `vendor-prefix/name` string "
f"(reverse-DNS prefix required), got {identifier!r}"
)