From e66cce25388af6db2898b2a2bd27393356b9bb97 Mon Sep 17 00:00:00 2001 From: Rishav Mitra Date: Mon, 13 Apr 2026 10:59:09 -0700 Subject: [PATCH] fix: task.wait() hangs indefinitely when task enters input_required (#3798) * fix: resolve OpenAPI 3.x server variables in _create_default_client When an OpenAPI spec defines server variables (e.g. `https://{region}.api.example.com/v1`), the default values are now substituted before constructing the httpx client base URL. Previously, the URL was used as-is, causing all requests to fail for specs that use server variable templating. Fixes #1681 * fix: use str.replace instead of format_map for server variable substitution format_map applies Python string formatting rules, so variable names like {api.version} would be treated as attribute access and raise errors. Literal token replacement handles all valid OpenAPI variable names safely. * fix: task.wait() now returns on input_required instead of hanging Previously, wait() used a terminal-state allowlist (completed, failed, cancelled), so tasks entering input_required would hang until timeout. Replaced with inverse logic: return whenever the task exits the 'working' state. This handles input_required and any future blocking states without needing to update the allowlist. Fixes #3779 * fix: include submitted in in_progress_states to avoid premature return * fix: revert submitted, update state docstring to match MCP spec * fix: add _wait_terminal() so result() waits for completed/failed/cancelled wait() correctly returns on input_required for human-in-the-loop use cases, but result() needs to wait until the task fully resolves. Add a private _wait_terminal() helper that loops through non-terminal states and use it in all result() implementations. --- src/fastmcp/client/tasks.py | 29 ++++++++++++++----- .../tasks/test_client_task_notifications.py | 26 +++++++++++++++++ 2 files changed, 48 insertions(+), 7 deletions(-) diff --git a/src/fastmcp/client/tasks.py b/src/fastmcp/client/tasks.py index ae6b0ad98..60166c5ca 100644 --- a/src/fastmcp/client/tasks.py +++ b/src/fastmcp/client/tasks.py @@ -216,8 +216,8 @@ class Task(abc.ABC, Generic[TaskResultT]): on status changes when server sends notifications/tasks/status. Args: - state: Desired state ('submitted', 'working', 'completed', 'failed'). - If None, waits for any terminal state (completed/failed) + state: Desired state ('working', 'input_required', 'completed', 'failed', 'cancelled'). + If None, waits until the task exits the 'working' state (completed, failed, cancelled, input_required, etc.) timeout: Maximum time to wait in seconds Returns: @@ -237,7 +237,7 @@ class Task(abc.ABC, Generic[TaskResultT]): self._status_event = asyncio.Event() start = time.time() - terminal_states = {"completed", "failed", "cancelled"} + in_progress_states = {"working"} poll_interval = 0.5 # Fallback polling interval (500ms) while True: @@ -245,7 +245,7 @@ class Task(abc.ABC, Generic[TaskResultT]): if self._status_cache: current = self._status_cache.status if state is None: - if current in terminal_states: + if current not in in_progress_states: return self._status_cache elif current == state: return self._status_cache @@ -269,6 +269,21 @@ class Task(abc.ABC, Generic[TaskResultT]): # Fallback: poll server (notification didn't arrive in time) self._status_cache = await self._client.get_task_status(self._task_id) + async def _wait_terminal(self, timeout: float = 300.0) -> GetTaskResult: + """Wait until task reaches a terminal state (completed, failed, cancelled). + + Unlike wait(), this will not return on input_required — it continues + waiting until the task fully resolves. Used internally by result(). + """ + terminal_states = {"completed", "failed", "cancelled"} + status = await self.wait(timeout=timeout) + while status.status not in terminal_states: + # Task is in a non-terminal state (e.g. input_required) — reset + # cache so the next wait() call blocks instead of returning immediately. + self._status_cache = None + status = await self.wait(timeout=timeout) + return status + async def cancel(self) -> None: """Cancel this task, transitioning it to cancelled state. @@ -354,7 +369,7 @@ class ToolTask(Task["CallToolResult"]): self._check_client_connected() # Wait for completion using event-based wait (respects notifications) - await self.wait() + await self._wait_terminal() # Get the raw result (dict or CallToolResult) raw_result = await self._client.get_task_result(self._task_id) @@ -445,7 +460,7 @@ class PromptTask(Task[mcp.types.GetPromptResult]): self._check_client_connected() # Wait for completion using event-based wait (respects notifications) - await self.wait() + await self._wait_terminal() # Get the raw MCP result mcp_result = await self._client.get_task_result(self._task_id) @@ -517,7 +532,7 @@ class ResourceTask( self._check_client_connected() # Wait for completion using event-based wait (respects notifications) - await self.wait() + await self._wait_terminal() # Get the raw MCP result mcp_result = await self._client.get_task_result(self._task_id) diff --git a/tests/client/tasks/test_client_task_notifications.py b/tests/client/tasks/test_client_task_notifications.py index 8fba3aad2..b02d149fc 100644 --- a/tests/client/tasks/test_client_task_notifications.py +++ b/tests/client/tasks/test_client_task_notifications.py @@ -7,6 +7,7 @@ and invoke user callbacks. import asyncio import time +from datetime import datetime, timezone import pytest from mcp.types import GetTaskResult @@ -207,3 +208,28 @@ async def test_notification_with_failed_task(task_notification_server): assert ( status.statusMessage is not None ) # Error details in statusMessage per spec + + +async def test_wait_returns_on_input_required(task_notification_server): + """wait() should return immediately when task enters input_required, not hang.""" + async with Client(task_notification_server) as client: + task = await client.call_tool("quick_task", {"value": 1}, task=True) + + # Directly inject an input_required status into the cache and signal the event + now = datetime.now(timezone.utc) + input_required_status = GetTaskResult( + taskId=task._task_id, + status="input_required", + statusMessage="Waiting for user input", + createdAt=now, + lastUpdatedAt=now, + ttl=None, + ) + task._status_cache = input_required_status + if task._status_event is None: + task._status_event = asyncio.Event() + task._status_event.set() + + # Should return immediately with input_required, not hang for 300s + status = await task.wait(timeout=2.0) + assert status.status == "input_required"