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:
parent
3748011d35
commit
5269b7ddd5
2 changed files with 49 additions and 12 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue