Move ClientSession onto JSONRPCDispatcher and delete BaseSession

ClientSession keeps its public surface (constructor, typed methods, manual
initialize, context-manager lifecycle) but now owns a JSONRPCDispatcher
instead of inheriting the v1 BaseSession receive loop. Server-initiated
requests are answered through the existing callbacks via the closed-union
parse; notifications validate-or-drop and tee to message_handler; transport
exceptions reach message_handler through the dispatcher's stream-exception
observer. A from_dispatcher constructor accepts a pre-built dispatcher for
in-process embedding.

mcp.shared.session shrinks to the surviving names: the ProgressFnT re-export
and a typing-only RequestResponder stub for MessageHandlerFnT annotations.

Behavior changes (deliberate, to be covered in the migration guide):
- request ids count from 1; the progress token follows
- timeouts use the dispatcher error text and send notifications/cancelled,
  so a timed-out server handler is interrupted instead of running on
- responses with unknown ids are ignored per spec instead of surfacing a
  RuntimeError to message_handler
- a raising request callback is answered with code 0 and the exception text
- notification callbacks run concurrently (no completion-before-response)

Three interaction-suite divergence entries are resolved and deleted, and the
server-to-client cancellation requirement is now pinned by a passing test.

Author: maxisbey

src/mcp/client/session.py

@@ -1,14 +1,13 @@
-"""Request context for MCP handlers."""
+"""Request context for MCP client handlers."""
 
 from dataclasses import dataclass
 from typing import Any, Generic
 
 from typing_extensions import TypeVar
 
-from mcp.shared.session import BaseSession
 from mcp.types import RequestId, RequestParamsMeta
 
-SessionT = TypeVar("SessionT", bound=BaseSession[Any, Any, Any, Any, Any])
+SessionT = TypeVar("SessionT", default=Any)
 
 
 @dataclass(kw_only=True)

src/mcp/shared/_context.py

@@ -17,6 +17,7 @@
 from mcp.types import (
     INVALID_PARAMS,
     LATEST_PROTOCOL_VERSION,
+    METHOD_NOT_FOUND,
     CallToolResult,
     Implementation,
     InitializedNotification,
@@ -751,8 +752,34 @@ async def test_receive_loop_answers_malformed_inbound_request_with_invalid_param
 
 
 @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 test_receive_loop_answers_unknown_request_method_with_method_not_found():
+    """A server request whose method is not in the ServerRequest union gets -32601
+    (METHOD_NOT_FOUND) on the wire, not a validation failure (-32602)."""
+    async with raw_client_session() as (_session, to_client, from_client):
+        await to_client.send(SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=7, method="x/unknown")))
+        out = await from_client.receive()
+    assert isinstance(out.message, JSONRPCError)
+    assert out.message.id == 7
+    assert out.message.error == types.ErrorData(code=METHOD_NOT_FOUND, message="Method not found", data="x/unknown")
+
+
+@pytest.mark.anyio
+async def test_receive_loop_drops_unknown_notification_method_without_response():
+    """An unknown notification method is dropped silently: JSON-RPC forbids
+    responses to notifications, and the receive loop keeps serving."""
+    async with raw_client_session() as (_session, to_client, from_client):
+        await to_client.send(SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="x/unknown")))
+        # The next wire output must be the answer to this follow-up ping,
+        # proving the notification produced no response and the loop survived.
+        await to_client.send(SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")))
+        out = await from_client.receive()
+    assert isinstance(out.message, JSONRPCResponse)
+    assert out.message.id == 1
+
+
+@pytest.mark.anyio
+async def test_raising_sampling_callback_answers_with_code_zero():
+    """A raising request callback is answered through the dispatcher's exception boundary."""
 
     async def boom(ctx: object, params: object) -> types.CreateMessageResult:
         raise RuntimeError("sampling boom")
@@ -767,7 +794,7 @@ async def boom(ctx: object, params: object) -> types.CreateMessageResult:
         )
         out = await from_client.receive()
     assert isinstance(out.message, JSONRPCError)
-    assert out.message.error.code == INVALID_PARAMS
+    assert out.message.error == types.ErrorData(code=0, message="sampling boom")
 
 
 @pytest.mark.anyio
@@ -841,23 +868,68 @@ async def handler(msg: object) -> None:
 
 
 @pytest.mark.anyio
-async def test_receive_loop_swallows_progress_callback_exception(caplog: pytest.LogCaptureFixture):
+async def test_progress_callback_exception_is_swallowed(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()
+        if isinstance(msg, types.ProgressNotification):
+            delivered.set()
+
+    async with raw_client_session(message_handler=handler) as (session, to_client, from_client):
+        async with anyio.create_task_group() as tg:
+
+            async def call() -> None:
+                await session.send_request(types.PingRequest(), types.EmptyResult, progress_callback=boom)
+
+            tg.start_soon(call)
+            request = await from_client.receive()
+            assert isinstance(request.message, JSONRPCRequest)
+            # The request id doubles as the progress token.
+            params = {"progressToken": request.message.id, "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; the
+            # raising callback was swallowed and logged.
+            await delivered.wait()
+            await to_client.send(SessionMessage(JSONRPCResponse(jsonrpc="2.0", id=request.message.id, result={})))
+    assert "progress callback raised" in caplog.text
 
-    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
+
+@pytest.mark.anyio
+async def test_from_dispatcher_runs_over_direct_dispatch():
+    """A session built with from_dispatcher works without a stream pair (in-process embedding)."""
+    from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair
+    from mcp.shared.dispatcher import DispatchContext
+    from mcp.shared.transport_context import TransportContext
+
+    client_side, server_side = create_direct_dispatcher_pair()
+
+    async def server_on_request(
+        ctx: DispatchContext[TransportContext], method: str, params: dict[str, object] | None
+    ) -> dict[str, object]:
+        assert method == "ping"
+        return {}
+
+    notified: list[str] = []
+
+    async def server_on_notify(
+        ctx: DispatchContext[TransportContext], method: str, params: dict[str, object] | None
+    ) -> None:
+        notified.append(method)
+
+    session = ClientSession.from_dispatcher(client_side)
+    results: list[types.EmptyResult] = []
+    async with anyio.create_task_group() as tg:
+        await tg.start(server_side.run, server_on_request, server_on_notify)
+        async with session:
+            results.append(await session.send_ping(meta=None))
+            # related_request_id routing is JSON-RPC plumbing; on other
+            # dispatchers the notification is sent without it.
+            await session.send_notification(types.RootsListChangedNotification(), related_request_id=7)
+        server_side.close()
+    assert results == [types.EmptyResult()]
+    assert notified == ["notifications/roots/list_changed"]

src/mcp/shared/session.py

@@ -287,14 +287,6 @@ def __post_init__(self) -> None:
             "A response that arrives after the sender issued notifications/cancelled is ignored; the "
             "request stays failed and no error is raised."
         ),
-        divergence=Divergence(
-            note=(
-                "A response whose id matches no in-flight request is delivered to the message handler "
-                "as a RuntimeError rather than being silently ignored. The post-cancellation case is the "
-                "same code path; tested in its unknown-id form because that is deterministic without the "
-                "client-side cancellation API the SDK does not yet provide."
-            ),
-        ),
     ),
     "protocol:cancel:server-survives": Requirement(
         source="sdk",
@@ -306,19 +298,6 @@ def __post_init__(self) -> None:
             "A server that abandons an in-flight server-initiated request (sampling, elicitation, roots) "
             "cancels it, and the client stops processing the cancelled request."
         ),
-        divergence=Divergence(
-            note=(
-                "Abandoning a server-side send_request emits no cancellation notification, and the client "
-                "could not act on one anyway: client callbacks run inline in the receive loop, so a "
-                "cancellation is not even read until the callback has finished."
-            ),
-        ),
-        deferred=(
-            "Not implemented in the SDK: abandoning a server-side send_request emits no cancellation "
-            "notification (the same sender-side gap recorded on protocol:timeout:sends-cancellation), and "
-            "the client could not act on one anyway because client callbacks run inline in the receive "
-            "loop, so a cancellation would not even be read until the callback had already finished."
-        ),
     ),
     "protocol:cancel:unknown-id-ignored": Requirement(
         source=f"{SPEC_BASE_URL}/basic/utilities/cancellation#error-handling",
@@ -466,13 +445,6 @@ def __post_init__(self) -> None:
             "When a request times out, the sender issues notifications/cancelled for that request before "
             "failing the local call."
         ),
-        divergence=Divergence(
-            note=(
-                "Client seat only: the client raises locally and sends nothing on timeout, so the server keeps "
-                "running the handler. The server seat conforms: a timed-out server-initiated request is followed "
-                "by notifications/cancelled on the wire."
-            ),
-        ),
     ),
     "protocol:timeout:session-survives": Requirement(
         source=f"{SPEC_BASE_URL}/basic/lifecycle#timeouts",

tests/client/test_resource_cleanup.py

@@ -11,7 +11,7 @@
 from inline_snapshot import snapshot
 
 from mcp import MCPError, types
-from mcp.client import ClientSession
+from mcp.client import ClientRequestContext, ClientSession
 from mcp.server import Server, ServerRequestContext
 from mcp.shared.memory import MessageStream, create_client_server_memory_streams
 from mcp.shared.message import SessionMessage
@@ -155,14 +155,70 @@ async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestPara
     assert result == snapshot(CallToolResult(content=[TextContent(text="unbothered")]))
 
 
+@requirement("protocol:cancel:server-to-client")
+async def test_abandoned_server_request_cancels_the_client_callback(connect: Connect) -> None:
+    """A server that abandons a sampling request cancels it, interrupting the client's callback.
+
+    The handler gives up on its sampling request by cancelling the scope around it; the courtesy
+    notifications/cancelled that follows interrupts the client's sampling callback mid-await.
+    """
+    callback_started = anyio.Event()
+    callback_cancelled = anyio.Event()
+
+    async def sampling_callback(
+        context: ClientRequestContext, params: types.CreateMessageRequestParams
+    ) -> types.CreateMessageResult:
+        callback_started.set()
+        try:
+            await anyio.Event().wait()  # blocks until the cancellation interrupts it
+        except anyio.get_cancelled_exc_class():
+            callback_cancelled.set()
+            raise
+        raise NotImplementedError  # unreachable
+
+    async def list_tools(
+        ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
+    ) -> types.ListToolsResult:
+        return types.ListToolsResult(tools=[types.Tool(name="impatient", input_schema={"type": "object"})])
+
+    async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
+        assert params.name == "impatient"
+        request = types.CreateMessageRequest(
+            params=types.CreateMessageRequestParams(
+                messages=[types.SamplingMessage(role="user", content=TextContent(text="Say hello."))],
+                max_tokens=8,
+            )
+        )
+        async with anyio.create_task_group() as abandon_scope:
+
+            async def sample() -> None:
+                await ctx.session.send_request(request, types.CreateMessageResult)
+                raise NotImplementedError  # unreachable: the scope is cancelled
+
+            abandon_scope.start_soon(sample)
+            await callback_started.wait()
+            abandon_scope.cancel_scope.cancel()
+        with anyio.fail_after(5):
+            await callback_cancelled.wait()
+        return CallToolResult(content=[TextContent(text="abandoned")])
+
+    server = Server("abandoner", on_list_tools=list_tools, on_call_tool=call_tool)
+
+    async with connect(server, sampling_callback=sampling_callback) as client:
+        result = await client.call_tool("impatient", {})
+
+    assert result == snapshot(CallToolResult(content=[TextContent(text="abandoned")]))
+    assert callback_cancelled.is_set()
+
+
 @requirement("protocol:cancel:late-response-ignored")
-async def test_a_response_for_an_unknown_request_id_surfaces_to_the_message_handler() -> None:
-    """A response whose id matches no in-flight request is surfaced to the message handler as a RuntimeError.
+async def test_a_response_for_an_unknown_request_id_is_ignored() -> None:
+    """A response whose id matches no in-flight request is ignored, as the spec asks.
 
     The spec says a sender SHOULD ignore a response that arrives after it issued a cancellation;
     that is the same client-side code path as any response with an unknown id, and that form is
-    deterministic to test without depending on the cancellation API the SDK does not yet provide.
-    See the divergence note on the requirement.
+    deterministic to test without depending on a client-side cancellation API. Nothing reaches
+    the message handler and the session keeps serving.
 
     A real Server cannot be made to answer with a fabricated id, so the test plays the server's
     side of the wire by hand. Reserve this pattern for behaviour no real server can produce. The
@@ -228,7 +284,6 @@ async def message_handler(message: IncomingMessage) -> None:
             pong = await session.send_request(PingRequest(), EmptyResult)
 
         assert pong == snapshot(EmptyResult())
-        assert len(incoming) == 1
-        assert isinstance(incoming[0], RuntimeError)
-        # The full message embeds the response object's repr; only the prefix is stable.
-        assert str(incoming[0]).startswith("Received response with an unknown request ID:")
+        # The fabricated response was dropped silently: the ping after it still
+        # round-tripped, and nothing was surfaced to the message handler.
+        assert incoming == []

tests/client/test_session.py

@@ -1,12 +1,15 @@
 """Logging interactions against the low-level Server, driven through the public Client API.
 
 Notification ordering: the in-memory transport delivers every server-to-client message on one
-ordered stream, and the client's receive loop dispatches each incoming message to completion
-before reading the next one. Over streamable HTTP that ordered single-stream guarantee holds
-only for messages that carry a ``related_request_id`` (they ride the originating request's POST
-stream); without it the message routes to the standalone GET stream and may arrive after the
-response. These tests pass ``related_request_id`` so they can collect into a plain list and
-assert after the request completes on every transport leg -- no events, no waiting.
+ordered stream, and the client starts notification callbacks in arrival order. Callbacks run
+concurrently with the rest of the session (no completion-before-response guarantee), but a
+callback with no internal awaits runs to completion as soon as it starts, which keeps
+plain-list collection deterministic here. Over streamable HTTP the ordered single-stream
+guarantee holds only for messages that carry a ``related_request_id`` (they ride the
+originating request's POST stream); without it the message routes to the standalone GET stream
+and may arrive after the response. These tests pass ``related_request_id`` and use await-free
+callbacks so they can collect into a plain list and assert after the request completes on
+every transport leg -- no events, no waiting.
 """
 
 import pytest

tests/interaction/_requirements.py

@@ -87,8 +87,8 @@ async def ignore(progress: float, total: float | None, message: str | None) -> N
     async with connect(server) as client:
         result = await client.call_tool("inspect", {}, progress_callback=ignore)
 
-    # The token is the request id of the tools/call request itself (initialize is request 0).
-    assert result == snapshot(CallToolResult(content=[TextContent(text="1")]))
+    # The token is the request id of the tools/call request itself (initialize is request 1).
+    assert result == snapshot(CallToolResult(content=[TextContent(text="2")]))
 
 
 @requirement("protocol:progress:no-token")