add timeout to request_call_tool to enable clients to unblock if server doesn't send an event in a reasonable time period

Author: davemssavage

src/mcp/client/session.py

@@ -291,7 +291,9 @@ async def request_call_tool(
         name: str,
         arguments: dict[str, Any] | None = None,
         progress_callback: ProgressFnT | None = None,
-    ) -> types.RequestId:
+        timeout: float | None = None,
+        cancel_if_not_resumable: bool = False,
+    ) -> types.RequestId | None:
         if self._resumable:
             captured_token = None
             captured = anyio.Event()
@@ -317,12 +319,20 @@ async def capture_token(token: str):
                 metadata=metadata,
             )
 
-            while captured_token is None:
-                await captured.wait()
-
-            await self._request_state_manager.update_resume_token(request_id, captured_token)
-
-            return request_id
+            try:
+                with anyio.fail_after(timeout):
+                    while captured_token is None:
+                        await captured.wait()
+
+                await self._request_state_manager.update_resume_token(request_id, captured_token)
+
+                return request_id
+            except TimeoutError:
+                if cancel_if_not_resumable:
+                    with anyio.CancelScope(shield=True):
+                        with anyio.move_on_after(timeout):
+                            await self.cancel_call_tool(request_id=request_id)
+                return None
         else:
             return await self.start_request(
                 types.ClientRequest(

tests/client/test_session.py

@@ -561,6 +561,7 @@ async def message_handler(
         tg.start_soon(mock_server)
 
         request_id = await session.request_call_tool("hello", {"name": "world"})
+        assert request_id is not None
         with anyio.fail_after(1):
             result = await session.join_call_tool(request_id)
 
@@ -633,7 +634,7 @@ async def message_handler(
         tg.start_soon(mock_server)
 
         request_id = await session.request_call_tool("hello", {"name": "world"})
-
+        assert request_id is not None
         with anyio.fail_after(3):
             result = await session.join_call_tool(
                 request_id, request_read_timeout_seconds=timedelta(seconds=0.5), done_on_timeout=False
@@ -763,6 +764,7 @@ async def progress_callback1(progress: float, total: float | None, message: str
                 raise RuntimeError("Unexpected progress value")
 
         request_id = await session.request_call_tool("hello", {"name": "world"}, progress_callback1)
+        assert request_id is not None
 
         with anyio.fail_after(3):
             await progress_1.wait()
@@ -910,6 +912,8 @@ async def progress_callback2(progress: float, total: float | None, message: str
                 raise RuntimeError("Unexpected progress value")
 
         request_id = await session1.request_call_tool("hello", {"name": "world"}, progress_callback1)
+        assert request_id is not None
+
         with anyio.fail_after(1):
             await progress_1_1.wait()
 
@@ -986,6 +990,8 @@ async def progress_callback(progress: float, total: float | None, message: str |
             pass
 
         request_id = await session.request_call_tool("hello", {"name": "world"}, progress_callback)
+        assert request_id is not None
+
         assert await session.cancel_call_tool(request_id)
         with anyio.fail_after(1):
             await cancelled.wait()

tests/shared/test_streamable_http.py

@@ -1346,6 +1346,73 @@ async def run_tool():
         assert len(request_state_manager_2._response_streams) == 0
 
 
+@pytest.mark.anyio
+async def test_streamablehttp_client_non_blocking_timeout(event_server: tuple[SimpleEventStore, str]):
+    """Test client session start timeout."""
+    _, server_url = event_server
+
+    with anyio.fail_after(10):
+        # Variables to track the state
+        captured_notifications: list[types.ServerNotification] = []
+        tool_started = anyio.Event()
+        tool_cancelled = anyio.Event()
+        
+        request_state_manager = InMemoryRequestStateManager[types.ClientRequest, types.ClientResult]()
+
+        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.set()
+                    else:
+                        await tool_cancelled.wait()
+                    
+                if isinstance(message.root, types.CancelledNotification):
+                    nonlocal tool_cancelled
+                    tool_cancelled.set()
+
+
+        # 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,
+            _,
+        ):
+            async with ClientSession(
+                read_stream,
+                write_stream,
+                message_handler=message_handler,
+                request_state_manager=request_state_manager,
+            ) as session:
+                # Initialize the session
+                result = await session.initialize()
+                assert isinstance(result, InitializeResult)
+
+                # Start a long-running tool in a task
+                async with anyio.create_task_group() as tg:
+                    async def run_tool():
+                        request_id = await session.request_call_tool(
+                            "long_running_with_checkpoints", arguments={},
+                            timeout=0.01,
+                            cancel_if_not_resumable=True
+                        )
+                        assert request_id is None
+
+                    tg.start_soon(run_tool)
+
+                    await tool_started.wait()
+                    await tool_cancelled.wait()
+
+        assert tool_started.is_set() and tool_cancelled.is_set()
+        assert len(request_state_manager._progress_callbacks) == 0
+        assert len(request_state_manager._response_streams) == 0
+
+
 @pytest.mark.anyio
 async def test_streamablehttp_client_resumption_timeout(event_server: tuple[SimpleEventStore, str]):
     """Test client session to resume a long running tool via non blocking api with timeout."""
@@ -1401,6 +1468,7 @@ async def run_tool():
                         captured_request_id = await session.request_call_tool(
                             "long_running_with_checkpoints", arguments={}
                         )
+                        assert captured_request_id is not None
 
                         result = await session.join_call_tool(
                             captured_request_id,