feat: PeerMixin and Peer wrapper

maxisbey · maxisbey · commit 870cb088c768 · 2026-04-16T21:53:11.000Z
PeerMixin defines the typed server-to-client request methods (sample with
overloads, elicit_form, elicit_url, list_roots, ping) once. Each method
constrains `self: Outbound` so any class with send_request/notify can mix it
in — pyright checks the host structurally at the call site. The mixin does no
capability gating; that's the host's send_request's job.

Peer is a trivial standalone wrapper for when you have a bare Outbound (e.g.
a dispatcher) and want the typed sugar without writing your own host class.

6 tests over DirectDispatcher, 0.03s.
diff --git a/src/mcp/shared/peer.py b/src/mcp/shared/peer.py
@@ -0,0 +1,194 @@
+"""Typed MCP request sugar over an `Outbound`.
+
+`PeerMixin` defines the server-to-client request methods (sampling, elicitation,
+roots, ping) once. Any class that satisfies `Outbound` (i.e. has `send_request`
+and `notify`) can mix it in and get the typed methods for free — `Context`,
+`Connection`, `Client`, or the bare `Peer` wrapper below.
+
+The mixin does no capability gating: it builds the params, calls
+``self.send_request(method, params)``, and parses the result into the typed
+model. Gating (and `NoBackChannelError`) is the host's `send_request`'s job.
+"""
+
+from collections.abc import Mapping
+from typing import Any, overload
+
+from pydantic import BaseModel
+
+from mcp.shared.dispatcher import CallOptions, Outbound
+from mcp.types import (
+    CreateMessageRequestParams,
+    CreateMessageResult,
+    CreateMessageResultWithTools,
+    ElicitRequestedSchema,
+    ElicitRequestFormParams,
+    ElicitRequestURLParams,
+    ElicitResult,
+    IncludeContext,
+    ListRootsResult,
+    ModelPreferences,
+    SamplingMessage,
+    Tool,
+    ToolChoice,
+)
+
+__all__ = ["Peer", "PeerMixin"]
+
+
+def _dump(model: BaseModel) -> dict[str, Any]:
+    return model.model_dump(by_alias=True, mode="json", exclude_none=True)
+
+
+class PeerMixin:
+    """Typed server-to-client request methods.
+
+    Each method constrains ``self`` to `Outbound` so the mixin can be applied
+    to anything with ``send_request``/``notify`` — pyright checks the host
+    class structurally at the call site.
+    """
+
+    @overload
+    async def sample(
+        self: Outbound,
+        messages: list[SamplingMessage],
+        *,
+        max_tokens: int,
+        system_prompt: str | None = None,
+        include_context: IncludeContext | None = None,
+        temperature: float | None = None,
+        stop_sequences: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        model_preferences: ModelPreferences | None = None,
+        tools: None = None,
+        tool_choice: ToolChoice | None = None,
+        opts: CallOptions | None = None,
+    ) -> CreateMessageResult: ...
+    @overload
+    async def sample(
+        self: Outbound,
+        messages: list[SamplingMessage],
+        *,
+        max_tokens: int,
+        system_prompt: str | None = None,
+        include_context: IncludeContext | None = None,
+        temperature: float | None = None,
+        stop_sequences: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        model_preferences: ModelPreferences | None = None,
+        tools: list[Tool],
+        tool_choice: ToolChoice | None = None,
+        opts: CallOptions | None = None,
+    ) -> CreateMessageResultWithTools: ...
+    async def sample(
+        self: Outbound,
+        messages: list[SamplingMessage],
+        *,
+        max_tokens: int,
+        system_prompt: str | None = None,
+        include_context: IncludeContext | None = None,
+        temperature: float | None = None,
+        stop_sequences: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        model_preferences: ModelPreferences | None = None,
+        tools: list[Tool] | None = None,
+        tool_choice: ToolChoice | None = None,
+        opts: CallOptions | None = None,
+    ) -> CreateMessageResult | CreateMessageResultWithTools:
+        """Send a ``sampling/createMessage`` request to the peer.
+
+        Raises:
+            MCPError: The peer responded with an error.
+            NoBackChannelError: The host's transport context has no
+                back-channel for server-initiated requests.
+        """
+        params = CreateMessageRequestParams(
+            messages=messages,
+            system_prompt=system_prompt,
+            include_context=include_context,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            stop_sequences=stop_sequences,
+            metadata=metadata,
+            model_preferences=model_preferences,
+            tools=tools,
+            tool_choice=tool_choice,
+        )
+        result = await self.send_request("sampling/createMessage", _dump(params), opts)
+        if tools is not None:
+            return CreateMessageResultWithTools.model_validate(result)
+        return CreateMessageResult.model_validate(result)
+
+    async def elicit_form(
+        self: Outbound,
+        message: str,
+        requested_schema: ElicitRequestedSchema,
+        opts: CallOptions | None = None,
+    ) -> ElicitResult:
+        """Send a form-mode ``elicitation/create`` request.
+
+        Raises:
+            MCPError: The peer responded with an error.
+            NoBackChannelError: No back-channel for server-initiated requests.
+        """
+        params = ElicitRequestFormParams(message=message, requested_schema=requested_schema)
+        result = await self.send_request("elicitation/create", _dump(params), opts)
+        return ElicitResult.model_validate(result)
+
+    async def elicit_url(
+        self: Outbound,
+        message: str,
+        url: str,
+        elicitation_id: str,
+        opts: CallOptions | None = None,
+    ) -> ElicitResult:
+        """Send a URL-mode ``elicitation/create`` request.
+
+        Raises:
+            MCPError: The peer responded with an error.
+            NoBackChannelError: No back-channel for server-initiated requests.
+        """
+        params = ElicitRequestURLParams(message=message, url=url, elicitation_id=elicitation_id)
+        result = await self.send_request("elicitation/create", _dump(params), opts)
+        return ElicitResult.model_validate(result)
+
+    async def list_roots(self: Outbound, opts: CallOptions | None = None) -> ListRootsResult:
+        """Send a ``roots/list`` request.
+
+        Raises:
+            MCPError: The peer responded with an error.
+            NoBackChannelError: No back-channel for server-initiated requests.
+        """
+        result = await self.send_request("roots/list", None, opts)
+        return ListRootsResult.model_validate(result)
+
+    async def ping(self: Outbound, opts: CallOptions | None = None) -> None:
+        """Send a ``ping`` request and ignore the result.
+
+        Raises:
+            MCPError: The peer responded with an error.
+            NoBackChannelError: No back-channel for server-initiated requests.
+        """
+        await self.send_request("ping", None, opts)
+
+
+class Peer(PeerMixin):
+    """Standalone wrapper that gives any `Outbound` the `PeerMixin` sugar.
+
+    `Context` and `Connection` mix `PeerMixin` in directly; use `Peer` when
+    you have a bare dispatcher (or any `Outbound`) and want the typed methods
+    without writing your own host class.
+    """
+
+    def __init__(self, outbound: Outbound) -> None:
+        self._outbound = outbound
+
+    async def send_request(
+        self,
+        method: str,
+        params: Mapping[str, Any] | None,
+        opts: CallOptions | None = None,
+    ) -> dict[str, Any]:
+        return await self._outbound.send_request(method, params, opts)
+
+    async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
+        await self._outbound.notify(method, params)
diff --git a/tests/shared/test_peer.py b/tests/shared/test_peer.py
@@ -0,0 +1,128 @@
+"""Tests for `PeerMixin` and `Peer`.
+
+Each PeerMixin method is tested by wrapping a `DirectDispatcher` in `Peer`,
+calling the typed method, 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.shared.dispatcher import DispatchContext
+from mcp.shared.peer import Peer
+from mcp.shared.transport_context import TransportContext
+from mcp.types import (
+    CreateMessageResult,
+    CreateMessageResultWithTools,
+    ElicitResult,
+    ListRootsResult,
+    SamplingMessage,
+    TextContent,
+    Tool,
+)
+
+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 = Peer(client)
+        with anyio.fail_after(5):
+            result = await peer.sample(
+                [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_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 = Peer(client)
+        with anyio.fail_after(5):
+            result = await peer.sample(
+                [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_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 = Peer(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 = Peer(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 = Peer(client)
+        with anyio.fail_after(5):
+            result = await peer.list_roots()
+    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_ping_sends_ping_and_returns_none():
+    rec = _Recorder({})
+    async with running_pair(direct_pair, server_on_request=rec.on_request) as (client, *_):
+        peer = Peer(client)
+        with anyio.fail_after(5):
+            result = await peer.ping()
+    method, _ = rec.seen[0]
+    assert method == "ping"
+    assert result is None
