fix(exec): retain process trees after root exit

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)

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()

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,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:
+        owner = ExecTool._process_tree_owner(process)
-            _reap_pid(process.pid)
-            return
         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

@@ -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

@@ -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