fastmcp/tests/client/test_client_extensions.py

"""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 a second tool emits a custom notification the user extension
    observes, proving the two coexist on one live connection. Pinned to
    ``mode="legacy"`` because FastMCP task submission is a legacy-era feature.
    """
    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)], mode="legacy")

    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]