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:
Rishav Mitra 2026-04-13 10:59:09 -07:00 committed by GitHub
commit e66cce2538
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 48 additions and 7 deletions

View file

@ -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)

View file

@ -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"