fix(ingest): compose the MCP timeout with anyio's own cancel scope, not asyncio.wait_for (kø-z)

asyncio.wait_for cancelling stdio_call_tool's run() from outside the anyio
task groups it awaits (stdio_client, ClientSession) never surfaced a
TimeoutError: measured against a real hanging server, the mismatch produced
an anyio.BrokenResourceError wrapped in a BaseExceptionGroup instead. Moving
the deadline to anyio.fail_after, nested inside both task groups, lets
anyio tear down its own structure cleanly and raise a plain TimeoutError,
which is now translated into IngestError(code="mcp_timeout") alongside the
mcp_tool_error/mcp_non_text_content family.
This commit is contained in:
Kjell Tore Guttormsen 2026-08-03 21:02:52 +02:00
commit 5269b7ddd5
2 changed files with 49 additions and 12 deletions

View file

@ -140,6 +140,7 @@ def stdio_call_tool(
"""
import asyncio
import anyio
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
@ -157,20 +158,34 @@ def stdio_call_tool(
env["MCP_CREDENTIAL"] = credential
async def run() -> str:
params = StdioServerParameters(command=command, args=list(args), env=env)
async with stdio_client(params) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
result = await session.call_tool(tool, {})
if result.isError:
raise IngestError(
f"MCP tool {tool!r} on {server_ref!r} returned an error result",
code="mcp_tool_error",
)
return _body_from_content(result.content, tool=tool, server_ref=server_ref)
# The deadline is an anyio cancel scope, not `asyncio.wait_for`, and it wraps BOTH
# nested task groups (`stdio_client`, `ClientSession`) from the INSIDE. `wait_for`
# cancels from outside a structure anyio itself owns, and the two cancellation
# mechanisms do not compose: measured (2026-08-03), that mismatch surfaced as an
# `anyio.BrokenResourceError` wrapped in a `BaseExceptionGroup` — never a `TimeoutError`
# — because a background reader task lost its write end mid-teardown. anyio's own scope
# is what the task groups already coordinate cancellation through, so nesting inside it
# tears down cleanly and raises a plain `TimeoutError` at the `with` statement.
with anyio.fail_after(timeout_seconds):
params = StdioServerParameters(command=command, args=list(args), env=env)
async with stdio_client(params) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
result = await session.call_tool(tool, {})
if result.isError:
raise IngestError(
f"MCP tool {tool!r} on {server_ref!r} returned an error result",
code="mcp_tool_error",
)
return _body_from_content(result.content, tool=tool, server_ref=server_ref)
try:
return asyncio.run(asyncio.wait_for(run(), timeout=timeout_seconds))
return asyncio.run(run())
except TimeoutError as exc:
raise IngestError(
f"MCP tool {tool!r} on {server_ref!r} did not respond within {timeout_seconds}s",
code="mcp_timeout",
) from exc
except BaseException as exc: # noqa: BLE001 — re-raised unchanged unless we own it
owned = _unwrap_ingest_error(exc)
if owned is None:

View file

@ -120,6 +120,28 @@ def test_tool_error_reaches_the_caller_as_ingest_error(
assert excinfo.value.code == "mcp_tool_error"
def test_tool_timeout_reaches_the_caller_as_ingest_error(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A tool that outlives ``timeout_seconds`` must surface as a typed ``IngestError``.
``asyncio.wait_for`` wraps ``run()`` itself two nested anyio task groups
(``stdio_client``, ``ClientSession``). Only a REAL hanging server proves what the timeout
cancellation actually surfaces as: a canned-tool test never enters a task group and cannot
observe this at all (the same reason (x)'s ``mcp_tool_error`` unwrap needed a real subprocess).
RED until the timeout is caught and re-raised as an owned ``IngestError``.
"""
script = _write_server(
tmp_path,
"import time\n@server.tool()\ndef cost_docs() -> str:\n time.sleep(5)\n return 'late'\n",
)
get = _transport(script, monkeypatch, timeout_seconds=0.2)
with pytest.raises(IngestError) as excinfo:
get(f"mcp+stdio://{_SERVER_REF}/cost_docs", None)
assert excinfo.value.code == "mcp_timeout"
def test_missing_server_ref_fails_before_any_process_is_spawned(
monkeypatch: pytest.MonkeyPatch,
) -> None: