diff --git a/fastmcp_slim/fastmcp/client/client.py b/fastmcp_slim/fastmcp/client/client.py index 2106c38c7..1d2bfaea5 100644 --- a/fastmcp_slim/fastmcp/client/client.py +++ b/fastmcp_slim/fastmcp/client/client.py @@ -8,7 +8,7 @@ import secrets import ssl import uuid import weakref -from collections.abc import Callable, Coroutine +from collections.abc import Callable, Coroutine, Mapping, Sequence from contextlib import AsyncExitStack, asynccontextmanager, suppress from dataclasses import dataclass, field from pathlib import Path @@ -31,7 +31,11 @@ from mcp.client.caching import ( ClientResponseCache, InMemoryResponseCacheStore, ) -from mcp.client.extension import NotificationBinding +from mcp.client.extension import ( + ClientExtension, + NotificationBinding, + ResultClaim, +) from mcp.client.session import ClientRequestContext, MessageHandlerFnT from mcp_types import ( GetTaskResult, @@ -148,6 +152,86 @@ def _synthesize_discover(protocol_version: str) -> mcp_types.DiscoverResult: ) +@dataclass +class _FoldedExtensions: + """`Client(extensions=...)` folded into the shapes `ClientSession` consumes. + + `ad` maps each extension identifier to its advertised settings (the SEP-2133 + capability ad), `claims` maps each identifier to its `ResultClaim`s, and + `bindings` is the flat list of `NotificationBinding`s the extensions observe. + """ + + ad: dict[str, dict[str, Any]] + claims: dict[str, tuple[ResultClaim[Any], ...]] + bindings: list[NotificationBinding[Any]] + + +def _fold_extensions( + extensions: Sequence[ClientExtension] | None, +) -> _FoldedExtensions: + """Decompose `ClientExtension` instances into `ClientSession` kwargs. + + Mirrors the SDK Client's own folding, using only the public `ClientExtension` + surface (`settings()`, `claims()`, `notifications()`). Duplicate identifiers, + result-type tags, or notification methods across extensions raise here rather + than at session construction, naming both owners. + """ + folded = _FoldedExtensions(ad={}, claims={}, bindings=[]) + if not extensions: + return folded + if isinstance(extensions, Mapping): + raise TypeError( + "extensions= takes a sequence of ClientExtension instances; use " + "mcp.client.advertise(identifier, settings) for an advertise-only entry" + ) + 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" + ) + if identifier in folded.ad: + raise ValueError( + f"extension identifier {identifier!r} is passed more than once" + ) + folded.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 + if extension_claims: + folded.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 + folded.bindings.append(binding) + return folded + + def _evicting_message_handler( cache: ClientResponseCache, user_handler: MessageHandlerFnT | None ) -> MessageHandlerFnT: @@ -275,6 +359,18 @@ class Client( modern-only, so a cache is inert on legacy connections. A custom `CacheConfig` store requires `target_id`, since FastMCP transports expose no server URL to derive a shared-store identity from. + extensions: Opt-in client extensions (SEP-2133), a sequence of + `mcp.client.extension.ClientExtension` instances. Each contributes its + capability advertisement, its result claims, and its notification bindings, + all of which are threaded into the underlying session. User-supplied + notification bindings compose with FastMCP's internal task-status binding + rather than replacing it. For an advertise-only entry, use + `mcp.client.advertise(identifier, settings)`. + result_claims: Additional `ResultClaim`s (SEP-2133) keyed by the identifier of + an extension already advertised through `extensions`, merged with that + extension's own claims. Rarely needed directly; prefer declaring claims on + the extension itself. Claimed shapes are modern-only and inert on a legacy + connection. Examples: ```python @@ -368,6 +464,8 @@ class Client( prior_discover: mcp_types.DiscoverResult | None = None, input_required_max_rounds: int = DEFAULT_INPUT_REQUIRED_MAX_ROUNDS, cache: CacheConfig | bool | None = None, + extensions: Sequence[ClientExtension] | None = None, + result_claims: Mapping[str, Sequence[ResultClaim[Any]]] | None = None, ) -> None: self.name = name or self.generate_name() @@ -452,6 +550,11 @@ class Client( self._response_cache, effective_message_handler ) + # Opt-in client extensions (SEP-2133) and their result claims. Retained so + # `new()` can rebuild an independent set of session kwargs per clone. + self._extensions_arg = extensions + self._result_claims_arg = result_claims + self._session_kwargs: SessionKwargs = { "sampling_callback": None, "list_roots_callback": None, @@ -459,10 +562,7 @@ class Client( "message_handler": effective_message_handler, "read_timeout_seconds": read_timeout_seconds, "client_info": client_info, - # SDK v2 does not carry `notifications/tasks/status` in any protocol - # version's core notification tables, so it is never tee'd to the - # message_handler; a binding routes it to Task objects instead. - "notification_bindings": [self._task_status_binding()], + **self._build_extension_kwargs(), } if roots is not None: @@ -684,10 +784,10 @@ class Client( ) else: new_client._session_kwargs["message_handler"] = base_handler - # Rebind the task-status notification binding so it routes to the clone. - new_client._session_kwargs["notification_bindings"] = [ - new_client._task_status_binding() - ] + # Rebuild the extension-contributed kwargs (capability ad, result claims, + # notification bindings) so the clone's task-status binding routes to the + # clone while user extensions still compose with it. + new_client._session_kwargs.update(new_client._build_extension_kwargs()) new_client.name += f":{secrets.token_hex(2)}" @@ -1160,6 +1260,37 @@ class Client( status = GetTaskResult.model_validate(params.model_dump()) task._handle_status_notification(status) + def _build_extension_kwargs(self) -> SessionKwargs: + """Session kwargs contributed by `extensions=` / `result_claims=`. + + Folds the user's `ClientExtension` instances into the capability ad, result + claims, and notification bindings the SDK `ClientSession` consumes, then + merges in any explicitly-passed `result_claims`. The internal task-status + binding is always prepended to the folded bindings so user extensions + *compose* with it rather than clobbering it; a user extension that binds the + same `notifications/tasks/status` method surfaces a duplicate-method error + from the SDK rather than silently replacing FastMCP's routing. + """ + folded = _fold_extensions(self._extensions_arg) + + claims: dict[str, tuple[ResultClaim[Any], ...]] = dict(folded.claims) + for identifier, extra in (self._result_claims_arg or {}).items(): + existing = claims.get(identifier, ()) + claims[identifier] = (*existing, *extra) + + kwargs: SessionKwargs = { + # The internal task binding must lead so user bindings extend it. + "notification_bindings": [ + self._task_status_binding(), + *folded.bindings, + ], + } + if folded.ad: + kwargs["extensions"] = folded.ad + if claims: + kwargs["result_claims"] = claims + return kwargs + def _task_status_binding(self) -> NotificationBinding[TaskStatusNotificationParams]: """Build a binding routing `notifications/tasks/status` to Task objects. diff --git a/fastmcp_slim/fastmcp/client/transports/base.py b/fastmcp_slim/fastmcp/client/transports/base.py index 87326413c..8151f8040 100644 --- a/fastmcp_slim/fastmcp/client/transports/base.py +++ b/fastmcp_slim/fastmcp/client/transports/base.py @@ -1,12 +1,12 @@ import abc import contextlib -from collections.abc import AsyncIterator, Sequence +from collections.abc import AsyncIterator, Mapping, Sequence from typing import Any, Literal, TypeVar import httpx import mcp_types from mcp import ClientSession -from mcp.client.extension import NotificationBinding +from mcp.client.extension import NotificationBinding, ResultClaim from mcp.client.session import ( ElicitationFnT, ListRootsFnT, @@ -32,6 +32,8 @@ class SessionKwargs(TypedDict, total=False): message_handler: MessageHandlerFnT | None client_info: mcp_types.Implementation | None notification_bindings: Sequence[NotificationBinding[Any]] | None + extensions: dict[str, dict[str, Any]] | None + result_claims: Mapping[str, Sequence[ResultClaim[Any]]] | None class ClientTransport(abc.ABC): diff --git a/tests/client/test_client_extensions.py b/tests/client/test_client_extensions.py new file mode 100644 index 000000000..1f2fb72ba --- /dev/null +++ b/tests/client/test_client_extensions.py @@ -0,0 +1,217 @@ +"""Tests for surfacing SEP-2133 client extensions on ``fastmcp.Client``. + +Covers that ``extensions=`` / ``result_claims=`` are folded into the underlying +``ClientSession`` kwargs on construction, that user-supplied notification +bindings *compose* with FastMCP's internal task-status binding rather than +clobbering it, and that both bindings actually fire against a live server. +""" + +import asyncio +from typing import Any, Literal + +import pytest +from mcp.client.extension import ( + ClaimContext, + ClientExtension, + NotificationBinding, + ResultClaim, +) +from mcp_types import CallToolResult, Result +from pydantic import BaseModel + +from fastmcp import FastMCP +from fastmcp.client import Client +from fastmcp.server.dependencies import get_context + +CUSTOM_METHOD = "notifications/x-test/ping" +TASK_STATUS_METHOD = "notifications/tasks/status" +EXTENSION_ID = "test.example.com/demo" + + +class PingParams(BaseModel): + value: int = 0 + + +class ClaimedResult(Result): + result_type: Literal["x-test/claimed"] + payload: str = "" + + +async def _resolve_claimed(result: ClaimedResult, ctx: ClaimContext) -> CallToolResult: + return CallToolResult(content=[]) + + +def _make_claim() -> ResultClaim[ClaimedResult]: + return ResultClaim( + result_type="x-test/claimed", + model=ClaimedResult, + resolve=_resolve_claimed, + ) + + +class _DemoExtension(ClientExtension): + """Extension contributing a settings ad, a result claim, and a binding.""" + + identifier = EXTENSION_ID + + def __init__(self, received: list[PingParams] | None = None) -> None: + self._received = received if received is not None else [] + + def settings(self) -> dict[str, Any]: + return {"enabled": True} + + def claims(self): + return (_make_claim(),) + + def notifications(self): + async def _handler(params: PingParams) -> None: + self._received.append(params) + + return ( + NotificationBinding( + method=CUSTOM_METHOD, + params_type=PingParams, + handler=_handler, + ), + ) + + +def _binding_methods(client: Client) -> list[str]: + bindings = client._session_kwargs.get("notification_bindings") or [] + return [b.method for b in bindings] + + +def test_extension_folds_into_session_kwargs(): + """A ClientExtension's ad, claim, and binding reach the session kwargs.""" + client = Client(FastMCP("srv"), extensions=[_DemoExtension()]) + + assert client._session_kwargs.get("extensions") == {EXTENSION_ID: {"enabled": True}} + result_claims = client._session_kwargs.get("result_claims") + assert result_claims is not None + assert [c.result_type for c in result_claims[EXTENSION_ID]] == ["x-test/claimed"] + + +def test_binding_composes_with_internal_task_binding(): + """User binding is appended to (not replacing) the task-status binding.""" + client = Client(FastMCP("srv"), extensions=[_DemoExtension()]) + + methods = _binding_methods(client) + assert TASK_STATUS_METHOD in methods + assert CUSTOM_METHOD in methods + # The internal task binding must lead so user bindings extend it. + assert methods[0] == TASK_STATUS_METHOD + + +def test_no_extensions_leaves_only_task_binding(): + """Without extensions, only the internal task-status binding is registered.""" + client = Client(FastMCP("srv")) + + assert _binding_methods(client) == [TASK_STATUS_METHOD] + assert "extensions" not in client._session_kwargs + assert "result_claims" not in client._session_kwargs + + +def test_new_preserves_extension_composition(): + """new() rebuilds the clone with both the task binding and user bindings.""" + client = Client(FastMCP("srv"), extensions=[_DemoExtension()]) + clone = client.new() + + methods = _binding_methods(clone) + assert methods[0] == TASK_STATUS_METHOD + assert CUSTOM_METHOD in methods + assert clone._session_kwargs.get("extensions") == {EXTENSION_ID: {"enabled": True}} + + +def test_result_claims_merge_with_extension_claims(): + """Explicit result_claims merge with an advertised extension's own claims.""" + + class ExtraClaimed(Result): + result_type: Literal["x-test/extra"] + + async def _resolve_extra(result: ExtraClaimed, ctx: ClaimContext) -> CallToolResult: + return CallToolResult(content=[]) + + extra_claim = ResultClaim( + result_type="x-test/extra", + model=ExtraClaimed, + resolve=_resolve_extra, + ) + + client = Client( + FastMCP("srv"), + extensions=[_DemoExtension()], + result_claims={EXTENSION_ID: [extra_claim]}, + ) + + result_claims = client._session_kwargs.get("result_claims") + assert result_claims is not None + tags = {c.result_type for c in result_claims[EXTENSION_ID]} + assert tags == {"x-test/claimed", "x-test/extra"} + + +async def test_user_binding_clobbering_task_method_is_rejected(): + """A user extension binding the task-status method cannot silently replace it. + + Composition means the internal task binding always leads; a user extension + that binds the same method collides with it, and the SDK session rejects the + duplicate at connect time rather than letting one silently win. + """ + + class TaskClobberExtension(ClientExtension): + identifier = "test.example.com/clobber" + + def notifications(self): + async def _handler(params: PingParams) -> None: ... + + return ( + NotificationBinding( + method=TASK_STATUS_METHOD, + params_type=PingParams, + handler=_handler, + ), + ) + + client = Client(FastMCP("srv"), extensions=[TaskClobberExtension()]) + with pytest.raises(RuntimeError, match="duplicate notification binding"): + async with client: + pass + + +async def test_both_bindings_fire_against_live_server(): + """The internal task binding and a user extension binding both fire. + + A ``task=True`` tool drives ``notifications/tasks/status`` (the internal + binding) while the same tool emits a custom notification the user extension + observes, proving the two coexist on one live connection. + """ + received: list[PingParams] = [] + mcp = FastMCP("compose-server") + + @mcp.tool + async def emit(value: int) -> int: + ctx = get_context() + # Emit a custom (non-core) notification straight onto the outbound + # channel; unknown methods route to the client's notification bindings. + await ctx.session._connection.notify(CUSTOM_METHOD, {"value": value}) + return value + + @mcp.tool(task=True) + async def background(value: int) -> int: + await asyncio.sleep(0.02) + return value * 2 + + client = Client(mcp, extensions=[_DemoExtension(received)]) + + async with client: + # The user extension binding fires on the custom notification. + await client.call_tool("emit", {"value": 21}) + # The internal task binding fires on the task-status notification. + task = await client.call_tool("background", {"value": 5}, task=True) + status = await task.wait(timeout=2.0) + # Give the custom-notification queue a moment to drain. + await asyncio.sleep(0.1) + + # Internal task binding fired: the task completed via a status notification. + assert status.status == "completed" + # User extension binding fired: it observed the custom notification. + assert [p.value for p in received] == [21]