Python: Kill the shell process group on stateless shell timeout

python/packages/tools/agent_framework_tools/shell/_killtree.py

@@ -52,8 +52,13 @@ async def kill_process_tree(
     *,
     grace: float = 2.0,
 ) -> None:
-    """Terminate ``proc`` and all of its descendants. Best-effort, never raises."""
-    if proc.returncode is not None:
+    """Terminate ``proc`` and all of its descendants. Best-effort, never raises.
+
+    On POSIX ``proc`` must have been started with ``start_new_session=True``.
+    Its process group is signalled even after ``proc`` has exited, which
+    reaches background and double-forked jobs that are no longer descendants.
+    """
+    if proc.returncode is not None and sys.platform == "win32":
         return
     if _has_psutil:
         await _kill_via_psutil(proc, grace=grace)
@@ -68,26 +73,28 @@ async def _kill_via_psutil(
 ) -> None:
     if psutil is None:
         raise RuntimeError("_kill_via_psutil called without psutil available")
-    try:
-        parent = psutil.Process(proc.pid)
-    except psutil.NoSuchProcess:
+    victims: list[psutil.Process] = []
+    if proc.returncode is None:
+        with contextlib.suppress(psutil.NoSuchProcess):
+            parent = psutil.Process(proc.pid)
+            victims.append(parent)
+            victims.extend(parent.children(recursive=True))
+    if not victims and sys.platform == "win32":
         return
-    try:
-        descendants = parent.children(recursive=True)
-    except psutil.NoSuchProcess:
-        descendants = []
-    victims = [parent, *descendants]
 
     # Phase 1: SIGTERM (or terminate() on Windows, which also asks nicely).
+    if sys.platform != "win32":
+        _signal_group(proc, signal.SIGTERM)
     for v in victims:
         with contextlib.suppress(psutil.NoSuchProcess, psutil.AccessDenied):
             v.terminate()
 
     # Wait briefly for graceful exit.
-    with contextlib.suppress(asyncio.TimeoutError):
-        await asyncio.wait_for(proc.wait(), timeout=grace)
+    await _wait_for_exit(proc, grace=grace)
 
     # Phase 2: SIGKILL anything still alive.
+    if sys.platform != "win32":
+        _signal_group(proc, signal.SIGKILL)
     for v in victims:
         with contextlib.suppress(psutil.NoSuchProcess, psutil.AccessDenied):
             if v.is_running():
@@ -122,11 +129,40 @@ async def _kill_via_stdlib(
         with contextlib.suppress(ProcessLookupError, OSError):
             proc.kill()
         return
+    _signal_group(proc, signal.SIGTERM)
+    await _wait_for_exit(proc, grace=grace)
+    _signal_group(proc, signal.SIGKILL)
+
+
+def _signal_group(proc: asyncio.subprocess.Process, sig: int) -> None:
+    """POSIX only. ``proc.pid`` is the group id under ``start_new_session``, also after ``proc`` is reaped."""
+    with contextlib.suppress(ProcessLookupError, PermissionError):
+        os.killpg(proc.pid, sig)
+
+
+def _group_alive(proc: asyncio.subprocess.Process) -> bool:
     try:
-        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
-        with contextlib.suppress(asyncio.TimeoutError):
-            await asyncio.wait_for(proc.wait(), timeout=grace)
-            return
-        os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
+        os.killpg(proc.pid, 0)
     except (ProcessLookupError, PermissionError):
-        pass
+        return False
+    if psutil is None:
+        return True
+    # killpg also succeeds for a group of zombies that their new parent has not reaped yet.
+    for p in psutil.process_iter(["status"]):
+        with contextlib.suppress(ProcessLookupError):
+            if p.info["status"] != psutil.STATUS_ZOMBIE and os.getpgid(p.pid) == proc.pid:
+                return True
+    return False
+
+
+async def _wait_for_exit(proc: asyncio.subprocess.Process, *, grace: float) -> None:
+    """Wait up to ``grace`` for ``proc`` and, on POSIX, the rest of its process group to exit."""
+
+    async def exited() -> None:
+        await proc.wait()
+        # A process group has no handle to await, so poll it.
+        while sys.platform != "win32" and _group_alive(proc):  # ruff:ignore[async-busy-wait]
+            await asyncio.sleep(0.05)
+
+    with contextlib.suppress(asyncio.TimeoutError):
+        await asyncio.wait_for(exited(), timeout=grace)

python/packages/tools/tests/test_local_shell_tool.py

@@ -1,13 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import asyncio
+import ctypes
 import json
 import os
 import shlex
 import shutil
 import signal
 import sys
-from collections.abc import Awaitable, Mapping, Sequence
+from collections.abc import Awaitable, Iterator, Mapping, Sequence
 from contextlib import nullcontext, suppress
 from pathlib import Path
 from types import SimpleNamespace
@@ -415,6 +416,84 @@ async def test_stateless_timeout_kills_long_command() -> None:
     assert result.timed_out is True
 
 
+def _pid_alive(pid: int) -> bool:
+    try:
+        return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE
+    except psutil.NoSuchProcess:
+        return False
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="POSIX process-group regression")
+@pytest.mark.parametrize(
+    "template",
+    ["sleep 5 & echo $! > {pidfile}", "(sleep 5 & echo $! > {pidfile}); sleep 30"],
+    ids=["background", "double_forked"],
+)
+async def test_stateless_timeout_kills_background_jobs_holding_pipes(tmp_path: Path, template: str) -> None:
+    pidfile = tmp_path / "bg.pid"
+    tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True, timeout=0.5)
+
+    started = asyncio.get_running_loop().time()
+    result = await asyncio.wait_for(tool.run(template.format(pidfile=pidfile)), timeout=10.0)
+    elapsed = asyncio.get_running_loop().time() - started
+
+    pid = int(pidfile.read_text().strip())
+    alive = _pid_alive(pid)
+    if alive:
+        os.kill(pid, signal.SIGKILL)
+    assert result.timed_out is True
+    assert elapsed < 2.0
+    assert not alive
+
+
+@pytest.mark.skipif(sys.platform != "linux" or shutil.which("setsid") is None, reason="needs /proc and setsid")
+async def test_stateless_timeout_closes_pipes_held_by_detached_child(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    from agent_framework_tools.shell import _capture
+
+    monkeypatch.setattr(_capture, "_DRAIN_GRACE", 0.2)
+    pidfile = tmp_path / "detached.pid"
+    tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True, timeout=0.5)
+
+    fds_before = len(os.listdir("/proc/self/fd"))
+    result = await tool.run(f"(setsid sh -c 'echo $$ > {pidfile}; exec sleep 5' &); sleep 30")
+    await asyncio.sleep(0)  # pipe transports close their descriptors in connection_lost, one loop pass later
+    fds_after = len(os.listdir("/proc/self/fd"))
+
+    os.kill(int(pidfile.read_text().strip()), signal.SIGKILL)
+    assert result.timed_out is True
+    assert fds_after == fds_before
+
+
+@pytest.fixture
+def unreaped_orphans() -> Iterator[None]:
+    """Make this process the subreaper for orphans and leave them as zombies, like a slow container init."""
+    libc = ctypes.CDLL(None, use_errno=True)
+    pr_set_child_subreaper = 36
+    assert libc.prctl(pr_set_child_subreaper, 1, 0, 0, 0) == 0
+    try:
+        yield
+    finally:
+        libc.prctl(pr_set_child_subreaper, 0, 0, 0, 0)
+        for child in psutil.Process().children():
+            with suppress(psutil.NoSuchProcess, ChildProcessError):
+                if child.status() == psutil.STATUS_ZOMBIE:
+                    os.waitpid(child.pid, 0)
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="Linux child subreaper")
+async def test_stateless_timeout_does_not_wait_for_unreaped_zombies(unreaped_orphans: None) -> None:
+    tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True, timeout=0.5)
+
+    started = asyncio.get_running_loop().time()
+    result = await tool.run("sleep 30 & sleep 30")
+    elapsed = asyncio.get_running_loop().time() - started
+
+    assert result.timed_out is True
+    assert elapsed < 2.0
+
+
 async def test_policy_denies_before_execution() -> None:
     tool = LocalShellTool(
         mode="stateless",

python/packages/tools/tests/test_shell_killtree.py

@@ -58,9 +58,14 @@ def test_resolve_taskkill_uses_systemroot_and_caches(monkeypatch) -> None:
     assert _resolve_taskkill() == expected_path
 
 
-async def test_kill_process_tree_short_circuits_or_delegates() -> None:
+async def test_kill_process_tree_short_circuits_or_delegates(monkeypatch) -> None:
+    import agent_framework_tools.shell._killtree as killtree_module
+
+    monkeypatch.setattr(killtree_module.sys, "platform", "win32")
     proc = cast(asyncio.subprocess.Process, _FakeAsyncProcess(returncode=0))
-    await kill_process_tree(proc)
+    with patch("agent_framework_tools.shell._killtree._kill_via_psutil", AsyncMock()) as via_psutil:
+        await kill_process_tree(proc)
+    via_psutil.assert_not_awaited()
 
     live = cast(asyncio.subprocess.Process, _FakeAsyncProcess(returncode=None))
     with (
@@ -72,7 +77,7 @@ async def test_kill_process_tree_short_circuits_or_delegates() -> None:
     via_psutil.assert_awaited_once_with(live, grace=2.0)
 
 
-async def test_kill_via_psutil_terminates_parent_and_children() -> None:
+async def test_kill_via_psutil_terminates_parent_and_children(monkeypatch) -> None:
     import agent_framework_tools.shell._killtree as killtree_module
 
     no_such_process = type("NoSuchProcess", (Exception,), {})
@@ -85,6 +90,7 @@ async def test_kill_via_psutil_terminates_parent_and_children() -> None:
         AccessDenied=access_denied,
     )
     proc = cast(asyncio.subprocess.Process, _FakeAsyncProcess(pid=4321, returncode=None))
+    monkeypatch.setattr(killtree_module.os, "killpg", MagicMock(), raising=False)
 
     with patch.object(killtree_module, "psutil", fake_psutil):
         await _kill_via_psutil(proc, grace=0.01)
@@ -95,12 +101,13 @@ async def test_kill_via_psutil_terminates_parent_and_children() -> None:
     child.kill.assert_called_once()
 
 
-async def test_kill_via_psutil_handles_missing_parent_process() -> None:
+async def test_kill_via_psutil_handles_missing_parent_process(monkeypatch) -> None:
     import agent_framework_tools.shell._killtree as killtree_module
 
     no_such_process = type("NoSuchProcess", (Exception,), {})
     fake_psutil = MagicMock(Process=MagicMock(side_effect=no_such_process()), NoSuchProcess=no_such_process)
     proc = cast(asyncio.subprocess.Process, _FakeAsyncProcess(pid=9999, returncode=None))
+    monkeypatch.setattr(killtree_module.os, "killpg", MagicMock(), raising=False)
 
     with patch.object(killtree_module, "psutil", fake_psutil):
         await _kill_via_psutil(proc, grace=0.01)
@@ -127,7 +134,6 @@ async def test_kill_via_stdlib_posix_escalates_to_sigkill(monkeypatch) -> None:
 
     monkeypatch.setattr(killtree_module.sys, "platform", "darwin")
     killpg = MagicMock()
-    monkeypatch.setattr(killtree_module.os, "getpgid", lambda pid: 99, raising=False)
     monkeypatch.setattr(killtree_module.os, "killpg", killpg, raising=False)
     monkeypatch.setattr(killtree_module.signal, "SIGKILL", 9, raising=False)
 
@@ -146,5 +152,31 @@ async def fake_wait_for(awaitable: Any, timeout: float) -> None:
     with patch("agent_framework_tools.shell._killtree.asyncio.wait_for", side_effect=fake_wait_for):
         await _kill_via_stdlib(proc, grace=0.01)
 
-    assert killpg.call_args_list[0].args == (99, killtree_module.signal.SIGTERM)
-    assert killpg.call_args_list[1].args == (99, killtree_module.signal.SIGKILL)
+    assert killpg.call_args_list[0].args == (12, killtree_module.signal.SIGTERM)
+    assert killpg.call_args_list[1].args == (12, killtree_module.signal.SIGKILL)
+
+
+async def test_kill_via_psutil_signals_process_group_after_shell_exited(monkeypatch) -> None:
+    import agent_framework_tools.shell._killtree as killtree_module
+
+    monkeypatch.setattr(killtree_module.sys, "platform", "linux")
+
+    def fake_killpg(pgid: int, sig: int) -> None:
+        if sig == 0:
+            raise ProcessLookupError
+
+    killpg = MagicMock(side_effect=fake_killpg)
+    monkeypatch.setattr(killtree_module.os, "killpg", killpg, raising=False)
+    monkeypatch.setattr(killtree_module.signal, "SIGKILL", 9, raising=False)
+    fake_psutil = MagicMock(NoSuchProcess=type("NoSuchProcess", (Exception,), {}))
+    proc = cast(asyncio.subprocess.Process, _FakeAsyncProcess(pid=77, returncode=0))
+
+    with patch.object(killtree_module, "psutil", fake_psutil):
+        await kill_process_tree(proc, grace=0.01)
+
+    fake_psutil.Process.assert_not_called()
+    assert [c.args for c in killpg.call_args_list] == [
+        (77, killtree_module.signal.SIGTERM),
+        (77, 0),
+        (77, killtree_module.signal.SIGKILL),
+    ]