[v2] ClientSession runs on JSONRPCDispatcher; BaseSession removed (#2838)

tests/interaction/lowlevel/test_cancellation.py

@@ -11,11 +11,12 @@ import pytest
 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
 from mcp.types import (
+    REQUEST_TIMEOUT,
     CallToolResult,
     EmptyResult,
     ErrorData,
@@ -155,14 +156,71 @@ async def test_cancellation_for_unknown_request_is_ignored(connect: Connect) ->
     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 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)
+            with anyio.fail_after(5):
+                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 a client-side cancellation API.
+
+    "Ignored" is proved in two halves: the pong round-trip proves the read loop survived the
+    fabricated response (the ordered in-memory stream routed it first), and `surfaced` holding
+    only the control notification proves the fabricated response was never delivered to
+    `message_handler` (v1 surfaced it there as a RuntimeError).
 
     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
@@ -208,14 +266,18 @@ async def test_a_response_for_an_unknown_request_id_surfaces_to_the_message_hand
         assert isinstance(ping, SessionMessage)
         assert isinstance(ping.message, JSONRPCRequest)
         assert ping.message.method == "ping"
-        # First answer with a fabricated id that matches nothing in flight, then the real id.
+        # First a fabricated id that matches nothing in flight, then a control notification that
+        # is surfaced to message_handler (proving the handler is live), then the real id.
         await server_write.send(respond(9999, EmptyResult()))
+        await server_write.send(
+            SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/tools/list_changed"))
+        )
         await server_write.send(respond(ping.message.id, EmptyResult()))
 
-    incoming: list[IncomingMessage] = []
+    surfaced: list[IncomingMessage] = []
 
     async def message_handler(message: IncomingMessage) -> None:
-        incoming.append(message)
+        surfaced.append(message)
 
     async with (
         create_client_server_memory_streams() as ((client_read, client_write), server_streams),
@@ -228,7 +290,56 @@ async def test_a_response_for_an_unknown_request_id_surfaces_to_the_message_hand
             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 stream is ordered, so the fabricated response was routed before the control
+        # notification: only the control surfaced, so the unknown-id response was dropped.
+        assert surfaced == snapshot([types.ToolListChangedNotification()])
+
+
+@requirement("protocol:cancel:initialize-not-cancellable")
+async def test_timed_out_initialize_sends_no_cancellation() -> None:
+    """An abandoned initialize is not followed by notifications/cancelled on the wire (spec-mandated).
+
+    A real Server always answers initialize, so the test plays a stalling server by hand.
+    """
+    received_methods: list[str] = []
+
+    async def scripted_server(streams: MessageStream) -> None:
+        server_read, server_write = streams
+
+        # Hold the initialize request unanswered until the client's read timeout fires.
+        init = await server_read.receive()
+        assert isinstance(init, SessionMessage)
+        assert isinstance(init.message, JSONRPCRequest)
+        received_methods.append(init.message.method)
+
+        follow_up = await server_read.receive()
+        assert isinstance(follow_up, SessionMessage)
+        assert isinstance(follow_up.message, JSONRPCRequest)
+        received_methods.append(follow_up.message.method)
+        await server_write.send(
+            SessionMessage(
+                JSONRPCResponse(
+                    jsonrpc="2.0",
+                    id=follow_up.message.id,
+                    result=EmptyResult().model_dump(by_alias=True, mode="json", exclude_none=True),
+                )
+            )
+        )
+
+    async with (
+        create_client_server_memory_streams() as ((client_read, client_write), server_streams),
+        anyio.create_task_group() as task_group,
+        # The session-level read timeout is the only public pathway that abandons initialize.
+        ClientSession(client_read, client_write, read_timeout_seconds=0.000001) as session,
+    ):
+        task_group.start_soon(scripted_server, server_streams)
+        with anyio.fail_after(5):
+            with pytest.raises(MCPError) as exc_info:
+                await session.initialize()
+            assert exc_info.value.error.code == REQUEST_TIMEOUT
+            # Override the session-level timeout: this ping must round-trip normally.
+            pong = await session.send_request(PingRequest(), EmptyResult, request_read_timeout_seconds=5)
+
+        assert pong == snapshot(EmptyResult())
+        # The stream is ordered, so a courtesy cancel would have arrived ahead of the ping.
+        assert received_methods == snapshot(["initialize", "ping"])

tests/interaction/lowlevel/test_logging.py

@@ -1,12 +1,8 @@
 """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.
+Notification ordering: await-free callbacks finish in arrival order, and passing
+``related_request_id`` keeps each notification on the originating request's POST stream over
+streamable HTTP, so plain-list collection is deterministic on every transport leg.
 """
 
 import pytest

tests/interaction/lowlevel/test_progress.py

@@ -87,8 +87,8 @@ async def test_progress_token_visible_to_handler(connect: Connect) -> None:
     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")

tests/interaction/lowlevel/test_timeouts.py

@@ -3,8 +3,9 @@
 The handler blocks on an event that is never set, so the awaited response can never arrive and
 any positive timeout fires deterministically on the next event-loop pass. Per-request timeouts are
 set to an effectively-zero duration; the session-level test runs on trio's virtual clock instead
-(see the comment there). Either way the tests add no wall-clock time to the suite. (Zero itself
-cannot be used: a falsy read_timeout_seconds is silently treated as "no timeout".)
+(see the comment there). Either way the tests add no wall-clock time to the suite. (Zero would
+also time out immediately, but a tiny positive value keeps the duration visible in the
+cancellation reason these tests snapshot.)
 """
 
 import anyio
@@ -13,9 +14,13 @@ from inline_snapshot import snapshot
 from trio.testing import MockClock
 
 from mcp import MCPError, types
+from mcp.client import ClientRequestContext
+from mcp.client._memory import InMemoryTransport
 from mcp.client.client import Client
 from mcp.server import Server, ServerRequestContext
-from mcp.types import REQUEST_TIMEOUT, CallToolResult, ErrorData, TextContent
+from mcp.shared.message import SessionMessage
+from mcp.types import REQUEST_TIMEOUT, CallToolResult, ErrorData, JSONRPCNotification, TextContent
+from tests.interaction._helpers import RecordingTransport
 from tests.interaction._requirements import requirement
 
 pytestmark = pytest.mark.anyio
@@ -26,16 +31,19 @@ pytestmark = pytest.mark.anyio
 async def test_request_timeout_fails_the_pending_call() -> None:
     """A request whose response does not arrive within its read timeout fails with a timeout error.
 
-    No cancellation is sent to the server (see the divergence note on the requirement): the handler
-    starts and is still running after the caller has already given up. The test waits for the
-    handler to have started only after the timeout has fired, so the timeout itself races nothing.
+    The timeout is followed by notifications/cancelled, which interrupts the server's handler.
     """
     handler_started = anyio.Event()
+    handler_cancelled = anyio.Event()
 
     async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
         assert params.name == "block"
         handler_started.set()
-        await anyio.Event().wait()  # blocks until the session is torn down
+        try:
+            await anyio.Event().wait()  # blocks until the courtesy cancellation interrupts it
+        except anyio.get_cancelled_exc_class():
+            handler_cancelled.set()
+            raise
         raise NotImplementedError  # unreachable
 
     server = Server("blocker", on_call_tool=call_tool)
@@ -44,18 +52,79 @@ async def test_request_timeout_fails_the_pending_call() -> None:
         with pytest.raises(MCPError) as exc_info:
             await client.call_tool("block", {}, read_timeout_seconds=0.000001)
 
-        # The request was already on the wire: the handler still runs even though the caller gave up.
+        # The request was already on the wire: the handler started and was then cancelled.
         with anyio.fail_after(5):
             await handler_started.wait()
+            await handler_cancelled.wait()
 
     assert exc_info.value.error == snapshot(
         ErrorData(
             code=REQUEST_TIMEOUT,
-            message="Timed out while waiting for response to CallToolRequest. Waited 1e-06 seconds.",
+            message="Request 'tools/call' timed out",
         )
     )
 
 
+@requirement("protocol:timeout:basic")
+@requirement("protocol:timeout:sends-cancellation")
+async def test_server_request_timeout_sends_cancellation_to_the_client() -> None:
+    """A server-initiated request that times out fails server-side and cancels the client's work.
+
+    The sampling callback answers only after the server gave up; the late response is discarded.
+    """
+    release = anyio.Event()
+    callback_started = anyio.Event()
+    errors: list[ErrorData] = []
+
+    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,
+            )
+        )
+        with pytest.raises(MCPError) as exc_info:
+            await ctx.session.send_request(request, types.CreateMessageResult, request_read_timeout_seconds=0.000001)
+        errors.append(exc_info.value.error)
+        release.set()
+        return CallToolResult(content=[TextContent(text="gave up")])
+
+    server = Server("impatient", on_list_tools=list_tools, on_call_tool=call_tool)
+    recording = RecordingTransport(InMemoryTransport(server))
+
+    async def sampling_callback(
+        context: ClientRequestContext, params: types.CreateMessageRequestParams
+    ) -> types.CreateMessageResult:
+        callback_started.set()
+        with anyio.fail_after(5):
+            await release.wait()
+        return types.CreateMessageResult(role="assistant", content=TextContent(text="too late"), model="test-model")
+
+    async with Client(recording, sampling_callback=sampling_callback) as client:
+        result = await client.call_tool("impatient", {})
+
+    assert result == snapshot(CallToolResult(content=[TextContent(text="gave up")]))
+    assert callback_started.is_set()
+    assert errors == snapshot([ErrorData(code=REQUEST_TIMEOUT, message="Request 'sampling/createMessage' timed out")])
+    cancellations = [
+        item.message
+        for item in recording.received
+        if isinstance(item, SessionMessage)
+        and isinstance(item.message, JSONRPCNotification)
+        and item.message.method == "notifications/cancelled"
+    ]
+    # requestId 1 is the sampling request, the server's first outbound request.
+    assert [notification.params for notification in cancellations] == snapshot(
+        [{"requestId": 1, "reason": "timed out after 1e-06s"}]
+    )
+
+
 @requirement("protocol:timeout:session-survives")
 async def test_session_serves_requests_after_timeout() -> None:
     """A timed-out request does not poison the session: the next request succeeds."""
@@ -73,7 +142,7 @@ async def test_session_serves_requests_after_timeout() -> None:
     async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
         if params.name == "echo":
             return CallToolResult(content=[TextContent(text="still alive")])
-        await anyio.Event().wait()  # blocks until the session is torn down
+        await anyio.Event().wait()  # blocks until the courtesy cancellation interrupts it
         raise NotImplementedError  # unreachable
 
     server = Server("blocker", on_list_tools=list_tools, on_call_tool=call_tool)
@@ -105,7 +174,7 @@ async def test_session_level_timeout_applies_to_every_request() -> None:
 
     async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
         assert params.name == "block"
-        await anyio.Event().wait()  # blocks until the session is torn down
+        await anyio.Event().wait()  # blocks until the courtesy cancellation interrupts it
         raise NotImplementedError  # unreachable
 
     server = Server("blocker", on_call_tool=call_tool)
@@ -117,6 +186,6 @@ async def test_session_level_timeout_applies_to_every_request() -> None:
     assert exc_info.value.error == snapshot(
         ErrorData(
             code=REQUEST_TIMEOUT,
-            message="Timed out while waiting for response to CallToolRequest. Waited 0.05 seconds.",
+            message="Request 'tools/call' timed out",
         )
     )

tests/interaction/lowlevel/test_wire.py

@@ -61,7 +61,7 @@ def _echo_server() -> Server:
 async def test_request_ids_are_unique_and_never_null() -> None:
     """Every request the client sends carries a distinct, non-null id.
 
-    The id sequence is pinned: sequential integers from zero, in send order.
+    The id sequence is pinned: sequential integers from one, in send order.
     """
     recording = RecordingTransport(InMemoryTransport(_echo_server()))
 
@@ -77,7 +77,7 @@ async def test_request_ids_are_unique_and_never_null() -> None:
     assert len(request_ids) == len(set(request_ids))
     # initialize, tools/list, tools/call, tools/call, ping -- the client does not issue a
     # schema-cache refresh here because the explicit tools/list already populated the cache.
-    assert request_ids == snapshot([0, 1, 2, 3, 4])
+    assert request_ids == snapshot([1, 2, 3, 4, 5])
 
 
 @requirement("protocol:notifications:no-response")