modelcontextprotocol--python-sdk/tests/shared/test_peer.py

"""Tests for `ClientPeer`.

Each typed method is tested by wrapping a `DirectDispatcher` in `ClientPeer`,
calling it, and asserting (a) the right method+params went out and (b) the
return value is the typed result model.
"""

from collections.abc import Mapping
from typing import Any

import anyio
import pytest
from mcp_types import (
    CreateMessageResult,
    CreateMessageResultWithTools,
    ElicitResult,
    ListRootsResult,
    SamplingMessage,
    TextContent,
    Tool,
    ToolChoice,
)

from mcp.shared.dispatcher import DispatchContext
from mcp.shared.exceptions import MCPDeprecationWarning
from mcp.shared.peer import ClientPeer, dump_params
from mcp.shared.transport_context import TransportContext

from .conftest import direct_pair
from .test_dispatcher import running_pair

DCtx = DispatchContext[TransportContext]


class _Recorder:
    def __init__(self, result: dict[str, Any]) -> None:
        self.result = result
        self.seen: list[tuple[str, Mapping[str, Any] | None]] = []

    async def on_request(self, ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
        self.seen.append((method, params))
        return self.result


@pytest.mark.anyio
async def test_peer_sample_sends_create_message_and_returns_typed_result():
    rec = _Recorder({"role": "assistant", "content": {"type": "text", "text": "hi"}, "model": "m"})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.sample(  # pyright: ignore[reportDeprecated]
                [SamplingMessage(role="user", content=TextContent(type="text", text="hello"))],
                max_tokens=10,
            )
        method, params = rec.seen[0]
        assert method == "sampling/createMessage"
        assert params is not None and params["maxTokens"] == 10
        assert isinstance(result, CreateMessageResult)
        assert result.model == "m"


@pytest.mark.anyio
async def test_peer_sample_validates_result_alias_only():
    """Peer results validate alias-only; a snake_case key from the wire is
    ignored as extra, not populated by Python field name."""
    snake = {"role": "assistant", "content": {"type": "text", "text": "x"}, "model": "m", "stop_reason": "endTurn"}
    rec = _Recorder(snake)
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.sample(  # pyright: ignore[reportDeprecated]
                [SamplingMessage(role="user", content=TextContent(type="text", text="q"))], max_tokens=1
            )
        assert isinstance(result, CreateMessageResult)
        assert result.stop_reason is None


@pytest.mark.anyio
async def test_peer_sample_with_tools_returns_with_tools_result():
    rec = _Recorder({"role": "assistant", "content": [{"type": "text", "text": "x"}], "model": "m"})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.sample(  # pyright: ignore[reportDeprecated]
                [SamplingMessage(role="user", content=TextContent(type="text", text="q"))],
                max_tokens=5,
                tools=[Tool(name="t", input_schema={"type": "object"})],
            )
        method, params = rec.seen[0]
        assert method == "sampling/createMessage"
        assert params is not None and params["tools"][0]["name"] == "t"
        assert isinstance(result, CreateMessageResultWithTools)


@pytest.mark.anyio
async def test_peer_sample_with_tool_choice_only_returns_with_tools_result():
    # tool_choice alone is tools-mode: the answer may carry array content.
    rec = _Recorder({"role": "assistant", "content": [{"type": "text", "text": "x"}], "model": "m"})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.sample(  # pyright: ignore[reportDeprecated]
                [SamplingMessage(role="user", content=TextContent(type="text", text="q"))],
                max_tokens=5,
                tool_choice=ToolChoice(mode="none"),
            )
        assert isinstance(result, CreateMessageResultWithTools)


@pytest.mark.anyio
async def test_peer_elicit_form_sends_elicitation_create_with_form_params():
    rec = _Recorder({"action": "accept", "content": {"name": "Max"}})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.elicit_form("Your name?", requested_schema={"type": "object", "properties": {}})
        method, params = rec.seen[0]
        assert method == "elicitation/create"
        assert params is not None and params["mode"] == "form"
        assert params["message"] == "Your name?"
        assert isinstance(result, ElicitResult)


@pytest.mark.anyio
async def test_peer_elicit_url_sends_elicitation_create_with_url_params():
    rec = _Recorder({"action": "accept"})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.elicit_url("Auth needed", url="https://example.com/auth", elicitation_id="e1")
        method, params = rec.seen[0]
        assert method == "elicitation/create"
        assert params is not None and params["mode"] == "url"
        assert params["url"] == "https://example.com/auth"
        assert isinstance(result, ElicitResult)


@pytest.mark.anyio
async def test_peer_list_roots_sends_roots_list_and_returns_typed_result():
    rec = _Recorder({"roots": [{"uri": "file:///workspace"}]})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.list_roots()  # pyright: ignore[reportDeprecated]
        method, _ = rec.seen[0]
        assert method == "roots/list"
        assert isinstance(result, ListRootsResult)
        assert len(result.roots) == 1
        assert str(result.roots[0].uri) == "file:///workspace"


@pytest.mark.anyio
async def test_peer_list_roots_with_meta_sends_meta_in_params():
    rec = _Recorder({"roots": []})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            await peer.list_roots(meta={"traceId": "t1"})  # pyright: ignore[reportDeprecated]
        method, params = rec.seen[0]
        assert method == "roots/list"
        assert params == {"_meta": {"traceId": "t1"}}


@pytest.mark.anyio
async def test_peer_list_roots_is_deprecated_sep_2577():
    rec = _Recorder({"roots": []})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            with pytest.warns(
                MCPDeprecationWarning, match=r"The roots capability is deprecated as of 2026-07-28 \(SEP-2577\)\."
            ):
                await peer.list_roots()  # pyright: ignore[reportDeprecated]
        assert rec.seen[0][0] == "roots/list"


def test_dump_params_merges_meta_over_model_meta():
    out = dump_params(None, None)
    assert out is None
    out = dump_params(None, {"k": 1})
    assert out == {"_meta": {"k": 1}}


def test_dump_params_serializes_meta_by_alias():
    """`progress_token` (the Python key an inbound `ctx.meta` carries) emits
    its wire alias `progressToken`; undeclared keys pass through unchanged."""
    out = dump_params(None, {"progress_token": 7, "traceparent": "00-abc"})
    assert out == {"_meta": {"progressToken": 7, "traceparent": "00-abc"}}
    # The wire spelling is already canonical and survives as-is.
    out = dump_params(None, {"progressToken": "tok"})
    assert out == {"_meta": {"progressToken": "tok"}}


@pytest.mark.anyio
async def test_peer_notify_forwards_to_wrapped_outbound():
    sent: list[tuple[str, Mapping[str, Any] | None]] = []

    class _Out:
        async def send_raw_request(
            self, method: str, params: Mapping[str, Any] | None, opts: Any = None
        ) -> dict[str, Any]:
            raise NotImplementedError

        async def notify(self, method: str, params: Mapping[str, Any] | None, opts: Any = None) -> None:
            sent.append((method, params))

    await ClientPeer(_Out()).notify("n", {"x": 1})
    assert sent == [("n", {"x": 1})]


@pytest.mark.anyio
async def test_peer_ping_sends_ping_and_returns_none():
    rec = _Recorder({})
    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
        peer = ClientPeer(client)
        with anyio.fail_after(5):
            result = await peer.ping()
        method, _ = rec.seen[0]
        assert method == "ping"
        assert result is None