Surface extensions= and result_claims= on fastmcp.Client

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.
 

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):

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]