Add a client extension API (#3034)

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
4 changed files with 3410 additions and 180 deletions

View file

@ -7,7 +7,7 @@ server's real Starlette app through the in-process streaming bridge, so the full
(session ids, SSE encoding, session management) runs with no sockets, threads, or subprocesses.
"""
from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from functools import partial
from typing import Any, Protocol
@ -30,6 +30,7 @@ from starlette.responses import Response
from starlette.routing import Mount, Route
from mcp.client.client import Client
from mcp.client.extension import ClientExtension
from mcp.client.session import ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
from mcp.client.sse import sse_client
from mcp.client.streamable_http import streamable_http_client
@ -70,6 +71,7 @@ class Connect(Protocol):
message_handler: MessageHandlerFnT | None = None,
client_info: Implementation | None = None,
elicitation_callback: ElicitationFnT | None = None,
extensions: Sequence[ClientExtension] | None = None,
spec_version: str = LATEST_HANDSHAKE_VERSION,
) -> AbstractAsyncContextManager[Client]: ...
@ -85,6 +87,7 @@ async def connect_in_memory(
message_handler: MessageHandlerFnT | None = None,
client_info: Implementation | None = None,
elicitation_callback: ElicitationFnT | None = None,
extensions: Sequence[ClientExtension] | None = None,
spec_version: str = LATEST_HANDSHAKE_VERSION,
) -> AsyncIterator[Client]:
"""Yield a Client connected to the server over the in-memory transport.
@ -103,6 +106,7 @@ async def connect_in_memory(
message_handler=message_handler,
client_info=client_info,
elicitation_callback=elicitation_callback,
extensions=extensions,
) as client:
yield client
@ -122,6 +126,7 @@ async def connect_over_streamable_http(
message_handler: MessageHandlerFnT | None = None,
client_info: Implementation | None = None,
elicitation_callback: ElicitationFnT | None = None,
extensions: Sequence[ClientExtension] | None = None,
spec_version: str = LATEST_HANDSHAKE_VERSION,
) -> AsyncIterator[Client]:
"""Yield a Client connected to the server's streamable HTTP app, entirely in process.
@ -156,6 +161,7 @@ async def connect_over_streamable_http(
message_handler=message_handler,
client_info=client_info,
elicitation_callback=elicitation_callback,
extensions=extensions,
) as client,
):
yield client
@ -357,6 +363,7 @@ async def connect_over_sse(
message_handler: MessageHandlerFnT | None = None,
client_info: Implementation | None = None,
elicitation_callback: ElicitationFnT | None = None,
extensions: Sequence[ClientExtension] | None = None,
spec_version: str = LATEST_HANDSHAKE_VERSION,
) -> AsyncIterator[Client]:
"""Yield a Client connected to the server's legacy SSE transport, entirely in process."""
@ -390,5 +397,6 @@ async def connect_over_sse(
message_handler=message_handler,
client_info=client_info,
elicitation_callback=elicitation_callback,
extensions=extensions,
) as client:
yield client

View file

@ -2384,6 +2384,69 @@ REQUIREMENTS: dict[str, Requirement] = {
),
),
# ═══════════════════════════════════════════════════════════════════════════
# Extensions (SEP-2133): client-side result claims and the capability ad
# ═══════════════════════════════════════════════════════════════════════════
"extensions:client:claimed-result-resolved": Requirement(
source=f"{SPEC_2026_BASE_URL}/basic#resulttype",
behavior=(
"A tools/call answered with an extension-claimed resultType is finished by the owning "
"ClientExtension's claim resolver, and Client.call_tool returns the resolver's ordinary "
"CallToolResult. The resolver may send follow-up requests through the session it is handed."
),
added_in="2026-07-28",
),
"extensions:client:claimed-result-undeclared-invalid": Requirement(
source=f"{SPEC_2026_BASE_URL}/basic#resulttype",
behavior=(
"A resultType unrecognized by the client is invalid: a claimed shape delivered to a client that "
"did not construct the owning extension fails result validation (the supported set is core plus "
"declared claims, never more)."
),
added_in="2026-07-28",
note=(
"Known leniency: the monolith result surface still accepts an unknown tag when the payload "
"also parses as a complete core result (open result_type, extras ignored). Rejecting tags "
"outside core plus active claims is a tracked follow-up ruling."
),
),
"extensions:client:capability-ad:gates-server-behaviour": Requirement(
source=f"{SPEC_2026_BASE_URL}/basic#resulttype",
behavior=(
"The per-request _meta capability ad carries each declared extension's identifier and settings, "
"and is what entitles the server to substitute that extension's claimed shapes: a server "
"extension gating on the ad sees the declared settings, and refuses a non-declaring client with "
"-32021 (missing required client capability)."
),
added_in="2026-07-28",
),
"extensions:client:capability-ad:legacy-omits-claimed": Requirement(
source=f"{SPEC_2026_BASE_URL}/basic#resulttype",
behavior=(
"On a legacy connection no claim can activate, and the initialize capability ad omits "
"claim-bearing identifiers in the same breath (claim-less identifiers still advertise), so the "
"client never advertises an extension whose claimed shapes it would reject."
),
removed_in="2026-07-28",
note=(
"The legacy-era half of the ad/claims coupling: only a handshake connection can exhibit it, so "
"the version window ends where the modern era begins."
),
arm_exclusions=(ArmExclusion(reason="requires-session", transport="streamable-http-stateless"),),
),
"extensions:client:notification-binding-delivery": Requirement(
source=f"{SPEC_2026_BASE_URL}/basic#resulttype",
behavior=(
"A vendor server notification bound by a ClientExtension's NotificationBinding is validated "
"against the binding's params type and delivered to its handler serially, in dispatch order."
),
added_in="2026-07-28",
deferred=(
"Covered at session tier by tests/client/test_session_notification_bindings.py: no public "
"server-side surface emits vendor-method notifications (ServerNotification is a closed union), "
"and HTTP-modern arrival additionally needs the subscriptions/listen client runtime."
),
),
# ═══════════════════════════════════════════════════════════════════════════
# Transports (in-suite coverage)
# ═══════════════════════════════════════════════════════════════════════════
"transport:streamable-http:stateful": Requirement(
@ -3341,6 +3404,19 @@ REQUIREMENTS: dict[str, Requirement] = {
transports=("streamable-http",),
note="Only observable over streamable HTTP: headers are derived from the cached tool schema at the seam.",
),
"client-transport:http:vendor-name-param-header": Requirement(
source="sdk",
behavior=(
"A vendor request type declaring name_param mirrors that wire-params key into the Mcp-Name "
"header of its outgoing HTTP request, with no client-side registration of the method."
),
added_in="2026-07-28",
transports=("streamable-http",),
note=(
"SDK mechanism honouring the per-extension Mcp-Name requirements (e.g. SEP-2663 mandates the "
"header for tasks/*); only observable over streamable HTTP, where headers exist."
),
),
"client-transport:http:stateless-ignores-session-id": Requirement(
source=f"{SPEC_2026_BASE_URL}/basic/transports#stateless-request-headers",
behavior=(

View file

@ -0,0 +1,174 @@
"""Client extensions (SEP-2133) over the full client-server loop: a server extension
substitutes a claimed `tools/call` shape and the declaring client's `ClientExtension` resolves it."""
from collections.abc import Awaitable, Callable, Sequence
from typing import Any, Literal
import mcp_types as types
import pytest
from inline_snapshot import snapshot
from mcp_types import MISSING_REQUIRED_CLIENT_CAPABILITY, CallToolResult, Result, TextContent
from pydantic import ValidationError
from mcp import MCPError
from mcp.client import ClaimContext, ClientExtension, ResultClaim, advertise
from mcp.server.context import CallNext, HandlerResult, ServerRequestContext
from mcp.server.extension import Extension
from mcp.server.mcpserver import Context, MCPServer, require_client_extension
from tests.interaction._connect import Connect
from tests.interaction._requirements import requirement
pytestmark = pytest.mark.anyio
_RECEIPTS = "com.example/receipts"
_FLAGS = "com.example/flags"
class ReceiptResult(Result):
result_type: Literal["receipt"] = "receipt"
receipt_token: str
settings_echo: dict[str, Any] | None = None
_Resolver = Callable[[ReceiptResult, ClaimContext], Awaitable[CallToolResult]]
class Receipts(ClientExtension):
"""Client half: claims the `receipt` shape with the test's resolver and settings."""
identifier = _RECEIPTS
def __init__(self, resolve: _Resolver, settings: dict[str, Any] | None = None) -> None:
self._resolve = resolve
self._settings = {} if settings is None else settings
def settings(self) -> dict[str, Any]:
return self._settings
def claims(self) -> Sequence[ResultClaim[Any]]:
return [ResultClaim(result_type="receipt", model=ReceiptResult, resolve=self._resolve)]
class _ReceiptIssuer(Extension):
"""Server half: answers `buy` with the claimed shape; every other tool passes through."""
identifier = _RECEIPTS
async def intercept_tool_call(
self, params: types.CallToolRequestParams, ctx: ServerRequestContext[Any, Any], call_next: CallNext
) -> HandlerResult:
if params.name != "buy":
return await call_next(ctx)
return {"resultType": "receipt", "receiptToken": "r-117"}
def _receipt_shop(issuer: Extension) -> MCPServer:
server = MCPServer("shop", extensions=[issuer])
@server.tool()
def buy(item: str) -> CallToolResult:
"""Buy an item."""
raise NotImplementedError # the server extension answers `buy` before the tool runs
@server.tool()
def redeem(token: str) -> str:
"""Exchange a receipt token for the goods."""
return f"goods for {token}"
return server
@requirement("extensions:client:claimed-result-resolved")
async def test_claimed_result_is_finished_by_the_owning_extensions_resolver(connect: Connect) -> None:
"""The owning extension's claim resolver redeems the substituted `receipt` through
`ctx.session`, and `call_tool` returns the resolver's plain `CallToolResult`."""
received: list[ReceiptResult] = []
async def redeem_receipt(claimed: ReceiptResult, ctx: ClaimContext) -> CallToolResult:
received.append(claimed)
return await ctx.session.call_tool("redeem", {"token": claimed.receipt_token})
async with connect(_receipt_shop(_ReceiptIssuer()), extensions=[Receipts(redeem_receipt)]) as client:
result = await client.call_tool("buy", {"item": "lamp"})
assert [claimed.receipt_token for claimed in received] == ["r-117"]
assert result == snapshot(
CallToolResult(content=[TextContent(text="goods for r-117")], structured_content={"result": "goods for r-117"})
)
@requirement("extensions:client:claimed-result-undeclared-invalid")
async def test_claimed_shape_fails_validation_for_a_client_without_the_extension(connect: Connect) -> None:
"""Spec-mandated: an unrecognized `resultType` is invalid, so a client without the
owning extension fails to parse the claimed shape."""
async with connect(_receipt_shop(_ReceiptIssuer())) as client:
with pytest.raises(ValidationError):
await client.call_tool("buy", {"item": "lamp"})
class _SettingsEchoIssuer(Extension):
"""Server half: requires the declaring client, then echoes its declared settings."""
identifier = _RECEIPTS
async def intercept_tool_call(
self, params: types.CallToolRequestParams, ctx: ServerRequestContext[Any, Any], call_next: CallNext
) -> HandlerResult:
require_client_extension(ctx, _RECEIPTS)
client_params = ctx.session.client_params
assert client_params is not None
extensions = client_params.capabilities.extensions
assert extensions is not None
return {"resultType": "receipt", "receiptToken": "echo", "settingsEcho": extensions[_RECEIPTS]}
@requirement("extensions:client:capability-ad:gates-server-behaviour")
async def test_per_request_ad_carries_settings_and_gates_the_claimed_substitution(connect: Connect) -> None:
"""The per-request `_meta` capability ad gates the claimed substitution: declared
settings reach the resolver and a non-declaring client is refused with -32021."""
server = MCPServer("shop", extensions=[_SettingsEchoIssuer()])
@server.tool()
def buy(item: str) -> CallToolResult:
"""Buy an item."""
raise NotImplementedError # the server extension answers `buy` before the tool runs
received: list[ReceiptResult] = []
async def keep(claimed: ReceiptResult, ctx: ClaimContext) -> CallToolResult:
received.append(claimed)
return CallToolResult(content=[TextContent(text="done")])
async with connect(server, extensions=[Receipts(keep, settings={"tier": "gold"})]) as client:
result = await client.call_tool("buy", {"item": "lamp"})
assert result.content == [TextContent(text="done")]
assert [claimed.settings_echo for claimed in received] == [{"tier": "gold"}]
async with connect(server) as client:
with pytest.raises(MCPError) as exc_info:
await client.call_tool("buy", {"item": "lamp"})
assert exc_info.value.code == MISSING_REQUIRED_CLIENT_CAPABILITY
async def _unreachable_resolve(claimed: ReceiptResult, ctx: ClaimContext) -> CallToolResult:
raise NotImplementedError # no claimed shape can be delivered on a legacy wire
@requirement("extensions:client:capability-ad:legacy-omits-claimed")
async def test_legacy_ad_omits_claim_bearing_identifiers_but_keeps_claim_less_ones(connect: Connect) -> None:
"""On a legacy connection the claim-bearing identifier drops out of the initialize
capability ad while an ad-only identifier still advertises."""
server = MCPServer("introspector")
@server.tool()
def declared(ctx: Context) -> list[str]:
"""Report the extension identifiers the client advertised."""
capabilities = ctx.client_capabilities
assert capabilities is not None
return sorted(capabilities.extensions or {})
client_extensions = [Receipts(_unreachable_resolve), advertise(_FLAGS)]
async with connect(server, extensions=client_extensions) as client:
result = await client.call_tool("declared", {})
assert result.structured_content == {"result": [_FLAGS]}

View file

@ -9,7 +9,7 @@ result-envelope shape, so every assertion here is necessarily wire-level.
import json
from collections.abc import Callable
from typing import Any
from typing import Any, Literal
import anyio
import httpx
@ -30,7 +30,9 @@ from mcp_types import (
JSONRPCResponse,
ListToolsResult,
PaginatedRequestParams,
Request,
RequestParams,
Result,
ServerCapabilities,
TextContent,
Tool,
@ -551,3 +553,56 @@ async def test_modern_client_stops_mirroring_after_a_re_list_drops_the_tool() ->
before, after = tool_calls
assert before.headers.get("mcp-param-region") == "x"
assert not any(k.startswith("mcp-param-") for k in after.headers)
class _JobParams(RequestParams):
job_id: str
class _JobStatusRequest(Request[_JobParams, Literal["com.example/jobs.status"]]):
method: Literal["com.example/jobs.status"] = "com.example/jobs.status"
name_param = "jobId"
class _JobStatusResult(Result):
status: str
@requirement("client-transport:http:vendor-name-param-header")
async def test_vendor_request_with_name_param_carries_mcp_name_on_the_wire() -> None:
"""`send_request` mirrors an unregistered vendor request's `name_param` value into the
`Mcp-Name` header while the body keeps the params key unchanged."""
async def job_status(ctx: ServerRequestContext, params: _JobParams) -> _JobStatusResult:
assert params.job_id == "job-7"
return _JobStatusResult(status="running")
server = _server()
server.add_request_handler("com.example/jobs.status", _JobParams, job_status)
requests: list[httpx.Request] = []
async def on_request(request: httpx.Request) -> None:
requests.append(request)
discover = DiscoverResult(
supported_versions=[LATEST_MODERN_VERSION],
capabilities=ServerCapabilities(),
server_info=Implementation(name="srv", version="0"),
)
with anyio.fail_after(5):
async with (
mounted_app(server, on_request=on_request) as (http, _),
Client(
streamable_http_client(f"{BASE_URL}/mcp", http_client=http),
mode=LATEST_MODERN_VERSION,
prior_discover=discover,
) as client,
):
request = _JobStatusRequest(params=_JobParams(job_id="job-7"))
result = await client.session.send_request(request, _JobStatusResult)
assert result.status == "running"
[wire_request] = requests
assert wire_request.headers["mcp-name"] == "job-7"
assert json.loads(wire_request.content)["params"]["jobId"] == "job-7"