diff --git a/python/packages/tools/agent_framework_tools/shell/_killtree.py b/python/packages/tools/agent_framework_tools/shell/_killtree.py index 8719db6dcf..ad3ae485fa 100644 --- a/python/packages/tools/agent_framework_tools/shell/_killtree.py +++ b/python/packages/tools/agent_framework_tools/shell/_killtree.py @@ -23,6 +23,10 @@ import os import signal import sys +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from psutil import Process as PsutilProcess try: # pragma: no cover - importable on every platform we ship import psutil @@ -52,8 +56,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 +77,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[PsutilProcess] = [] + 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 +133,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) diff --git a/python/packages/tools/tests/test_local_shell_tool.py b/python/packages/tools/tests/test_local_shell_tool.py index b24c99a85f..83b6396d23 100644 --- a/python/packages/tools/tests/test_local_shell_tool.py +++ b/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", diff --git a/python/packages/tools/tests/test_shell_killtree.py b/python/packages/tools/tests/test_shell_killtree.py index 304b85ffe4..e98a807feb 100644 --- a/python/packages/tools/tests/test_shell_killtree.py +++ b/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), + ]