modelcontextprotocol--python-sdk/tests/interaction/mcpserver/test_extensions.py

"""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 anyio
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, NotificationBinding, 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._stamp import Unstamp
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, unstamped: Unstamp
) -> 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 unstamped(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]}


class ReceiptIssuedParams(types.NotificationParams):
    seq: int


class ReceiptIssued(types.Notification[ReceiptIssuedParams, Literal["notifications/receipts"]]):
    """Vendor notification the Receipts extension observes; off-spec, so only a binding sees it."""

    method: Literal["notifications/receipts"] = "notifications/receipts"
    params: ReceiptIssuedParams


class ReceiptFeed(ClientExtension):
    """Client half: observes the vendor `notifications/receipts` stream."""

    identifier = _RECEIPTS

    def __init__(self, handler: Callable[[ReceiptIssuedParams], Awaitable[None]]) -> None:
        self._handler = handler

    def notifications(self) -> Sequence[NotificationBinding[Any]]:
        return [
            NotificationBinding(method="notifications/receipts", params_type=ReceiptIssuedParams, handler=self._handler)
        ]


@requirement("extensions:client:notification-binding-delivery")
async def test_vendor_notifications_reach_the_owning_extensions_binding_in_order(connect: Connect) -> None:
    """A tool emits vendor notifications through `ctx.session.send_notification`; the declaring
    client's binding receives the validated params serially, in dispatch order."""
    received: list[int] = []
    done = anyio.Event()

    async def on_receipt(params: ReceiptIssuedParams) -> None:
        received.append(params.seq)
        if params.seq == 3:
            done.set()

    server = MCPServer("issuer")

    @server.tool()
    async def buy_many(ctx: Context) -> str:
        """Issue three receipts."""
        for seq in (1, 2, 3):
            await ctx.session.send_notification(
                ReceiptIssued(params=ReceiptIssuedParams(seq=seq)), related_request_id=ctx.request_id
            )
        return "issued"

    async with connect(server, extensions=[ReceiptFeed(on_receipt)]) as client:
        await client.call_tool("buy_many", {})
        with anyio.fail_after(5):
            await done.wait()

    assert received == [1, 2, 3]