From a84a4b52aa90eeed5bb6819e8c4d69542b8b52db Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 24 Aug 2026 15:10:02 -0700 Subject: [PATCH] fix(cli): clean up pytest subprocesses on test-client disconnect The dev-server test-run endpoint spawned pytest inside a fire-and-forget asyncio.create_task() and piped its output through an unbounded asyncio.Queue. Nothing owned that task, so a client that disconnected mid-run left pytest, its descendants, and the output pump running until the server itself exited, and the queue could grow without bound while no consumer was draining it. The response iterator now owns the subprocess for its whole lifetime: it spawns pytest, reads bounded chunks straight off the pipe so the client applies natural backpressure, and terminates the process tree in a finally block. Termination reaches descendants rather than just the direct child - on POSIX pytest is started as its own process-group leader and signalled with os.killpg, and on Windows it runs in a new process group torn down with taskkill /T. Cleanup escalates from a graceful signal to a forced kill after a bounded wait, and falls back to signalling the direct child if the process group turns out not to exist. Cleanup runs under an anyio shield, so the cancel scope the server cancels on client disconnect cannot interrupt it partway. The shield covers the common case, where the disconnect arrives while the iterator is parked reading pytest output or awaiting process exit. It is not a guarantee on every path: if the disconnect lands while the iterator is suspended at a yield, the async generator is dropped rather than cancelled, and its finally block runs at async-generator finalization instead. That finalization does happen under CPython, but its timing is not deterministic. Behavior change: disconnecting from the test-output stream now aborts the in-flight pytest run. Previously the run continued to completion in the background after the client went away. Nothing persists the result of a run - the output is only streamed - so a background completion was unobservable, but a caller that relied on starting a run and hanging up must now keep the response stream open until it ends. The endpoint path, its parameters, and the streamed byte content are unchanged. Co-authored-by: George Weale PiperOrigin-RevId: 970107514 --- src/google/adk/cli/dev_server.py | 186 +++++++--- .../cli/test_adk_web_server_tests.py | 343 +++++++++++++++++- 2 files changed, 469 insertions(+), 60 deletions(-) diff --git a/src/google/adk/cli/dev_server.py b/src/google/adk/cli/dev_server.py index cffa28ea..67c2e097 100644 --- a/src/google/adk/cli/dev_server.py +++ b/src/google/adk/cli/dev_server.py @@ -31,17 +31,21 @@ network, and never use it for a production or multi-user deployment. from __future__ import annotations import asyncio +from collections.abc import AsyncIterator import json import logging import os from pathlib import Path import shutil +import signal +import subprocess import sys import time from typing import Any from typing import Iterator from typing import Optional +import anyio from fastapi import FastAPI from fastapi import HTTPException from fastapi import Request as FastAPIRequest @@ -86,6 +90,10 @@ from .utils.state import create_empty_state logger = logging.getLogger("google_adk." + __name__) _EVAL_SET_FILE_EXTENSION = ".evalset.json" +_PROCESS_TERMINATION_GRACE_SECONDS = 0.5 +_PROCESS_TERMINATOR_TIMEOUT_SECONDS = 1.0 +_TEST_OUTPUT_CHUNK_BYTES = 64 * 1024 +_IS_WINDOWS = os.name == "nt" TAG_DEBUG = "Debug" TAG_EVALUATION = "Evaluation" @@ -286,6 +294,126 @@ def _check_code_reference( ) +async def _signal_process_tree( + process: asyncio.subprocess.Process, + *, + force: bool, +) -> None: + """Requests termination of a subprocess and its descendants.""" + if _IS_WINDOWS: + command = ["taskkill", "/PID", str(process.pid), "/T"] + if force: + command.append("/F") + try: + terminator = await asyncio.create_subprocess_exec( + *command, + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + ) + try: + await asyncio.wait_for( + terminator.wait(), timeout=_PROCESS_TERMINATOR_TIMEOUT_SECONDS + ) + except asyncio.TimeoutError: + terminator.kill() + try: + await asyncio.wait_for( + terminator.wait(), timeout=_PROCESS_TERMINATION_GRACE_SECONDS + ) + except asyncio.TimeoutError: + logger.warning("taskkill did not exit for process %d", process.pid) + except OSError: + logger.warning("Unable to run taskkill for process %d", process.pid) + if force and process.returncode is None: + process.kill() + return + + kill_process_group = getattr(os, "killpg", None) + if kill_process_group is not None: + try: + kill_process_group( + process.pid, + ( + getattr(signal, "SIGKILL", signal.SIGTERM) + if force + else signal.SIGTERM + ), + ) + return + except ProcessLookupError: + # No such group: everything already exited, or the child never led one. + # The returncode check below tells those apart. + pass + except OSError: + logger.warning("Unable to signal process group %d", process.pid) + if process.returncode is None: + (process.kill if force else process.terminate)() + + +async def _terminate_process_tree( + process: asyncio.subprocess.Process, +) -> None: + """Terminates a process tree within bounded waits.""" + if process.returncode is not None: + return + + for force in (False, True): + await _signal_process_tree(process, force=force) + try: + await asyncio.wait_for( + process.wait(), timeout=_PROCESS_TERMINATION_GRACE_SECONDS + ) + return + except asyncio.TimeoutError: + pass + + logger.error("Process tree %d did not terminate cleanly", process.pid) + + +async def _stream_test_output( + *, + agent_dir: str, + test_name: str | None, +) -> AsyncIterator[bytes]: + """Runs pytest and yields bounded output chunks until completion.""" + cmd_args = [ + sys.executable, + "-m", + "pytest", + os.path.join(os.path.dirname(__file__), "agent_test_runner.py"), + "-s", + "-vv", + ] + if test_name: + name_to_use = test_name[:-5] if test_name.endswith(".json") else test_name + cmd_args.extend(["-k", name_to_use]) + + env = os.environ.copy() + env["ADK_TEST_FOLDER"] = agent_dir + process = await asyncio.create_subprocess_exec( + *cmd_args, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + env=env, + creationflags=( + getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) + if _IS_WINDOWS + else 0 + ), + start_new_session=not _IS_WINDOWS, + ) + + try: + if process.stdout is None: + raise RuntimeError("pytest output pipe was not created") + while chunk := await process.stdout.read(_TEST_OUTPUT_CHUNK_BYTES): + yield chunk + await process.wait() + finally: + with anyio.CancelScope(shield=True): + await _terminate_process_tree(process) + + class DevServer(ApiServer): """Development server that extends ApiServer with dev-only endpoints. @@ -801,60 +929,10 @@ class DevServer(ApiServer): ) -> StreamingResponse: """Runs tests and streams pytest output.""" agent_dir = self._get_agent_dir(app_name) - - import subprocess - - queue: asyncio.Queue[str | None] = asyncio.Queue() - - async def run_pytest_subprocess(): - cmd_args = [ - sys.executable, - "-m", - "pytest", - os.path.join(os.path.dirname(__file__), "agent_test_runner.py"), - "-s", - "-vv", - ] - if test_name: - name_to_use = ( - test_name[:-5] if test_name.endswith(".json") else test_name - ) - cmd_args.extend(["-k", name_to_use]) - - # Ensure environment variable is set - env = os.environ.copy() - env["ADK_TEST_FOLDER"] = agent_dir - - try: - process = await asyncio.create_subprocess_exec( - *cmd_args, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - env=env, - ) - - while True: - line = await process.stdout.readline() - if not line: - break - await queue.put(line.decode("utf-8")) - - await process.wait() - finally: - # Signal completion to generator - await queue.put(None) - - # Start pytest in a background task - asyncio.create_task(run_pytest_subprocess()) - - async def generate(): - while True: - item = await queue.get() - if item is None: - break - yield item.encode("utf-8") - - return StreamingResponse(generate(), media_type="text/plain") + return StreamingResponse( + _stream_test_output(agent_dir=agent_dir, test_name=test_name), + media_type="text/plain", + ) @app.put("/dev/apps/{app_name}/tests/{test_name}") async def create_test( diff --git a/tests/unittests/cli/test_adk_web_server_tests.py b/tests/unittests/cli/test_adk_web_server_tests.py index 7e09d53e..62781ba2 100644 --- a/tests/unittests/cli/test_adk_web_server_tests.py +++ b/tests/unittests/cli/test_adk_web_server_tests.py @@ -15,16 +15,79 @@ from __future__ import annotations import asyncio -import json +import functools import os +import signal +import subprocess +import sys from unittest.mock import AsyncMock +from unittest.mock import call from unittest.mock import MagicMock from unittest.mock import patch +import anyio from fastapi.testclient import TestClient +from google.adk.cli import dev_server from google.adk.cli.fast_api import get_fast_api_app import pytest +_LIFECYCLE_TEST_TIMEOUT_SECONDS = 10 + +# Two processes that ignore SIGTERM. The grandchild announces itself on an +# inherited pipe, so EOF on that pipe means the whole group is gone. +_SIGTERM_PROOF_GRANDCHILD = ( + "import os, signal, sys, time\n" + "signal.signal(signal.SIGTERM, signal.SIG_IGN)\n" + "os.write(int(sys.argv[1]), b'up')\n" + "time.sleep(30)\n" +) +_SIGTERM_PROOF_PARENT = ( + "import os, signal, subprocess, sys, time\n" + "signal.signal(signal.SIGTERM, signal.SIG_IGN)\n" + "fd = int(sys.argv[1])\n" + "subprocess.Popen([sys.executable, '-c', sys.argv[2], str(fd)]," + " pass_fds=(fd,))\n" + "os.close(fd)\n" + "time.sleep(30)\n" +) + + +def _bounded(test_fn): + """Fails a deadlocked lifecycle test instead of hanging the suite. + + asyncio.wait_for is not enough: a task parked inside a shielded cancel scope + never sees the cancellation, and wait_for then waits on it forever. This + abandons the task instead of awaiting it. + """ + + @functools.wraps(test_fn) + async def wrapper(*args, **kwargs): + task = asyncio.ensure_future(test_fn(*args, **kwargs)) + done, _ = await asyncio.wait( + {task}, timeout=_LIFECYCLE_TEST_TIMEOUT_SECONDS + ) + if not done: + task.cancel() + pytest.fail( + f"{test_fn.__name__} did not finish within" + f" {_LIFECYCLE_TEST_TIMEOUT_SECONDS}s" + ) + await task + + return wrapper + + +async def _read_pipe(read_fd: int, size: int, timeout: float) -> bytes | None: + """Reads from a non-blocking pipe; b"" means every writer is gone.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + try: + return os.read(read_fd, size) + except BlockingIOError: + await asyncio.sleep(0.01) + return None + @pytest.fixture def test_client(tmp_path): @@ -156,12 +219,9 @@ def test_rebuild_single_test(test_client): def test_run_tests(test_client): - from unittest.mock import AsyncMock - from unittest.mock import MagicMock - from unittest.mock import patch - mock_process = MagicMock() - mock_process.stdout.readline = AsyncMock( + mock_process.returncode = 0 + mock_process.stdout.read = AsyncMock( side_effect=[b"line1\n", b"line2\n", b""] ) mock_process.wait = AsyncMock(return_value=0) @@ -179,3 +239,274 @@ def test_run_tests(test_client): content = response.content assert b"line1\n" in content assert b"line2\n" in content + + +@pytest.mark.asyncio +@_bounded +async def test_stream_test_output_starts_new_session_on_posix(monkeypatch): + monkeypatch.setattr(dev_server, "_IS_WINDOWS", False) + mock_process = MagicMock() + mock_process.returncode = 0 + mock_process.stdout.read = AsyncMock(return_value=b"") + mock_process.wait = AsyncMock(return_value=0) + mock_create_subprocess = AsyncMock(return_value=mock_process) + + with ( + patch( + "google.adk.cli.dev_server.asyncio.create_subprocess_exec", + new=mock_create_subprocess, + ), + patch( + "google.adk.cli.dev_server._terminate_process_tree", + new=AsyncMock(), + ), + ): + async for _ in dev_server._stream_test_output( + agent_dir="agent", test_name=None + ): + pass + + # Signalling the group only reaches descendants if the child leads its own. + assert mock_create_subprocess.await_args.kwargs["start_new_session"] is True + + +@pytest.mark.asyncio +@_bounded +async def test_stream_test_output_shields_cleanup_from_cancel_scope(): + read_started = anyio.Event() + cleanup_finished = anyio.Event() + mock_process = MagicMock() + mock_process.returncode = None + + async def read_output(_size): + read_started.set() + await anyio.sleep_forever() + + async def cleanup(_process): + await anyio.sleep(0) + cleanup_finished.set() + + mock_process.stdout.read = AsyncMock(side_effect=read_output) + + with ( + patch( + "google.adk.cli.dev_server.asyncio.create_subprocess_exec", + new=AsyncMock(return_value=mock_process), + ), + patch( + "google.adk.cli.dev_server._terminate_process_tree", + new=AsyncMock(side_effect=cleanup), + ), + ): + + async def consume_output(): + async for _ in dev_server._stream_test_output( + agent_dir="agent", test_name=None + ): + pass + + async with anyio.create_task_group() as task_group: + task_group.start_soon(consume_output) + await read_started.wait() + task_group.cancel_scope.cancel() + + assert cleanup_finished.is_set() + + +@pytest.mark.asyncio +@_bounded +async def test_stream_test_output_cleans_up_after_partial_output(): + mock_process = MagicMock() + mock_process.returncode = None + mock_process.stdout.read = AsyncMock(return_value=b"partial output") + mock_cleanup = AsyncMock() + + with ( + patch( + "google.adk.cli.dev_server.asyncio.create_subprocess_exec", + new=AsyncMock(return_value=mock_process), + ), + patch( + "google.adk.cli.dev_server._terminate_process_tree", + new=mock_cleanup, + ), + ): + output = dev_server._stream_test_output(agent_dir="agent", test_name=None) + assert await anext(output) == b"partial output" + await output.aclose() + + mock_cleanup.assert_awaited_once_with(mock_process) + + +@pytest.mark.asyncio +@_bounded +async def test_stream_test_output_shields_cleanup_during_process_wait(): + wait_started = anyio.Event() + cleanup_finished = anyio.Event() + mock_process = MagicMock() + mock_process.returncode = None + mock_process.stdout.read = AsyncMock(return_value=b"") + + async def wait_for_exit(): + wait_started.set() + await anyio.sleep_forever() + + async def cleanup(_process): + await anyio.sleep(0) + cleanup_finished.set() + + mock_process.wait = AsyncMock(side_effect=wait_for_exit) + + with ( + patch( + "google.adk.cli.dev_server.asyncio.create_subprocess_exec", + new=AsyncMock(return_value=mock_process), + ), + patch( + "google.adk.cli.dev_server._terminate_process_tree", + new=AsyncMock(side_effect=cleanup), + ), + ): + + async def consume_output(): + async for _ in dev_server._stream_test_output( + agent_dir="agent", test_name=None + ): + pass + + async with anyio.create_task_group() as task_group: + task_group.start_soon(consume_output) + await wait_started.wait() + task_group.cancel_scope.cancel() + + assert cleanup_finished.is_set() + + +@pytest.mark.asyncio +@_bounded +async def test_terminate_process_tree_stops_after_graceful_exit(): + mock_process = MagicMock() + mock_process.returncode = None + mock_process.wait = AsyncMock(return_value=0) + mock_signal = AsyncMock() + + with patch("google.adk.cli.dev_server._signal_process_tree", new=mock_signal): + await dev_server._terminate_process_tree(mock_process) + + mock_signal.assert_awaited_once_with(mock_process, force=False) + mock_process.wait.assert_awaited_once() + + +@pytest.mark.asyncio +@_bounded +async def test_terminate_process_tree_escalates_after_timeout(): + mock_process = MagicMock() + mock_process.returncode = None + mock_process.wait = AsyncMock(side_effect=[asyncio.TimeoutError, 0]) + mock_signal = AsyncMock() + + with patch("google.adk.cli.dev_server._signal_process_tree", new=mock_signal): + await dev_server._terminate_process_tree(mock_process) + + assert mock_signal.await_args_list == [ + call(mock_process, force=False), + call(mock_process, force=True), + ] + + +@pytest.mark.skipif(os.name == "nt", reason="POSIX process groups") +@pytest.mark.asyncio +@_bounded +async def test_terminate_process_tree_kills_real_grandchild(): + read_fd, write_fd = os.pipe() + os.set_blocking(read_fd, False) + process = await asyncio.create_subprocess_exec( + sys.executable, + "-c", + _SIGTERM_PROOF_PARENT, + str(write_fd), + _SIGTERM_PROOF_GRANDCHILD, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + pass_fds=(write_fd,), + start_new_session=True, + ) + os.close(write_fd) + try: + assert await _read_pipe(read_fd, 2, 5) == b"up" + + await dev_server._terminate_process_tree(process) + + # Both processes ignore SIGTERM, so only the forced escalation can land. + assert process.returncode == -signal.SIGKILL + # The grandchild is the last holder of the write end; EOF means it is gone. + assert await _read_pipe(read_fd, 1, 3) == b"" + finally: + os.close(read_fd) + if process.returncode is None: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + await process.wait() + + +@pytest.mark.asyncio +@_bounded +async def test_signal_process_tree_falls_back_to_direct_child(monkeypatch): + mock_process = MagicMock() + mock_process.pid = 123 + mock_process.returncode = None + monkeypatch.setattr(dev_server, "_IS_WINDOWS", False) + monkeypatch.setattr( + dev_server.os, + "killpg", + MagicMock(side_effect=ProcessLookupError), + raising=False, + ) + + await dev_server._signal_process_tree(mock_process, force=True) + + mock_process.kill.assert_called_once_with() + + +@pytest.mark.asyncio +@_bounded +async def test_signal_process_tree_targets_posix_process_group(monkeypatch): + mock_process = MagicMock() + mock_process.pid = 123 + mock_killpg = MagicMock() + monkeypatch.setattr(dev_server, "_IS_WINDOWS", False) + monkeypatch.setattr(dev_server.os, "killpg", mock_killpg, raising=False) + + await dev_server._signal_process_tree(mock_process, force=False) + + mock_killpg.assert_called_once_with(123, dev_server.signal.SIGTERM) + + +@pytest.mark.asyncio +@_bounded +async def test_signal_process_tree_targets_windows_descendants(monkeypatch): + mock_process = MagicMock() + mock_process.pid = 123 + mock_terminator = MagicMock() + mock_terminator.wait = AsyncMock(return_value=0) + mock_create_subprocess = AsyncMock(return_value=mock_terminator) + monkeypatch.setattr(dev_server, "_IS_WINDOWS", True) + monkeypatch.setattr( + dev_server.asyncio, + "create_subprocess_exec", + mock_create_subprocess, + ) + + await dev_server._signal_process_tree(mock_process, force=True) + + mock_create_subprocess.assert_awaited_once_with( + "taskkill", + "/PID", + "123", + "/T", + "/F", + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + )