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

tests/client/test_session.py

@@ -1,6 +1,11 @@
 from __future__ import annotations
 
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any
+
 import anyio
+import anyio.streams.memory
 import pytest
 
 from mcp import types
@@ -10,12 +15,14 @@ from mcp.shared.message import SessionMessage
 from mcp.shared.session import RequestResponder
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 from mcp.types import (
+    INVALID_PARAMS,
     LATEST_PROTOCOL_VERSION,
     CallToolResult,
     Implementation,
     InitializedNotification,
     InitializeRequest,
     InitializeResult,
+    JSONRPCError,
     JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
@@ -26,6 +33,29 @@ from mcp.types import (
     client_request_adapter,
 )
 
+_SendToClient = anyio.streams.memory.MemoryObjectSendStream[SessionMessage | Exception]
+_RecvFromClient = anyio.streams.memory.MemoryObjectReceiveStream[SessionMessage]
+
+
+@asynccontextmanager
+async def raw_client_session(
+    **kwargs: Any,
+) -> AsyncIterator[tuple[ClientSession, _SendToClient, _RecvFromClient]]:
+    """Yield `(session, send_to_client, recv_from_client)` with the receive loop running.
+
+    `send_to_client` accepts `SessionMessage | Exception` so tests can inject
+    transport-level exceptions. No initialize handshake is performed.
+    """
+    s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
+    c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage](32)
+    async with ClientSession(s2c_recv, c2s_send, **kwargs) as session:
+        try:
+            with anyio.fail_after(5):
+                yield session, s2c_send, c2s_recv
+        finally:
+            s2c_send.close()
+            c2s_recv.close()
+
 
 @pytest.mark.anyio
 async def test_client_session_initialize():
@@ -705,3 +735,129 @@ async def test_client_tool_call_with_meta(meta: RequestParamsMeta | None):
         await session.initialize()
 
         await session.call_tool(name=mocked_tool.name, arguments={"foo": "bar"}, meta=meta)
+
+
+@pytest.mark.anyio
+async def test_receive_loop_answers_malformed_inbound_request_with_invalid_params():
+    """A request that fails ServerRequest validation gets an INVALID_PARAMS error response."""
+    async with raw_client_session() as (_session, to_client, from_client):
+        await to_client.send(
+            SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=7, method="sampling/createMessage", params={"broken": 1}))
+        )
+        out = await from_client.receive()
+    assert isinstance(out.message, JSONRPCError)
+    assert out.message.id == 7
+    assert out.message.error.code == INVALID_PARAMS
+
+
+@pytest.mark.anyio
+async def test_receive_loop_answers_invalid_params_when_sampling_callback_raises():
+    """Same boundary catches exceptions from the request handler itself."""
+
+    async def boom(ctx: object, params: object) -> types.CreateMessageResult:
+        raise RuntimeError("sampling boom")
+
+    params = types.CreateMessageRequestParams(
+        messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="hi"))],
+        max_tokens=10,
+    ).model_dump(by_alias=True, mode="json", exclude_none=True)
+    async with raw_client_session(sampling_callback=boom) as (_session, to_client, from_client):
+        await to_client.send(
+            SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=8, method="sampling/createMessage", params=params))
+        )
+        out = await from_client.receive()
+    assert isinstance(out.message, JSONRPCError)
+    assert out.message.error.code == INVALID_PARAMS
+
+
+@pytest.mark.anyio
+async def test_receive_loop_logs_and_drops_malformed_notification(caplog: pytest.LogCaptureFixture):
+    """A notification that fails ServerNotification validation is logged and dropped."""
+    seen: list[object] = []
+    delivered = anyio.Event()
+
+    async def handler(msg: object) -> None:
+        seen.append(msg)
+        delivered.set()
+
+    async with raw_client_session(message_handler=handler) as (_session, to_client, _):
+        await to_client.send(SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="not/a/spec/notification")))
+        # Follow with a valid notification so we know the loop is still alive.
+        await to_client.send(
+            SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/tools/list_changed"))
+        )
+        await delivered.wait()
+    assert isinstance(seen[0], types.ToolListChangedNotification)
+    assert "Failed to validate notification" in caplog.text
+
+
+@pytest.mark.anyio
+async def test_receive_loop_forwards_transport_exception_to_message_handler():
+    seen: list[object] = []
+    delivered = anyio.Event()
+
+    async def handler(msg: object) -> None:
+        seen.append(msg)
+        delivered.set()
+
+    async with raw_client_session(message_handler=handler) as (_session, to_client, _):
+        exc = ValueError("bad bytes")
+        await to_client.send(exc)
+        await delivered.wait()
+    assert seen == [exc]
+
+
+@pytest.mark.anyio
+async def test_receive_loop_consumes_server_cancelled_without_reaching_message_handler():
+    """A server-sent notifications/cancelled is swallowed, matching the pre-swap contract.
+
+    The server dispatcher now emits this on sampling/elicitation timeout, but
+    ClientSession has no in-flight tracking to act on it, so surfacing it would
+    only break user handlers that exhaustively match ServerNotification.
+    """
+    seen: list[object] = []
+    delivered = anyio.Event()
+
+    async def handler(msg: object) -> None:
+        seen.append(msg)
+        delivered.set()
+
+    async with raw_client_session(message_handler=handler) as (_session, to_client, _):
+        await to_client.send(
+            SessionMessage(
+                JSONRPCNotification(
+                    jsonrpc="2.0", method="notifications/cancelled", params={"requestId": 1, "reason": "timed out"}
+                )
+            )
+        )
+        # Follow with a notification that does reach the handler so we can
+        # assert ordering deterministically.
+        await to_client.send(
+            SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/tools/list_changed"))
+        )
+        await delivered.wait()
+    assert len(seen) == 1
+    assert isinstance(seen[0], types.ToolListChangedNotification)
+
+
+@pytest.mark.anyio
+async def test_receive_loop_swallows_progress_callback_exception(caplog: pytest.LogCaptureFixture):
+    delivered = anyio.Event()
+
+    async def boom(progress: float, total: float | None, message: str | None) -> None:
+        raise RuntimeError("progress boom")
+
+    async def handler(msg: object) -> None:
+        delivered.set()
+
+    async with raw_client_session(message_handler=handler) as (session, to_client, _):
+        # Register the callback under a known token without sending a request.
+        session._progress_callbacks[42] = boom  # pyright: ignore[reportPrivateUsage]
+        params = {"progressToken": 42, "progress": 0.5}
+        await to_client.send(
+            SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/progress", params=params))
+        )
+        # The progress notification also reaches the message handler after the
+        # callback runs, so this fires once the callback's exception is handled.
+        await delivered.wait()
+    assert "Progress callback raised an exception" in caplog.text

tests/client/transports/test_memory.py

@@ -1,8 +1,15 @@
 """Tests for InMemoryTransport."""
 
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any
+
+import anyio
+import anyio.lowlevel
 import pytest
 
 from mcp import Client, types
+from mcp.client import _memory
 from mcp.client._memory import InMemoryTransport
 from mcp.server import Server, ServerRequestContext
 from mcp.server.mcpserver import MCPServer
@@ -95,3 +102,47 @@ async def test_raise_exceptions(mcpserver_server: MCPServer):
     transport = InMemoryTransport(mcpserver_server, raise_exceptions=True)
     async with transport as (read_stream, _write_stream):
         assert read_stream is not None
+
+
+async def test_aexit_with_well_behaved_lifespan_runs_teardown_without_cancel():
+    """A lifespan that finishes promptly on EOF should run to completion.
+
+    The transport closes the streams first and waits for the server to exit
+    naturally, so teardown observes no cancellation.
+    """
+    teardown_ran = anyio.Event()
+
+    @asynccontextmanager
+    async def lifespan(_: Server[Any]) -> AsyncIterator[dict[str, Any]]:
+        yield {}
+        await anyio.lowlevel.checkpoint()
+        teardown_ran.set()
+
+    server = Server(name="test_server", lifespan=lifespan)
+    with anyio.fail_after(5):
+        async with InMemoryTransport(server):
+            pass
+    assert teardown_ran.is_set()
+
+
+async def test_aexit_with_blocking_lifespan_is_bounded(monkeypatch: pytest.MonkeyPatch):
+    """A lifespan that never returns must not hang `__aexit__` forever.
+
+    After EOFing the server the transport waits `SERVER_SHUTDOWN_GRACE` for a
+    natural exit, then cancels the server task as a backstop so the
+    task-group join completes.
+    """
+    monkeypatch.setattr(_memory, "SERVER_SHUTDOWN_GRACE", 0.05)
+    teardown_started = anyio.Event()
+
+    @asynccontextmanager
+    async def blocking_lifespan(_: Server[Any]) -> AsyncIterator[dict[str, Any]]:
+        yield {}
+        teardown_started.set()
+        await anyio.Event().wait()
+
+    server = Server(name="test_server", lifespan=blocking_lifespan)
+    with anyio.fail_after(5):
+        async with InMemoryTransport(server):
+            pass
+    assert teardown_started.is_set()

tests/conftest.py

@@ -1,6 +1,44 @@
+import os
+from collections.abc import Iterator
+
 import pytest
 
+# OpenTelemetry's `set_tracer_provider` is set-once per process, so the suite
+# uses a single span-capture mechanism: logfire's `capfire` fixture (its
+# `configure()` swaps span processors on repeat calls rather than re-setting
+# the provider). Logfire's default `distributed_tracing=None` emits a
+# RuntimeWarning + diagnostic span when incoming W3C trace context is
+# extracted; several tests exercise that propagation deliberately, so opt in
+# suite-wide. Set before logfire is imported anywhere.
+os.environ.setdefault("LOGFIRE_DISTRIBUTED_TRACING", "true")
+
+import opentelemetry.trace  # noqa: E402  (env var must be set before logfire import below)
+from logfire.testing import CaptureLogfire  # noqa: E402
+
+import mcp.shared._otel  # noqa: E402
+
 
 @pytest.fixture
 def anyio_backend():
     return "asyncio"
+
+
+@pytest.fixture(name="capfire")
+def _capfire_isolated(capfire: CaptureLogfire) -> Iterator[CaptureLogfire]:
+    """Override of logfire's `capfire` that scopes the MCP tracer to the test.
+
+    `capfire` installs a real tracer provider, and logfire's proxy machinery
+    mutates the cached `mcp.shared._otel._tracer` to delegate to it for the
+    rest of the process. Without isolation, every subsequent test in the same
+    worker would emit real spans, and `send_raw_request` would inject a real
+    `traceparent` into outbound `_meta`, breaking the interaction-suite
+    snapshots that pin `_meta={}` under a no-op tracer.
+
+    Setup points `_tracer` at the now-live provider so MCP spans record;
+    teardown replaces it with a `NoOpTracer`.
+    """
+    mcp.shared._otel._tracer = opentelemetry.trace.get_tracer_provider().get_tracer("mcp-python-sdk")
+    try:
+        yield capfire
+    finally:
+        mcp.shared._otel._tracer = opentelemetry.trace.NoOpTracer()

tests/interaction/lowlevel/test_roots.py

@@ -163,4 +163,4 @@ async def test_roots_list_changed_reaches_server_handler(connect: Connect) -> No
         with anyio.fail_after(5):
             await delivered.wait()
 
-    assert received == snapshot([None])
+    assert received == snapshot([types.NotificationParams()])

tests/interaction/transports/_bridge.py

@@ -35,6 +35,8 @@ import httpx
 from anyio.streams.memory import MemoryObjectReceiveStream
 from starlette.types import ASGIApp, Message, Scope
 
+from mcp.shared._compat import resync_tracer
+
 
 class _StreamingResponseBody(httpx.AsyncByteStream):
     """A response body that yields chunks as the application produces them.
@@ -88,6 +90,7 @@ class StreamingASGITransport(httpx.AsyncBaseTransport):
         if self._cancel_on_close:
             self._task_group.cancel_scope.cancel()
         await self._task_group.__aexit__(exc_type, exc_value, traceback)
+        await resync_tracer()
 
     async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
         assert isinstance(request.stream, httpx.AsyncByteStream)

tests/issues/test_malformed_input.py

@@ -1,151 +0,0 @@
-# Claude Debug
-"""Test for HackerOne vulnerability report #3156202 - malformed input DOS."""
-
-import anyio
-import pytest
-
-from mcp.server.models import InitializationOptions
-from mcp.server.session import ServerSession
-from mcp.shared.message import SessionMessage
-from mcp.types import INVALID_PARAMS, JSONRPCError, JSONRPCMessage, JSONRPCRequest, ServerCapabilities
-
-
-@pytest.mark.anyio
-async def test_malformed_initialize_request_does_not_crash_server():
-    """Test that malformed initialize requests return proper error responses
-    instead of crashing the server (HackerOne #3156202).
-    """
-    # Create in-memory streams for testing
-    read_send_stream, read_receive_stream = anyio.create_memory_object_stream[SessionMessage | Exception](10)
-    write_send_stream, write_receive_stream = anyio.create_memory_object_stream[SessionMessage](10)
-
-    try:
-        # Create a malformed initialize request (missing required params field)
-        malformed_request = JSONRPCRequest(
-            jsonrpc="2.0",
-            id="f20fe86132ed4cd197f89a7134de5685",
-            method="initialize",
-            # params=None  # Missing required params field
-        )
-
-        # Wrap in session message
-        request_message = SessionMessage(message=malformed_request)
-
-        # Start a server session
-        async with ServerSession(
-            read_stream=read_receive_stream,
-            write_stream=write_send_stream,
-            init_options=InitializationOptions(
-                server_name="test_server",
-                server_version="1.0.0",
-                capabilities=ServerCapabilities(),
-            ),
-        ):
-            # Send the malformed request
-            await read_send_stream.send(request_message)
-
-            # Give the session time to process the request
-            await anyio.sleep(0.1)
-
-            # Check that we received an error response instead of a crash
-            try:
-                response_message = write_receive_stream.receive_nowait()
-                response = response_message.message
-
-                # Verify it's a proper JSON-RPC error response
-                assert isinstance(response, JSONRPCError)
-                assert response.jsonrpc == "2.0"
-                assert response.id == "f20fe86132ed4cd197f89a7134de5685"
-                assert response.error.code == INVALID_PARAMS
-                assert "Invalid request parameters" in response.error.message
-
-                # Verify the session is still alive and can handle more requests
-                # Send another malformed request to confirm server stability
-                another_malformed_request = JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id="test_id_2",
-                    method="tools/call",
-                    # params=None  # Missing required params
-                )
-                another_request_message = SessionMessage(message=another_malformed_request)
-
-                await read_send_stream.send(another_request_message)
-                await anyio.sleep(0.1)
-
-                # Should get another error response, not a crash
-                second_response_message = write_receive_stream.receive_nowait()
-                second_response = second_response_message.message
-
-                assert isinstance(second_response, JSONRPCError)
-                assert second_response.id == "test_id_2"
-                assert second_response.error.code == INVALID_PARAMS
-
-            except anyio.WouldBlock:  # pragma: no cover
-                pytest.fail("No response received - server likely crashed")
-    finally:  # pragma: lax no cover
-        # Close all streams to ensure proper cleanup
-        await read_send_stream.aclose()
-        await write_send_stream.aclose()
-        await read_receive_stream.aclose()
-        await write_receive_stream.aclose()
-
-
-@pytest.mark.anyio
-async def test_multiple_concurrent_malformed_requests():
-    """Test that multiple concurrent malformed requests don't crash the server."""
-    # Create in-memory streams for testing
-    read_send_stream, read_receive_stream = anyio.create_memory_object_stream[SessionMessage | Exception](100)
-    write_send_stream, write_receive_stream = anyio.create_memory_object_stream[SessionMessage](100)
-
-    try:
-        # Start a server session
-        async with ServerSession(
-            read_stream=read_receive_stream,
-            write_stream=write_send_stream,
-            init_options=InitializationOptions(
-                server_name="test_server",
-                server_version="1.0.0",
-                capabilities=ServerCapabilities(),
-            ),
-        ):
-            # Send multiple malformed requests concurrently
-            malformed_requests: list[SessionMessage] = []
-            for i in range(10):
-                malformed_request = JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id=f"malformed_{i}",
-                    method="initialize",
-                    # params=None  # Missing required params
-                )
-                request_message = SessionMessage(message=malformed_request)
-                malformed_requests.append(request_message)
-
-            # Send all requests
-            for request in malformed_requests:
-                await read_send_stream.send(request)
-
-            # Give time to process
-            await anyio.sleep(0.2)
-
-            # Verify we get error responses for all requests
-            error_responses: list[JSONRPCMessage] = []
-            try:
-                while True:
-                    response_message = write_receive_stream.receive_nowait()
-                    error_responses.append(response_message.message)
-            except anyio.WouldBlock:
-                pass  # No more messages
-
-            # Should have received 10 error responses
-            assert len(error_responses) == 10
-
-            for i, response in enumerate(error_responses):
-                assert isinstance(response, JSONRPCError)
-                assert response.id == f"malformed_{i}"
-                assert response.error.code == INVALID_PARAMS
-    finally:  # pragma: lax no cover
-        # Close all streams to ensure proper cleanup
-        await read_send_stream.aclose()
-        await write_send_stream.aclose()
-        await read_receive_stream.aclose()
-        await write_receive_stream.aclose()

tests/server/conftest.py

@@ -0,0 +1,46 @@
+"""Shared fixtures for server-side tests."""
+
+from collections.abc import Iterator
+
+import pytest
+from logfire.testing import CaptureLogfire, TestExporter
+from opentelemetry.sdk.trace import ReadableSpan
+
+
+class SpanCapture:
+    """Thin adapter over logfire's `TestExporter` for asserting on MCP spans.
+
+    `finished()` returns the raw `ReadableSpan` objects emitted by the
+    `mcp-python-sdk` instrumentation scope, filtered to exclude logfire's
+    synthetic `pending_span` markers, so tests can assert directly on
+    `.name`, `.kind`, `.status`, `.attributes`, `.parent`, `.events`.
+    """
+
+    def __init__(self, exporter: TestExporter) -> None:
+        self._exporter = exporter
+
+    def clear(self) -> None:
+        self._exporter.clear()
+
+    def finished(self) -> list[ReadableSpan]:
+        return [
+            s
+            for s in self._exporter.exported_spans
+            if s.instrumentation_scope is not None
+            and s.instrumentation_scope.name == "mcp-python-sdk"
+            and not (s.attributes and s.attributes.get("logfire.span_type") == "pending_span")
+        ]
+
+
+@pytest.fixture
+def spans(capfire: CaptureLogfire) -> Iterator[SpanCapture]:
+    """In-memory MCP span capture, cleared before and after each test.
+
+    Backed by the project-level `capfire` override (see `tests/conftest.py`),
+    which scopes `mcp.shared._otel._tracer` to the test so the real tracer
+    doesn't leak into later tests in the same worker.
+    """
+    capture = SpanCapture(capfire.exporter)
+    capture.clear()
+    yield capture
+    capture.clear()

tests/server/test_connection.py

@@ -0,0 +1,254 @@
+"""Tests for `Connection`.
+
+`Connection` wraps an `Outbound` (the standalone stream). Its `notify` is
+best-effort (never raises); `send_raw_request` is gated on
+`has_standalone_channel`. Tested with a stub `Outbound` so we can assert wire
+shape and inject failures.
+"""
+
+import logging
+from collections.abc import Mapping
+from typing import Any
+
+import anyio
+import pytest
+from pydantic import ValidationError
+
+from mcp.server.connection import Connection
+from mcp.shared.dispatcher import CallOptions
+from mcp.shared.exceptions import NoBackChannelError
+from mcp.types import (
+    LATEST_PROTOCOL_VERSION,
+    ClientCapabilities,
+    CreateMessageRequest,
+    CreateMessageRequestParams,
+    ElicitationCapability,
+    EmptyResult,
+    Implementation,
+    InitializeRequestParams,
+    ListRootsRequest,
+    ListRootsResult,
+    PingRequest,
+    RootsCapability,
+    SamplingCapability,
+    SamplingContextCapability,
+    SamplingToolsCapability,
+)
+
+
+def _client_params(capabilities: ClientCapabilities) -> InitializeRequestParams:
+    return InitializeRequestParams(
+        protocol_version=LATEST_PROTOCOL_VERSION,
+        capabilities=capabilities,
+        client_info=Implementation(name="t", version="0"),
+    )
+
+
+class StubOutbound:
+    def __init__(
+        self, *, result: dict[str, Any] | None = None, raise_on_send: type[BaseException] | None = None
+    ) -> None:
+        self.requests: list[tuple[str, Mapping[str, Any] | None]] = []
+        self.notifications: list[tuple[str, Mapping[str, Any] | None]] = []
+        self._result = result if result is not None else {}
+        self._raise_on_send = raise_on_send
+
+    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))
+        return self._result
+
+    async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
+        if self._raise_on_send is not None:
+            raise self._raise_on_send()
+        self.notifications.append((method, params))
+
+
+@pytest.mark.anyio
+async def test_connection_notify_forwards_to_outbound():
+    out = StubOutbound()
+    conn = Connection(out, has_standalone_channel=True)
+    await conn.notify("notifications/message", {"level": "info", "data": "hi"})
+    assert out.notifications == [("notifications/message", {"level": "info", "data": "hi"})]
+
+
+@pytest.mark.anyio
+async def test_connection_notify_swallows_broken_stream_and_debug_logs(caplog: pytest.LogCaptureFixture):
+    caplog.set_level(logging.DEBUG, logger="mcp.server.connection")
+    out = StubOutbound(raise_on_send=anyio.BrokenResourceError)
+    conn = Connection(out, has_standalone_channel=True)
+    await conn.notify("notifications/message", {"data": "x"})  # must not raise
+    assert "stream closed" in caplog.text.lower()
+
+
+@pytest.mark.anyio
+async def test_connection_notify_drops_when_no_standalone_channel(caplog: pytest.LogCaptureFixture):
+    caplog.set_level(logging.DEBUG, logger="mcp.server.connection")
+    out = StubOutbound()
+    conn = Connection(out, has_standalone_channel=False)
+    await conn.notify("notifications/message", {"data": "x"})  # must not raise
+    assert out.notifications == []
+    assert "no standalone channel" in caplog.text.lower()
+
+
+@pytest.mark.anyio
+async def test_connection_send_raw_request_raises_nobackchannel_when_no_standalone_channel():
+    conn = Connection(StubOutbound(), has_standalone_channel=False)
+    with pytest.raises(NoBackChannelError):
+        await conn.send_raw_request("ping", None)
+
+
+@pytest.mark.anyio
+async def test_connection_send_raw_request_forwards_when_standalone_channel_present():
+    out = StubOutbound()
+    conn = Connection(out, has_standalone_channel=True)
+    result = await conn.send_raw_request("ping", None)
+    assert out.requests == [("ping", None)]
+    assert result == {}
+
+
+@pytest.mark.anyio
+async def test_connection_send_request_with_spec_type_infers_result_type():
+    out = StubOutbound(result={"roots": [{"uri": "file:///ws"}]})
+    conn = Connection(out, has_standalone_channel=True)
+    result = await conn.send_request(ListRootsRequest())
+    method, _ = out.requests[0]
+    assert method == "roots/list"
+    assert isinstance(result, ListRootsResult)
+    assert str(result.roots[0].uri) == "file:///ws"
+
+
+@pytest.mark.anyio
+async def test_connection_send_request_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"}
+    conn = Connection(StubOutbound(result=snake), has_standalone_channel=True)
+    result = await conn.send_request(CreateMessageRequest(params=CreateMessageRequestParams(messages=[], max_tokens=1)))
+    assert result.stop_reason is None
+
+
+@pytest.mark.anyio
+async def test_connection_send_request_with_result_type_kwarg_validates_custom_type():
+    out = StubOutbound(result={})
+    conn = Connection(out, has_standalone_channel=True)
+    result = await conn.send_request(PingRequest(), result_type=EmptyResult)
+    assert isinstance(result, EmptyResult)
+
+
+@pytest.mark.anyio
+async def test_connection_send_request_nonconforming_result_raises_validation_error():
+    conn = Connection(StubOutbound(result={"bogus": 1}), has_standalone_channel=True)
+    with pytest.raises(ValidationError):
+        await conn.send_request(ListRootsRequest())
+
+
+@pytest.mark.anyio
+async def test_connection_ping_sends_ping_on_standalone():
+    out = StubOutbound()
+    conn = Connection(out, has_standalone_channel=True)
+    await conn.ping()
+    assert out.requests == [("ping", None)]
+
+
+@pytest.mark.anyio
+async def test_connection_log_sends_logging_message_notification():
+    out = StubOutbound()
+    conn = Connection(out, has_standalone_channel=True)
+    await conn.log("info", {"k": "v"}, logger="my.logger")
+    method, params = out.notifications[0]
+    assert method == "notifications/message"
+    assert params is not None
+    assert params["level"] == "info"
+    assert params["data"] == {"k": "v"}
+    assert params["logger"] == "my.logger"
+
+
+@pytest.mark.anyio
+async def test_connection_log_with_meta_includes_meta_in_params():
+    out = StubOutbound()
+    conn = Connection(out, has_standalone_channel=True)
+    await conn.log("info", "x", meta={"traceId": "abc"})
+    _, params = out.notifications[0]
+    assert params is not None
+    assert params["_meta"] == {"traceId": "abc"}
+
+
+@pytest.mark.anyio
+async def test_connection_list_changed_notifications_send_correct_methods():
+    out = StubOutbound()
+    conn = Connection(out, has_standalone_channel=True)
+    await conn.send_tool_list_changed()
+    await conn.send_prompt_list_changed()
+    await conn.send_resource_list_changed()
+    await conn.send_resource_updated("file:///workspace/a.txt")
+    methods = [m for m, _ in out.notifications]
+    assert methods == [
+        "notifications/tools/list_changed",
+        "notifications/prompts/list_changed",
+        "notifications/resources/list_changed",
+        "notifications/resources/updated",
+    ]
+    assert out.notifications[-1][1] == {"uri": "file:///workspace/a.txt"}
+
+
+@pytest.mark.anyio
+async def test_connection_send_tool_list_changed_with_meta_includes_meta_only_params():
+    out = StubOutbound()
+    conn = Connection(out, has_standalone_channel=True)
+    await conn.send_tool_list_changed(meta={"k": 1})
+    assert out.notifications == [("notifications/tools/list_changed", {"_meta": {"k": 1}})]
+
+
+def test_connection_check_capability_false_before_initialized():
+    conn = Connection(StubOutbound(), has_standalone_channel=True)
+    assert conn.check_capability(ClientCapabilities(sampling=SamplingCapability())) is False
+
+
+@pytest.mark.parametrize(
+    ("have", "want", "expected"),
+    [
+        (ClientCapabilities(roots=None), ClientCapabilities(roots=RootsCapability()), False),
+        (
+            ClientCapabilities(roots=RootsCapability(list_changed=False)),
+            ClientCapabilities(roots=RootsCapability(list_changed=True)),
+            False,
+        ),
+        (ClientCapabilities(sampling=None), ClientCapabilities(sampling=SamplingCapability()), False),
+        (
+            ClientCapabilities(sampling=SamplingCapability()),
+            ClientCapabilities(sampling=SamplingCapability(context=SamplingContextCapability())),
+            False,
+        ),
+        (
+            ClientCapabilities(sampling=SamplingCapability()),
+            ClientCapabilities(sampling=SamplingCapability(tools=SamplingToolsCapability())),
+            False,
+        ),
+        (
+            ClientCapabilities(sampling=SamplingCapability(tools=SamplingToolsCapability())),
+            ClientCapabilities(sampling=SamplingCapability(tools=SamplingToolsCapability())),
+            True,
+        ),
+        (ClientCapabilities(experimental=None), ClientCapabilities(experimental={"a": {}}), False),
+        (ClientCapabilities(experimental={"a": {}}), ClientCapabilities(experimental={"b": {}}), False),
+        (ClientCapabilities(experimental={"a": {"x": 1}}), ClientCapabilities(experimental={"a": {"x": 2}}), False),
+        (ClientCapabilities(experimental={"a": {}}), ClientCapabilities(experimental={"a": {}}), True),
+    ],
+)
+def test_check_capability_per_field_branches(have: ClientCapabilities, want: ClientCapabilities, expected: bool):
+    conn = Connection(StubOutbound(), has_standalone_channel=True)
+    conn.client_params = _client_params(have)
+    assert conn.check_capability(want) is expected
+
+
+def test_connection_check_capability_true_when_client_declares_it():
+    conn = Connection(StubOutbound(), has_standalone_channel=True)
+    conn.client_params = _client_params(
+        ClientCapabilities(sampling=SamplingCapability(), roots=RootsCapability(list_changed=True))
+    )
+    conn.initialized.set()
+    assert conn.check_capability(ClientCapabilities(sampling=SamplingCapability())) is True
+    assert conn.check_capability(ClientCapabilities(roots=RootsCapability(list_changed=True))) is True
+    assert conn.check_capability(ClientCapabilities(elicitation=ElicitationCapability())) is False

tests/server/test_lowlevel_exception_handling.py

@@ -1,64 +1,8 @@
-from unittest.mock import AsyncMock, Mock
-
 import anyio
 import pytest
 
-from mcp import types
 from mcp.server.lowlevel.server import Server
-from mcp.server.session import ServerSession
 from mcp.shared.message import SessionMessage
-from mcp.shared.session import RequestResponder
-
-
-@pytest.mark.anyio
-async def test_exception_handling_with_raise_exceptions_true():
-    """Transport exceptions are re-raised when raise_exceptions=True."""
-    server = Server("test-server")
-    session = Mock(spec=ServerSession)
-
-    test_exception = RuntimeError("Test error")
-
-    with pytest.raises(RuntimeError, match="Test error"):
-        await server._handle_message(test_exception, session, {}, raise_exceptions=True)
-
-
-@pytest.mark.anyio
-async def test_exception_handling_with_raise_exceptions_false():
-    """Transport exceptions are logged locally but not sent to the client.
-
-    The transport that reported the error is likely broken; writing back
-    through it races with stream closure (#1967, #2064). The TypeScript,
-    Go, and C# SDKs all log locally only.
-    """
-    server = Server("test-server")
-    session = Mock(spec=ServerSession)
-    session.send_log_message = AsyncMock()
-
-    await server._handle_message(RuntimeError("Test error"), session, {}, raise_exceptions=False)
-
-    session.send_log_message.assert_not_called()
-
-
-@pytest.mark.anyio
-async def test_normal_message_handling_not_affected():
-    """Test that normal messages still work correctly"""
-    server = Server("test-server")
-    session = Mock(spec=ServerSession)
-
-    # Create a mock RequestResponder
-    responder = Mock(spec=RequestResponder)
-    responder.request = types.PingRequest(method="ping")
-    responder.__enter__ = Mock(return_value=responder)
-    responder.__exit__ = Mock(return_value=None)
-
-    # Mock the _handle_request method to avoid complex setup
-    server._handle_request = AsyncMock()
-
-    # Should handle normally without any exception handling
-    await server._handle_message(responder, session, {}, raise_exceptions=False)
-
-    # Verify _handle_request was called
-    server._handle_request.assert_called_once()
 
 
 @pytest.mark.anyio
@@ -71,23 +15,21 @@ async def test_server_run_exits_cleanly_when_transport_yields_exception_then_clo
     1. Transport yields an Exception into the read stream
        (streamable_http.py does this in its broad POST-handler except).
     2. Transport closes the read stream (terminate() in stateless mode).
-    3. _receive_loop exits its `async with read_stream, write_stream:` block,
-       closing the write stream.
-    4. Meanwhile _handle_message(exc) was spawned via tg.start_soon and runs
-       after the write stream is closed.
+    3. The read loop exits and closes the write stream.
 
-    Before the fix, _handle_message tried to send_log_message through the
-    closed write stream, raising ClosedResourceError inside the TaskGroup
-    and crashing server.run(). After the fix, it only logs locally.
+    Before the fix, the message handler tried to send_log_message through the
+    closed write stream, raising ClosedResourceError and crashing server.run().
+    After the fix (and now in the dispatcher), the exception is only logged
+    locally.
     """
     server = Server("test-server")
 
     read_send, read_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
     # Zero-buffer on the write stream forces send() to block until received.
-    # With no receiver, a send() sits blocked until _receive_loop exits its
-    # `async with self._read_stream, self._write_stream:` block and closes the
-    # stream, at which point the blocked send raises ClosedResourceError.
-    # This deterministically reproduces the race without sleeps.
+    # With no receiver, a send() sits blocked until the read loop exits its
+    # `async with read_stream, write_stream:` block and closes the stream, at
+    # which point the blocked send raises ClosedResourceError. This
+    # deterministically reproduces the race without sleeps.
     write_send, write_recv = anyio.create_memory_object_stream[SessionMessage](0)
 
     # What the streamable HTTP transport does: push the exception, then close.
@@ -96,11 +38,11 @@ async def test_server_run_exits_cleanly_when_transport_yields_exception_then_clo
 
     with anyio.fail_after(5):
         # stateless=True so server.run doesn't wait for initialize handshake.
-        # Before this fix, this raised ExceptionGroup(ClosedResourceError).
+        # Before the fix, this raised ExceptionGroup(ClosedResourceError).
         await server.run(read_recv, write_send, server.create_initialization_options(), stateless=True)
 
-    # write_send was closed inside _receive_loop's `async with`; receive_nowait
-    # raises EndOfStream iff the buffer is empty (i.e., server wrote nothing).
+    # write_send was closed inside run's `async with`; receive_nowait raises
+    # EndOfStream iff the buffer is empty (i.e., server wrote nothing).
     with pytest.raises(anyio.EndOfStream):
         write_recv.receive_nowait()
     write_recv.close()

tests/server/test_runner.py

@@ -0,0 +1,952 @@
+"""Tests for `ServerRunner`.
+
+End-to-end over `JSONRPCDispatcher` with a real lowlevel `Server` as the
+registry. The `connected_runner` helper starts both sides and (by default)
+performs the initialize handshake, so each test exercises only the behaviour
+under test.
+"""
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any, cast
+
+import anyio
+import pytest
+from opentelemetry.trace import SpanKind, StatusCode
+
+import mcp.server.runner
+from mcp.server.context import ServerRequestContext
+from mcp.server.lowlevel.server import NotificationOptions, Server
+from mcp.server.models import InitializationOptions
+from mcp.server.runner import ServerRunner, _extract_meta, otel_middleware
+from mcp.server.session import ServerSession
+from mcp.shared.dispatcher import DispatchContext, DispatchMiddleware, OnRequest
+from mcp.shared.exceptions import MCPError
+from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
+from mcp.shared.peer import dump_params
+from mcp.shared.transport_context import TransportContext
+from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
+from mcp.types import (
+    INTERNAL_ERROR,
+    INVALID_PARAMS,
+    LATEST_PROTOCOL_VERSION,
+    METHOD_NOT_FOUND,
+    CallToolRequestParams,
+    ClientCapabilities,
+    ErrorData,
+    Implementation,
+    InitializeRequestParams,
+    ListToolsResult,
+    NotificationParams,
+    PaginatedRequestParams,
+    ProgressNotificationParams,
+    RequestParams,
+    SetLevelRequestParams,
+    Tool,
+)
+
+from ..shared.conftest import jsonrpc_pair
+from ..shared.test_dispatcher import Recorder, echo_handlers
+from .conftest import SpanCapture
+
+Ctx = ServerRequestContext[dict[str, Any], Any]
+
+
+def _initialize_params() -> dict[str, Any]:
+    return InitializeRequestParams(
+        protocol_version=LATEST_PROTOCOL_VERSION,
+        capabilities=ClientCapabilities(),
+        client_info=Implementation(name="test-client", version="1.0"),
+    ).model_dump(by_alias=True, exclude_none=True)
+
+
+_seen_ctx: list[Ctx] = []
+SrvT = Server[dict[str, Any]]
+
+
+@pytest.fixture
+def server() -> SrvT:
+    """A lowlevel Server with one tools/list handler registered."""
+    _seen_ctx.clear()
+
+    async def list_tools(ctx: Ctx, params: PaginatedRequestParams | None) -> ListToolsResult:
+        _seen_ctx.append(ctx)
+        return ListToolsResult(tools=[Tool(name="t", input_schema={"type": "object"})])
+
+    return Server(name="test-server", version="0.0.1", on_list_tools=list_tools)
+
+
+@asynccontextmanager
+async def connected_runner(
+    server: SrvT,
+    *,
+    initialized: bool = True,
+    stateless: bool = False,
+    has_standalone_channel: bool = True,
+    init_options: InitializationOptions | None = None,
+    session_id: str | None = None,
+    dispatch_middleware: list[DispatchMiddleware] | None = None,
+) -> AsyncIterator[tuple[JSONRPCDispatcher[TransportContext], ServerRunner[dict[str, Any]]]]:
+    """Yield `(client, runner)` running over an in-memory JSON-RPC dispatcher pair.
+
+    Starts the client (echo handlers) and `runner.run()` in a task group, wraps
+    the body in `anyio.fail_after(5)`, and cancels on exit. When
+    `initialized` is true the helper performs the real `initialize` request
+    before yielding, so tests start past the init-gate via the public path.
+    """
+    client, server_d, close = jsonrpc_pair()
+    assert isinstance(client, JSONRPCDispatcher) and isinstance(server_d, JSONRPCDispatcher)
+    runner = ServerRunner(
+        server=server,
+        dispatcher=server_d,
+        lifespan_state={},
+        has_standalone_channel=has_standalone_channel,
+        init_options=init_options,
+        session_id=session_id,
+        stateless=stateless,
+        dispatch_middleware=dispatch_middleware or [],
+    )
+    c_req, c_notify = echo_handlers(Recorder())
+    body_exc: BaseException | None = None
+    async with anyio.create_task_group() as tg:
+        await tg.start(client.run, c_req, c_notify)
+        await tg.start(runner.run)
+        try:
+            with anyio.fail_after(5):
+                if initialized:
+                    await client.send_raw_request("initialize", _initialize_params())
+                yield client, runner
+        except BaseException as e:
+            # Capture and re-raise outside the task group so test failures
+            # surface as the original exception, not an ExceptionGroup wrapper.
+            body_exc = e
+        close()
+    if body_exc is not None:
+        raise body_exc
+
+
+@pytest.mark.anyio
+async def test_connected_runner_propagates_body_exception_unwrapped(server: SrvT):
+    """The harness re-raises body exceptions as-is, not as `ExceptionGroup`."""
+    with pytest.raises(RuntimeError, match="boom"):
+        async with connected_runner(server):
+            raise RuntimeError("boom")
+
+
+@pytest.mark.anyio
+async def test_runner_handles_initialize_and_populates_connection(server: SrvT):
+    async with connected_runner(server, initialized=False) as (client, runner):
+        result = await client.send_raw_request("initialize", _initialize_params())
+    assert result["serverInfo"]["name"] == "test-server"
+    assert "tools" in result["capabilities"]
+    assert runner.connection.client_params is not None
+    assert runner.connection.client_params.client_info.name == "test-client"
+    assert runner.connection.protocol_version == LATEST_PROTOCOL_VERSION
+    assert runner.connection.initialize_accepted is True
+
+
+@pytest.mark.anyio
+async def test_runner_initialize_opens_gate_but_event_fires_only_after_initialized_notification(server: SrvT):
+    """`initialize` commits the gate flag and peer info, but the public
+    `connection.initialized` event waits for `notifications/initialized` (the
+    point from which the spec permits server-initiated requests)."""
+    async with connected_runner(server, initialized=False) as (client, runner):
+        await client.send_raw_request("initialize", _initialize_params())
+        assert runner.connection.initialize_accepted is True
+        assert not runner.connection.initialized.is_set()
+        await client.notify("notifications/initialized", None)
+        await runner.connection.initialized.wait()
+
+
+@pytest.mark.anyio
+async def test_runner_gates_requests_before_initialize(server: SrvT):
+    async with connected_runner(server, initialized=False) as (client, _):
+        with 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="")
+        # ping is exempt from the gate
+        assert await client.send_raw_request("ping", None) == {}
+
+
+@pytest.mark.anyio
+async def test_runner_routes_to_handler_and_builds_context(server: SrvT):
+    async with connected_runner(server) as (client, runner):
+        result = await client.send_raw_request("tools/list", None)
+    assert result["tools"][0]["name"] == "t"
+    ctx = _seen_ctx[0]
+    assert isinstance(ctx, ServerRequestContext)
+    assert ctx.lifespan_context == {}
+    assert isinstance(ctx.session, ServerSession)
+    assert ctx.session is runner.session
+    assert ctx.request_id is not None
+
+
+@pytest.mark.anyio
+async def test_runner_spec_method_with_no_handler_raises_method_not_found(server: SrvT):
+    async with connected_runner(server) as (client, _):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("resources/list", None)
+    assert exc.value.error.code == METHOD_NOT_FOUND
+
+
+@pytest.mark.anyio
+async def test_runner_non_spec_method_with_no_handler_raises_method_not_found(server: SrvT):
+    """Upfront validation is gated to spec methods, so a non-spec method
+    skips it and reaches handler lookup."""
+    async with connected_runner(server) as (client, _):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("nonexistent/method", None)
+    assert exc.value.error.code == METHOD_NOT_FOUND
+
+
+@pytest.mark.anyio
+async def test_runner_malformed_params_for_unregistered_spec_method_raises_invalid_params(server: SrvT):
+    """A spec method with malformed params is INVALID_PARAMS even with no handler."""
+    async with connected_runner(server) as (client, _):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("tools/call", {"name": 123})
+    assert exc.value.error == ErrorData(code=INVALID_PARAMS, message="Invalid request parameters", data="")
+
+
+@pytest.mark.anyio
+async def test_runner_rejects_snake_case_initialize_params(server: SrvT):
+    """Inbound wire payloads validate alias-only; Python field names are not
+    accepted (`protocol_version` must arrive as `protocolVersion`)."""
+    snake = {
+        "protocol_version": LATEST_PROTOCOL_VERSION,
+        "capabilities": {},
+        "client_info": {"name": "c", "version": "0"},
+    }
+    async with connected_runner(server, initialized=False) as (client, _):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("initialize", snake)
+    assert exc.value.error.code == INVALID_PARAMS
+
+
+@pytest.mark.anyio
+async def test_runner_initialize_with_absent_params_returns_invalid_params_and_stays_alive(server: SrvT):
+    """Re-covers what the old `tests/issues/test_malformed_input.py` pinned: a
+    malformed `initialize` is rejected and the runner keeps serving."""
+    async with connected_runner(server, initialized=False) as (client, _):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("initialize", None)
+        assert exc.value.error.code == INVALID_PARAMS
+        result = await client.send_raw_request("initialize", _initialize_params())
+    assert result["serverInfo"]["name"] == "test-server"
+
+
+@pytest.mark.anyio
+async def test_runner_rejects_snake_case_params_for_custom_handler(server: SrvT):
+    """Custom-method handlers (which skip the spec-method gate) still validate
+    alias-only at the per-handler boundary."""
+
+    async def handler(ctx: Ctx, params: ProgressNotificationParams) -> dict[str, Any]:
+        return {"ok": True}
+
+    server.add_request_handler("custom/progress", ProgressNotificationParams, handler)
+    async with connected_runner(server) as (client, _):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("custom/progress", {"progress_token": 1, "progress": 0.5})
+        assert exc.value.error.code == INVALID_PARAMS
+        result = await client.send_raw_request("custom/progress", {"progressToken": 1, "progress": 0.5})
+    assert result == {"ok": True}
+
+
+@pytest.mark.anyio
+async def test_runner_on_notify_drops_snake_case_params(server: SrvT, caplog: pytest.LogCaptureFixture):
+    """Notification params validate alias-only; snake_case is dropped as malformed."""
+
+    async def handler(ctx: Ctx, params: ProgressNotificationParams) -> None:
+        raise NotImplementedError
+
+    server.add_notification_handler("notifications/roots/list_changed", ProgressNotificationParams, handler)
+    async with connected_runner(server) as (client, _):
+        await client.notify("notifications/roots/list_changed", {"progress_token": 1, "progress": 0.5})
+        await client.send_raw_request("tools/list", None)
+    assert "dropped 'notifications/roots/list_changed': malformed params" in caplog.text
+
+
+@pytest.mark.anyio
+async def test_runner_on_notify_initialized_sets_flag_and_connection_event(server: SrvT):
+    async with connected_runner(server, initialized=False) as (client, runner):
+        await client.notify("notifications/initialized", None)
+        await runner.connection.initialized.wait()
+    assert runner.connection.initialize_accepted is True
+
+
+@pytest.mark.anyio
+async def test_runner_on_notify_malformed_initialized_does_not_initialize(
+    server: SrvT, caplog: pytest.LogCaptureFixture
+):
+    """A malformed `notifications/initialized` drops like any other malformed
+    notification and leaves the connection uninitialized."""
+    async with connected_runner(server, initialized=False) as (client, runner):
+        await client.notify("notifications/initialized", {"_meta": 42})
+        await anyio.wait_all_tasks_blocked()
+        assert runner.connection.initialize_accepted is False
+        assert not runner.connection.initialized.is_set()
+    assert "dropped 'notifications/initialized': malformed params" in caplog.text
+
+
+@pytest.mark.anyio
+async def test_runner_on_notify_initialized_routes_to_registered_handler_after_state_set(server: SrvT):
+    """A handler registered for `notifications/initialized` fires after the
+    runner flips the init state, so it observes an initialized connection."""
+    seen: list[bool] = []
+    delivered = anyio.Event()
+
+    async def on_initialized(ctx: Ctx, params: NotificationParams | None) -> None:
+        seen.append(runner.connection.initialize_accepted and runner.connection.initialized.is_set())
+        delivered.set()
+
+    server.add_notification_handler("notifications/initialized", NotificationParams, on_initialized)
+    async with connected_runner(server, initialized=False) as (client, runner):
+        await client.notify("notifications/initialized", {"_meta": {"k": "v"}})
+        await delivered.wait()
+    assert seen == [True]
+
+
+def test_server_add_request_handler_rejects_initialize():
+    async def handler(ctx: Ctx, params: InitializeRequestParams) -> dict[str, Any]:
+        raise NotImplementedError
+
+    server: SrvT = Server(name="s")
+    with pytest.raises(ValueError, match="Server.middleware"):
+        server.add_request_handler("initialize", InitializeRequestParams, handler)
+    assert server.get_request_handler("initialize") is None
+
+
+@pytest.mark.anyio
+async def test_runner_on_notify_routes_to_registered_handler(server: SrvT):
+    seen: list[tuple[Any, Any]] = []
+    delivered = anyio.Event()
+
+    async def on_roots_changed(ctx: Ctx, params: NotificationParams | None) -> None:
+        seen.append((ctx, params))
+        if len(seen) == 2:
+            delivered.set()
+
+    server.add_notification_handler("notifications/roots/list_changed", NotificationParams, on_roots_changed)
+    async with connected_runner(server) as (client, _):
+        await client.notify("notifications/roots/list_changed", None)
+        await client.notify("notifications/roots/list_changed", {})
+        await delivered.wait()
+    assert isinstance(seen[0][0], ServerRequestContext)
+    # Absent and present-but-empty wire params both validate to the defaults model.
+    assert seen[0][1] == NotificationParams()
+    assert isinstance(seen[1][1], NotificationParams)
+
+
+@pytest.mark.anyio
+async def test_runner_on_notify_handler_exception_is_swallowed_and_logged(
+    server: SrvT, caplog: pytest.LogCaptureFixture
+):
+    """A notification handler crashing must not tear down the connection."""
+
+    async def boom(ctx: Ctx, params: NotificationParams | None) -> None:
+        raise RuntimeError("notification handler boom")
+
+    server.add_notification_handler("notifications/roots/list_changed", NotificationParams, boom)
+    async with connected_runner(server) as (client, _):
+        await client.notify("notifications/roots/list_changed", None)
+        # Connection still alive: a request after the crashing handler succeeds.
+        result = await client.send_raw_request("tools/list", None)
+    assert result["tools"][0]["name"] == "t"
+    assert "notification handler for 'notifications/roots/list_changed' raised" in caplog.text
+
+
+@pytest.mark.anyio
+async def test_runner_on_notify_drops_malformed_params(server: SrvT, caplog: pytest.LogCaptureFixture):
+    """Malformed notification params are logged and dropped, not raised."""
+
+    async def on_level(ctx: Ctx, params: SetLevelRequestParams) -> None:
+        raise NotImplementedError
+
+    server.add_notification_handler("notifications/roots/list_changed", SetLevelRequestParams, on_level)
+    async with connected_runner(server) as (client, _):
+        await client.notify("notifications/roots/list_changed", {"level": "not-a-level"})
+        result = await client.send_raw_request("tools/list", None)
+    assert result["tools"][0]["name"] == "t"
+    assert "dropped 'notifications/roots/list_changed': malformed params" in caplog.text
+
+
+@pytest.mark.anyio
+async def test_runner_on_notify_drops_absent_params_when_model_requires_them(
+    server: SrvT, caplog: pytest.LogCaptureFixture
+):
+    """A params-less progress notification is dropped, not delivered as None.
+
+    `on_progress` is typed to receive a non-Optional `ProgressNotificationParams`;
+    the previous server validated the full notification union and dropped this
+    as malformed before dispatch.
+    """
+
+    async def on_progress(ctx: Ctx, params: ProgressNotificationParams) -> None:
+        raise NotImplementedError
+
+    server.add_notification_handler("notifications/progress", ProgressNotificationParams, on_progress)
+    async with connected_runner(server) as (client, _):
+        await client.notify("notifications/progress", None)
+        result = await client.send_raw_request("tools/list", None)
+    assert result["tools"][0]["name"] == "t"
+    assert "dropped 'notifications/progress': malformed params" in caplog.text
+    assert "notification handler for" not in caplog.text
+
+
+@pytest.mark.anyio
+async def test_runner_absent_wire_params_reaches_request_handler_as_defaults_model():
+    """A request with no `params` member on the wire reaches the handler as
+    the params model with its defaults, never `None`.
+
+    The in-SDK client always attaches `_meta`, so a dispatch middleware
+    forwards `params=None` to model what an external client sends.
+    """
+    seen: list[PaginatedRequestParams | None] = []
+
+    async def list_tools(ctx: Ctx, params: PaginatedRequestParams | None) -> ListToolsResult:
+        seen.append(params)
+        return ListToolsResult(tools=[])
+
+    def drop_params(next_on_request: OnRequest) -> OnRequest:
+        async def wrapped(dctx: DispatchContext[Any], method: str, params: Any) -> dict[str, Any]:
+            return await next_on_request(dctx, method, None if method == "tools/list" else params)
+
+        return wrapped
+
+    server: SrvT = Server(name="s", on_list_tools=list_tools)
+    async with connected_runner(server, dispatch_middleware=[drop_params]) as (client, _):
+        await client.send_raw_request("tools/list", None)
+    assert seen == [PaginatedRequestParams()]
+
+
+@pytest.mark.anyio
+async def test_runner_absent_wire_params_for_required_params_custom_method_is_invalid_params():
+    """A custom method whose `params_type` has required fields rejects absent
+    wire params as INVALID_PARAMS rather than invoking the handler with None."""
+
+    class GreetParams(RequestParams):
+        name: str
+
+    async def greet(ctx: Ctx, params: GreetParams) -> dict[str, Any]:
+        raise NotImplementedError
+
+    def drop_params(next_on_request: OnRequest) -> OnRequest:
+        async def wrapped(dctx: DispatchContext[Any], method: str, params: Any) -> dict[str, Any]:
+            return await next_on_request(dctx, method, None if method == "custom/greet" else params)
+
+        return wrapped
+
+    server: SrvT = Server(name="s")
+    server.add_request_handler("custom/greet", GreetParams, greet)
+    async with connected_runner(server, dispatch_middleware=[drop_params]) as (client, _):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("custom/greet", {"name": "x"})
+    assert exc.value.error.code == INVALID_PARAMS
+
+
+@pytest.mark.anyio
+async def test_runner_on_notify_drops_before_init_and_unknown_methods(server: SrvT):
+    seen: list[Any] = []
+
+    async def on_roots(ctx: Ctx, params: NotificationParams | None) -> None:
+        seen.append(params)
+
+    server.add_notification_handler("notifications/roots/list_changed", NotificationParams, on_roots)
+    async with connected_runner(server, initialized=False) as (client, _):
+        await client.notify("notifications/roots/list_changed", None)  # before init: dropped
+        await client.notify("notifications/initialized", None)
+        await client.notify("notifications/unknown", None)  # no handler: dropped
+        await client.notify("notifications/roots/list_changed", None)  # post-init: delivered
+        await anyio.wait_all_tasks_blocked()
+    assert seen == [NotificationParams()]  # only the post-init one reached the handler
+
+
+@pytest.mark.anyio
+async def test_runner_dispatch_middleware_wraps_everything_including_initialize(server: SrvT):
+    seen_methods: list[str] = []
+
+    def trace_mw(next_on_request: Any) -> Any:
+        async def wrapped(dctx: Any, method: str, params: Any) -> Any:
+            seen_methods.append(method)
+            return await next_on_request(dctx, method, params)
+
+        return wrapped
+
+    async with connected_runner(server, dispatch_middleware=[trace_mw]) as (client, _):
+        await client.send_raw_request("tools/list", None)
+    assert seen_methods == ["initialize", "tools/list"]
+
+
+@pytest.mark.anyio
+async def test_runner_server_middleware_wraps_every_request_including_initialize(server: SrvT):
+    seen: list[tuple[str, Any]] = []
+
+    async def ctx_mw(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
+        seen.append((method, params))
+        return await call_next()
+
+    server.middleware.append(ctx_mw)
+    async with connected_runner(server) as (client, _):
+        await client.send_raw_request("ping", None)
+        await client.send_raw_request("tools/list", {"_meta": {"k": "v"}})
+    assert [m for m, _ in seen] == ["initialize", "ping", "tools/list"]
+    # params arrive raw (Mapping), not as a validated model
+    assert seen[2][1] == {"_meta": {"k": "v"}}
+
+
+@pytest.mark.anyio
+async def test_runner_middleware_raise_after_call_next_on_initialize_leaves_connection_uninitialized(server: SrvT):
+    """A middleware failure after `call_next()` on `initialize` reaches the
+    client as an error and skips the state commit: the pre-init gate stays
+    closed and `connection.initialized` never fires."""
+
+    async def reject_initialize(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
+        result = await call_next()
+        if method == "initialize":
+            raise MCPError(code=INTERNAL_ERROR, message="rejected by middleware")
+        return result
+
+    server.middleware.append(reject_initialize)
+    async with connected_runner(server, initialized=False) as (client, runner):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("initialize", _initialize_params())
+        assert exc.value.error.message == "rejected by middleware"
+        with pytest.raises(MCPError) as gate_exc:
+            await client.send_raw_request("tools/list", None)
+        assert gate_exc.value.error == ErrorData(code=INVALID_PARAMS, message="Invalid request parameters", data="")
+        # ping passes through the middleware untouched
+        assert await client.send_raw_request("ping", None) == {}
+    assert runner.connection.initialize_accepted is False
+    assert runner.connection.client_params is None
+    assert runner.connection.protocol_version is None
+    assert not runner.connection.initialized.is_set()
+
+
+@pytest.mark.anyio
+async def test_runner_server_middleware_observes_method_not_found_via_call_next_raise(server: SrvT):
+    seen: list[tuple[str, type[BaseException] | None]] = []
+
+    async def observe(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
+        try:
+            return await call_next()
+        except MCPError as e:
+            seen.append((method, type(e)))
+            raise
+
+    server.middleware.append(observe)
+    async with connected_runner(server) as (client, _):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("nonexistent/method", None)
+    assert exc.value.error.code == METHOD_NOT_FOUND
+    assert seen == [("nonexistent/method", MCPError)]
+
+
+@pytest.mark.anyio
+async def test_runner_server_middleware_wraps_notifications(server: SrvT):
+    """The same chain wraps `_on_notify`: it sees `notifications/initialized`,
+    pre-init drops, and registered notification handlers, with
+    `ctx.request_id is None`."""
+    seen: list[tuple[str, bool]] = []
+
+    async def observe(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
+        seen.append((method, ctx.request_id is None))
+        return await call_next()
+
+    async def on_roots(ctx: Ctx, params: NotificationParams | None) -> None:
+        return None
+
+    server.add_notification_handler("notifications/roots/list_changed", NotificationParams, on_roots)
+    server.middleware.append(observe)
+    async with connected_runner(server, initialized=False) as (client, _):
+        await client.notify("notifications/roots/list_changed", None)  # pre-init drop, still observed
+        await client.notify("notifications/initialized", None)
+        await client.notify("notifications/roots/list_changed", None)
+        await anyio.wait_all_tasks_blocked()
+    assert seen == [
+        ("notifications/roots/list_changed", True),
+        ("notifications/initialized", True),
+        ("notifications/roots/list_changed", True),
+    ]
+
+
+def test_extract_meta_returns_none_for_absent_or_malformed():
+    """Context construction is independent of `_meta` validity; the params
+    validation inside `call_next()` is what surfaces the error."""
+    assert _extract_meta(None) is None
+    assert _extract_meta({}) is None
+    assert _extract_meta({"_meta": "not-a-dict"}) is None
+    assert _extract_meta({"_meta": {"progressToken": []}}) is None
+    assert _extract_meta({"_meta": {"progressToken": "x", "k": 1}}) == {"progress_token": "x", "k": 1}
+
+
+def test_extract_meta_round_trips_through_dump_params():
+    """Forwarding an inbound `ctx.meta` outbound (`meta=ctx.meta`) re-emits the
+    wire key `progressToken`, not the Python field name `_extract_meta`
+    validation produced."""
+    meta = _extract_meta({"_meta": {"progressToken": 7, "k": 1}})
+    assert meta is not None
+    assert dump_params(None, dict(meta)) == {"_meta": {"progressToken": 7, "k": 1}}
+
+
+@pytest.mark.anyio
+async def test_runner_server_middleware_runs_outermost_first(server: SrvT):
+    order: list[str] = []
+
+    def make_mw(tag: str) -> Any:
+        async def mw(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
+            order.append(f"{tag}-in")
+            result = await call_next()
+            order.append(f"{tag}-out")
+            return result
+
+        return mw
+
+    server.middleware.extend([make_mw("a"), make_mw("b")])
+    async with connected_runner(server) as (client, _):
+        order.clear()  # drop the wrap of the helper's `initialize`
+        await client.send_raw_request("tools/list", None)
+    assert order == ["a-in", "b-in", "b-out", "a-out"]
+
+
+@pytest.mark.anyio
+async def test_runner_handler_returning_none_yields_empty_result(server: SrvT):
+    async def set_level(ctx: Ctx, params: SetLevelRequestParams) -> None:
+        return None
+
+    server.add_request_handler("logging/setLevel", SetLevelRequestParams, set_level)
+    async with connected_runner(server) as (client, _):
+        result = await client.send_raw_request("logging/setLevel", {"level": "info"})
+    assert result == {}
+
+
+@pytest.mark.anyio
+async def test_runner_handler_returning_error_data_produces_jsonrpc_error(server: SrvT):
+    """A handler returning `ErrorData` reaches the client as a JSON-RPC error,
+    not a success result, matching `BaseSession._send_response`."""
+
+    async def set_level(ctx: Ctx, params: SetLevelRequestParams) -> ErrorData:
+        return ErrorData(code=INVALID_PARAMS, message="bad level", data={"got": params.level})
+
+    server.add_request_handler("logging/setLevel", SetLevelRequestParams, set_level)
+    async with connected_runner(server) as (client, _):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("logging/setLevel", {"level": "info"})
+    assert exc.value.error == ErrorData(code=INVALID_PARAMS, message="bad level", data={"got": "info"})
+
+
+@pytest.mark.anyio
+async def test_runner_server_middleware_observes_handler_error_data_as_mcp_error(server: SrvT):
+    """A handler returning `ErrorData` raises `MCPError` inside `call_next()`,
+    so observation middleware records the failure instead of seeing a
+    successful-looking `ErrorData` return."""
+    seen: list[MCPError] = []
+
+    async def observe(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
+        try:
+            return await call_next()
+        except MCPError as e:
+            seen.append(e)
+            raise
+
+    async def set_level(ctx: Ctx, params: SetLevelRequestParams) -> ErrorData:
+        return ErrorData(code=INVALID_PARAMS, message="bad level")
+
+    server.middleware.append(observe)
+    server.add_request_handler("logging/setLevel", SetLevelRequestParams, set_level)
+    async with connected_runner(server) as (client, _):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("logging/setLevel", {"level": "info"})
+    assert exc.value.error == ErrorData(code=INVALID_PARAMS, message="bad level")
+    assert [e.error.message for e in seen] == ["bad level"]
+
+
+@pytest.mark.anyio
+async def test_runner_middleware_returning_error_data_produces_jsonrpc_error(server: SrvT):
+    """A middleware that short-circuits with an `ErrorData` return gets the
+    same treatment as a handler return: the wire sees a JSON-RPC error."""
+
+    async def short_circuit(ctx: Ctx, method: str, params: Any, call_next: Any) -> Any:
+        return ErrorData(code=INVALID_PARAMS, message="denied")
+
+    server.middleware.append(short_circuit)
+    async with connected_runner(server, initialized=False) as (client, _):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("tools/list", None)
+    assert exc.value.error == ErrorData(code=INVALID_PARAMS, message="denied")
+
+
+@pytest.mark.anyio
+async def test_runner_handler_returning_unsupported_type_surfaces_as_error(server: SrvT):
+    async def bad_return(ctx: Ctx, params: PaginatedRequestParams | None) -> int:
+        return 42
+
+    # cast: deliberately registering a handler with a bad return type to
+    # exercise the runtime check; pyright would (correctly) reject it otherwise.
+    server.add_request_handler("tools/list", PaginatedRequestParams, cast(Any, bad_return))
+    async with connected_runner(server) as (client, _):
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("tools/list", None)
+    assert exc.value.error.code == 0
+    assert "int" in exc.value.error.message
+
+
+@pytest.mark.anyio
+async def test_runner_stateless_skips_init_gate(server: SrvT):
+    async with connected_runner(server, initialized=False, stateless=True, has_standalone_channel=False) as (client, _):
+        result = await client.send_raw_request("tools/list", None)
+    assert result["tools"][0]["name"] == "t"
+
+
+@pytest.mark.anyio
+async def test_runner_stateless_connection_initialized_event_set_on_construction(server: SrvT):
+    """`connection.initialized` mirrors the gate flag in stateless mode so
+    `await connection.initialized.wait()` does not hang when no handshake
+    arrives."""
+    async with connected_runner(server, initialized=False, stateless=True, has_standalone_channel=False) as (_, runner):
+        assert runner.connection.initialize_accepted is True
+        assert runner.connection.initialized.is_set()
+        await runner.connection.initialized.wait()
+
+
+@pytest.mark.anyio
+async def test_server_add_request_handler_routes_custom_method_with_validated_params(server: SrvT):
+    """Custom methods outside the spec `ClientRequest` union skip upfront
+    validation and route to the registered handler."""
+
+    class GreetParams(RequestParams):
+        name: str
+
+    received: list[GreetParams] = []
+
+    async def greet(ctx: Ctx, params: GreetParams) -> dict[str, Any]:
+        received.append(params)
+        return {"greeting": f"hello {params.name}"}
+
+    server.add_request_handler("custom/greet", GreetParams, greet)
+    async with connected_runner(server) as (client, _):
+        result = await client.send_raw_request("custom/greet", {"name": "world"})
+    assert result == {"greeting": "hello world"}
+    assert isinstance(received[0], GreetParams)
+    assert received[0].name == "world"
+
+
+@pytest.mark.anyio
+async def test_runner_initialize_result_reflects_init_options():
+    async def list_tools(ctx: Ctx, params: PaginatedRequestParams | None) -> ListToolsResult:
+        raise NotImplementedError
+
+    server: SrvT = Server(name="caps-test", on_list_tools=list_tools, instructions="be nice")
+    init_options = server.create_initialization_options(NotificationOptions(tools_changed=True), {"ext": {"k": "v"}})
+    async with connected_runner(server, initialized=False, init_options=init_options) as (client, _):
+        result = await client.send_raw_request("initialize", _initialize_params())
+    assert result["capabilities"]["tools"]["listChanged"] is True
+    assert result["capabilities"]["experimental"] == {"ext": {"k": "v"}}
+    assert result["serverInfo"]["name"] == "caps-test"
+    assert result["instructions"] == "be nice"
+
+
+@pytest.mark.anyio
+async def test_runner_initialize_echoes_supported_version_and_falls_back_to_latest(server: SrvT):
+    oldest = SUPPORTED_PROTOCOL_VERSIONS[0]
+    async with connected_runner(server, initialized=False) as (client, _):
+        params = {**_initialize_params(), "protocolVersion": oldest}
+        result = await client.send_raw_request("initialize", params)
+        assert result["protocolVersion"] == oldest
+    async with connected_runner(server, initialized=False) as (client, _):
+        params = {**_initialize_params(), "protocolVersion": "1999-01-01"}
+        result = await client.send_raw_request("initialize", params)
+        assert result["protocolVersion"] == LATEST_PROTOCOL_VERSION
+
+
+@pytest.mark.anyio
+async def test_otel_middleware_emits_server_span_with_method_and_target(server: SrvT, spans: SpanCapture):
+    async def call_tool(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
+        return {"content": [], "isError": False}
+
+    server.add_request_handler("tools/call", CallToolRequestParams, call_tool)
+    async with connected_runner(server, dispatch_middleware=[otel_middleware]) as (client, _):
+        spans.clear()
+        result = await client.send_raw_request("tools/call", {"name": "mytool", "arguments": {}})
+    assert result == {"content": [], "isError": False}
+    finished = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    [span] = finished
+    assert span.name == "MCP handle tools/call mytool"
+    assert span.attributes is not None
+    assert span.attributes["mcp.method.name"] == "tools/call"
+    assert isinstance(span.attributes["jsonrpc.request.id"], str)
+    assert span.status.status_code == StatusCode.UNSET
+
+
+@pytest.mark.anyio
+async def test_otel_trace_context_propagates_client_to_server(server: SrvT, spans: SpanCapture):
+    """The client dispatcher injects traceparent into `_meta`; the server's
+    `otel_middleware` extracts it, so client and server spans share a trace."""
+    async with connected_runner(server, dispatch_middleware=[otel_middleware]) as (client, _):
+        spans.clear()
+        await client.send_raw_request("tools/list", None)
+    [client_span] = [s for s in spans.finished() if s.kind == SpanKind.CLIENT]
+    [server_span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert server_span.parent is not None
+    assert client_span.context is not None and server_span.context is not None
+    assert server_span.parent.span_id == client_span.context.span_id
+    assert server_span.context.trace_id == client_span.context.trace_id
+    assert client_span.attributes is not None and server_span.attributes is not None
+    assert client_span.attributes["jsonrpc.request.id"] == server_span.attributes["jsonrpc.request.id"]
+
+
+@pytest.mark.anyio
+async def test_otel_middleware_malformed_traceparent_degrades_to_no_parent(server: SrvT, spans: SpanCapture):
+    """A non-string traceparent in `_meta` must not fail the request; the
+    server span simply gets no parent."""
+
+    def break_traceparent(next_on_request: OnRequest) -> OnRequest:
+        async def wrapped(dctx: DispatchContext[Any], method: str, params: Any) -> dict[str, Any]:
+            mangled = {"_meta": {"traceparent": 123}} if method == "tools/list" else params
+            return await next_on_request(dctx, method, mangled)
+
+        return wrapped
+
+    async with connected_runner(server, dispatch_middleware=[break_traceparent, otel_middleware]) as (client, _):
+        spans.clear()
+        await client.send_raw_request("tools/list", None)
+    [server_span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert server_span.parent is None
+
+
+@pytest.mark.anyio
+async def test_otel_middleware_validation_failure_sets_sanitized_status(server: SrvT, spans: SpanCapture):
+    """Malformed params set the sanitized wire message as span status and do
+    not record the pydantic exception (it carries client input)."""
+    async with connected_runner(server, dispatch_middleware=[otel_middleware]) as (client, _):
+        spans.clear()
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("tools/call", {"name": 123})
+    assert exc.value.error.code == INVALID_PARAMS
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.status.status_code == StatusCode.ERROR
+    assert span.status.description == "Invalid request parameters"
+    assert not span.events
+
+
+@pytest.mark.anyio
+async def test_otel_middleware_records_error_status_on_mcp_error(server: SrvT, spans: SpanCapture):
+    async with connected_runner(server, dispatch_middleware=[otel_middleware]) as (client, _):
+        spans.clear()
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("resources/list", None)
+        assert exc.value.error.code == METHOD_NOT_FOUND
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.status.status_code == StatusCode.ERROR
+    assert span.status.description == "Method not found"
+    # MCPError is a protocol-level response, not a crash - no traceback event.
+    assert not [e for e in span.events if e.name == "exception"]
+
+
+@pytest.mark.anyio
+async def test_otel_middleware_records_error_status_on_handler_exception(server: SrvT, spans: SpanCapture):
+    async def failing(ctx: Ctx, params: PaginatedRequestParams | None) -> Any:
+        raise ValueError("handler blew up")
+
+    server.add_request_handler("tools/list", PaginatedRequestParams, failing)
+    async with connected_runner(server, dispatch_middleware=[otel_middleware]) as (client, _):
+        spans.clear()
+        with pytest.raises(MCPError) as exc:
+            await client.send_raw_request("tools/list", None)
+        assert exc.value.error.code == 0
+    [span] = [s for s in spans.finished() if s.kind == SpanKind.SERVER]
+    assert span.status.status_code == StatusCode.ERROR
+    assert span.status.description == "handler blew up"
+    [event] = [e for e in span.events if e.name == "exception"]
+    assert event.attributes is not None
+    assert event.attributes["exception.type"] == "ValueError"
+
+
+@pytest.mark.anyio
+async def test_runner_connection_exit_stack_unwinds_after_run_returns(server: SrvT) -> None:
+    """`runner.connection.exit_stack` is closed when the dispatcher loop ends."""
+    cleaned: list[int] = []
+
+    async def _append(i: int) -> None:
+        cleaned.append(i)
+
+    async with connected_runner(server) as (client, runner):
+        for i in (1, 2, 3):
+            runner.connection.exit_stack.push_async_callback(_append, i)
+        await client.send_raw_request("tools/list", None)
+        assert cleaned == []
+    assert cleaned == [3, 2, 1]
+
+
+@pytest.mark.anyio
+async def test_runner_exit_stack_cleanup_exception_is_logged_not_propagated(
+    server: SrvT, caplog: pytest.LogCaptureFixture
+) -> None:
+    """A raising cleanup callback is caught and logged; `run()` exits cleanly."""
+    cleaned: list[str] = []
+
+    async def _ok() -> None:
+        cleaned.append("ok")
+
+    async def _boom() -> None:
+        raise RuntimeError("cleanup failed")
+
+    async with connected_runner(server) as (client, runner):
+        runner.connection.exit_stack.push_async_callback(_ok)
+        runner.connection.exit_stack.push_async_callback(_boom)
+        await client.send_raw_request("tools/list", None)
+    assert cleaned == ["ok"]
+    assert "connection exit_stack cleanup raised" in caplog.text
+
+
+@pytest.mark.anyio
+async def test_runner_exit_stack_blocking_cleanup_abandoned_after_grace(
+    server: SrvT, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    """A cleanup callback that never returns is abandoned once the grace period
+    elapses: `run()` exits, later callbacks in the unwind are cancelled at
+    their first checkpoint, and a warning is logged. Grace 0 means the deadline
+    is already expired on entry, so the abandonment is immediate."""
+    monkeypatch.setattr(mcp.server.runner, "_EXIT_STACK_CLOSE_TIMEOUT", 0)
+    ran: list[str] = []
+    release = anyio.Event()
+
+    async def _abandoned() -> None:
+        # LIFO unwind: pushed first, so it runs after the blocker. By then the
+        # deadline has fired, so this checkpoint raises and the line below is
+        # unreachable (if abandonment broke, the missing warning fails the
+        # caplog assert).
+        await anyio.sleep(0)
+        raise NotImplementedError
+
+    async def _blocker() -> None:
+        ran.append("blocker started")
+        await release.wait()
+        raise NotImplementedError
+
+    async with connected_runner(server) as (client, runner):
+        runner.connection.exit_stack.push_async_callback(_abandoned)
+        runner.connection.exit_stack.push_async_callback(_blocker)
+        await client.send_raw_request("tools/list", None)
+    assert ran == ["blocker started"]
+    assert "abandoning remaining callbacks" in caplog.text
+
+
+@pytest.mark.anyio
+async def test_runner_exit_stack_fast_cleanup_completes_within_grace(
+    server: SrvT, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Well-behaved cleanup callbacks run to completion under the bounded
+    unwind and no abandonment warning is logged. Uses the production grace;
+    the deadline never delays a fast unwind, it only bounds a hung one."""
+    cleaned: list[int] = []
+
+    async def _append(i: int) -> None:
+        await anyio.sleep(0)
+        cleaned.append(i)
+
+    async with connected_runner(server) as (client, runner):
+        for i in (1, 2):
+            runner.connection.exit_stack.push_async_callback(_append, i)
+        await client.send_raw_request("tools/list", None)
+    assert cleaned == [2, 1]
+    assert "abandoning remaining callbacks" not in caplog.text

tests/server/test_server_context.py

@@ -0,0 +1,101 @@
+"""Tests for the server-side `Context`.
+
+`Context` extends `BaseContext` (forwarding to a `DispatchContext`) with
+`lifespan`, `connection`, and request-scoped `log`. End-to-end tested over
+`DirectDispatcher`.
+"""
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any
+
+import anyio
+import pytest
+
+from mcp.server.connection import Connection
+from mcp.server.context import Context
+from mcp.shared.dispatcher import DispatchContext
+from mcp.shared.transport_context import TransportContext
+
+from ..shared.conftest import direct_pair
+from ..shared.test_dispatcher import Recorder, echo_handlers, running_pair
+
+DCtx = DispatchContext[TransportContext]
+
+
+@dataclass
+class _Lifespan:
+    name: str
+
+
+@pytest.mark.anyio
+async def test_context_exposes_lifespan_and_connection_and_forwards_base_context():
+    captured: list[Context[_Lifespan]] = []
+    conn = Connection.__new__(Connection)  # placeholder until running_pair gives us the dispatcher
+
+    async def server_on_request(dctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
+        ctx: Context[_Lifespan] = Context(dctx, lifespan=_Lifespan("app"), connection=conn)
+        captured.append(ctx)
+        return {}
+
+    async with running_pair(direct_pair, server_on_request=server_on_request) as (client, server, *_):
+        # Now we have the server dispatcher; build the real Connection bound to it.
+        conn.__init__(server, has_standalone_channel=True, session_id="sess-1")
+        with anyio.fail_after(5):
+            await client.send_raw_request("t", None)
+        ctx = captured[0]
+        assert ctx.lifespan.name == "app"
+        assert ctx.connection is conn
+        assert ctx.transport.kind == "direct"
+        assert ctx.can_send_request is True
+        assert ctx.session_id == "sess-1"
+        assert ctx.headers is None
+
+
+@pytest.mark.anyio
+async def test_context_log_sends_request_scoped_message_notification():
+    crec = Recorder()
+    _, c_notify = echo_handlers(crec)
+
+    async def server_on_request(dctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
+        ctx: Context[_Lifespan] = Context(
+            dctx, lifespan=_Lifespan("app"), connection=Connection(dctx, has_standalone_channel=True)
+        )
+        await ctx.log("debug", "hello")
+        return {}
+
+    async with running_pair(direct_pair, server_on_request=server_on_request, client_on_notify=c_notify) as (
+        client,
+        *_,
+    ):
+        with anyio.fail_after(5):
+            await client.send_raw_request("t", None)
+            await crec.notified.wait()
+        method, params = crec.notifications[0]
+        assert method == "notifications/message"
+        assert params is not None and params["level"] == "debug" and params["data"] == "hello"
+
+
+@pytest.mark.anyio
+async def test_context_log_includes_logger_and_meta_when_supplied():
+    crec = Recorder()
+    _, c_notify = echo_handlers(crec)
+
+    async def server_on_request(dctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
+        ctx: Context[_Lifespan] = Context(
+            dctx, lifespan=_Lifespan("app"), connection=Connection(dctx, has_standalone_channel=True)
+        )
+        await ctx.log("info", "x", logger="my.log", meta={"traceId": "t"})
+        return {}
+
+    async with running_pair(direct_pair, server_on_request=server_on_request, client_on_notify=c_notify) as (
+        client,
+        *_,
+    ):
+        with anyio.fail_after(5):
+            await client.send_raw_request("t", None)
+            await crec.notified.wait()
+        _, params = crec.notifications[0]
+        assert params is not None
+        assert params["logger"] == "my.log"
+        assert params["_meta"] == {"traceId": "t"}

tests/server/test_session.py

@@ -1,535 +1,160 @@
-from typing import Any
+"""Tests for `ServerSession`.
+
+`ServerSession` is a thin proxy over a dispatcher and a `Connection`. Tested
+with a stub dispatcher so we can assert what reaches the wire (method, params,
+`CallOptions`, related-request-id) without standing up a full transport.
+"""
+
+from collections.abc import Mapping
+from typing import Any, cast
 
-import anyio
 import pytest
 
 from mcp import 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.connection import Connection
 from mcp.server.session import ServerSession
-from mcp.shared.exceptions import MCPError
-from mcp.shared.message import SessionMessage
-from mcp.shared.session import RequestResponder
+from mcp.shared.dispatcher import CallOptions
+from mcp.shared.exceptions import NoBackChannelError
+from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
+from mcp.shared.message import ServerMessageMetadata
 from mcp.types import (
-    ClientNotification,
-    CompletionsCapability,
-    InitializedNotification,
-    PromptsCapability,
-    ResourcesCapability,
-    ServerCapabilities,
+    LATEST_PROTOCOL_VERSION,
+    ClientCapabilities,
+    Implementation,
+    InitializeRequestParams,
+    SamplingCapability,
+    SamplingToolsCapability,
 )
 
 
-@pytest.mark.anyio
-async def test_server_session_initialize():
-    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)
+class StubDispatcher:
+    """Records `send_raw_request` / `notify` calls and returns a canned result."""
 
-    # Create a message handler to catch exceptions
-    async def message_handler(  # pragma: no cover
-        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
-    ) -> None:
-        if isinstance(message, Exception):
-            raise message
+    def __init__(self, result: dict[str, Any] | None = None) -> None:
+        self.requests: list[tuple[str, Mapping[str, Any] | None, CallOptions | None, Any]] = []
+        self.result = result if result is not None else {}
 
-    received_initialized = False
+    async def send_raw_request(
+        self,
+        method: str,
+        params: Mapping[str, Any] | None,
+        opts: CallOptions | None = None,
+        *,
+        _related_request_id: Any = None,
+    ) -> dict[str, Any]:
+        self.requests.append((method, params, opts, _related_request_id))
+        return self.result
 
-    async def run_server():
-        nonlocal received_initialized
+    async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
+        raise NotImplementedError
 
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            InitializationOptions(
-                server_name="mcp",
-                server_version="0.1.0",
-                capabilities=ServerCapabilities(),
-            ),
-        ) as server_session:
-            async for message in server_session.incoming_messages:  # pragma: no branch
-                if isinstance(message, Exception):  # pragma: no cover
-                    raise message
 
-                if isinstance(message, ClientNotification) and isinstance(
-                    message, InitializedNotification
-                ):  # pragma: no branch
-                    received_initialized = True
-                    return
-
-    try:
-        async with (
-            ClientSession(
-                server_to_client_receive,
-                client_to_server_send,
-                message_handler=message_handler,
-            ) as client_session,
-            anyio.create_task_group() as tg,
-        ):
-            tg.start_soon(run_server)
-
-            await client_session.initialize()
-    except anyio.ClosedResourceError:  # pragma: no cover
-        pass
-
-    assert received_initialized
+def _make_session(
+    dispatcher: StubDispatcher,
+    *,
+    capabilities: ClientCapabilities | None = None,
+    has_standalone_channel: bool = True,
+) -> ServerSession:
+    conn = Connection(dispatcher, has_standalone_channel=has_standalone_channel)
+    if capabilities is not None:
+        conn.client_params = InitializeRequestParams(
+            protocol_version=LATEST_PROTOCOL_VERSION,
+            capabilities=capabilities,
+            client_info=Implementation(name="c", version="0"),
+        )
+    # cast: `ServerSession` is typed to take `JSONRPCDispatcher` but only ever
+    # calls `send_raw_request` / `notify`, so the stub is structurally sufficient.
+    return ServerSession(cast("JSONRPCDispatcher[Any]", dispatcher), conn)
 
 
 @pytest.mark.anyio
-async def test_check_client_capability():
-    """check_client_capability reflects the capabilities sent by the client at initialize."""
-    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)
+async def test_send_request_forwards_timeout_and_progress_callback_as_call_options():
+    dispatcher = StubDispatcher(result={"roots": []})
+    session = _make_session(dispatcher)
 
-    initialized = anyio.Event()
-
-    async def list_roots_callback(context: Any) -> types.ListRootsResult:  # pragma: no cover
-        return types.ListRootsResult(roots=[])
-
-    async def run_server(server_session: ServerSession):
-        async for message in server_session.incoming_messages:  # pragma: no branch
-            if isinstance(message, ClientNotification) and isinstance(
-                message, InitializedNotification
-            ):  # pragma: no branch
-                initialized.set()
-                return
-
-    async with (
-        ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            InitializationOptions(server_name="mcp", server_version="0.1.0", capabilities=ServerCapabilities()),
-        ) as server_session,
-        ClientSession(
-            server_to_client_receive,
-            client_to_server_send,
-            list_roots_callback=list_roots_callback,
-        ) as client_session,
-        anyio.create_task_group() as tg,
-    ):
-        tg.start_soon(run_server, server_session)
-        await client_session.initialize()
-        with anyio.fail_after(5):
-            await initialized.wait()
-
-        # ClientSession advertises roots when a list_roots_callback is provided.
-        assert server_session.check_client_capability(types.ClientCapabilities(roots=types.RootsCapability()))
-        # ClientSession does not advertise sampling without a sampling_callback.
-        assert not server_session.check_client_capability(types.ClientCapabilities(sampling=types.SamplingCapability()))
-
-
-@pytest.mark.anyio
-async def test_server_capabilities():
-    notification_options = NotificationOptions()
-    experimental_capabilities: dict[str, Any] = {}
-
-    async def noop_list_prompts(
-        ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
-    ) -> types.ListPromptsResult:
+    async def on_progress(progress: float, total: float | None, message: str | None) -> None:
         raise NotImplementedError
 
-    async def noop_list_resources(
-        ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
-    ) -> types.ListResourcesResult:
-        raise NotImplementedError
-
-    async def noop_completion(ctx: ServerRequestContext, params: types.CompleteRequestParams) -> types.CompleteResult:
-        raise NotImplementedError
-
-    # No capabilities
-    server = Server("test")
-    caps = server.get_capabilities(notification_options, experimental_capabilities)
-    assert caps.prompts is None
-    assert caps.resources is None
-    assert caps.completions is None
-
-    # With prompts handler
-    server = Server("test", on_list_prompts=noop_list_prompts)
-    caps = server.get_capabilities(notification_options, experimental_capabilities)
-    assert caps.prompts == PromptsCapability(list_changed=False)
-    assert caps.resources is None
-    assert caps.completions is None
-
-    # With prompts + resources handlers
-    server = Server("test", on_list_prompts=noop_list_prompts, on_list_resources=noop_list_resources)
-    caps = server.get_capabilities(notification_options, experimental_capabilities)
-    assert caps.prompts == PromptsCapability(list_changed=False)
-    assert caps.resources == ResourcesCapability(subscribe=False, list_changed=False)
-    assert caps.completions is None
-
-    # With prompts + resources + completion handlers
-    server = Server(
-        "test",
-        on_list_prompts=noop_list_prompts,
-        on_list_resources=noop_list_resources,
-        on_completion=noop_completion,
+    result = await session.send_request(
+        types.ListRootsRequest(),
+        types.ListRootsResult,
+        request_read_timeout_seconds=2.5,
+        metadata=ServerMessageMetadata(related_request_id=7),
+        progress_callback=on_progress,
     )
-    caps = server.get_capabilities(notification_options, experimental_capabilities)
-    assert caps.prompts == PromptsCapability(list_changed=False)
-    assert caps.resources == ResourcesCapability(subscribe=False, list_changed=False)
-    assert caps.completions == CompletionsCapability()
+    assert isinstance(result, types.ListRootsResult)
+    method, _params, opts, related = dispatcher.requests[0]
+    assert method == "roots/list"
+    assert opts == {"timeout": 2.5, "on_progress": on_progress}
+    assert related == 7
 
 
 @pytest.mark.anyio
-async def test_server_session_initialize_with_older_protocol_version():
-    """Test that server accepts and responds with older protocol (2024-11-05)."""
-    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
-
-    received_initialized = False
-    received_protocol_version = None
-
-    async def run_server():
-        nonlocal received_initialized
-
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            InitializationOptions(
-                server_name="mcp",
-                server_version="0.1.0",
-                capabilities=ServerCapabilities(),
-            ),
-        ) as server_session:
-            async for message in server_session.incoming_messages:  # pragma: no branch
-                if isinstance(message, Exception):  # pragma: no cover
-                    raise message
-
-                if isinstance(message, types.ClientNotification) and isinstance(
-                    message, InitializedNotification
-                ):  # pragma: no branch
-                    received_initialized = True
-                    return
-
-    async def mock_client():
-        nonlocal received_protocol_version
-
-        # Send initialization request with older protocol version (2024-11-05)
-        await client_to_server_send.send(
-            SessionMessage(
-                types.JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id=1,
-                    method="initialize",
-                    params=types.InitializeRequestParams(
-                        protocol_version="2024-11-05",
-                        capabilities=types.ClientCapabilities(),
-                        client_info=types.Implementation(name="test-client", version="1.0.0"),
-                    ).model_dump(by_alias=True, mode="json", exclude_none=True),
-                )
-            )
-        )
-
-        # Wait for the initialize response
-        init_response_message = await server_to_client_receive.receive()
-        assert isinstance(init_response_message.message, types.JSONRPCResponse)
-        result_data = init_response_message.message.result
-        init_result = types.InitializeResult.model_validate(result_data)
-
-        # Check that the server responded with the requested protocol version
-        received_protocol_version = init_result.protocol_version
-        assert received_protocol_version == "2024-11-05"
-
-        # Send initialized notification
-        await client_to_server_send.send(
-            SessionMessage(types.JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized"))
-        )
-
-    async with (
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-        server_to_client_receive,
-        anyio.create_task_group() as tg,
-    ):
-        tg.start_soon(run_server)
-        tg.start_soon(mock_client)
-
-    assert received_initialized
-    assert received_protocol_version == "2024-11-05"
+async def test_send_request_omits_call_options_when_none_given():
+    dispatcher = StubDispatcher(result={"roots": []})
+    session = _make_session(dispatcher)
+    await session.send_request(types.ListRootsRequest(), types.ListRootsResult)
+    _method, _params, opts, related = dispatcher.requests[0]
+    assert opts is None
+    assert related is None
 
 
 @pytest.mark.anyio
-async def test_ping_request_before_initialization():
-    """Test that ping requests are allowed before initialization is complete."""
-    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
-
-    ping_response_received = False
-    ping_response_id = None
-
-    async def run_server():
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            InitializationOptions(
-                server_name="mcp",
-                server_version="0.1.0",
-                capabilities=ServerCapabilities(),
-            ),
-        ) as server_session:
-            async for message in server_session.incoming_messages:  # pragma: no branch
-                if isinstance(message, Exception):  # pragma: no cover
-                    raise message
-
-                # We should receive a ping request before initialization
-                if isinstance(message, RequestResponder) and isinstance(
-                    message.request, types.PingRequest
-                ):  # pragma: no branch
-                    # Respond to the ping
-                    with message:
-                        await message.respond(types.EmptyResult())
-                    return
-
-    async def mock_client():
-        nonlocal ping_response_received, ping_response_id
-
-        # Send ping request before any initialization
-        await client_to_server_send.send(SessionMessage(types.JSONRPCRequest(jsonrpc="2.0", id=42, method="ping")))
-
-        # Wait for the ping response
-        ping_response_message = await server_to_client_receive.receive()
-        assert isinstance(ping_response_message.message, types.JSONRPCResponse)
-
-        ping_response_received = True
-        ping_response_id = ping_response_message.message.id
-
-    async with (
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-        server_to_client_receive,
-        anyio.create_task_group() as tg,
-    ):
-        tg.start_soon(run_server)
-        tg.start_soon(mock_client)
-
-    assert ping_response_received
-    assert ping_response_id == 42
+async def test_send_request_timeout_zero_means_no_timeout():
+    """0 falls through BaseSession's `or`-fallback, so it has always meant
+    "no timeout"; ClientSession still reads it that way."""
+    dispatcher = StubDispatcher(result={})
+    session = _make_session(dispatcher)
+    await session.send_request(types.PingRequest(), types.EmptyResult, request_read_timeout_seconds=0)
+    assert dispatcher.requests[0][2] is None
 
 
 @pytest.mark.anyio
-async def test_create_message_tool_result_validation():
-    """Test tool_use/tool_result validation in create_message."""
-    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
-
-    async with (
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-        server_to_client_receive,
-    ):
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            InitializationOptions(
-                server_name="test",
-                server_version="0.1.0",
-                capabilities=ServerCapabilities(),
-            ),
-        ) as session:
-            # Set up client params with sampling.tools capability for the test
-            session._client_params = types.InitializeRequestParams(
-                protocol_version=types.LATEST_PROTOCOL_VERSION,
-                capabilities=types.ClientCapabilities(
-                    sampling=types.SamplingCapability(tools=types.SamplingToolsCapability())
-                ),
-                client_info=types.Implementation(name="test", version="1.0"),
-            )
-
-            tool = types.Tool(name="test_tool", input_schema={"type": "object"})
-            text = types.TextContent(type="text", text="hello")
-            tool_use = types.ToolUseContent(type="tool_use", id="call_1", name="test_tool", input={})
-            tool_result = types.ToolResultContent(type="tool_result", tool_use_id="call_1", content=[])
-
-            # Case 1: tool_result mixed with other content
-            with pytest.raises(ValueError, match="only tool_result content"):
-                await session.create_message(
-                    messages=[
-                        types.SamplingMessage(role="user", content=text),
-                        types.SamplingMessage(role="assistant", content=tool_use),
-                        types.SamplingMessage(role="user", content=[tool_result, text]),  # mixed!
-                    ],
-                    max_tokens=100,
-                    tools=[tool],
-                )
-
-            # Case 2: tool_result without previous message
-            with pytest.raises(ValueError, match="requires a previous message"):
-                await session.create_message(
-                    messages=[types.SamplingMessage(role="user", content=tool_result)],
-                    max_tokens=100,
-                    tools=[tool],
-                )
-
-            # Case 3: tool_result without previous tool_use
-            with pytest.raises(ValueError, match="do not match any tool_use"):
-                await session.create_message(
-                    messages=[
-                        types.SamplingMessage(role="user", content=text),
-                        types.SamplingMessage(role="user", content=tool_result),
-                    ],
-                    max_tokens=100,
-                    tools=[tool],
-                )
-
-            # Case 4: mismatched tool IDs
-            with pytest.raises(ValueError, match="ids of tool_result blocks and tool_use blocks"):
-                await session.create_message(
-                    messages=[
-                        types.SamplingMessage(role="user", content=text),
-                        types.SamplingMessage(role="assistant", content=tool_use),
-                        types.SamplingMessage(
-                            role="user",
-                            content=types.ToolResultContent(type="tool_result", tool_use_id="wrong_id", content=[]),
-                        ),
-                    ],
-                    max_tokens=100,
-                    tools=[tool],
-                )
-
-            # Case 5: text-only message with tools (no tool_results) - passes validation
-            # Covers has_tool_results=False branch.
-            # We use move_on_after because validation happens synchronously before
-            # send_request, which would block indefinitely waiting for a response.
-            # The timeout lets validation pass, then cancels the blocked send.
-            with anyio.move_on_after(0.01):
-                await session.create_message(
-                    messages=[types.SamplingMessage(role="user", content=text)],
-                    max_tokens=100,
-                    tools=[tool],
-                )
-
-            # Case 6: valid matching tool_result/tool_use IDs - passes validation
-            # Covers tool_use_ids == tool_result_ids branch.
-            # (see Case 5 comment for move_on_after explanation)
-            with anyio.move_on_after(0.01):
-                await session.create_message(
-                    messages=[
-                        types.SamplingMessage(role="user", content=text),
-                        types.SamplingMessage(role="assistant", content=tool_use),
-                        types.SamplingMessage(role="user", content=tool_result),
-                    ],
-                    max_tokens=100,
-                    tools=[tool],
-                )
-
-            # Case 7: validation runs even without `tools` parameter
-            # (tool loop continuation may omit tools while containing tool_result)
-            with pytest.raises(ValueError, match="do not match any tool_use"):
-                await session.create_message(
-                    messages=[
-                        types.SamplingMessage(role="user", content=text),
-                        types.SamplingMessage(role="user", content=tool_result),
-                    ],
-                    max_tokens=100,
-                    # Note: no tools parameter
-                )
-
-            # Case 8: empty messages list - skips validation entirely
-            # Covers the `if messages:` branch (line 280->302)
-            with anyio.move_on_after(0.01):  # pragma: no branch
-                await session.create_message(messages=[], max_tokens=100)
+async def test_send_request_without_back_channel_or_related_id_fails_fast():
+    """No standalone channel and no related request to ride on: raise instead
+    of parking forever on a response that cannot arrive."""
+    dispatcher = StubDispatcher(result={})
+    session = _make_session(dispatcher, has_standalone_channel=False)
+    with pytest.raises(NoBackChannelError):
+        await session.send_request(types.PingRequest(), types.EmptyResult)
+    assert dispatcher.requests == []
+    # With a related request id the message rides that request's stream.
+    await session.send_request(
+        types.PingRequest(), types.EmptyResult, metadata=ServerMessageMetadata(related_request_id=3)
+    )
+    assert dispatcher.requests[0][3] == 3
 
 
 @pytest.mark.anyio
-async def test_create_message_without_tools_capability():
-    """Test that create_message raises MCPError when tools are provided without capability."""
-    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
-
-    async with (
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-        server_to_client_receive,
-    ):
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            InitializationOptions(
-                server_name="test",
-                server_version="0.1.0",
-                capabilities=ServerCapabilities(),
-            ),
-        ) as session:
-            # Set up client params WITHOUT sampling.tools capability
-            session._client_params = types.InitializeRequestParams(
-                protocol_version=types.LATEST_PROTOCOL_VERSION,
-                capabilities=types.ClientCapabilities(sampling=types.SamplingCapability()),
-                client_info=types.Implementation(name="test", version="1.0"),
-            )
-
-            tool = types.Tool(name="test_tool", input_schema={"type": "object"})
-            text = types.TextContent(type="text", text="hello")
-
-            # Should raise MCPError when tools are provided but client lacks capability
-            with pytest.raises(MCPError) as exc_info:
-                await session.create_message(
-                    messages=[types.SamplingMessage(role="user", content=text)],
-                    max_tokens=100,
-                    tools=[tool],
-                )
-            assert "does not support sampling tools capability" in exc_info.value.error.message
-
-            # Should also raise MCPError when tool_choice is provided
-            with pytest.raises(MCPError) as exc_info:
-                await session.create_message(
-                    messages=[types.SamplingMessage(role="user", content=text)],
-                    max_tokens=100,
-                    tool_choice=types.ToolChoice(mode="auto"),
-                )
-            assert "does not support sampling tools capability" in exc_info.value.error.message
+async def test_send_request_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"}
+    session = _make_session(StubDispatcher(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_other_requests_blocked_before_initialization():
-    """Test that non-ping requests are still blocked before initialization."""
-    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
+async def test_create_message_with_tools_returns_with_tools_result():
+    dispatcher = StubDispatcher(result={"role": "assistant", "content": [{"type": "text", "text": "ok"}], "model": "m"})
+    session = _make_session(
+        dispatcher, capabilities=ClientCapabilities(sampling=SamplingCapability(tools=SamplingToolsCapability()))
+    )
+    result = await session.create_message(
+        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, _related = dispatcher.requests[0]
+    assert method == "sampling/createMessage"
+    assert params is not None and params["tools"][0]["name"] == "t"
 
-    error_response_received = False
-    error_code = None
 
-    async def run_server():
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            InitializationOptions(
-                server_name="mcp",
-                server_version="0.1.0",
-                capabilities=ServerCapabilities(),
-            ),
-        ):
-            # Server should handle the request and send an error response
-            # No need to process incoming_messages since the error is handled automatically
-            await anyio.sleep(0.1)  # Give time for the request to be processed
-
-    async def mock_client():
-        nonlocal error_response_received, error_code
-
-        # Try to send a non-ping request before initialization
-        await client_to_server_send.send(
-            SessionMessage(types.JSONRPCRequest(jsonrpc="2.0", id=1, method="prompts/list"))
-        )
-
-        # Wait for the error response
-        error_message = await server_to_client_receive.receive()
-        if isinstance(error_message.message, types.JSONRPCError):  # pragma: no branch
-            error_response_received = True
-            error_code = error_message.message.error.code
-
-    async with (
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-        server_to_client_receive,
-        anyio.create_task_group() as tg,
-    ):
-        tg.start_soon(run_server)
-        tg.start_soon(mock_client)
-
-    assert error_response_received
-    assert error_code == types.INVALID_PARAMS
+def test_check_client_capability_delegates_to_connection():
+    dispatcher = StubDispatcher()
+    session = _make_session(dispatcher, capabilities=ClientCapabilities(sampling=SamplingCapability()))
+    assert session.check_client_capability(ClientCapabilities(sampling=SamplingCapability())) is True
+    assert session.check_client_capability(ClientCapabilities(experimental={"x": {}})) is False

tests/server/test_session_race_condition.py

@@ -1,132 +0,0 @@
-"""Test for race condition fix in initialization flow.
-
-This test verifies that requests can be processed immediately after
-responding to InitializeRequest, without waiting for InitializedNotification.
-
-This is critical for HTTP transport where requests can arrive in any order.
-"""
-
-import anyio
-import pytest
-
-from mcp import types
-from mcp.server.models import InitializationOptions
-from mcp.server.session import ServerSession
-from mcp.shared.message import SessionMessage
-from mcp.shared.session import RequestResponder
-from mcp.types import ServerCapabilities, Tool
-
-
-@pytest.mark.anyio
-async def test_request_immediately_after_initialize_response():
-    """Test that requests are accepted immediately after initialize response.
-
-    This reproduces the race condition in stateful HTTP mode where:
-    1. Client sends InitializeRequest
-    2. Server responds with InitializeResult
-    3. Client immediately sends tools/list (before server receives InitializedNotification)
-    4. Without fix: Server rejects with "Received request before initialization was complete"
-    5. With fix: Server accepts and processes the request
-
-    This test simulates the HTTP transport behavior where InitializedNotification
-    may arrive in a separate POST request after other requests.
-    """
-    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](10)
-
-    tools_list_success = False
-    error_received = None
-
-    async def run_server():
-        nonlocal tools_list_success
-
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            InitializationOptions(
-                server_name="test-server",
-                server_version="1.0.0",
-                capabilities=ServerCapabilities(
-                    tools=types.ToolsCapability(list_changed=False),
-                ),
-            ),
-        ) as server_session:
-            async for message in server_session.incoming_messages:  # pragma: no branch
-                if isinstance(message, Exception):  # pragma: no cover
-                    raise message
-
-                # Handle tools/list request
-                if isinstance(message, RequestResponder):
-                    if isinstance(message.request, types.ListToolsRequest):  # pragma: no branch
-                        tools_list_success = True
-                        # Respond with a tool list
-                        with message:
-                            await message.respond(
-                                types.ListToolsResult(
-                                    tools=[
-                                        Tool(
-                                            name="example_tool",
-                                            description="An example tool",
-                                            input_schema={"type": "object", "properties": {}},
-                                        )
-                                    ]
-                                )
-                            )
-
-                # Handle InitializedNotification
-                if isinstance(message, types.ClientNotification):
-                    if isinstance(message, types.InitializedNotification):  # pragma: no branch
-                        # Done - exit gracefully
-                        return
-
-    async def mock_client():
-        nonlocal error_received
-
-        # Step 1: Send InitializeRequest
-        await client_to_server_send.send(
-            SessionMessage(
-                types.JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id=1,
-                    method="initialize",
-                    params=types.InitializeRequestParams(
-                        protocol_version=types.LATEST_PROTOCOL_VERSION,
-                        capabilities=types.ClientCapabilities(),
-                        client_info=types.Implementation(name="test-client", version="1.0.0"),
-                    ).model_dump(by_alias=True, mode="json", exclude_none=True),
-                )
-            )
-        )
-
-        # Step 2: Wait for InitializeResult
-        init_msg = await server_to_client_receive.receive()
-        assert isinstance(init_msg.message, types.JSONRPCResponse)
-
-        # Step 3: Immediately send tools/list BEFORE InitializedNotification
-        # This is the race condition scenario
-        await client_to_server_send.send(SessionMessage(types.JSONRPCRequest(jsonrpc="2.0", id=2, method="tools/list")))
-
-        # Step 4: Check the response
-        tools_msg = await server_to_client_receive.receive()
-        if isinstance(tools_msg.message, types.JSONRPCError):  # pragma: no cover
-            error_received = tools_msg.message.error.message
-
-        # Step 5: Send InitializedNotification
-        await client_to_server_send.send(
-            SessionMessage(types.JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized"))
-        )
-
-    async with (
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-        server_to_client_receive,
-        anyio.create_task_group() as tg,
-    ):
-        tg.start_soon(run_server)
-        tg.start_soon(mock_client)
-
-    # With the PR fix: tools_list_success should be True, error_received should be None
-    # Without the fix: error_received would contain "Received request before initialization was complete"
-    assert tools_list_success, f"tools/list should have succeeded. Error received: {error_received}"
-    assert error_received is None, f"Expected no error, but got: {error_received}"

tests/server/test_stateless_mode.py

@@ -7,45 +7,35 @@ that appropriate errors are raised when attempting to use unsupported features.
 See: https://github.com/modelcontextprotocol/python-sdk/issues/1097
 """
 
-from collections.abc import AsyncGenerator
 from typing import Any
+from unittest.mock import Mock
 
 import anyio
 import pytest
 
 from mcp import types
-from mcp.server.models import InitializationOptions
+from mcp.server.connection import Connection
+from mcp.server.context import ServerRequestContext
+from mcp.server.lowlevel.server import Server
 from mcp.server.session import ServerSession
-from mcp.shared.exceptions import StatelessModeNotSupported
+from mcp.shared.exceptions import NoBackChannelError, StatelessModeNotSupported
+from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
 from mcp.shared.message import SessionMessage
-from mcp.types import ServerCapabilities
+from mcp.types import JSONRPCRequest, JSONRPCResponse, ListToolsResult, PaginatedRequestParams
+
+
+def _make_session(*, stateless: bool) -> ServerSession:
+    """A `ServerSession` with a mock dispatcher; the stateless guard fires before any send."""
+    return ServerSession(
+        Mock(spec=JSONRPCDispatcher),
+        Connection(Mock(), has_standalone_channel=False),
+        stateless=stateless,
+    )
 
 
 @pytest.fixture
-async def stateless_session() -> AsyncGenerator[ServerSession, None]:
-    """Create a stateless ServerSession for testing."""
-    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
-
-    init_options = InitializationOptions(
-        server_name="test",
-        server_version="0.1.0",
-        capabilities=ServerCapabilities(),
-    )
-
-    async with (
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-        server_to_client_receive,
-    ):
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            init_options,
-            stateless=True,
-        ) as session:
-            yield session
+def stateless_session() -> ServerSession:
+    return _make_session(stateless=True)
 
 
 @pytest.mark.anyio
@@ -126,30 +116,8 @@ async def test_exception_has_method_attribute(stateless_session: ServerSession):
 
 
 @pytest.fixture
-async def stateful_session() -> AsyncGenerator[ServerSession, None]:
-    """Create a stateful ServerSession for testing."""
-    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
-
-    init_options = InitializationOptions(
-        server_name="test",
-        server_version="0.1.0",
-        capabilities=ServerCapabilities(),
-    )
-
-    async with (
-        client_to_server_send,
-        client_to_server_receive,
-        server_to_client_send,
-        server_to_client_receive,
-    ):
-        async with ServerSession(
-            client_to_server_receive,
-            server_to_client_send,
-            init_options,
-            stateless=False,
-        ) as session:
-            yield session
+def stateful_session() -> ServerSession:
+    return _make_session(stateless=False)
 
 
 @pytest.mark.anyio
@@ -175,3 +143,45 @@ async def test_stateful_mode_does_not_raise_stateless_error(
 
     assert send_request_called
     assert isinstance(result, types.ListRootsResult)
+
+
+@pytest.mark.anyio
+async def test_server_run_stateless_wires_no_standalone_channel():
+    """`Server.run(stateless=True)` must wire `Connection.has_standalone_channel=False`.
+
+    Stateless HTTP has no standalone GET stream, so server-initiated requests on
+    the connection must fail fast with `NoBackChannelError` rather than write to
+    a channel that will never deliver a response. The `ServerSession` typed
+    helpers carry their own stateless guard (tested above); this pins the
+    `Connection` wiring that `Server.run` produces.
+    """
+    captured: list[Connection] = []
+
+    async def list_tools(ctx: ServerRequestContext[Any], params: PaginatedRequestParams | None) -> ListToolsResult:
+        # `ServerRequestContext` doesn't expose `connection` directly yet (it
+        # will after the Context rework); reach it via the session for now.
+        captured.append(ctx.session._connection)  # pyright: ignore[reportPrivateUsage]
+        return ListToolsResult(tools=[])
+
+    server: Server[Any] = Server("test", on_list_tools=list_tools)
+
+    to_server, server_read = anyio.create_memory_object_stream[SessionMessage | Exception](10)
+    server_write, from_server = anyio.create_memory_object_stream[SessionMessage](10)
+
+    async def run_server() -> None:
+        await server.run(server_read, server_write, server.create_initialization_options(), stateless=True)
+
+    async with anyio.create_task_group() as tg, to_server, server_read, server_write, from_server:
+        tg.start_soon(run_server)
+        # stateless=True skips the init gate, so tools/list routes immediately.
+        await to_server.send(SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=1, method="tools/list")))
+        with anyio.fail_after(5):
+            response = (await from_server.receive()).message
+        assert isinstance(response, JSONRPCResponse)
+        tg.cancel_scope.cancel()
+
+    assert len(captured) == 1
+    conn = captured[0]
+    assert conn.has_standalone_channel is False
+    with pytest.raises(NoBackChannelError):
+        await conn.ping()

tests/server/test_validation.py

@@ -120,6 +120,16 @@ def test_validate_tool_use_result_messages_raises_when_tool_result_without_previ
         validate_tool_use_result_messages(messages)
 
 
+def test_validate_tool_use_result_messages_raises_when_previous_message_has_no_tool_use() -> None:
+    """Raises when tool_result follows a message that has content but no tool_use."""
+    messages = [
+        SamplingMessage(role="assistant", content=TextContent(type="text", text="just text")),
+        SamplingMessage(role="user", content=ToolResultContent(type="tool_result", tool_use_id="tool-1")),
+    ]
+    with pytest.raises(ValueError, match="do not match any tool_use in the previous message"):
+        validate_tool_use_result_messages(messages)
+
+
 def test_validate_tool_use_result_messages_raises_when_tool_result_ids_dont_match_tool_use() -> None:
     """Raises when tool_result IDs don't match tool_use IDs."""
     messages = [

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")