[v2] Dispatcher/ServerRunner receive-path swap — replaces BaseSession (#2710)

tests/shared/conftest.py

@@ -0,0 +1,61 @@
+"""Shared fixtures for `Dispatcher` contract tests.
+
+The `pair_factory` fixture parametrizes contract tests over every `Dispatcher`
+implementation, so the same behavioral assertions run against `DirectDispatcher`
+(in-memory) and `JSONRPCDispatcher` (over crossed anyio memory streams).
+"""
+
+from collections.abc import Callable
+
+import anyio
+import pytest
+
+from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair
+from mcp.shared.dispatcher import Dispatcher
+from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
+from mcp.shared.message import SessionMessage
+from mcp.shared.transport_context import TransportContext
+
+DispatcherTriple = tuple[Dispatcher[TransportContext], Dispatcher[TransportContext], Callable[[], None]]
+PairFactory = Callable[..., DispatcherTriple]
+
+
+def direct_pair(*, can_send_request: bool = True) -> DispatcherTriple:
+    client, server = create_direct_dispatcher_pair(can_send_request=can_send_request)
+
+    def close() -> None:
+        client.close()
+        server.close()
+
+    return client, server, close
+
+
+def jsonrpc_pair(*, can_send_request: bool = True) -> DispatcherTriple:
+    """Two `JSONRPCDispatcher`s wired over crossed in-memory streams."""
+    c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+
+    def builder(_meta: object) -> TransportContext:
+        return TransportContext(kind="jsonrpc", can_send_request=can_send_request)
+
+    client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send, transport_builder=builder)
+    server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send, transport_builder=builder)
+
+    def close() -> None:
+        for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
+            s.close()
+
+    return client, server, close
+
+
+@pytest.fixture(
+    params=[
+        pytest.param(direct_pair, id="direct"),
+        pytest.param(jsonrpc_pair, id="jsonrpc"),
+    ]
+)
+def pair_factory(request: pytest.FixtureRequest) -> PairFactory:
+    return request.param
+
+
+__all__ = ["PairFactory", "direct_pair", "jsonrpc_pair"]

tests/shared/test_context.py

@@ -0,0 +1,133 @@
+"""Tests for `BaseContext`.
+
+`BaseContext` is composition over a `DispatchContext` - it forwards
+`transport`/`cancel_requested`/`send_raw_request`/`notify`/`progress`
+and adds `meta`. It must satisfy `Outbound` so `ClientPeer` can wrap it.
+"""
+
+from collections.abc import Mapping
+from typing import Any
+
+import anyio
+import pytest
+
+from mcp.shared.context import BaseContext
+from mcp.shared.dispatcher import DispatchContext
+from mcp.shared.peer import ClientPeer
+from mcp.shared.transport_context import TransportContext
+
+from .conftest import direct_pair, jsonrpc_pair
+from .test_dispatcher import Recorder, echo_handlers, running_pair
+
+DCtx = DispatchContext[TransportContext]
+
+
+@pytest.mark.anyio
+async def test_base_context_forwards_transport_and_cancel_requested():
+    captured: list[BaseContext[TransportContext]] = []
+
+    async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
+        bctx = BaseContext(ctx)
+        captured.append(bctx)
+        return {}
+
+    async with running_pair(direct_pair, server_on_request=server_on_request) as (client, *_):
+        with anyio.fail_after(5):
+            await client.send_raw_request("t", None)
+        bctx = captured[0]
+        assert bctx.transport.kind == "direct"
+        assert isinstance(bctx.cancel_requested, anyio.Event)
+        assert bctx.can_send_request is True
+        assert bctx.meta is None
+
+
+@pytest.mark.anyio
+async def test_base_context_can_send_request_reflects_dispatch_context_closed_state():
+    """`can_send_request` must track the dctx, not the static transport flag,
+    so it agrees with whether `send_raw_request` would raise."""
+    captured: list[BaseContext[TransportContext]] = []
+
+    async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
+        captured.append(BaseContext(ctx))
+        return {}
+
+    async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, *_):
+        with anyio.fail_after(5):
+            await client.send_raw_request("t", None)
+        bctx = captured[0]
+        assert bctx.transport.can_send_request is True
+        assert bctx.can_send_request is False
+
+
+@pytest.mark.anyio
+async def test_base_context_send_raw_request_and_notify_forward_to_dispatch_context():
+    crec = Recorder()
+    c_req, c_notify = echo_handlers(crec)
+
+    async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
+        bctx = BaseContext(ctx)
+        sample = await bctx.send_raw_request("sampling/createMessage", {"x": 1})
+        await bctx.notify("notifications/message", {"level": "info"})
+        return {"sample": sample}
+
+    async with running_pair(
+        direct_pair,
+        server_on_request=server_on_request,
+        client_on_request=c_req,
+        client_on_notify=c_notify,
+    ) as (client, *_):
+        with anyio.fail_after(5):
+            result = await client.send_raw_request("tools/call", None)
+            await crec.notified.wait()
+        assert crec.requests == [("sampling/createMessage", {"x": 1})]
+        assert crec.notifications == [("notifications/message", {"level": "info"})]
+        assert result["sample"] == {"echoed": "sampling/createMessage", "params": {"x": 1}}
+
+
+@pytest.mark.anyio
+async def test_base_context_report_progress_invokes_caller_on_progress():
+    received: list[tuple[float, float | None, str | None]] = []
+
+    async def on_progress(progress: float, total: float | None, message: str | None) -> None:
+        received.append((progress, total, message))
+
+    async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
+        bctx = BaseContext(ctx)
+        await bctx.report_progress(0.5, total=1.0, message="halfway")
+        return {}
+
+    async with running_pair(direct_pair, server_on_request=server_on_request) as (client, *_):
+        with anyio.fail_after(5):
+            await client.send_raw_request("t", None, {"on_progress": on_progress})
+        assert received == [(0.5, 1.0, "halfway")]
+
+
+@pytest.mark.anyio
+async def test_base_context_satisfies_outbound_so_peer_mixin_works():
+    """Wrapping a BaseContext in ClientPeer proves it satisfies Outbound structurally."""
+
+    async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
+        bctx = BaseContext(ctx)
+        await ClientPeer(bctx).ping()
+        return {}
+
+    crec = Recorder()
+    c_req, c_notify = echo_handlers(crec)
+    async with running_pair(
+        direct_pair, server_on_request=server_on_request, client_on_request=c_req, client_on_notify=c_notify
+    ) as (client, *_):
+        with anyio.fail_after(5):
+            await client.send_raw_request("t", None)
+        assert crec.requests == [("ping", None)]
+
+
+@pytest.mark.anyio
+async def test_base_context_meta_holds_supplied_request_params_meta():
+    async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
+        bctx = BaseContext(ctx, meta={"progressToken": "abc"})
+        assert bctx.meta is not None and bctx.meta.get("progressToken") == "abc"
+        return {}
+
+    async with running_pair(direct_pair, server_on_request=server_on_request) as (client, *_):
+        with anyio.fail_after(5):
+            await client.send_raw_request("t", None)

tests/shared/test_dispatcher.py

@@ -0,0 +1,299 @@
+"""Behavioral tests for the Dispatcher Protocol.
+
+The contract tests are parametrized over every `Dispatcher` implementation via
+the `pair_factory` fixture (see `conftest.py`); they must pass for both
+`DirectDispatcher` and `JSONRPCDispatcher`. Implementation-specific tests pass
+a concrete factory directly.
+"""
+
+from collections.abc import AsyncIterator, Mapping
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any
+
+import anyio
+import pytest
+
+from mcp.shared.direct_dispatcher import DirectDispatcher, create_direct_dispatcher_pair
+from mcp.shared.dispatcher import DispatchContext, Dispatcher, OnNotify, OnRequest, Outbound
+from mcp.shared.exceptions import MCPError
+from mcp.shared.transport_context import TransportContext
+from mcp.types import INTERNAL_ERROR, INVALID_PARAMS, INVALID_REQUEST, REQUEST_TIMEOUT, ErrorData, Tool
+
+from .conftest import PairFactory, direct_pair
+
+
+class Recorder:
+    def __init__(self) -> None:
+        self.requests: list[tuple[str, Mapping[str, Any] | None]] = []
+        self.notifications: list[tuple[str, Mapping[str, Any] | None]] = []
+        self.contexts: list[DispatchContext[TransportContext]] = []
+        self.notified = anyio.Event()
+
+
+def echo_handlers(recorder: Recorder) -> tuple[OnRequest, OnNotify]:
+    async def on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        # Strip `_meta` so JSON-RPC and direct dispatch record identically:
+        # the JSON-RPC outbound path always attaches `_meta` (otel injection).
+        recorded = {k: v for k, v in (params or {}).items() if k != "_meta"} if params is not None else None
+        recorder.requests.append((method, recorded))
+        recorder.contexts.append(ctx)
+        return {"echoed": method, "params": recorded or {}}
+
+    async def on_notify(ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None) -> None:
+        recorder.notifications.append((method, params))
+        recorder.notified.set()
+
+    return on_request, on_notify
+
+
+@asynccontextmanager
+async def running_pair(
+    factory: PairFactory,
+    *,
+    server_on_request: OnRequest | None = None,
+    server_on_notify: OnNotify | None = None,
+    client_on_request: OnRequest | None = None,
+    client_on_notify: OnNotify | None = None,
+    can_send_request: bool = True,
+) -> AsyncIterator[tuple[Dispatcher[TransportContext], Dispatcher[TransportContext], Recorder, Recorder]]:
+    """Yield `(client, server, client_recorder, server_recorder)` with both `run()` loops live."""
+    client, server, close = factory(can_send_request=can_send_request)
+    client_rec, server_rec = Recorder(), Recorder()
+    c_req, c_notify = echo_handlers(client_rec)
+    s_req, s_notify = echo_handlers(server_rec)
+    try:
+        async with anyio.create_task_group() as tg:
+            await tg.start(client.run, client_on_request or c_req, client_on_notify or c_notify)
+            await tg.start(server.run, server_on_request or s_req, server_on_notify or s_notify)
+            try:
+                yield client, server, client_rec, server_rec
+            finally:
+                tg.cancel_scope.cancel()
+    finally:
+        close()
+
+
+@pytest.mark.anyio
+async def test_send_raw_request_returns_result_from_peer_on_request(pair_factory: PairFactory):
+    async with running_pair(pair_factory) as (client, _server, _crec, srec):
+        with anyio.fail_after(5):
+            result = await client.send_raw_request("tools/list", {"cursor": "abc"})
+    assert result == {"echoed": "tools/list", "params": {"cursor": "abc"}}
+    assert srec.requests == [("tools/list", {"cursor": "abc"})]
+
+
+@pytest.mark.anyio
+async def test_send_raw_request_reraises_mcperror_from_handler_unchanged(pair_factory: PairFactory):
+    async def on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        raise MCPError(code=INVALID_PARAMS, message="bad cursor")
+
+    async with running_pair(pair_factory, server_on_request=on_request) as (client, *_):
+        with anyio.fail_after(5), pytest.raises(MCPError) as exc:
+            await client.send_raw_request("tools/list", {})
+    assert exc.value.error.code == INVALID_PARAMS
+    assert exc.value.error.message == "bad cursor"
+
+
+@pytest.mark.anyio
+async def test_send_raw_request_maps_validation_error_to_invalid_params(pair_factory: PairFactory):
+    """A pydantic `ValidationError` from the handler surfaces as the
+    normalized INVALID_PARAMS shape on every dispatcher."""
+
+    async def on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        Tool.model_validate({"name": 123})  # raises ValidationError
+        raise NotImplementedError
+
+    async with running_pair(pair_factory, server_on_request=on_request) as (client, *_):
+        with anyio.fail_after(5), pytest.raises(MCPError) as exc:
+            await client.send_raw_request("tools/list", None)
+    assert exc.value.error == ErrorData(code=INVALID_PARAMS, message="Invalid request parameters", data="")
+
+
+@pytest.mark.anyio
+async def test_send_raw_request_with_timeout_raises_mcperror_request_timeout(pair_factory: PairFactory):
+    async def on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        await anyio.sleep_forever()
+        raise NotImplementedError
+
+    async with running_pair(pair_factory, server_on_request=on_request) as (client, *_):
+        with anyio.fail_after(5), pytest.raises(MCPError) as exc:
+            await client.send_raw_request("slow", None, {"timeout": 0})
+    assert exc.value.error.code == REQUEST_TIMEOUT
+
+
+@pytest.mark.anyio
+async def test_notify_invokes_peer_on_notify(pair_factory: PairFactory):
+    async with running_pair(pair_factory) as (client, _server, _crec, srec):
+        with anyio.fail_after(5):
+            await client.notify("notifications/initialized", {"v": 1})
+            await srec.notified.wait()
+    assert srec.notifications == [("notifications/initialized", {"v": 1})]
+
+
+@pytest.mark.anyio
+async def test_ctx_send_raw_request_round_trips_to_calling_side(pair_factory: PairFactory):
+    """A handler's ctx.send_raw_request reaches the side that made the inbound request."""
+
+    async def server_on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        sample = await ctx.send_raw_request("sampling/createMessage", {"prompt": "hi"})
+        return {"sampled": sample}
+
+    async with running_pair(pair_factory, server_on_request=server_on_request) as (client, _server, crec, _srec):
+        with anyio.fail_after(5):
+            result = await client.send_raw_request("tools/call", None)
+    assert crec.requests == [("sampling/createMessage", {"prompt": "hi"})]
+    assert result == {"sampled": {"echoed": "sampling/createMessage", "params": {"prompt": "hi"}}}
+
+
+@pytest.mark.anyio
+async def test_ctx_send_raw_request_raises_nobackchannelerror_when_transport_disallows(pair_factory: PairFactory):
+    async def server_on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        return await ctx.send_raw_request("sampling/createMessage", None)
+
+    async with running_pair(pair_factory, server_on_request=server_on_request, can_send_request=False) as (client, *_):
+        with anyio.fail_after(5), pytest.raises(MCPError) as exc:
+            await client.send_raw_request("tools/call", None)
+    assert exc.value.error.code == INVALID_REQUEST
+
+
+@pytest.mark.anyio
+async def test_ctx_notify_invokes_calling_side_on_notify(pair_factory: PairFactory):
+    async def server_on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        await ctx.notify("notifications/message", {"level": "info"})
+        return {}
+
+    async with running_pair(pair_factory, server_on_request=server_on_request) as (client, _server, crec, _srec):
+        with anyio.fail_after(5):
+            await client.send_raw_request("tools/call", None)
+            await crec.notified.wait()
+    assert crec.notifications == [("notifications/message", {"level": "info"})]
+
+
+@pytest.mark.anyio
+async def test_ctx_progress_invokes_caller_on_progress_callback(pair_factory: PairFactory):
+    async def server_on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        await ctx.progress(0.5, total=1.0, message="halfway")
+        return {}
+
+    received: list[tuple[float, float | None, str | None]] = []
+
+    async def on_progress(progress: float, total: float | None, message: str | None) -> None:
+        received.append((progress, total, message))
+
+    async with running_pair(pair_factory, server_on_request=server_on_request) as (client, *_):
+        with anyio.fail_after(5):
+            await client.send_raw_request("tools/call", None, {"on_progress": on_progress})
+    assert received == [(0.5, 1.0, "halfway")]
+
+
+@pytest.mark.anyio
+async def test_ctx_progress_is_noop_when_caller_supplied_no_callback(pair_factory: PairFactory):
+    async def server_on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        await ctx.progress(0.5)
+        return {"ok": True}
+
+    async with running_pair(pair_factory, server_on_request=server_on_request) as (client, *_):
+        with anyio.fail_after(5):
+            result = await client.send_raw_request("tools/call", None)
+    assert result == {"ok": True}
+
+
+@pytest.mark.anyio
+async def test_ctx_message_metadata_is_none_when_transport_attaches_nothing(pair_factory: PairFactory):
+    """Plain requests carry no transport metadata, so handlers see `None`."""
+    async with running_pair(pair_factory) as (client, _server, _crec, srec):
+        with anyio.fail_after(5):
+            await client.send_raw_request("tools/call", None)
+    assert len(srec.contexts) == 1
+    assert srec.contexts[0].message_metadata is None
+
+
+@pytest.mark.anyio
+async def test_ctx_request_id_exposes_inbound_id(pair_factory: PairFactory):
+    """JSON-RPC carries the wire id through; direct dispatch has none."""
+    async with running_pair(pair_factory) as (client, _server, _crec, srec):
+        with anyio.fail_after(5):
+            await client.send_raw_request("tools/call", None)
+            await client.send_raw_request("tools/call", None)
+    a, b = (ctx.request_id for ctx in srec.contexts)
+    assert (a is None and b is None) or (isinstance(a, int) and isinstance(b, int) and a != b)
+
+
+@pytest.mark.anyio
+async def test_direct_send_raw_request_wraps_non_mcperror_exception_as_internal_error_with_cause():
+    """DirectDispatcher-specific: the original exception is chained via __cause__."""
+
+    async def on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
+    ) -> dict[str, Any]:
+        raise ValueError("oops")
+
+    async with running_pair(direct_pair, server_on_request=on_request) as (client, *_):
+        with anyio.fail_after(5), pytest.raises(MCPError) as exc:
+            await client.send_raw_request("tools/list", {})
+        assert exc.value.error.code == INTERNAL_ERROR
+        assert isinstance(exc.value.__cause__, ValueError)
+
+
+@pytest.mark.anyio
+async def test_direct_send_raw_request_issued_before_peer_run_blocks_until_peer_ready():
+    client, server = create_direct_dispatcher_pair()
+    s_req, s_notify = echo_handlers(Recorder())
+    c_req, c_notify = echo_handlers(Recorder())
+
+    async def late_start():
+        await anyio.sleep(0)
+        await server.run(s_req, s_notify)
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(client.run, c_req, c_notify)
+        tg.start_soon(late_start)
+        with anyio.fail_after(5):
+            result = await client.send_raw_request("ping", None)
+        assert result == {"echoed": "ping", "params": {}}
+        client.close()
+        server.close()
+
+
+@pytest.mark.anyio
+async def test_direct_send_raw_request_and_notify_raise_runtimeerror_when_no_peer_connected():
+    d = DirectDispatcher(TransportContext(kind="direct", can_send_request=True))
+    with pytest.raises(RuntimeError, match="no peer"):
+        await d.send_raw_request("ping", None)
+    with pytest.raises(RuntimeError, match="no peer"):
+        await d.notify("ping", None)
+
+
+@pytest.mark.anyio
+async def test_direct_close_makes_run_return():
+    client, server = create_direct_dispatcher_pair()
+    on_request, on_notify = echo_handlers(Recorder())
+    with anyio.fail_after(5):
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(server.run, on_request, on_notify)
+            tg.start_soon(client.run, on_request, on_notify)
+            client.close()
+            server.close()
+
+
+if TYPE_CHECKING:
+    _d: Dispatcher[TransportContext] = DirectDispatcher(TransportContext(kind="direct", can_send_request=True))
+    _o: Outbound = _d

tests/shared/test_otel.py

@@ -10,9 +10,6 @@ from mcp.server.mcpserver import MCPServer
 pytestmark = pytest.mark.anyio
 
 
-# Logfire warns about propagated trace context by default (distributed_tracing=None).
-# This is expected here since we're testing cross-boundary context propagation.
-@pytest.mark.filterwarnings("ignore::RuntimeWarning")
 async def test_client_and_server_spans(capfire: CaptureLogfire):
     """Verify that calling a tool produces client and server spans with correct attributes."""
     server = MCPServer("test")

tests/shared/test_peer.py

@@ -0,0 +1,190 @@
+"""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.shared.dispatcher import DispatchContext
+from mcp.shared.peer import ClientPeer, dump_params
+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 = ClientPeer(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_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(
+                [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(
+                [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 = 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()
+        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"})
+        method, params = rec.seen[0]
+        assert method == "roots/list"
+        assert params == {"_meta": {"traceId": "t1"}}
+
+
+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) -> 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

tests/shared/test_progress_notifications.py

@@ -1,264 +0,0 @@
-from typing import Any
-from unittest.mock import patch
-
-import anyio
-import pytest
-
-from mcp import Client, types
-from mcp.client.session import ClientSession
-from mcp.server import Server, ServerRequestContext
-from mcp.server.lowlevel import NotificationOptions
-from mcp.server.models import InitializationOptions
-from mcp.server.session import ServerSession
-from mcp.shared.message import SessionMessage
-from mcp.shared.session import RequestResponder
-
-
-@pytest.mark.anyio
-async def test_bidirectional_progress_notifications():
-    """Test that both client and server can send progress notifications."""
-    # Create memory streams for client/server
-    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](5)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](5)
-
-    # Run a server session so we can send progress updates in tool
-    async def run_server():
-        # Create a server session
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            InitializationOptions(
-                server_name="ProgressTestServer",
-                server_version="0.1.0",
-                capabilities=server.get_capabilities(NotificationOptions(), {}),
-            ),
-        ) as server_session:
-            async for message in server_session.incoming_messages:
-                try:
-                    await server._handle_message(message, server_session, {})
-                except Exception as e:  # pragma: no cover
-                    raise e
-
-    # Track progress updates
-    server_progress_updates: list[dict[str, Any]] = []
-    client_progress_updates: list[dict[str, Any]] = []
-
-    # Progress tokens
-    server_progress_token = "server_token_123"
-    client_progress_token = "client_token_456"
-
-    # Register progress handler
-    async def handle_progress(ctx: ServerRequestContext, params: types.ProgressNotificationParams) -> None:
-        server_progress_updates.append(
-            {
-                "token": params.progress_token,
-                "progress": params.progress,
-                "total": params.total,
-                "message": params.message,
-            }
-        )
-
-    # Register list tool handler
-    async def handle_list_tools(
-        ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
-    ) -> types.ListToolsResult:
-        return types.ListToolsResult(
-            tools=[
-                types.Tool(
-                    name="test_tool",
-                    description="A tool that sends progress notifications <o/",
-                    input_schema={},
-                )
-            ]
-        )
-
-    # Register tool handler
-    async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.CallToolResult:
-        # Make sure we received a progress token
-        if params.name == "test_tool":
-            assert params.meta is not None
-            progress_token = params.meta.get("progress_token")
-            assert progress_token is not None
-            assert progress_token == client_progress_token
-
-            # Send progress notifications using ctx.session
-            await ctx.session.send_progress_notification(
-                progress_token=progress_token,
-                progress=0.25,
-                total=1.0,
-                message="Server progress 25%",
-            )
-
-            await ctx.session.send_progress_notification(
-                progress_token=progress_token,
-                progress=0.5,
-                total=1.0,
-                message="Server progress 50%",
-            )
-
-            await ctx.session.send_progress_notification(
-                progress_token=progress_token,
-                progress=1.0,
-                total=1.0,
-                message="Server progress 100%",
-            )
-
-            return types.CallToolResult(content=[types.TextContent(type="text", text="Tool executed successfully")])
-
-        raise ValueError(f"Unknown tool: {params.name}")  # pragma: no cover
-
-    # Create a server with progress capability
-    server = Server(
-        name="ProgressTestServer",
-        on_progress=handle_progress,
-        on_list_tools=handle_list_tools,
-        on_call_tool=handle_call_tool,
-    )
-
-    # Client message handler to store progress notifications
-    async def handle_client_message(
-        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
-    ) -> None:
-        if isinstance(message, Exception):  # pragma: no cover
-            raise message
-
-        if isinstance(message, types.ServerNotification):  # pragma: no branch
-            if isinstance(message, types.ProgressNotification):  # pragma: no branch
-                params = message.params
-                client_progress_updates.append(
-                    {
-                        "token": params.progress_token,
-                        "progress": params.progress,
-                        "total": params.total,
-                        "message": params.message,
-                    }
-                )
-
-    # Test using client
-    async with (
-        ClientSession(
-            server_to_client_receive,
-            client_to_server_send,
-            message_handler=handle_client_message,
-        ) as client_session,
-        anyio.create_task_group() as tg,
-    ):
-        # Start the server in a background task
-        tg.start_soon(run_server)
-
-        # Initialize the client connection
-        await client_session.initialize()
-
-        # Call list_tools with progress token
-        await client_session.list_tools()
-
-        # Call test_tool with progress token
-        await client_session.call_tool("test_tool", meta={"progress_token": client_progress_token})
-
-        # Send progress notifications from client to server
-        await client_session.send_progress_notification(
-            progress_token=server_progress_token,
-            progress=0.33,
-            total=1.0,
-            message="Client progress 33%",
-        )
-
-        await client_session.send_progress_notification(
-            progress_token=server_progress_token,
-            progress=0.66,
-            total=1.0,
-            message="Client progress 66%",
-        )
-
-        await client_session.send_progress_notification(
-            progress_token=server_progress_token,
-            progress=1.0,
-            total=1.0,
-            message="Client progress 100%",
-        )
-
-        # Wait and exit
-        await anyio.sleep(0.5)
-        tg.cancel_scope.cancel()
-
-    # Verify client received progress updates from server
-    assert len(client_progress_updates) == 3
-    assert client_progress_updates[0]["token"] == client_progress_token
-    assert client_progress_updates[0]["progress"] == 0.25
-    assert client_progress_updates[0]["message"] == "Server progress 25%"
-    assert client_progress_updates[2]["progress"] == 1.0
-
-    # Verify server received progress updates from client
-    assert len(server_progress_updates) == 3
-    assert server_progress_updates[0]["token"] == server_progress_token
-    assert server_progress_updates[0]["progress"] == 0.33
-    assert server_progress_updates[0]["message"] == "Client progress 33%"
-    assert server_progress_updates[2]["progress"] == 1.0
-
-
-@pytest.mark.anyio
-async def test_progress_callback_exception_logging():
-    """Test that exceptions in progress callbacks are logged and \
-        don't crash the session."""
-    # Track logged warnings
-    logged_errors: list[str] = []
-
-    def mock_log_exception(msg: str, *args: Any, **kwargs: Any) -> None:
-        logged_errors.append(msg % args if args else msg)
-
-    # Create a progress callback that raises an exception
-    async def failing_progress_callback(progress: float, total: float | None, message: str | None) -> None:
-        raise ValueError("Progress callback failed!")
-
-    # Create a server with a tool that sends progress notifications
-    async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.CallToolResult:
-        if params.name == "progress_tool":
-            assert ctx.request_id is not None
-            # Send a progress notification
-            await ctx.session.send_progress_notification(
-                progress_token=ctx.request_id,
-                progress=50.0,
-                total=100.0,
-                message="Halfway done",
-            )
-            return types.CallToolResult(content=[types.TextContent(type="text", text="progress_result")])
-        raise ValueError(f"Unknown tool: {params.name}")  # pragma: no cover
-
-    async def handle_list_tools(
-        ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
-    ) -> types.ListToolsResult:
-        return types.ListToolsResult(
-            tools=[
-                types.Tool(
-                    name="progress_tool",
-                    description="A tool that sends progress notifications",
-                    input_schema={},
-                )
-            ]
-        )
-
-    server = Server(
-        name="TestProgressServer",
-        on_call_tool=handle_call_tool,
-        on_list_tools=handle_list_tools,
-    )
-
-    # Test with mocked logging
-    with patch("mcp.shared.session.logging.exception", side_effect=mock_log_exception):
-        async with Client(server) as client:
-            # Call tool with a failing progress callback
-            result = await client.call_tool(
-                "progress_tool",
-                arguments={},
-                progress_callback=failing_progress_callback,
-            )
-
-            # Verify the request completed successfully despite the callback failure
-            assert len(result.content) == 1
-            content = result.content[0]
-            assert isinstance(content, types.TextContent)
-            assert content.text == "progress_result"
-
-            # Check that a warning was logged for the progress callback exception
-            assert len(logged_errors) > 0
-            assert any("Progress callback raised an exception" in warning for warning in logged_errors)

tests/shared/test_session.py

@@ -23,19 +23,6 @@ from mcp.types import (
 )
 
 
-@pytest.mark.anyio
-async def test_in_flight_requests_cleared_after_completion():
-    """Verify that _in_flight is empty after all requests complete."""
-    server = Server(name="test server")
-    async with Client(server) as client:
-        # Send a request and wait for response
-        response = await client.send_ping()
-        assert isinstance(response, EmptyResult)
-
-        # Verify _in_flight is empty
-        assert len(client.session._in_flight) == 0
-
-
 @pytest.mark.anyio
 async def test_request_cancellation():
     """Test that requests can be cancelled while in-flight."""

tests/shared/test_streamable_http.py

@@ -40,6 +40,7 @@ from mcp.server.streamable_http import (
 )
 from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
 from mcp.server.transport_security import TransportSecuritySettings
+from mcp.shared._compat import resync_tracer
 from mcp.shared._context import RequestContext
 from mcp.shared._context_streams import create_context_streams
 from mcp.shared.message import ClientMessageMetadata, ServerMessageMetadata, SessionMessage
@@ -1187,6 +1188,8 @@ async def test_streamable_http_client_resumption(event_app: tuple[SimpleEventSto
                     # Kill the client session while tool is waiting on lock
                     tg.cancel_scope.cancel()
 
+    await resync_tracer()
+
     async with make_client(app, headers=headers) as httpx_client2:
         async with streamable_http_client(f"{BASE_URL}/mcp", http_client=httpx_client2) as (
             read_stream,