diff --git a/python/packages/tools/agent_framework_tools/shell/_executor.py b/python/packages/tools/agent_framework_tools/shell/_executor.py index 446d227ae1c..cb1871a33fc 100644 --- a/python/packages/tools/agent_framework_tools/shell/_executor.py +++ b/python/packages/tools/agent_framework_tools/shell/_executor.py @@ -22,6 +22,9 @@ from ._truncate import truncate_head_tail as _truncate from ._types import ShellResult +# Bounds the post-kill drain: a process that left the session can still hold the pipes open. +_DRAIN_GRACE = 2.0 + def _popen_kwargs_for_group() -> dict[str, object]: """Platform-specific process-group isolation so we can kill children too.""" @@ -75,9 +78,11 @@ async def run_stateless( await kill_process_tree(proc) # Drain any queued output. try: - stdout_bytes, stderr_bytes = await proc.communicate() + stdout_bytes, stderr_bytes = await asyncio.wait_for(proc.communicate(), timeout=_DRAIN_GRACE) except Exception: stdout_bytes, stderr_bytes = b"", b"" + # Release the pipes a detached process still holds instead of leaving them to its exit. + proc._transport.close() # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue, reportPrivateUsage] duration_ms = int((time.monotonic() - started) * 1000) stdout_str, stdout_truncated = _truncate(stdout_bytes or b"", max_output_bytes) diff --git a/python/packages/tools/agent_framework_tools/shell/_killtree.py b/python/packages/tools/agent_framework_tools/shell/_killtree.py index 8719db6dcfa..f40664097b0 100644 --- a/python/packages/tools/agent_framework_tools/shell/_killtree.py +++ b/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) diff --git a/python/packages/tools/tests/test_local_shell_tool.py b/python/packages/tools/tests/test_local_shell_tool.py index f9dcbdf172c..e2a3d2b9da5 100644 --- a/python/packages/tools/tests/test_local_shell_tool.py +++ b/python/packages/tools/tests/test_local_shell_tool.py @@ -1,19 +1,21 @@ # 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 contextlib import nullcontext +from collections.abc import Awaitable, Iterator, Mapping, Sequence +from contextlib import nullcontext, suppress from pathlib import Path from types import SimpleNamespace from typing import Any -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch +import psutil import pytest from agent_framework import Agent, AgentSession, BaseChatClient, ChatResponse, Content, FunctionInvocationLayer, Message from agent_framework.security import ( @@ -43,6 +45,7 @@ def __init__( self.returncode = returncode self.stdout = object() self.stderr = object() + self._transport = MagicMock() self._communicate_results = list(communicate_results or [(b"", b"")]) async def communicate(self) -> tuple[bytes, bytes]: @@ -142,6 +145,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: + import agent_framework_tools.shell._executor as executor_module + + monkeypatch.setattr(executor_module, "_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", @@ -292,6 +373,7 @@ async def test_run_stateless_timeout_returns_empty_output_if_drain_fails() -> No ) kill_tree.assert_awaited_once_with(proc) + proc._transport.close.assert_called_once_with() assert result.timed_out is True assert result.stdout == "" assert result.stderr == "" diff --git a/python/packages/tools/tests/test_shell_killtree.py b/python/packages/tools/tests/test_shell_killtree.py index 304b85ffe4d..e98a807febd 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), + ]