diff --git a/src/claude_agent_sdk/_internal/transport/subprocess_cli.py b/src/claude_agent_sdk/_internal/transport/subprocess_cli.py index 58abc438d..cf85bc64a 100644 --- a/src/claude_agent_sdk/_internal/transport/subprocess_cli.py +++ b/src/claude_agent_sdk/_internal/transport/subprocess_cli.py @@ -1,5 +1,6 @@ """Subprocess transport implementation using Claude Code CLI.""" +import asyncio import atexit import json import logging @@ -15,6 +16,7 @@ from typing import Any, cast import anyio +import sniffio from anyio.abc import Process from anyio.streams.text import TextReceiveStream, TextSendStream @@ -940,33 +942,35 @@ def emit(line: str) -> None: emit(framer.flush()) async def close(self) -> None: - """Close the transport and clean up resources. - - The whole body runs inside a shielded cancel scope. Cleanup is - routinely reached while the caller's task is being cancelled (e.g. - `async with ClaudeSDKClient()` unwinding on cancel), and an - unshielded close() would abort at the first await — before the - terminate/kill escalation ran — orphaning the CLI child, which then - surfaces as `` once nothing is left to wait() on it. - - Every await in *this* scope is bounded (~20s worst case), so an anyio - cancellation is delayed but never blocked: the stream `aclose()`s are a - non-blocking `close()` plus a checkpoint on both anyio backends (they - never await `wait_closed()`, so undrained stdin cannot wedge them), the - stderr task is cancelled before it is awaited, and the lock acquire and - every process `wait()` carry an explicit deadline. - - Caveat: an anyio shield only defers cancellation that *originates from - an anyio cancel scope*. A raw asyncio cancellation (`asyncio.wait_for` / - `asyncio.timeout` firing, a bare `task.cancel()`, loop shutdown) is - still delivered at the next await in here, and the escalation below only - catches `TimeoutError` — so it would be skipped. That is a pre-existing - limitation of the shield on the asyncio backend rather than something - this scope introduces, and it is still strictly better than before: the - `finally` keeps a still-running child in `_ACTIVE_CHILDREN` for the - atexit reaper instead of dropping it. Making the escalation robust to a - foreign `CancelledError` is a follow-up. + """Finish subprocess cleanup before cancellation reaches the caller. + + Raw asyncio cancellation bypasses AnyIO shields, so cleanup runs in a + separate task and the cancellation is re-raised after that task ends. + Caller timeouts do not bound this call because cleanup completes first; + full escalation may take about 20 seconds. """ + if sniffio.current_async_library() != "asyncio": + await self._close_impl() + return + + cleanup_task: asyncio.Task[None] = asyncio.create_task(self._close_impl()) + cancellation: asyncio.CancelledError | None = None + with anyio.CancelScope(shield=True): + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError as exc: + cancellation = exc + except Exception: + break + if cancellation is not None: + cleanup_error = cleanup_task.exception() + if cleanup_error is not None: + raise cancellation from cleanup_error + raise cancellation + cleanup_task.result() + + async def _close_impl(self) -> None: if not self._process: self._ready = False return diff --git a/tests/test_close_cancellation.py b/tests/test_close_cancellation.py index d3355de9f..b09e63829 100644 --- a/tests/test_close_cancellation.py +++ b/tests/test_close_cancellation.py @@ -6,16 +6,19 @@ child is left running -- an orphan that surfaces as `[claude] ` once nothing is left to wait() on it. -Every test here runs under both asyncio and trio (``anyio_backend`` in -conftest.py): the leak reproduces on both. +Cancellation-scope tests run under both asyncio and trio. The raw +``asyncio.wait_for`` regression runs only under asyncio. """ +import asyncio import json import os +import signal import subprocess import sys import textwrap from collections.abc import Iterator +from contextlib import suppress from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch @@ -110,15 +113,186 @@ async def test_close_under_cancellation_still_reaps_child(tmp_path: Path) -> Non assert process is not None pid = process.pid - with anyio.CancelScope() as scope: - scope.cancel() # every await inside close() would raise - await transport.close() + with patch.object( + transport, "_close_impl", wraps=transport._close_impl + ) as close_impl: + with anyio.CancelScope() as scope: + scope.cancel() # every await inside close() would raise + await transport.close() + close_impl.assert_awaited_once() assert process.returncode is not None assert _process_state(pid) == "gone" assert process not in _ACTIVE_CHILDREN +@posix_only +@pytest.mark.parametrize("anyio_backend", ["asyncio"]) +async def test_asyncio_timeout_still_reaps_child(tmp_path: Path) -> None: + script = _write_fake_cli(tmp_path) + script.write_text(FAKE_CLI + "\nimport time; time.sleep(60)\n") + transport = SubprocessCLITransport( + prompt="hi", + options=ClaudeAgentOptions(cli_path=str(script)), + ) + await transport.connect() + process = transport._process + assert process is not None + + try: + with ( + patch.object( + transport, "_close_impl", wraps=transport._close_impl + ) as close_impl, + pytest.raises(TimeoutError), + ): + await asyncio.wait_for(transport.close(), timeout=0.05) + close_impl.assert_awaited_once() + assert process.returncode == -signal.SIGTERM + assert _process_state(process.pid) == "gone" + assert process not in _ACTIVE_CHILDREN + finally: + await transport.close() + + +@posix_only +@pytest.mark.parametrize("anyio_backend", ["asyncio"]) +async def test_repeated_asyncio_cancellation_still_reaps_child( + tmp_path: Path, +) -> None: + script = _write_fake_cli(tmp_path) + script.write_text(FAKE_CLI + "\nimport time; time.sleep(60)\n") + transport = SubprocessCLITransport( + prompt="hi", + options=ClaudeAgentOptions(cli_path=str(script)), + ) + await transport.connect() + process = transport._process + assert process is not None + + with patch.object( + transport, "_close_impl", wraps=transport._close_impl + ) as close_impl: + close_task: asyncio.Task[None] = asyncio.create_task(transport.close()) + while close_impl.await_count == 0: + await anyio.sleep(0) + deadline = anyio.current_time() + 7 + try: + while not close_task.done() and anyio.current_time() < deadline: + close_task.cancel() + await anyio.sleep(0.02) + assert close_task.done(), ( + f"_close_impl restarted {close_impl.await_count} times" + ) + with pytest.raises(asyncio.CancelledError): + await close_task + finally: + if not close_task.done(): + with suppress(asyncio.CancelledError): + await close_task + + close_impl.assert_awaited_once() + assert process.returncode == -signal.SIGTERM + assert _process_state(process.pid) == "gone" + assert process not in _ACTIVE_CHILDREN + + +@pytest.mark.parametrize("anyio_backend", ["asyncio"]) +async def test_cleanup_failure_does_not_swallow_asyncio_cancellation() -> None: + transport = SubprocessCLITransport( + prompt="hi", + options=ClaudeAgentOptions(cli_path="/usr/bin/claude"), + ) + cleanup_started = asyncio.Event() + fail_cleanup = asyncio.Event() + cleanup_error = RuntimeError("cleanup failed") + + async def failing_close() -> None: + cleanup_started.set() + await fail_cleanup.wait() + raise cleanup_error + + # Python 3.14 reports a failed shielded task to the loop even when close() + # later retrieves and chains that failure. + with ( + patch.object(asyncio.get_running_loop(), "call_exception_handler") as handler, + patch.object(transport, "_close_impl", side_effect=failing_close) as close_impl, + ): + close_task: asyncio.Task[None] = asyncio.create_task(transport.close()) + await cleanup_started.wait() + close_task.cancel() + await anyio.sleep(0) + assert not close_task.done() + fail_cleanup.set() + with pytest.raises(asyncio.CancelledError) as raised: + await close_task + + close_impl.assert_awaited_once() + assert raised.value.__cause__ is cleanup_error + assert all( + call.args[0].get("exception") is cleanup_error + for call in handler.call_args_list + ) + + +@pytest.mark.parametrize("anyio_backend", ["asyncio"]) +async def test_cleanup_failure_does_not_break_asyncio_timeout() -> None: + transport = SubprocessCLITransport( + prompt="hi", + options=ClaudeAgentOptions(cli_path="/usr/bin/claude"), + ) + cleanup_started = asyncio.Event() + fail_cleanup = asyncio.Event() + cleanup_error = RuntimeError("cleanup failed") + + async def failing_close() -> None: + cleanup_started.set() + await fail_cleanup.wait() + raise cleanup_error + + async def release_cleanup() -> None: + await cleanup_started.wait() + await anyio.sleep(0.05) + fail_cleanup.set() + + release_task = asyncio.create_task(release_cleanup()) + try: + with ( + patch.object( + asyncio.get_running_loop(), "call_exception_handler" + ) as handler, + patch.object(transport, "_close_impl", side_effect=failing_close), + pytest.raises(TimeoutError) as raised, + ): + await asyncio.wait_for(transport.close(), timeout=0.01) + finally: + await release_task + assert isinstance(raised.value.__cause__, asyncio.CancelledError) + assert raised.value.__cause__.__cause__ is cleanup_error + assert all( + call.args[0].get("exception") is cleanup_error + for call in handler.call_args_list + ) + + +@pytest.mark.parametrize("anyio_backend", ["asyncio"]) +async def test_cleanup_failure_surfaces_without_cancellation() -> None: + transport = SubprocessCLITransport( + prompt="hi", + options=ClaudeAgentOptions(cli_path="/usr/bin/claude"), + ) + + with ( + patch.object( + transport, + "_close_impl", + side_effect=RuntimeError("cleanup failed"), + ), + pytest.raises(RuntimeError, match="cleanup failed"), + ): + await transport.close() + + @posix_only async def test_cancelled_client_context_leaves_no_child(tmp_path: Path) -> None: """A cancelled `async with ClaudeSDKClient()` must not leak the CLI child."""