diff --git a/src/acp/transports.py b/src/acp/transports.py index 31a704e..82a4de6 100644 --- a/src/acp/transports.py +++ b/src/acp/transports.py @@ -58,6 +58,7 @@ async def spawn_stdio_transport( This mirrors the defensive shutdown behaviour used by the MCP Python SDK: close stdin first, wait for graceful exit, then escalate to terminate/kill. + Stdin closure uses the same timeout as the waits before terminate/kill. """ merged_env = dict(default_environment()) if env: @@ -96,16 +97,9 @@ async def spawn_stdio_transport( finally: # Attempt graceful stdin shutdown first if process.stdin is not None: - try: - process.stdin.write_eof() - except (AttributeError, OSError, RuntimeError): - process.stdin.close() - with contextlib.suppress(Exception): - await process.stdin.drain() with contextlib.suppress(Exception): process.stdin.close() - with contextlib.suppress(Exception): - await process.stdin.wait_closed() + await asyncio.wait_for(process.stdin.wait_closed(), timeout=shutdown_timeout) try: await asyncio.wait_for(process.wait(), timeout=shutdown_timeout) diff --git a/tests/real_user/test_stdio_shutdown.py b/tests/real_user/test_stdio_shutdown.py new file mode 100644 index 0000000..8f3d8f3 --- /dev/null +++ b/tests/real_user/test_stdio_shutdown.py @@ -0,0 +1,69 @@ +import asyncio +import contextlib +import signal +import sys + +import pytest + +from acp.transports import spawn_stdio_transport + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ignore_terminate", [False, True]) +async def test_shutdown_with_unread_stdin(ignore_terminate: bool) -> None: + if ignore_terminate and sys.platform == "win32": + pytest.skip("Windows terminate() cannot be ignored") + script = "import signal, time; " + if ignore_terminate: + script += "signal.signal(signal.SIGTERM, signal.SIG_IGN); " + script += "print('ready', flush=True); time.sleep(60)" + transport = spawn_stdio_transport(sys.executable, "-c", script, shutdown_timeout=0.1) + reader, writer, process = await transport.__aenter__() + closing = None + try: + assert (await asyncio.wait_for(reader.readline(), timeout=5)).strip() == b"ready" + writer.write(b"x" * (1024 * 1024)) + assert writer.transport.get_write_buffer_size() > 0 + closing = asyncio.create_task(transport.__aexit__(None, None, None)) + await asyncio.wait_for(asyncio.shield(closing), timeout=5) + assert process.returncode is not None + if sys.platform != "win32": + assert process.returncode == -(signal.SIGKILL if ignore_terminate else signal.SIGTERM) + finally: + if process.returncode is None: + process.kill() + await process.wait() + if closing is not None: + await closing + else: + await transport.__aexit__(None, None, None) + + +@pytest.mark.asyncio +async def test_shutdown_flushes_stdin_before_eof() -> None: + script = "import sys; data = sys.stdin.buffer.read(); print(len(data), flush=True)" + async with spawn_stdio_transport(sys.executable, "-c", script) as (reader, writer, process): + writer.write(b"x" * (1024 * 1024)) + assert process.returncode == 0 + assert (await reader.readline()).strip() == b"1048576" + + +@pytest.mark.asyncio +async def test_shutdown_preserves_body_exception() -> None: + with pytest.raises(ValueError, match="body failure"): + async with spawn_stdio_transport(sys.executable, "-c", "import sys; sys.stdin.buffer.read()") as ( + _reader, + _writer, + process, + ): + raise ValueError("body failure") + assert process.returncode == 0 + + +@pytest.mark.asyncio +async def test_shutdown_after_child_exit() -> None: + async with spawn_stdio_transport(sys.executable, "-c", "pass") as (_reader, writer, process): + await asyncio.wait_for(process.wait(), timeout=5) + with contextlib.suppress(ConnectionError): + await writer.drain() + assert process.returncode == 0