fix(tools): scope file state by session

Made-with: Cursor

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)

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.

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)

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