Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(exec): clean up sessions on shutdown

Browse Source
This commit is contained in:
KDB
2026-07-21 13:48:51 +08:00
committed by Xubin Ren
parent 7cf3c71e3a
commit 8981995474
10 changed files with 390 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/loop.py
+16 -2
View File
@@ -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)."""
nanobot/agent/subagent.py
+9
View File
@@ -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)
nanobot/agent/tools/exec_session.py
+49 -9
View File
@@ -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,
)
nanobot/agent/tools/shell.py
+37
View File
@@ -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.
tests/agent/test_runner_injections.py
+1
View File
@@ -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
tests/agent/test_subagent_lifecycle.py
+18
View File
@@ -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
# ---------------------------------------------------------------------------
tests/cli/test_restart_command.py
+2 -1
View File
@@ -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
tests/tools/test_exec_platform.py
+16
View File
@@ -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:
tests/tools/test_exec_session_tools.py
+186
View File
@@ -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())
tests/tools/test_shell_reap.py
+56
View File
@@ -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)