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

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