From 995cc44e8995e7a0d936537a0e4b0c697f8acc05 Mon Sep 17 00:00:00 2001 From: yu-xin-c <2182712990@qq.com> Date: Thu, 9 Jul 2026 23:02:40 +0800 Subject: [PATCH] fix(exec): isolate exec session managers --- nanobot/agent/loop.py | 3 ++ nanobot/agent/subagent.py | 3 ++ nanobot/agent/tools/context.py | 1 + nanobot/agent/tools/exec_session.py | 14 ++---- nanobot/agent/tools/shell.py | 1 + tests/agent/test_exec_session_isolation.py | 40 +++++++++++++++++ tests/tools/test_exec_session_tools.py | 51 ++++++++++++++++++++++ tests/tools/test_tool_loader.py | 31 ++++++++++++- 8 files changed, 132 insertions(+), 12 deletions(-) create mode 100644 tests/agent/test_exec_session_isolation.py diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 17507ba6..fbbadfcd 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -28,6 +28,7 @@ from nanobot.agent.model_runtime import ModelRuntimeResolver from nanobot.agent.runner import _MAX_INJECTIONS_PER_TURN, AgentRunner, AgentRunSpec from nanobot.agent.subagent import SubagentManager from nanobot.agent.tools.context import RequestContext, bind_request_context, reset_request_context +from nanobot.agent.tools.exec_session import ExecSessionManager from nanobot.agent.tools.file_state import FileStateStore, bind_file_states, reset_file_states from nanobot.agent.tools.message import MessageTool from nanobot.agent.tools.registry import ToolRegistry @@ -355,6 +356,7 @@ class AgentLoop: # One file-read/write tracker per logical session. The tool registry is # shared by this loop, so tools resolve the active state via contextvars. self._file_state_store = FileStateStore() + self._exec_session_manager = ExecSessionManager() self.runner = AgentRunner() self.subagents = SubagentManager( workspace=workspace, @@ -540,6 +542,7 @@ class AgentLoop: bus=self.bus, subagent_manager=self.subagents, cron_service=self.cron_service, + exec_session_manager=self._exec_session_manager, sessions=self.sessions, provider_snapshot_loader=provider_snapshot_loader, image_generation_provider_configs=self._image_generation_provider_configs, diff --git a/nanobot/agent/subagent.py b/nanobot/agent/subagent.py index cf4a1ecd..f4d1fe6a 100644 --- a/nanobot/agent/subagent.py +++ b/nanobot/agent/subagent.py @@ -19,6 +19,7 @@ from nanobot.agent.tools.context import ( bind_request_context, reset_request_context, ) +from nanobot.agent.tools.exec_session import ExecSessionManager from nanobot.agent.tools.file_state import FileStates from nanobot.agent.tools.loader import ToolLoader from nanobot.agent.tools.registry import ToolRegistry @@ -143,6 +144,7 @@ class SubagentManager: else defaults.fail_on_tool_error ) self.runner = AgentRunner() + self._exec_session_manager = ExecSessionManager() self._llm_wall_timeout_for_session = llm_wall_timeout_for_session self._running_tasks: dict[str, asyncio.Task[None]] = {} self._task_statuses: dict[str, SubagentStatus] = {} @@ -204,6 +206,7 @@ class SubagentManager: ctx = ToolContext( config=cfg, workspace=str(root.resolve()), + exec_session_manager=self._exec_session_manager, file_state_store=FileStates(), workspace_sandbox=workspace_sandbox_status( restrict_to_workspace=cfg.restrict_to_workspace, diff --git a/nanobot/agent/tools/context.py b/nanobot/agent/tools/context.py index c64da5f3..f6b09215 100644 --- a/nanobot/agent/tools/context.py +++ b/nanobot/agent/tools/context.py @@ -71,6 +71,7 @@ class ToolContext: bus: Any | None = None subagent_manager: Any | None = None cron_service: Any | None = None + exec_session_manager: Any | None = None sessions: Any | None = None file_state_store: Any = field(default=None) provider_snapshot_loader: Callable[[], Any] | None = None diff --git a/nanobot/agent/tools/exec_session.py b/nanobot/agent/tools/exec_session.py index d9db38c8..c93ed8d0 100644 --- a/nanobot/agent/tools/exec_session.py +++ b/nanobot/agent/tools/exec_session.py @@ -250,11 +250,7 @@ class ExecSessionManager: session = self._sessions.get(session_id) if session is None: raise KeyError(session_id) - if ( - owner_session_key - and session.owner_session_key - and session.owner_session_key != owner_session_key - ): + if session.owner_session_key and session.owner_session_key != owner_session_key: raise KeyError(session_id) if chars: @@ -296,9 +292,7 @@ class ExecSessionManager: owner_session_key=session.owner_session_key, ) for session_id, session in sorted(self._sessions.items()) - if not owner_session_key - or not session.owner_session_key - or session.owner_session_key == owner_session_key + if session.owner_session_key == owner_session_key ] async def _cleanup_locked(self) -> None: @@ -442,7 +436,7 @@ class WriteStdinTool(Tool): @classmethod def create(cls, ctx: Any) -> Tool: - return cls() + return cls(manager=getattr(ctx, "exec_session_manager", None)) @property def exclusive(self) -> bool: @@ -586,7 +580,7 @@ class ListExecSessionsTool(Tool): @classmethod def create(cls, ctx: Any) -> Tool: - return cls() + return cls(manager=getattr(ctx, "exec_session_manager", None)) @property def name(self) -> str: diff --git a/nanobot/agent/tools/shell.py b/nanobot/agent/tools/shell.py index 4aa36911..175d06df 100644 --- a/nanobot/agent/tools/shell.py +++ b/nanobot/agent/tools/shell.py @@ -188,6 +188,7 @@ class ExecTool(Tool): allowed_env_keys=cfg.allowed_env_keys, allow_patterns=cfg.allow_patterns, deny_patterns=cfg.deny_patterns, + session_manager=getattr(ctx, "exec_session_manager", None), ) def __init__( diff --git a/tests/agent/test_exec_session_isolation.py b/tests/agent/test_exec_session_isolation.py new file mode 100644 index 00000000..9d929ae3 --- /dev/null +++ b/tests/agent/test_exec_session_isolation.py @@ -0,0 +1,40 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +from nanobot.agent.loop import AgentLoop +from nanobot.bus.queue import MessageBus + + +def _provider() -> MagicMock: + provider = MagicMock() + provider.get_default_model.return_value = "test-model" + provider.generation = SimpleNamespace(max_tokens=4096, temperature=0.1, reasoning_effort=None) + return provider + + +def test_agent_loops_do_not_share_exec_session_managers(tmp_path): + loop_a = AgentLoop( + bus=MessageBus(), + provider=_provider(), + workspace=tmp_path / "a", + model="test-model", + context_window_tokens=4096, + ) + loop_b = AgentLoop( + bus=MessageBus(), + provider=_provider(), + workspace=tmp_path / "b", + model="test-model", + context_window_tokens=4096, + ) + + exec_a = loop_a.tools.get("exec") + stdin_a = loop_a.tools.get("write_stdin") + list_a = loop_a.tools.get("list_exec_sessions") + exec_b = loop_b.tools.get("exec") + + assert exec_a._session_manager is loop_a._exec_session_manager + assert stdin_a._manager is loop_a._exec_session_manager + assert list_a._manager is loop_a._exec_session_manager + assert exec_b._session_manager is loop_b._exec_session_manager + assert loop_a._exec_session_manager is not loop_b._exec_session_manager diff --git a/tests/tools/test_exec_session_tools.py b/tests/tools/test_exec_session_tools.py index b62dbbe5..cf20eb53 100644 --- a/tests/tools/test_exec_session_tools.py +++ b/tests/tools/test_exec_session_tools.py @@ -7,6 +7,7 @@ import subprocess import sys import time +from nanobot.agent.tools.context import RequestContext, bind_request_context, reset_request_context from nanobot.agent.tools.exec_session import ( ExecSessionManager, ListExecSessionsTool, @@ -398,6 +399,56 @@ def test_list_exec_sessions_reports_running_commands(tmp_path): assert "Session terminated." in cleanup +def test_exec_sessions_are_scoped_to_request_session_key(tmp_path): + async def run() -> tuple[str, str, str, str, str, str]: + manager = ExecSessionManager() + exec_tool = ExecTool(working_dir=str(tmp_path), timeout=5, session_manager=manager) + list_tool = ListExecSessionsTool(manager=manager) + stdin_tool = WriteStdinTool(manager=manager) + command = _python_command( + "import time; print('ready', flush=True); time.sleep(5)" + ) + + token_a = bind_request_context( + RequestContext(channel="cli", chat_id="a", session_key="cli:a") + ) + try: + initial = await exec_tool.execute(command=command, yield_time_ms=100) + sid = _session_id(initial) + owner_listing = await list_tool.execute() + finally: + reset_request_context(token_a) + + unbound_listing = await list_tool.execute() + + token_b = bind_request_context( + RequestContext(channel="cli", chat_id="b", session_key="cli:b") + ) + try: + other_listing = await list_tool.execute() + other_write = await stdin_tool.execute(session_id=sid, yield_time_ms=0) + finally: + reset_request_context(token_b) + + token_a = bind_request_context( + RequestContext(channel="cli", chat_id="a", session_key="cli:a") + ) + try: + cleanup = await stdin_tool.execute(session_id=sid, terminate=True, yield_time_ms=0) + finally: + reset_request_context(token_a) + + return sid, owner_listing, unbound_listing, other_listing, other_write, cleanup + + sid, owner_listing, unbound_listing, other_listing, other_write, cleanup = asyncio.run(run()) + + assert sid in owner_listing + assert unbound_listing == "No active exec sessions." + assert other_listing == "No active exec sessions." + assert other_write == f"Error: exec session not found: {sid!r}" + assert "Session terminated." in cleanup + + def test_list_exec_sessions_reports_empty_state(): result = asyncio.run(ListExecSessionsTool(manager=ExecSessionManager()).execute()) diff --git a/tests/tools/test_tool_loader.py b/tests/tools/test_tool_loader.py index c7ad9d69..f31e9dbd 100644 --- a/tests/tools/test_tool_loader.py +++ b/tests/tools/test_tool_loader.py @@ -57,8 +57,8 @@ def test_tool_context_has_required_fields(): field_names = {f.name for f in fields(ToolContext)} required = { "config", "workspace", "bus", "subagent_manager", - "cron_service", "file_state_store", "provider_snapshot_loader", - "image_generation_provider_configs", "timezone", + "cron_service", "exec_session_manager", "file_state_store", + "provider_snapshot_loader", "image_generation_provider_configs", "timezone", } assert required <= field_names @@ -68,6 +68,7 @@ def test_tool_context_defaults(): assert ctx.bus is None assert ctx.subagent_manager is None assert ctx.cron_service is None + assert ctx.exec_session_manager is None assert ctx.provider_snapshot_loader is None assert ctx.image_generation_provider_configs is None assert ctx.timezone == "UTC" @@ -137,6 +138,32 @@ def test_loader_registers_exec_with_real_tools_config(tmp_path): assert registry.has("exec") +def test_loader_wires_shared_exec_session_manager(tmp_path): + from types import SimpleNamespace + + from nanobot.agent.tools.exec_session import ExecSessionManager + from nanobot.agent.tools.registry import ToolRegistry + from nanobot.config.schema import ToolsConfig + + manager = ExecSessionManager() + ctx = ToolContext( + config=ToolsConfig(), + workspace=str(tmp_path), + subagent_manager=SimpleNamespace( + get_running_count=lambda: 0, + max_concurrent_subagents=4, + ), + exec_session_manager=manager, + timezone="UTC", + ) + registry = ToolRegistry() + ToolLoader().load(ctx, registry) + + assert registry.get("exec")._session_manager is manager + assert registry.get("write_stdin")._manager is manager + assert registry.get("list_exec_sessions")._manager is manager + + # --- Task 4: _FsTool.create() ---