mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 21:44:18 +02:00
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.
This commit is contained in:
parent
db6d7a8a61
commit
e66cce2538
2 changed files with 48 additions and 7 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue