modelcontextprotocol--python-sdk/tests/server/test_session.py

"""Tests for `ServerSession`, a thin proxy over two `Outbound` channels and a `Connection`.

Stub outbounds record what reaches the wire (method, params, `CallOptions`) and
which channel it routed to, without standing up a transport.
"""

from collections.abc import Mapping
from typing import Any

import mcp_types as types
import pytest
from mcp_types import (
    ClientCapabilities,
    Implementation,
    SamplingCapability,
    SamplingToolsCapability,
)
from mcp_types.version import LATEST_HANDSHAKE_VERSION, LATEST_MODERN_VERSION
from pydantic import ValidationError

from mcp.server.connection import Connection
from mcp.server.session import ServerSession
from mcp.shared.dispatcher import CallOptions
from mcp.shared.message import ServerMessageMetadata


class StubOutbound:
    """Structural `DispatchContext[Any]` stub: records calls and returns a canned result."""

    transport: Any = None
    can_send_request: bool = True
    request_id: Any = None
    message_metadata: Any = None
    cancel_requested: Any = None

    def __init__(self, result: dict[str, Any] | None = None) -> None:
        self.requests: list[tuple[str, Mapping[str, Any] | None, CallOptions | None]] = []
        self.notifications: list[tuple[str, Mapping[str, Any] | None]] = []
        self.progress_calls: list[tuple[float, float | None, str | None]] = []
        self.result = result if result is not None else {}

    async def send_raw_request(
        self,
        method: str,
        params: Mapping[str, Any] | None,
        opts: CallOptions | None = None,
    ) -> dict[str, Any]:
        self.requests.append((method, params, opts))
        return self.result

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

    async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None:
        self.progress_calls.append((progress, total, message))


def _make_session(
    outbound: StubOutbound,
    *,
    capabilities: ClientCapabilities | None = None,
    protocol_version: str = LATEST_HANDSHAKE_VERSION,
) -> ServerSession:
    """Single-channel session: the stub is both request and standalone outbound."""
    client_info = Implementation(name="c", version="0") if capabilities is not None else None
    conn = Connection.from_envelope(protocol_version, client_info, capabilities, outbound=outbound)
    return ServerSession(outbound, conn)


def _two_channel_session(request_ch: StubOutbound, standalone_ch: StubOutbound) -> ServerSession:
    """Distinct request/standalone outbounds so routing assertions can tell the channels apart."""
    conn = Connection.from_envelope(LATEST_HANDSHAKE_VERSION, None, None, outbound=standalone_ch)
    return ServerSession(request_ch, conn)


@pytest.mark.anyio
async def test_send_request_forwards_timeout_and_progress_callback_as_call_options():
    outbound = StubOutbound(result={"roots": []})
    session = _make_session(outbound)

    async def on_progress(progress: float, total: float | None, message: str | None) -> None:
        raise NotImplementedError

    result = await session.send_request(
        types.ListRootsRequest(),
        types.ListRootsResult,
        request_read_timeout_seconds=2.5,
        progress_callback=on_progress,
    )
    assert isinstance(result, types.ListRootsResult)
    method, _params, opts = outbound.requests[0]
    assert method == "roots/list"
    assert opts == {"timeout": 2.5, "on_progress": on_progress}


@pytest.mark.anyio
async def test_send_request_omits_call_options_when_none_given():
    outbound = StubOutbound(result={"roots": []})
    session = _make_session(outbound)
    await session.send_request(types.ListRootsRequest(), types.ListRootsResult)
    _method, _params, opts = outbound.requests[0]
    assert opts is None


@pytest.mark.anyio
async def test_send_request_timeout_zero_is_forwarded():
    """0 is a real timeout (`anyio.fail_after(0)` semantics); only `None` means no timeout."""
    outbound = StubOutbound(result={})
    session = _make_session(outbound)
    await session.send_request(types.PingRequest(), types.EmptyResult, request_read_timeout_seconds=0.0)
    assert outbound.requests[0][2] == {"timeout": 0.0}


@pytest.mark.anyio
async def test_send_request_without_related_id_routes_to_standalone_channel():
    """SDK-defined: no `related_request_id` routes onto the connection's standalone channel."""
    request_ch = StubOutbound()
    standalone_ch = StubOutbound(result={"roots": []})
    session = _two_channel_session(request_ch, standalone_ch)
    await session.send_request(types.ListRootsRequest(), types.ListRootsResult)
    assert request_ch.requests == []
    assert standalone_ch.requests[0][0] == "roots/list"


@pytest.mark.anyio
async def test_send_request_with_related_id_routes_to_request_channel():
    """SDK-defined: the request rides the originating POST's response stream over streamable HTTP."""
    request_ch = StubOutbound(result={"action": "cancel"})
    standalone_ch = StubOutbound()
    session = _two_channel_session(request_ch, standalone_ch)
    result = await session.send_request(
        types.ElicitRequest(params=types.ElicitRequestFormParams(message="q", requested_schema={})),
        types.ElicitResult,
        metadata=ServerMessageMetadata(related_request_id=7),
    )
    assert isinstance(result, types.ElicitResult)
    assert standalone_ch.requests == []
    assert request_ch.requests[0][0] == "elicitation/create"


@pytest.mark.anyio
async def test_send_notification_routes_by_related_request_id():
    """SDK-defined: notifications pick their channel by `related_request_id` exactly like requests."""
    request_ch = StubOutbound()
    standalone_ch = StubOutbound()
    session = _two_channel_session(request_ch, standalone_ch)
    await session.send_tool_list_changed()
    await session.send_progress_notification("tok", 0.5, related_request_id="req-1")
    assert [m for m, _ in standalone_ch.notifications] == ["notifications/tools/list_changed"]
    assert [m for m, _ in request_ch.notifications] == ["notifications/progress"]


@pytest.mark.anyio
async def test_report_progress_delegates_to_the_request_dispatch_context():
    """Delegates to the per-request `DispatchContext.progress` seam: token gating lives in the dispatcher."""
    request_ch = StubOutbound()
    standalone_ch = StubOutbound()
    session = _two_channel_session(request_ch, standalone_ch)
    await session.report_progress(0.5, total=1.0, message="halfway")
    assert request_ch.progress_calls == [(0.5, 1.0, "halfway")]
    assert standalone_ch.progress_calls == []
    assert request_ch.notifications == []


@pytest.mark.anyio
async def test_send_request_validates_the_client_result_against_the_surface_schema():
    """The surface gate rejects results the caller's `result_type` would accept."""
    session = _make_session(StubOutbound(result={"roots": "nope"}))
    with pytest.raises(ValidationError):
        await session.send_request(types.ListRootsRequest(), types.EmptyResult)


@pytest.mark.anyio
async def test_send_request_passes_a_spec_valid_client_result():
    session = _make_session(StubOutbound(result={"roots": [{"uri": "file:///ws"}]}))
    result = await session.send_request(types.ListRootsRequest(), types.ListRootsResult)
    assert isinstance(result, types.ListRootsResult)
    assert str(result.roots[0].uri) == "file:///ws"


@pytest.mark.anyio
async def test_send_request_skips_the_surface_gate_when_method_absent_at_version():
    session = _make_session(StubOutbound(result={}), protocol_version=LATEST_MODERN_VERSION)
    result = await session.send_request(types.PingRequest(), types.EmptyResult)
    assert isinstance(result, types.EmptyResult)


@pytest.mark.anyio
async def test_send_request_validates_result_alias_only():
    """A snake_case wire key is ignored as extra, not populated by Python field name."""
    snake = {"role": "assistant", "content": {"type": "text", "text": "x"}, "model": "m", "stop_reason": "endTurn"}
    session = _make_session(StubOutbound(result=snake))
    request = types.CreateMessageRequest(params=types.CreateMessageRequestParams(messages=[], max_tokens=1))
    result = await session.send_request(request, types.CreateMessageResult)
    assert result.stop_reason is None


@pytest.mark.anyio
async def test_create_message_with_tools_returns_with_tools_result():
    outbound = StubOutbound(result={"role": "assistant", "content": [{"type": "text", "text": "ok"}], "model": "m"})
    session = _make_session(
        outbound, capabilities=ClientCapabilities(sampling=SamplingCapability(tools=SamplingToolsCapability()))
    )
    result = await session.create_message(  # pyright: ignore[reportDeprecated]
        messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="hi"))],
        max_tokens=10,
        tools=[types.Tool(name="t", input_schema={"type": "object"})],
    )
    assert isinstance(result, types.CreateMessageResultWithTools)
    method, params, _opts = outbound.requests[0]
    assert method == "sampling/createMessage"
    assert params is not None and params["tools"][0]["name"] == "t"


def test_check_client_capability_delegates_to_connection():
    outbound = StubOutbound()
    session = _make_session(outbound, capabilities=ClientCapabilities(sampling=SamplingCapability()))
    assert session.check_client_capability(ClientCapabilities(sampling=SamplingCapability())) is True
    assert session.check_client_capability(ClientCapabilities(experimental={"x": {}})) is False


def test_protocol_version_proxies_connection():
    """SDK-defined: `session.protocol_version` reads through to the held `Connection`."""
    _ARBITRARY_VERSION = "sentinel-version"  # identity-only: any string the connection holds
    conn = Connection.from_envelope(_ARBITRARY_VERSION, None, None)
    session = ServerSession(StubOutbound(), conn)
    assert session.protocol_version == _ARBITRARY_VERSION
    assert session.client_params is None