From bcf5d8a6ed4f8871a601a7d512b689d758a413a2 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Wed, 12 Aug 2026 11:21:21 +0800 Subject: [PATCH] fix(exec): retain process trees after root exit --- nanobot/agent/tools/_windows_job.py | 219 ++++++++++++++++++++++++++++ nanobot/agent/tools/exec_session.py | 6 +- nanobot/agent/tools/shell.py | 148 +++++++++++++------ tests/tools/test_exec_platform.py | 20 +++ tests/tools/test_shell_reap.py | 61 ++++++++ 5 files changed, 411 insertions(+), 43 deletions(-) create mode 100644 nanobot/agent/tools/_windows_job.py diff --git a/nanobot/agent/tools/_windows_job.py b/nanobot/agent/tools/_windows_job.py new file mode 100644 index 000000000..83a7111f6 --- /dev/null +++ b/nanobot/agent/tools/_windows_job.py @@ -0,0 +1,219 @@ +"""Windows Job Object ownership for subprocess trees.""" + +from __future__ import annotations + +import ctypes +from ctypes import wintypes + +_CREATE_SUSPENDED = 0x00000004 +_PROCESS_SET_QUOTA = 0x0100 +_PROCESS_TERMINATE = 0x0001 +_TH32CS_SNAPTHREAD = 0x00000004 +_THREAD_SUSPEND_RESUME = 0x0002 +_JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x00002000 +_JOB_OBJECT_EXTENDED_LIMIT_INFORMATION = 9 +_INVALID_HANDLE_VALUE = ctypes.c_void_p(-1).value + + +class _IoCounters(ctypes.Structure): + _fields_ = [ + ("ReadOperationCount", ctypes.c_ulonglong), + ("WriteOperationCount", ctypes.c_ulonglong), + ("OtherOperationCount", ctypes.c_ulonglong), + ("ReadTransferCount", ctypes.c_ulonglong), + ("WriteTransferCount", ctypes.c_ulonglong), + ("OtherTransferCount", ctypes.c_ulonglong), + ] + + +class _BasicLimitInformation(ctypes.Structure): + _fields_ = [ + ("PerProcessUserTimeLimit", ctypes.c_longlong), + ("PerJobUserTimeLimit", ctypes.c_longlong), + ("LimitFlags", wintypes.DWORD), + ("MinimumWorkingSetSize", ctypes.c_size_t), + ("MaximumWorkingSetSize", ctypes.c_size_t), + ("ActiveProcessLimit", wintypes.DWORD), + ("Affinity", ctypes.c_size_t), + ("PriorityClass", wintypes.DWORD), + ("SchedulingClass", wintypes.DWORD), + ] + + +class _ExtendedLimitInformation(ctypes.Structure): + _fields_ = [ + ("BasicLimitInformation", _BasicLimitInformation), + ("IoInfo", _IoCounters), + ("ProcessMemoryLimit", ctypes.c_size_t), + ("JobMemoryLimit", ctypes.c_size_t), + ("PeakProcessMemoryUsed", ctypes.c_size_t), + ("PeakJobMemoryUsed", ctypes.c_size_t), + ] + + +class _ThreadEntry32(ctypes.Structure): + _fields_ = [ + ("dwSize", wintypes.DWORD), + ("cntUsage", wintypes.DWORD), + ("th32ThreadID", wintypes.DWORD), + ("th32OwnerProcessID", wintypes.DWORD), + ("tpBasePri", wintypes.LONG), + ("tpDeltaPri", wintypes.LONG), + ("dwFlags", wintypes.DWORD), + ] + + +_kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) +_kernel32.CreateJobObjectW.argtypes = [ctypes.c_void_p, wintypes.LPCWSTR] +_kernel32.CreateJobObjectW.restype = wintypes.HANDLE +_kernel32.SetInformationJobObject.argtypes = [ + wintypes.HANDLE, + ctypes.c_int, + ctypes.c_void_p, + wintypes.DWORD, +] +_kernel32.SetInformationJobObject.restype = wintypes.BOOL +_kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] +_kernel32.OpenProcess.restype = wintypes.HANDLE +_kernel32.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE] +_kernel32.AssignProcessToJobObject.restype = wintypes.BOOL +_kernel32.TerminateProcess.argtypes = [wintypes.HANDLE, wintypes.UINT] +_kernel32.TerminateProcess.restype = wintypes.BOOL +_kernel32.TerminateJobObject.argtypes = [wintypes.HANDLE, wintypes.UINT] +_kernel32.TerminateJobObject.restype = wintypes.BOOL +_kernel32.CreateToolhelp32Snapshot.argtypes = [wintypes.DWORD, wintypes.DWORD] +_kernel32.CreateToolhelp32Snapshot.restype = wintypes.HANDLE +_kernel32.Thread32First.argtypes = [wintypes.HANDLE, ctypes.POINTER(_ThreadEntry32)] +_kernel32.Thread32First.restype = wintypes.BOOL +_kernel32.Thread32Next.argtypes = [wintypes.HANDLE, ctypes.POINTER(_ThreadEntry32)] +_kernel32.Thread32Next.restype = wintypes.BOOL +_kernel32.OpenThread.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] +_kernel32.OpenThread.restype = wintypes.HANDLE +_kernel32.ResumeThread.argtypes = [wintypes.HANDLE] +_kernel32.ResumeThread.restype = wintypes.DWORD +_kernel32.CloseHandle.argtypes = [wintypes.HANDLE] +_kernel32.CloseHandle.restype = wintypes.BOOL + + +def _win_error(operation: str) -> OSError: + code = ctypes.get_last_error() + return OSError(code, f"{operation} failed (Windows error {code})") + + +def _close_handle(handle: int | None) -> None: + if handle: + _kernel32.CloseHandle(handle) + + +def _set_kill_on_close(handle: int, enabled: bool) -> None: + info = _ExtendedLimitInformation() + if enabled: + info.BasicLimitInformation.LimitFlags = _JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE + if not _kernel32.SetInformationJobObject( + handle, + _JOB_OBJECT_EXTENDED_LIMIT_INFORMATION, + ctypes.byref(info), + ctypes.sizeof(info), + ): + raise _win_error("SetInformationJobObject") + + +def _resume_primary_thread(pid: int) -> None: + snapshot = _kernel32.CreateToolhelp32Snapshot(_TH32CS_SNAPTHREAD, 0) + if snapshot == _INVALID_HANDLE_VALUE: + raise _win_error("CreateToolhelp32Snapshot") + try: + entry = _ThreadEntry32() + entry.dwSize = ctypes.sizeof(entry) + found = _kernel32.Thread32First(snapshot, ctypes.byref(entry)) + while found: + if entry.th32OwnerProcessID == pid: + thread = _kernel32.OpenThread( + _THREAD_SUSPEND_RESUME, + False, + entry.th32ThreadID, + ) + if not thread: + raise _win_error("OpenThread") + try: + if _kernel32.ResumeThread(thread) == 0xFFFFFFFF: + raise _win_error("ResumeThread") + return + finally: + _close_handle(thread) + found = _kernel32.Thread32Next(snapshot, ctypes.byref(entry)) + raise RuntimeError(f"suspended process {pid} has no resumable thread") + finally: + _close_handle(snapshot) + + +class WindowsJob: + """Own a process tree even after its root process exits.""" + + creation_flags = _CREATE_SUSPENDED + + def __init__(self, handle: int) -> None: + self._handle: int | None = handle + + @classmethod + def create(cls) -> WindowsJob: + handle = _kernel32.CreateJobObjectW(None, None) + if not handle: + raise _win_error("CreateJobObjectW") + try: + _set_kill_on_close(handle, True) + except Exception: + _close_handle(handle) + raise + return cls(handle) + + def assign_and_resume(self, pid: int) -> None: + """Atomically establish tree ownership before the root can spawn.""" + if self._handle is None: + raise RuntimeError("Windows job is already closed") + process = _kernel32.OpenProcess( + _PROCESS_SET_QUOTA | _PROCESS_TERMINATE, + False, + pid, + ) + if not process: + error = _win_error("OpenProcess") + self.close() + raise error + + assigned = False + try: + if not _kernel32.AssignProcessToJobObject(self._handle, process): + raise _win_error("AssignProcessToJobObject") + assigned = True + _resume_primary_thread(pid) + except Exception: + if assigned: + self.terminate() + else: + _kernel32.TerminateProcess(process, 1) + self.close() + raise + finally: + _close_handle(process) + + def release(self) -> None: + """Release ownership after successful output collection.""" + if self._handle is None: + return + _set_kill_on_close(self._handle, False) + self.close() + + def terminate(self) -> None: + """Terminate every process in the job and close its handle.""" + if self._handle is None: + return + try: + _kernel32.TerminateJobObject(self._handle, 1) + finally: + self.close() + + def close(self) -> None: + handle = self._handle + self._handle = None + _close_handle(handle) diff --git a/nanobot/agent/tools/exec_session.py b/nanobot/agent/tools/exec_session.py index 3a8d59bfc..16210409e 100644 --- a/nanobot/agent/tools/exec_session.py +++ b/nanobot/agent/tools/exec_session.py @@ -209,7 +209,11 @@ class _ExecSession: timeout=2.0, ) # Safety-net reap after normal exit. - from nanobot.agent.tools.shell import _reap_pid # pyright: ignore[reportPrivateUsage] + from nanobot.agent.tools.shell import ( # pyright: ignore[reportPrivateUsage] + ExecTool, + _reap_pid, # pyright: ignore[reportPrivateUsage] + ) + ExecTool._release_process_tree(self.process) # pyright: ignore[reportPrivateUsage] _reap_pid(self.process.pid) # pyright: ignore[reportPrivateUsage] elif yield_time_ms > 0: await self._wait_for_buffered_output() diff --git a/nanobot/agent/tools/shell.py b/nanobot/agent/tools/shell.py index 2c743c9dc..8583968d1 100644 --- a/nanobot/agent/tools/shell.py +++ b/nanobot/agent/tools/shell.py @@ -12,7 +12,7 @@ import sys from contextlib import suppress from dataclasses import dataclass from pathlib import Path, PureWindowsPath -from typing import Any +from typing import Any, Protocol, cast, runtime_checkable from loguru import logger from pydantic import Field @@ -42,6 +42,20 @@ from nanobot.security.workspace_access import current_scope_allows_loopback, cur from nanobot.security.workspace_policy import is_path_within _IS_WINDOWS = sys.platform == "win32" +_PROCESS_TREE_OWNER_ATTR = "_nanobot_process_tree_owner" + + +@runtime_checkable +class _ProcessTreeOwner(Protocol): + creation_flags: int + + def assign_and_resume(self, pid: int) -> None: ... + + def close(self) -> None: ... + + def release(self) -> None: ... + + def terminate(self) -> None: ... def _reap_pid(pid: int) -> None: @@ -369,6 +383,7 @@ class ExecTool(Tool): + result[-half:] ) + self._release_process_tree(process) return result except Exception as e: @@ -538,38 +553,55 @@ class ExecTool(Tool): ) -> asyncio.subprocess.Process: """Launch *command* in a platform-appropriate shell.""" if _IS_WINDOWS: + windows_job = None + creation_flags = 0 + if process_tree and sys.platform == "win32": + windows_job = ExecTool._create_windows_job() + creation_flags = windows_job.creation_flags # Default to PowerShell so single-line and multi-line commands # share the same shell semantics. cmd.exe is reachable via the # explicit shell="cmd" parameter (see _resolve_shell). default_program = shutil.which("pwsh") or shutil.which("powershell") or "powershell" program = shell_program or default_program program_name = PureWindowsPath(program).name.lower() - if program_name in ("cmd", "cmd.exe"): - cmd_env = {**env, "COMSPEC": program} - return await asyncio.create_subprocess_shell( - command, - stdin=stdin, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=cwd, - env=cmd_env, - ) - command = ExecTool._normalize_powershell_command(command) - command = ( - "[Console]::OutputEncoding = [System.Text.UTF8Encoding]::new($false)\n" - "if ($PSVersionTable.PSVersion.Major -lt 6) { $OutputEncoding = [Console]::OutputEncoding }\n" - "$PSDefaultParameterValues['Out-File:Encoding'] = 'utf8'\n" - f"{command}\n" - "if ($LASTEXITCODE -ne $null) { exit $LASTEXITCODE }" - ) - return await asyncio.create_subprocess_exec( - program, "-NoProfile", "-NonInteractive", "-Command", command, - stdin=stdin, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=cwd, - env=env, - ) + try: + if program_name in ("cmd", "cmd.exe"): + cmd_env = {**env, "COMSPEC": program} + process = await asyncio.create_subprocess_shell( + command, + stdin=stdin, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=cwd, + env=cmd_env, + creationflags=creation_flags, + ) + else: + command = ExecTool._normalize_powershell_command(command) + command = ( + "[Console]::OutputEncoding = [System.Text.UTF8Encoding]::new($false)\n" + "if ($PSVersionTable.PSVersion.Major -lt 6) { $OutputEncoding = [Console]::OutputEncoding }\n" + "$PSDefaultParameterValues['Out-File:Encoding'] = 'utf8'\n" + f"{command}\n" + "if ($LASTEXITCODE -ne $null) { exit $LASTEXITCODE }" + ) + process = await asyncio.create_subprocess_exec( + program, "-NoProfile", "-NonInteractive", "-Command", command, + stdin=stdin, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=cwd, + env=env, + creationflags=creation_flags, + ) + if windows_job is not None: + windows_job.assign_and_resume(process.pid) + setattr(process, _PROCESS_TREE_OWNER_ATTR, windows_job) + return process + except BaseException: + if windows_job is not None: + windows_job.terminate() + raise shell_program = shell_program or shutil.which("bash") or "/bin/bash" args: list[str] = [shell_program] shell_name = Path(shell_program).name.lower() @@ -688,22 +720,23 @@ class ExecTool(Tool): @staticmethod async def _kill_process_tree(process: asyncio.subprocess.Process) -> None: """Kill a session process and descendants, then reap the root process.""" - if process.returncode is not None: - _reap_pid(process.pid) - return + owner = ExecTool._process_tree_owner(process) try: - if _IS_WINDOWS: - with suppress(OSError, asyncio.TimeoutError): - await asyncio.wait_for( - asyncio.to_thread( - subprocess.run, - ["taskkill", "/PID", str(process.pid), "/T", "/F"], - check=False, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ), - timeout=5.0, - ) + if owner is not None: + owner.terminate() + elif _IS_WINDOWS: + if process.returncode is None: + with suppress(OSError, asyncio.TimeoutError): + await asyncio.wait_for( + asyncio.to_thread( + subprocess.run, + ["taskkill", "/PID", str(process.pid), "/T", "/F"], + check=False, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ), + timeout=5.0, + ) else: try: os.killpg(process.pid, signal.SIGKILL) @@ -716,8 +749,39 @@ class ExecTool(Tool): with suppress(asyncio.TimeoutError): await asyncio.wait_for(process.wait(), timeout=5.0) finally: + if owner is not None: + ExecTool._drop_process_tree_owner(process) _reap_pid(process.pid) + @staticmethod + def _process_tree_owner( + process: asyncio.subprocess.Process, + ) -> _ProcessTreeOwner | None: + attributes = getattr(process, "__dict__", None) + if not isinstance(attributes, dict): + return None + owner = cast(dict[str, object], attributes).get(_PROCESS_TREE_OWNER_ATTR) + return owner if isinstance(owner, _ProcessTreeOwner) else None + + @staticmethod + def _create_windows_job() -> _ProcessTreeOwner: + from nanobot.agent.tools._windows_job import WindowsJob + + return WindowsJob.create() + + @staticmethod + def _drop_process_tree_owner(process: asyncio.subprocess.Process) -> None: + with suppress(AttributeError): + delattr(process, _PROCESS_TREE_OWNER_ATTR) + + @staticmethod + def _release_process_tree(process: asyncio.subprocess.Process) -> None: + owner = ExecTool._process_tree_owner(process) + if owner is None: + return + owner.release() + ExecTool._drop_process_tree_owner(process) + def _build_env(self) -> dict[str, str]: """Build a minimal environment for subprocess execution. diff --git a/tests/tools/test_exec_platform.py b/tests/tools/test_exec_platform.py index c02638a17..49ddf7e75 100644 --- a/tests/tools/test_exec_platform.py +++ b/tests/tools/test_exec_platform.py @@ -21,6 +21,22 @@ _WINDOWS_ENV_KEYS = { } +class _FakeWindowsJob: + creation_flags = 0 + + def assign_and_resume(self, pid: int) -> None: + pass + + def close(self) -> None: + pass + + def release(self) -> None: + pass + + def terminate(self) -> None: + pass + + # --------------------------------------------------------------------------- # _build_env # --------------------------------------------------------------------------- @@ -661,6 +677,7 @@ class TestWindowsMultilineExec: with ( patch("nanobot.agent.tools.shell._IS_WINDOWS", True), patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_exec, + patch.object(ExecTool, "_create_windows_job", side_effect=_FakeWindowsJob), patch.object(ExecTool, "_guard_command", return_value=None), ): mock_exec.return_value = mock_proc @@ -682,6 +699,7 @@ class TestWindowsMultilineExec: with ( patch("nanobot.agent.tools.shell._IS_WINDOWS", True), patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_exec, + patch.object(ExecTool, "_create_windows_job", side_effect=_FakeWindowsJob), patch.object(ExecTool, "_guard_command", return_value=None), ): mock_exec.return_value = mock_proc @@ -745,6 +763,7 @@ class TestResolveShellWindows: with ( patch("nanobot.agent.tools.shell._IS_WINDOWS", True), patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_exec, + patch.object(ExecTool, "_create_windows_job", side_effect=_FakeWindowsJob), patch.object(ExecTool, "_guard_command", return_value=None), ): mock_exec.return_value = mock_proc @@ -766,6 +785,7 @@ class TestResolveShellWindows: with ( patch("nanobot.agent.tools.shell._IS_WINDOWS", True), patch("asyncio.create_subprocess_shell", new_callable=AsyncMock) as mock_shell, + patch.object(ExecTool, "_create_windows_job", side_effect=_FakeWindowsJob), patch.object(ExecTool, "_guard_command", return_value=None), ): mock_shell.return_value = mock_proc diff --git a/tests/tools/test_shell_reap.py b/tests/tools/test_shell_reap.py index 42a0eaa9e..1429475c8 100644 --- a/tests/tools/test_shell_reap.py +++ b/tests/tools/test_shell_reap.py @@ -3,7 +3,9 @@ from __future__ import annotations import asyncio +import base64 import shlex +import subprocess import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -13,6 +15,12 @@ from nanobot.agent.tools.exec_session import _ExecSession from nanobot.agent.tools.shell import ExecTool, _reap_pid +def _python_command(code: str) -> str: + if sys.platform == "win32": + return f"{subprocess.list2cmdline([sys.executable])} -u -c {subprocess.list2cmdline([code])}" + return f"{shlex.quote(sys.executable)} -u -c {shlex.quote(code)}" + + def test_reap_pid_noops_without_waitpid(): """On platforms (or test stubs) without waitpid, reaping is a no-op.""" with patch("nanobot.agent.tools.shell.os") as mock_os: @@ -193,6 +201,25 @@ async def test_execute_exception_during_communicate_kills_live_process(): kill_tree.assert_awaited_once_with(mock_proc) +@pytest.mark.asyncio +async def test_kill_process_tree_targets_group_after_root_exits(): + process = AsyncMock() + process.pid = 1006 + process.returncode = 0 + + with ( + patch("nanobot.agent.tools.shell._IS_WINDOWS", False), + patch("nanobot.agent.tools.shell.os.killpg", create=True) as kill_group, + patch("nanobot.agent.tools.shell.signal.SIGKILL", 9, create=True), + patch("nanobot.agent.tools.shell._reap_pid") as reap, + ): + await ExecTool._kill_process_tree(process) + + kill_group.assert_called_once_with(1006, 9) + process.kill.assert_not_called() + reap.assert_called_once_with(1006) + + @pytest.mark.skipif(sys.platform == "win32", reason="requires Unix process groups") @pytest.mark.asyncio async def test_execute_timeout_kills_background_process_tree(tmp_path): @@ -210,6 +237,40 @@ async def test_execute_timeout_kills_background_process_tree(tmp_path): assert not marker.exists() +@pytest.mark.asyncio +async def test_execute_timeout_kills_descendant_after_root_exits(tmp_path): + """Tree ownership must outlive a root shell that exits before timeout.""" + marker = tmp_path / "child-survived-root" + child_code = ( + "import pathlib,time; time.sleep(3.5); " + f"pathlib.Path({str(marker)!r}).write_text('alive')" + ) + child_payload = base64.b64encode(child_code.encode()).decode() + parent_code = ( + "import base64,subprocess,sys; " + f"child=base64.b64decode('{child_payload}').decode(); " + "subprocess.Popen([sys.executable, '-c', child])" + ) + spawned = [] + original_spawn = ExecTool._spawn + + async def capture_spawn(*args, **kwargs): + process = await original_spawn(*args, **kwargs) + spawned.append(process) + return process + + with patch.object(ExecTool, "_spawn", side_effect=capture_spawn): + result = await ExecTool(working_dir=str(tmp_path), timeout=2).execute( + command=_python_command(parent_code), + timeout=2, + ) + + assert "timed out" in result.lower() + assert spawned[0].returncode == 0 + await asyncio.sleep(2) + assert not marker.exists() + + def _mock_session_process(*, pid: int, returncode: int | None): process = AsyncMock() process.pid = pid