From ffed951f671c49a6a33b78f651860e85db4ae649 Mon Sep 17 00:00:00 2001 From: Justin Wilson Date: Mon, 24 Aug 2026 11:16:09 -0500 Subject: [PATCH] fix(client/stdio): allow FIFO cleanup of multiple transports on asyncio stdio_client wrapped its reader/writer in anyio.create_task_group() inside the async context manager. The task group binds its cancel scope to the task that opened the transport, and anyio then requires that task to close transports LIFO. Independent AsyncExitStacks, multi-client managers, and unordered pytest fixtures close FIFO and hit: RuntimeError: Attempted to exit cancel scope in a different task On asyncio, spawn the pipe tasks with asyncio.ensure_future so they are not stacked on the caller's cancel-scope. Shutdown order is unchanged. Trio still uses a task group and still requires LIFO. Regression tests pin oldest-first and newest-first teardown, and that an unhandled pipe-task failure still surfaces. Fixes #577 --- src/mcp/client/stdio.py | 74 ++++++++++++++++++++++++++++++++++---- tests/client/test_stdio.py | 60 +++++++++++++++++++++++++++++-- 2 files changed, 124 insertions(+), 10 deletions(-) diff --git a/src/mcp/client/stdio.py b/src/mcp/client/stdio.py index 3e03eef9ef..d103bf96bf 100644 --- a/src/mcp/client/stdio.py +++ b/src/mcp/client/stdio.py @@ -8,13 +8,14 @@ process nor hang on one. """ +import asyncio import logging import os import sys -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Callable, Coroutine from contextlib import asynccontextmanager, suppress from pathlib import Path -from typing import Literal, TextIO +from typing import Any, Literal, TextIO import anyio import anyio.lowlevel @@ -130,7 +131,7 @@ async def stdio_client( cwd=server.cwd, ) - # The spawn succeeded; no awaits until the task group is entered, or a + # The spawn succeeded; no awaits until the pipe tasks are running, or a # cancellation delivered in the gap would leak the live process. read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0) write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0) @@ -197,9 +198,7 @@ async def shutdown() -> None: # One pass so unblocked tasks exit via their except paths before the cancel. await anyio.lowlevel.checkpoint() - async with anyio.create_task_group() as tg: - tg.start_soon(stdout_reader) - tg.start_soon(stdin_writer) + async with _run_pipe_tasks(stdout_reader, stdin_writer) as cancel_pipe_tasks: try: yield read_stream, write_stream finally: @@ -210,11 +209,72 @@ async def shutdown() -> None: with anyio.CancelScope(shield=True): await shutdown() # Unstick pipe tasks a kill survivor's open pipe end could still block. - tg.cancel_scope.cancel() + cancel_pipe_tasks() # The cancel lands via throw(); one yield resyncs 3.11 coverage (gh-106749). await anyio.lowlevel.cancel_shielded_checkpoint() +@asynccontextmanager +async def _run_pipe_tasks( + *pipes: Callable[[], Coroutine[Any, Any, None]], +) -> AsyncGenerator[Callable[[], None], None]: + """Runs the pipe tasks for the duration of the body, yielding their canceller. + + On asyncio they are plain asyncio tasks rather than an anyio task group: a task + group binds its cancel scope to the task that opened the transport, and anyio then + requires that task to close its transports in the reverse of the order it opened + them. Callers holding several servers (a multi-server manager, independent exit + stacks, pytest fixtures) legitimately close them in other orders, and got a + "cancel scope" RuntimeError instead of a clean teardown -- see #577. Trio cannot + spawn a task outside a nursery, so it keeps the task group and still requires + LIFO closing. + """ + if not _on_asyncio(): + async with anyio.create_task_group() as tg: + for pipe in pipes: + tg.start_soon(pipe) + yield tg.cancel_scope.cancel + return + + tasks = [asyncio.ensure_future(pipe()) for pipe in pipes] + + def cancel_pipe_tasks() -> None: + for task in tasks: + task.cancel() + + errors: list[Exception] = [] + try: + yield cancel_pipe_tasks + finally: + cancel_pipe_tasks() + # Shielded, as the task group's own reaping was: a cancelled caller must + # still leave no pipe task behind. Every task is awaited before anything is + # re-raised, so a second failure cannot surface as an unretrieved exception. + with anyio.CancelScope(shield=True): + for task in tasks: + try: + await task + except asyncio.CancelledError: + pass # our own cancellation above, not a failure + except Exception as exc: # the pipe tasks' top-level handler + errors.append(exc) + # Outside the finally: a body that raised keeps its own exception. + if errors: + raise errors[0] + + +def _on_asyncio() -> bool: + """Whether the caller is running on anyio's asyncio backend. + + True exactly when an asyncio task is executing, which is when spawning further + asyncio tasks is meaningful; on trio there is no running loop to ask. + """ + try: + return asyncio.current_task() is not None + except RuntimeError: + return False + + def _parse_line(line: str) -> SessionMessage | Exception: """Parses one stdout line, returning parse errors as values for the session to surface.""" try: diff --git a/tests/client/test_stdio.py b/tests/client/test_stdio.py index 91f829ff98..7f27e6cd25 100644 --- a/tests/client/test_stdio.py +++ b/tests/client/test_stdio.py @@ -189,15 +189,16 @@ def pending_stdout_chunks(self) -> int: def install_fake_process( - monkeypatch: pytest.MonkeyPatch, process: FakeProcess, *, grace_period: float | None = 0.2 + monkeypatch: pytest.MonkeyPatch, *processes: FakeProcess, grace_period: float | None = 0.2 ) -> list[FakeProcess]: - """Route stdio_client's spawn and terminate seams to `process`. + """Route stdio_client's spawn and terminate seams to `processes`, one per spawn. Returns the list of processes the (fake) tree termination was invoked on. `grace_period=None` keeps the production stdin-close grace (affordable only on a virtual clock). """ terminated: list[FakeProcess] = [] + to_spawn = iter(processes) async def fake_spawn( command: str, @@ -206,7 +207,7 @@ async def fake_spawn( errlog: TextIO = sys.stderr, cwd: Path | str | None = None, ) -> FakeProcess: - return process + return next(to_spawn) async def fake_terminate_tree(proc: FakeProcess) -> None: terminated.append(proc) @@ -480,6 +481,59 @@ async def run_client_until_cancelled() -> None: assert terminated == [process] +@pytest.mark.anyio +@pytest.mark.parametrize("close_order", [(0, 1), (1, 0)], ids=["oldest-first", "newest-first"]) +async def test_two_transports_held_by_one_task_close_in_either_order( + monkeypatch: pytest.MonkeyPatch, close_order: tuple[int, int] +) -> None: + """One task holding two transports may close them oldest-first, not only newest-first. + + Pins issue #577: a multi-server manager, independent exit stacks, or unordered pytest + fixtures got a cancel-scope RuntimeError out of the oldest-first teardown. Asyncio + only (this module's backend): on trio the transport still borrows the caller's + nursery, so newest-first stays the requirement there. + """ + first = FakeProcess(on_stdin_close=lambda: first.exit(0)) + second = FakeProcess(on_stdin_close=lambda: second.exit(0)) + terminated = install_fake_process(monkeypatch, first, second) + + stacks = [AsyncExitStack(), AsyncExitStack()] + + with anyio.fail_after(5): + for stack in stacks: + await stack.enter_async_context(stdio_client(FAKE_PARAMS)) + for index in close_order: + await stacks[index].aclose() + + # Both servers went through the full shutdown, so neither needed terminating. + assert first.stdin_closed.is_set() + assert second.stdin_closed.is_set() + assert terminated == [] + + +@pytest.mark.anyio +async def test_an_unhandled_pipe_task_failure_surfaces_out_of_the_context_manager( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A pipe task failing in a way the transport does not handle reaches the caller. + + Undecodable server output crashes the reader (the encoding error handler defaults to + `strict`); exiting still shuts the server down cleanly, but the error is re-raised + rather than swallowed. + """ + process = FakeProcess(on_stdin_close=lambda: process.exit(0)) + terminated = install_fake_process(monkeypatch, process) + + with pytest.raises(UnicodeDecodeError): + with anyio.fail_after(5): + async with stdio_client(FAKE_PARAMS): + await process.feed(b"\xff\xfe not utf-8\n") + # Wait until the reader has actually decoded the bytes and died. + await anyio.wait_all_tasks_blocked() + + assert terminated == [] + + @pytest.mark.anyio async def test_writing_after_the_server_dies_reports_clean_closure(monkeypatch: pytest.MonkeyPatch) -> None: """A send racing the server's death must not surface a raw backend exception.