Add a client extension API (#3034)
This commit is contained in:
parent
7322ca56f4
commit
4df609119f
4 changed files with 3410 additions and 180 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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=(
|
||||
|
|
|
|||
174
tests/interaction/mcpserver/test_extensions.py
Normal file
174
tests/interaction/mcpserver/test_extensions.py
Normal 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]}
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Reference in a new issue