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.