diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index caa11832..9df71670 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -1277,11 +1277,25 @@ class AgentLoop: await self._publish_next_deferred_automation_turn(session_key) async def close_mcp(self) -> None: - """Drain pending background archives, then close MCP connections.""" + """Drain background work, stop exec sessions, then close MCP connections.""" if self._background_tasks: await asyncio.gather(*self._background_tasks, return_exceptions=True) self._background_tasks.clear() - await agent_context.close_mcp(self) + errors: list[BaseException] = [] + cleanup_steps = ( + self.subagents.close, + self._exec_session_manager.close_all, + lambda: agent_context.close_mcp(self), + ) + for cleanup in cleanup_steps: + try: + await cleanup() + except BaseException as exc: + errors.append(exc) + if len(errors) == 1: + raise errors[0] + if errors: + raise BaseExceptionGroup("failed to close agent resources", errors) def _schedule_background(self, coro) -> None: """Schedule a coroutine as a tracked background task (drained on shutdown).""" diff --git a/nanobot/agent/subagent.py b/nanobot/agent/subagent.py index f4d1fe6a..74d1eab4 100644 --- a/nanobot/agent/subagent.py +++ b/nanobot/agent/subagent.py @@ -459,6 +459,15 @@ class SubagentManager: await asyncio.gather(*tasks, return_exceptions=True) return len(tasks) + async def close(self) -> None: + """Cancel running subagents and close their shared exec sessions.""" + tasks = [task for task in self._running_tasks.values() if not task.done()] + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + await self._exec_session_manager.close_all() + def get_running_count(self) -> int: """Return the number of currently running subagents.""" return len(self._running_tasks) diff --git a/nanobot/agent/tools/exec_session.py b/nanobot/agent/tools/exec_session.py index c93ed8d0..24738dfc 100644 --- a/nanobot/agent/tools/exec_session.py +++ b/nanobot/agent/tools/exec_session.py @@ -61,12 +61,14 @@ class _ExecSession: cwd: str, timeout: int | None, owner_session_key: str | None = None, + process_tree: bool = False, ) -> None: self.session_id = session_id self.process = process self.command = command self.cwd = cwd self.owner_session_key = owner_session_key + self._process_tree = process_tree self.started_at = time.monotonic() # timeout None/0 means no limit; an infinite deadline is never reached. self.deadline = time.monotonic() + timeout if timeout else float("inf") @@ -171,17 +173,23 @@ class _ExecSession: ) async def kill(self) -> None: - if self.process.returncode is not None: - return - self.process.kill() + from nanobot.agent.tools.shell import ExecTool + try: - with suppress(asyncio.TimeoutError): - await asyncio.wait_for(self.process.wait(), timeout=5.0) + if self._process_tree: + await ExecTool._kill_process_tree(self.process) + else: + await ExecTool._kill_process(self.process) finally: - # Safety-net waitpid — prevent zombie if asyncio's child watcher - # did not reap the process (common in containers). - from nanobot.agent.tools.shell import _reap_pid - _reap_pid(self.process.pid) + with suppress(asyncio.TimeoutError): + await asyncio.wait_for( + asyncio.gather( + self._stdout_task, + self._stderr_task, + return_exceptions=True, + ), + timeout=2.0, + ) async def _wait_for_buffered_output(self) -> None: deadline = time.monotonic() + OUTPUT_DRAIN_GRACE_S @@ -198,6 +206,7 @@ class ExecSessionManager: self.idle_timeout = idle_timeout self._sessions: dict[str, _ExecSession] = {} self._lock = asyncio.Lock() + self._closed = False async def start( self, @@ -213,6 +222,8 @@ class ExecSessionManager: owner_session_key: str | None = None, ) -> tuple[str, _SessionPoll]: async with self._lock: + if self._closed: + raise RuntimeError("exec session manager is closed") await self._cleanup_locked() if len(self._sessions) >= self.max_sessions: raise RuntimeError(f"maximum exec sessions reached ({self.max_sessions})") @@ -225,6 +236,7 @@ class ExecSessionManager: cwd=cwd, timeout=timeout, owner_session_key=owner_session_key, + process_tree=True, ) self._sessions[session_id] = session @@ -295,6 +307,33 @@ class ExecSessionManager: if session.owner_session_key == owner_session_key ] + async def close_all(self) -> int: + """Terminate and remove all active sessions during shutdown.""" + async with self._lock: + self._closed = True + sessions = list(self._sessions.values()) + self._sessions.clear() + results = await asyncio.gather( + *(session.kill() for session in sessions), + return_exceptions=True, + ) + failures = [ + (session, result) + for session, result in zip(sessions, results, strict=True) + if isinstance(result, BaseException) + ] + if failures: + async with self._lock: + for session, _ in failures: + self._sessions[session.session_id] = session + if len(failures) == 1: + raise failures[0][1] + raise BaseExceptionGroup( + "failed to close exec sessions", + [result for _, result in failures], + ) + return len(sessions) + async def _cleanup_locked(self) -> None: now = time.monotonic() stale = [ @@ -319,6 +358,7 @@ class ExecSessionManager: return await ExecTool._spawn( command, cwd, env, shell_program, login, stdin=asyncio.subprocess.PIPE, + process_tree=True, ) diff --git a/nanobot/agent/tools/shell.py b/nanobot/agent/tools/shell.py index 175d06df..9fac8725 100644 --- a/nanobot/agent/tools/shell.py +++ b/nanobot/agent/tools/shell.py @@ -6,6 +6,8 @@ import asyncio import os import re import shutil +import signal +import subprocess import sys from contextlib import suppress from dataclasses import dataclass @@ -516,6 +518,7 @@ class ExecTool(Tool): login: bool = False, *, stdin: int = asyncio.subprocess.DEVNULL, + process_tree: bool = False, ) -> asyncio.subprocess.Process: """Launch *command* in a platform-appropriate shell.""" if _IS_WINDOWS: @@ -563,6 +566,7 @@ class ExecTool(Tool): stderr=asyncio.subprocess.PIPE, cwd=cwd, env=env, + **({"start_new_session": True} if process_tree else {}), ) @staticmethod @@ -655,6 +659,39 @@ class ExecTool(Tool): finally: _reap_pid(process.pid) + @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 + 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, + ) + else: + try: + os.killpg(process.pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError): + pass + + if process.returncode is None: + with suppress(ProcessLookupError): + process.kill() + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(process.wait(), timeout=5.0) + finally: + _reap_pid(process.pid) + def _build_env(self) -> dict[str, str]: """Build a minimal environment for subprocess execution. diff --git a/tests/agent/test_runner_injections.py b/tests/agent/test_runner_injections.py index e6522671..8a8ad912 100644 --- a/tests/agent/test_runner_injections.py +++ b/tests/agent/test_runner_injections.py @@ -37,6 +37,7 @@ def _make_loop(tmp_path): patch("nanobot.agent.loop.SessionManager"), \ patch("nanobot.agent.loop.SubagentManager") as mock_sub_mgr: mock_sub_mgr.return_value.cancel_by_session = AsyncMock(return_value=0) + mock_sub_mgr.return_value.close = AsyncMock() loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path) return loop diff --git a/tests/agent/test_subagent_lifecycle.py b/tests/agent/test_subagent_lifecycle.py index da7959e2..c85f37f1 100644 --- a/tests/agent/test_subagent_lifecycle.py +++ b/tests/agent/test_subagent_lifecycle.py @@ -65,6 +65,24 @@ async def _drain_subagent_tasks(sm: SubagentManager) -> None: await asyncio.sleep(0) +@pytest.mark.asyncio +async def test_close_cancels_tasks_before_closing_exec_sessions(tmp_path): + sm = _manager(tmp_path) + task = asyncio.create_task(asyncio.Event().wait()) + sm._running_tasks["t1"] = task + + async def close_exec_sessions() -> int: + assert task.done() + return 0 + + sm._exec_session_manager.close_all = AsyncMock(side_effect=close_exec_sessions) + + await sm.close() + + assert task.cancelled() + sm._exec_session_manager.close_all.assert_awaited_once() + + # --------------------------------------------------------------------------- # SubagentStatus defaults # --------------------------------------------------------------------------- diff --git a/tests/cli/test_restart_command.py b/tests/cli/test_restart_command.py index 2bd499bf..bdd5cc08 100644 --- a/tests/cli/test_restart_command.py +++ b/tests/cli/test_restart_command.py @@ -28,7 +28,8 @@ def _make_loop(): with patch("nanobot.agent.loop.ContextBuilder"), \ patch("nanobot.agent.loop.SessionManager"), \ - patch("nanobot.agent.loop.SubagentManager"): + patch("nanobot.agent.loop.SubagentManager") as mock_sub_mgr: + mock_sub_mgr.return_value.close = AsyncMock() loop = AgentLoop(bus=bus, provider=provider, workspace=workspace) return loop, bus diff --git a/tests/tools/test_exec_platform.py b/tests/tools/test_exec_platform.py index 3f1df3b4..514ccd95 100644 --- a/tests/tools/test_exec_platform.py +++ b/tests/tools/test_exec_platform.py @@ -114,6 +114,22 @@ class TestSpawnUnix: kwargs = mock_exec.call_args[1] assert kwargs["stdin"] == asyncio.subprocess.DEVNULL + @pytest.mark.asyncio + async def test_process_tree_starts_new_session(self): + with ( + patch("nanobot.agent.tools.shell._IS_WINDOWS", False), + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_exec, + ): + mock_exec.return_value = AsyncMock() + await ExecTool._spawn( + "echo hi", + "/tmp", + {"HOME": "/tmp"}, + process_tree=True, + ) + + assert mock_exec.call_args.kwargs["start_new_session"] is True + class TestSpawnWindows: diff --git a/tests/tools/test_exec_session_tools.py b/tests/tools/test_exec_session_tools.py index dfdddf1f..456d8f3e 100644 --- a/tests/tools/test_exec_session_tools.py +++ b/tests/tools/test_exec_session_tools.py @@ -1,12 +1,19 @@ from __future__ import annotations import asyncio +import base64 import re import shlex import subprocess import sys import time +from types import SimpleNamespace +from unittest.mock import AsyncMock +import pytest + +from nanobot.agent import context as agent_context +from nanobot.agent.loop import AgentLoop from nanobot.agent.tools.context import RequestContext, bind_request_context, reset_request_context from nanobot.agent.tools.exec_session import ( ExecSessionManager, @@ -454,3 +461,182 @@ def test_list_exec_sessions_reports_empty_state(): result = asyncio.run(ListExecSessionsTool(manager=ExecSessionManager()).execute()) assert result == "No active exec sessions." + + +def test_exec_session_manager_close_all_terminates_active_sessions(tmp_path): + async def run() -> None: + manager = ExecSessionManager() + tool = ExecTool(working_dir=str(tmp_path), timeout=30, session_manager=manager) + initial = await tool.execute( + command=_waiting_shell_command("ready"), + yield_time_ms=100, + ) + sid = _session_id(initial) + process = manager._sessions[sid].process + assert process.returncode is None + + closed = await manager.close_all() + + assert closed == 1 + assert process.returncode is not None + assert manager._sessions == {} + assert await manager.close_all() == 0 + + asyncio.run(run()) + + +def test_exec_session_manager_shutdown_terminates_child_processes(tmp_path): + async def run() -> None: + marker = tmp_path / "orphaned-child.txt" + child_code = ( + "import pathlib,time; time.sleep(2); " + f"pathlib.Path({str(marker)!r}).write_text('alive')" + ) + child_payload = base64.b64encode(child_code.encode()).decode() + parent_code = ( + "import base64,subprocess,sys,time; " + f"child=base64.b64decode('{child_payload}').decode(); " + "subprocess.Popen([sys.executable, '-c', child]); " + "print('ready', flush=True); time.sleep(4)" + ) + manager = ExecSessionManager() + tool = ExecTool(working_dir=str(tmp_path), timeout=30, session_manager=manager) + initial = await tool.execute(command=_python_command(parent_code), yield_time_ms=500) + assert "ready" in initial + assert "Process running" in initial + + await manager.close_all() + await asyncio.sleep(2.3) + + assert not marker.exists() + + asyncio.run(run()) + + +def test_exec_session_manager_rejects_new_sessions_after_shutdown(tmp_path): + async def run() -> str: + manager = ExecSessionManager() + await manager.close_all() + tool = ExecTool(working_dir=str(tmp_path), timeout=5, session_manager=manager) + return await tool.execute(command="echo should-not-run", yield_time_ms=0) + + result = asyncio.run(run()) + + assert result == "Error executing command: exec session manager is closed" + + +def test_exec_session_manager_retains_and_aggregates_failed_cleanup(): + async def run() -> None: + manager = ExecSessionManager() + first = SimpleNamespace( + session_id="first", + kill=AsyncMock(side_effect=OSError("first failed")), + ) + second = SimpleNamespace( + session_id="second", + kill=AsyncMock(side_effect=RuntimeError("second failed")), + ) + manager._sessions = {first.session_id: first, second.session_id: second} + + with pytest.raises(ExceptionGroup) as exc_info: + await manager.close_all() + + assert len(exc_info.value.exceptions) == 2 + assert manager._sessions == {first.session_id: first, second.session_id: second} + first.kill.assert_awaited_once() + second.kill.assert_awaited_once() + + first.kill.side_effect = None + second.kill.side_effect = None + assert await manager.close_all() == 2 + assert manager._sessions == {} + + asyncio.run(run()) + + +def test_exec_session_manager_preserves_single_cleanup_error(): + async def run() -> None: + manager = ExecSessionManager() + session = SimpleNamespace( + session_id="failed", + kill=AsyncMock(side_effect=OSError("cleanup failed")), + ) + manager._sessions = {session.session_id: session} + + with pytest.raises(OSError, match="cleanup failed"): + await manager.close_all() + + assert manager._sessions == {session.session_id: session} + + asyncio.run(run()) + + +def test_agent_loop_shutdown_closes_exec_sessions(tmp_path, monkeypatch): + async def run() -> None: + manager = ExecSessionManager() + tool = ExecTool(working_dir=str(tmp_path), timeout=30, session_manager=manager) + initial = await tool.execute( + command=_waiting_shell_command("ready"), + yield_time_ms=100, + ) + sid = _session_id(initial) + process = manager._sessions[sid].process + + monkeypatch.setattr(agent_context, "close_mcp", lambda _state: asyncio.sleep(0)) + loop = object.__new__(AgentLoop) + loop._background_tasks = [] + loop._exec_session_manager = manager + loop.subagents = SimpleNamespace(close=AsyncMock()) + + await loop.close_mcp() + await loop.close_mcp() + + assert process.returncode is not None + assert manager._sessions == {} + assert loop.subagents.close.await_count == 2 + + asyncio.run(run()) + + +def test_agent_loop_shutdown_attempts_all_cleanup_after_errors(monkeypatch): + async def run() -> None: + loop = object.__new__(AgentLoop) + loop._background_tasks = [] + loop.subagents = SimpleNamespace( + close=AsyncMock(side_effect=RuntimeError("subagent cleanup failed")), + ) + loop._exec_session_manager = SimpleNamespace( + close_all=AsyncMock(side_effect=OSError("exec cleanup failed")), + ) + close_mcp = AsyncMock() + monkeypatch.setattr(agent_context, "close_mcp", close_mcp) + + with pytest.raises(BaseExceptionGroup) as exc_info: + await loop.close_mcp() + + assert len(exc_info.value.exceptions) == 2 + loop.subagents.close.assert_awaited_once() + loop._exec_session_manager.close_all.assert_awaited_once() + close_mcp.assert_awaited_once_with(loop) + + asyncio.run(run()) + + +def test_agent_loop_shutdown_preserves_single_cleanup_error(monkeypatch): + async def run() -> None: + loop = object.__new__(AgentLoop) + loop._background_tasks = [] + loop.subagents = SimpleNamespace( + close=AsyncMock(side_effect=RuntimeError("subagent cleanup failed")), + ) + loop._exec_session_manager = SimpleNamespace(close_all=AsyncMock()) + close_mcp = AsyncMock() + monkeypatch.setattr(agent_context, "close_mcp", close_mcp) + + with pytest.raises(RuntimeError, match="subagent cleanup failed"): + await loop.close_mcp() + + loop._exec_session_manager.close_all.assert_awaited_once() + close_mcp.assert_awaited_once_with(loop) + + asyncio.run(run()) diff --git a/tests/tools/test_shell_reap.py b/tests/tools/test_shell_reap.py index 3f776e4c..98b38c39 100644 --- a/tests/tools/test_shell_reap.py +++ b/tests/tools/test_shell_reap.py @@ -211,6 +211,62 @@ async def test_exec_session_kill_reaps(): await asyncio.gather(session._stdout_task, session._stderr_task, return_exceptions=True) +@pytest.mark.asyncio +async def test_exec_session_kill_reaps_if_process_exits_before_kill(): + process = _mock_session_process(pid=2002, returncode=None) + process.kill.side_effect = ProcessLookupError("raced exit") + session = _ExecSession( + session_id="raced", + process=process, + command="sleep 1", + cwd="/tmp", + timeout=30, + owner_session_key=None, + ) + try: + with patch("nanobot.agent.tools.shell._reap_pid") as reap: + await session.kill() + reap.assert_called_once_with(2002) + finally: + await asyncio.gather(session._stdout_task, session._stderr_task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_exec_session_kill_waits_for_reader_tasks(): + process = _mock_session_process(pid=2003, returncode=None) + session = _ExecSession( + session_id="readers", + process=process, + command="sleep 1", + cwd="/tmp", + timeout=30, + owner_session_key=None, + ) + await asyncio.gather(session._stdout_task, session._stderr_task) + + release = asyncio.Event() + + async def reader() -> None: + await release.wait() + + async def wait_for_process() -> int: + release.set() + process.returncode = -9 + return -9 + + session._stdout_task = asyncio.create_task(reader()) + session._stderr_task = asyncio.create_task(reader()) + process.wait = AsyncMock(side_effect=wait_for_process) + try: + await session.kill() + + assert session._stdout_task.done() + assert session._stderr_task.done() + finally: + release.set() + await asyncio.gather(session._stdout_task, session._stderr_task) + + @pytest.mark.asyncio async def test_exec_session_poll_reaps_after_exit(): process = _mock_session_process(pid=2002, returncode=0)