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 <gweale@google.com>
PiperOrigin-RevId: 970107514

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(

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,
+  )