diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 75639199..490172af 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -28,7 +28,7 @@ from nanobot.agent.tools.ask import ( pending_ask_user_id, ) from nanobot.agent.tools.cron import CronTool -from nanobot.agent.tools.file_state import FileStates +from nanobot.agent.tools.file_state import FileStates, bind_file_states, reset_file_states from nanobot.agent.tools.filesystem import EditFileTool, ListDirTool, ReadFileTool, WriteFileTool from nanobot.agent.tools.message import MessageTool from nanobot.agent.tools.notebook import NotebookEditTool @@ -248,10 +248,9 @@ class AgentLoop: self.context = ContextBuilder(workspace, timezone=timezone, disabled_skills=disabled_skills) self.sessions = session_manager or SessionManager(workspace) self.tools = ToolRegistry() - # Per-session file-read/write tracker (issue #3571) — shared across - # the filesystem tools registered below so this AgentLoop does not - # leak read-dedup state into another loop's tools. - self._file_states = FileStates() + # 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_states_by_session: dict[str, FileStates] = {} self.runner = AgentRunner(provider) self.subagents = SubagentManager( provider=provider, @@ -313,6 +312,14 @@ class AgentLoop: self.commands = CommandRouter() register_builtin_commands(self.commands) + def _file_states_for_session(self, session_key: str | None) -> FileStates: + key = session_key or "__default__" + states = self._file_states_by_session.get(key) + if states is None: + states = FileStates() + self._file_states_by_session[key] = states + return states + def _sync_subagent_runtime_limits(self) -> None: """Keep subagent runtime limits aligned with mutable loop settings.""" self.subagents.max_iterations = self.max_iterations @@ -353,24 +360,19 @@ class AgentLoop: self.workspace if (self.restrict_to_workspace or self.exec_config.sandbox) else None ) extra_read = [BUILTIN_SKILLS_DIR] if allowed_dir else None - # One FileStates per session so read-dedup and read-before-edit - # tracking does not leak across sessions sharing this process - # (issue #3571). - file_states = self._file_states self.tools.register(AskUserTool()) self.tools.register( ReadFileTool( workspace=self.workspace, allowed_dir=allowed_dir, extra_allowed_dirs=extra_read, - file_states=file_states, ) ) for cls in (WriteFileTool, EditFileTool, ListDirTool): - self.tools.register(cls(workspace=self.workspace, allowed_dir=allowed_dir, file_states=file_states)) + self.tools.register(cls(workspace=self.workspace, allowed_dir=allowed_dir)) for cls in (GlobTool, GrepTool): - self.tools.register(cls(workspace=self.workspace, allowed_dir=allowed_dir, file_states=file_states)) - self.tools.register(NotebookEditTool(workspace=self.workspace, allowed_dir=allowed_dir, file_states=file_states)) + self.tools.register(cls(workspace=self.workspace, allowed_dir=allowed_dir)) + self.tools.register(NotebookEditTool(workspace=self.workspace, allowed_dir=allowed_dir)) if self.exec_config.enable: self.tools.register( ExecTool( @@ -630,25 +632,30 @@ class AgentLoop: return items - result = await self.runner.run(AgentRunSpec( - initial_messages=initial_messages, - tools=self.tools, - model=self.model, - max_iterations=self.max_iterations, - max_tool_result_chars=self.max_tool_result_chars, - hook=hook, - error_message="Sorry, I encountered an error calling the AI model.", - concurrent_tools=True, - workspace=self.workspace, - session_key=session.key if session else None, - context_window_tokens=self.context_window_tokens, - context_block_limit=self.context_block_limit, - provider_retry_mode=self.provider_retry_mode, - progress_callback=on_progress, - retry_wait_callback=on_retry_wait, - checkpoint_callback=_checkpoint, - injection_callback=_drain_pending, - )) + active_session_key = session.key if session else session_key + file_state_token = bind_file_states(self._file_states_for_session(active_session_key)) + try: + result = await self.runner.run(AgentRunSpec( + initial_messages=initial_messages, + tools=self.tools, + model=self.model, + max_iterations=self.max_iterations, + max_tool_result_chars=self.max_tool_result_chars, + hook=hook, + error_message="Sorry, I encountered an error calling the AI model.", + concurrent_tools=True, + workspace=self.workspace, + session_key=session.key if session else None, + context_window_tokens=self.context_window_tokens, + context_block_limit=self.context_block_limit, + provider_retry_mode=self.provider_retry_mode, + progress_callback=on_progress, + retry_wait_callback=on_retry_wait, + checkpoint_callback=_checkpoint, + injection_callback=_drain_pending, + )) + finally: + reset_file_states(file_state_token) self._last_usage = result.usage if result.stop_reason == "max_iterations": logger.warning("Max iterations ({}) reached", self.max_iterations) diff --git a/nanobot/agent/tools/file_state.py b/nanobot/agent/tools/file_state.py index 018581ac..4cf5af87 100644 --- a/nanobot/agent/tools/file_state.py +++ b/nanobot/agent/tools/file_state.py @@ -4,6 +4,7 @@ from __future__ import annotations import hashlib import os +from contextvars import ContextVar, Token from dataclasses import dataclass from pathlib import Path @@ -129,6 +130,26 @@ class FileStates: self._state.clear() +_current_file_states: ContextVar[FileStates | None] = ContextVar( + "nanobot_file_states", + default=None, +) + + +def current_file_states(default: FileStates) -> FileStates: + """Return the FileStates bound to the current agent task, or a fallback.""" + return _current_file_states.get() or default + + +def bind_file_states(file_states: FileStates) -> Token[FileStates | None]: + """Bind file read/write state for the current async task.""" + return _current_file_states.set(file_states) + + +def reset_file_states(token: Token[FileStates | None]) -> None: + _current_file_states.reset(token) + + # Module-level default instance, retained for backward compatibility with # tests and callers that reach in directly. Per-session callers should hold # their own FileStates instance instead of touching this one. diff --git a/nanobot/agent/tools/filesystem.py b/nanobot/agent/tools/filesystem.py index 500f0123..587a149f 100644 --- a/nanobot/agent/tools/filesystem.py +++ b/nanobot/agent/tools/filesystem.py @@ -9,7 +9,7 @@ from typing import Any from nanobot.agent.tools.base import Tool, tool_parameters from nanobot.agent.tools.schema import BooleanSchema, IntegerSchema, StringSchema, tool_parameters_schema -from nanobot.agent.tools.file_state import FileStates, _hash_file +from nanobot.agent.tools.file_state import FileStates, _hash_file, current_file_states from nanobot.utils.helpers import build_image_content_blocks, detect_image_mime from nanobot.config.paths import get_media_dir @@ -54,10 +54,17 @@ class _FsTool(Tool): self._workspace = workspace self._allowed_dir = allowed_dir self._extra_allowed_dirs = extra_allowed_dirs - # Read-dedup / read-before-edit state is scoped to one FileStates - # so it does not leak across sessions sharing this process - # (issue #3571). Bare constructions get a private instance. - self._file_states: FileStates = file_states if file_states is not None else FileStates() + # Explicit state is used by isolated runners like Dream/subagents. + # Main AgentLoop tools leave this unset and resolve state from the + # current async task, which keeps shared tool instances session-safe. + self._explicit_file_states = file_states + self._fallback_file_states = FileStates() + + @property + def _file_states(self) -> FileStates: + if self._explicit_file_states is not None: + return self._explicit_file_states + return current_file_states(self._fallback_file_states) def _resolve(self, path: str) -> Path: return _resolve_path(path, self._workspace, self._allowed_dir, self._extra_allowed_dirs) diff --git a/tests/tools/test_read_enhancements.py b/tests/tools/test_read_enhancements.py index b5577f05..490623d5 100644 --- a/tests/tools/test_read_enhancements.py +++ b/tests/tools/test_read_enhancements.py @@ -127,6 +127,36 @@ class TestReadDedupSessionIsolation: ) assert "line 0" in second + @pytest.mark.asyncio + async def test_shared_loop_tool_uses_bound_session_state(self, tmp_path): + f = tmp_path / "shared.txt" + f.write_text("\n".join(f"line {i}" for i in range(10)), encoding="utf-8") + + # AgentLoop registers one shared ReadFileTool instance. The session + # boundary is the task-local FileStates binding, not the tool object. + shared_tool = ReadFileTool(workspace=tmp_path) + session_a = file_state.FileStates() + session_b = file_state.FileStates() + + token = file_state.bind_file_states(session_a) + try: + first = await shared_tool.execute(path=str(f)) + repeat = await shared_tool.execute(path=str(f)) + finally: + file_state.reset_file_states(token) + + assert "line 0" in first + assert "unchanged" in repeat.lower() + + token = file_state.bind_file_states(session_b) + try: + second_session_read = await shared_tool.execute(path=str(f)) + finally: + file_state.reset_file_states(token) + + assert "unchanged" not in second_session_read.lower() + assert "line 0" in second_session_read + # --------------------------------------------------------------------------- # PDF support