mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-08-12 23:29:16 +03:00
Code Issues Packages Projects Releases Wiki Activity

fix(exec): retain process trees after root exit

Browse Source
This commit is contained in:
chengyongru
2026-08-12 14:37:27 +08:00
committed by chengyongru
parent d64b84604c
commit bcf5d8a6ed
5 changed files with 411 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/tools/_windows_job.py
+219
View File
@@ -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)
nanobot/agent/tools/exec_session.py
+5 -1
View File
@@ -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()
nanobot/agent/tools/shell.py
+71 -7
View File
@@ -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,22 +553,30 @@ 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()
try:
if program_name in ("cmd", "cmd.exe"):
cmd_env = {**env, "COMSPEC": program}
return await asyncio.create_subprocess_shell(
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"
@@ -562,14 +585,23 @@ class ExecTool(Tool):
f"{command}\n"
"if ($LASTEXITCODE -ne $null) { exit $LASTEXITCODE }"
)
return await asyncio.create_subprocess_exec(
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,11 +720,12 @@ 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:
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(
@@ -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.
tests/tools/test_exec_platform.py
+20
View File
@@ -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
tests/tools/test_shell_reap.py
+61
View File
@@ -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