simplify token capture using events rather than streams, add test for timeout on join and subsequent rejoin

Author: davemssavage

src/mcp/client/session.py

@@ -293,44 +293,36 @@ async def request_call_tool(
         progress_callback: ProgressFnT | None = None,
     ) -> types.RequestId:
         if self._resumable:
-            send_stream, receive_stream = anyio.create_memory_object_stream[str](1)
+            captured_token = None
+            captured = anyio.Event()
 
-            async def close() -> None:
-                await send_stream.aclose()
-                await receive_stream.aclose()
+            async def capture_token(token: str):
+                nonlocal captured_token
+                captured_token = token
+                captured.set()
 
-            self._exit_stack.push_async_callback(close)
+            metadata = ClientMessageMetadata(on_resumption_token_update=capture_token)
 
-            with send_stream, receive_stream:
-
-                async def send_token(token: str):
-                    try:
-                        await send_stream.send(token)
-                    except anyio.BrokenResourceError as e:
-                        raise e
-
-                metadata = ClientMessageMetadata(on_resumption_token_update=send_token)
-
-                request_id = await self.start_request(
-                    types.ClientRequest(
-                        types.CallToolRequest(
-                            method="tools/call",
-                            params=types.CallToolRequestParams(
-                                name=name,
-                                arguments=arguments,
-                            ),
-                        )
-                    ),
-                    progress_callback=progress_callback,
-                    metadata=metadata,
-                )
+            request_id = await self.start_request(
+                types.ClientRequest(
+                    types.CallToolRequest(
+                        method="tools/call",
+                        params=types.CallToolRequestParams(
+                            name=name,
+                            arguments=arguments,
+                        ),
+                    )
+                ),
+                progress_callback=progress_callback,
+                metadata=metadata,
+            )
 
-                await anyio.lowlevel.checkpoint()
+            while captured_token is None:
+                await captured.wait()
 
-                token = await receive_stream.receive()
-                await self._request_state_manager.update_resume_token(request_id, token)
+            await self._request_state_manager.update_resume_token(request_id, captured_token)
 
-                return request_id
+            return request_id
         else:
             return await self.start_request(
                 types.ClientRequest(

src/mcp/client/streamable_http.py

@@ -173,7 +173,7 @@ async def _handle_sse_event(
 
                 # Extract protocol version from initialization response
                 if is_initialization:
-                    self._maybe_extract_protocol_version_from_message(message)
+                    message = self._maybe_extract_protocol_version_from_message(message)
 
                 # If this is a response and we have original_request_id, replace it
                 if original_request_id is not None and isinstance(message.root, JSONRPCResponse | JSONRPCError):
@@ -192,10 +192,7 @@ async def _handle_sse_event(
 
             except Exception as exc:
                 logger.exception("Error parsing SSE message")
-                try:
-                    await read_stream_writer.send(exc)
-                except anyio.BrokenResourceError:
-                    pass
+                await read_stream_writer.send(exc)
                 return False
         else:
             logger.warning(f"Unknown SSE event: {sse.event}")
@@ -486,8 +483,8 @@ async def streamablehttp_client(
     read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
 
-    async with anyio.create_task_group() as tg:
-        try:
+    try:
+        async with anyio.create_task_group() as tg:
             logger.debug(f"Connecting to StreamableHTTP endpoint: {url}")
 
             async with httpx_client_factory(
@@ -519,6 +516,6 @@ def start_get_stream() -> None:
                     if transport.session_id and terminate_on_close:
                         await transport.terminate_session(client)
                     tg.cancel_scope.cancel()
-        finally:
-            await read_stream_writer.aclose()
-            await write_stream.aclose()
+    finally:
+        await read_stream_writer.aclose()
+        await write_stream.aclose()

tests/shared/test_streamable_http.py

@@ -9,6 +9,7 @@
 import socket
 import time
 from collections.abc import Generator
+from datetime import timedelta
 from typing import Any
 
 import anyio
@@ -1295,6 +1296,146 @@ async def run_tool():
         assert len(request_state_manager_2._response_streams) == 0
 
 
+@pytest.mark.anyio
+async def test_streamablehttp_client_resumption_timeout(event_server):
+    """Test client session to resume a long running tool via non blocking api."""
+    _, server_url = event_server
+
+    with anyio.fail_after(10):
+        # Variables to track the state
+        captured_session_id = None
+        captured_notifications = []
+        tool_started = False
+        captured_protocol_version = None
+        captured_request_id = None
+        request_state_manager_1 = InMemoryRequestStateManager()
+        request_state_manager_2 = InMemoryRequestStateManager()
+
+        async def message_handler(
+            message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+        ) -> None:
+            if isinstance(message, types.ServerNotification):
+                captured_notifications.append(message)
+                # Look for our special notification that indicates the tool is running
+                if isinstance(message.root, types.LoggingMessageNotification):
+                    if message.root.params.data == "Tool started":
+                        nonlocal tool_started
+                        tool_started = True
+
+        # First, start the client session and begin the long-running tool
+        async with streamablehttp_client(f"{server_url}/mcp", terminate_on_close=False) as (
+            read_stream,
+            write_stream,
+            get_session_id,
+        ):
+            async with ClientSession(
+                read_stream,
+                write_stream,
+                message_handler=message_handler,
+                request_state_manager=request_state_manager_1,
+            ) as session:
+                # Initialize the session
+                result = await session.initialize()
+                assert isinstance(result, InitializeResult)
+                captured_session_id = get_session_id()
+                assert captured_session_id is not None
+                # Capture the negotiated protocol version
+                captured_protocol_version = result.protocolVersion
+
+                # Start a long-running tool in a task
+                async with anyio.create_task_group() as tg:
+                    timed_out = anyio.Event()
+
+                    async def run_tool():
+                        nonlocal captured_request_id
+                        captured_request_id = await session.request_call_tool(
+                            "long_running_with_checkpoints", arguments={}
+                        )
+                        try:
+                            await session.join_call_tool(
+                                captured_request_id, request_read_timeout_seconds=timedelta(seconds=0.01)
+                            )
+                            raise RuntimeError("Expected timeout")
+                        except McpError as e:
+                            assert e.error.code == httpx.codes.REQUEST_TIMEOUT.value
+
+                        timed_out.set()
+
+                    tg.start_soon(run_tool)
+
+                    # Wait for the tool to start and at least one notification
+                    # and then kill the task group
+                    while (
+                        not tool_started or not captured_request_id or len(request_state_manager_1._resume_tokens) == 0
+                    ):
+                        await anyio.sleep(0.1)
+
+                    await timed_out.wait()
+
+                    tg.cancel_scope.cancel()
+
+        # Store pre notifications and clear the captured notifications
+        # for the post-resumption check
+        captured_notifications_pre = captured_notifications.copy()
+        captured_notifications = []
+
+        # Now resume the session with the same mcp-session-id and protocol version
+        headers = {}
+        if captured_session_id:
+            headers[MCP_SESSION_ID_HEADER] = captured_session_id
+        if captured_protocol_version:
+            headers[MCP_PROTOCOL_VERSION_HEADER] = captured_protocol_version
+
+        assert len(request_state_manager_1._requests) == 1, str(request_state_manager_1._requests)
+        assert len(request_state_manager_1._resume_tokens) == 1
+
+        request_state_manager_2._requests = request_state_manager_1._requests.copy()
+        request_state_manager_2._resume_tokens = request_state_manager_1._resume_tokens.copy()
+
+        async with streamablehttp_client(f"{server_url}/mcp", headers=headers) as (
+            read_stream,
+            write_stream,
+            _,
+        ):
+            async with ClientSession(
+                read_stream,
+                write_stream,
+                message_handler=message_handler,
+                request_state_manager=request_state_manager_2,
+            ) as session:
+                # Don't initialize - just use the existing session
+
+                # Resume the tool with the resumption token
+                assert captured_request_id is not None
+
+                result = await session.join_call_tool(captured_request_id)
+
+                # We should get a complete result
+                assert len(result.content) == 1
+                assert result.content[0].type == "text"
+                assert "Completed" in result.content[0].text
+
+                # We should have received the remaining notifications
+                assert len(captured_notifications) > 0
+
+                # Should not have the first notification
+                # Check that "Tool started" notification isn't repeated when resuming
+                assert not any(
+                    isinstance(n.root, types.LoggingMessageNotification) and n.root.params.data == "Tool started"
+                    for n in captured_notifications
+                )
+                # there is no intersection between pre and post notifications
+                assert not any(n in captured_notifications_pre for n in captured_notifications), (
+                    f"{captured_notifications_pre} -> {captured_notifications}"
+                )
+                
+        assert len(request_state_manager_1._progress_callbacks) == 0
+        assert len(request_state_manager_1._response_streams) == 0
+        assert len(request_state_manager_2._progress_callbacks) == 0
+        assert len(request_state_manager_2._resume_tokens) == 0
+        assert len(request_state_manager_2._response_streams) == 0
+
+
 @pytest.mark.anyio
 async def test_streamablehttp_server_sampling(basic_server, basic_server_url):
     """Test server-initiated sampling request through streamable HTTP transport."""