From 39c38b593f74e301937ee44085ff6ad0d570c423 Mon Sep 17 00:00:00 2001
From: Xubin Ren <xubinrencs@gmail.com>
Date: Fri, 1 May 2026 11:13:59 +0000
Subject: [PATCH] refactor(tools): move file state lookup out of loop

Made-with: Cursor
---
 nanobot/agent/loop.py             | 14 +++-----------
 nanobot/agent/tools/file_state.py | 20 ++++++++++++++++++++
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py
index 490172af6..b18bfe37c 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, bind_file_states, reset_file_states
+from nanobot.agent.tools.file_state import FileStateStore, 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
@@ -250,7 +250,7 @@ class AgentLoop:
         self.tools = ToolRegistry()
         # 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._file_state_store = FileStateStore()
         self.runner = AgentRunner(provider)
         self.subagents = SubagentManager(
             provider=provider,
@@ -312,14 +312,6 @@ 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
@@ -633,7 +625,7 @@ class AgentLoop:
             return items
 
         active_session_key = session.key if session else session_key
-        file_state_token = bind_file_states(self._file_states_for_session(active_session_key))
+        file_state_token = bind_file_states(self._file_state_store.for_session(active_session_key))
         try:
             result = await self.runner.run(AgentRunSpec(
                 initial_messages=initial_messages,
diff --git a/nanobot/agent/tools/file_state.py b/nanobot/agent/tools/file_state.py
index 4cf5af873..33673b3ef 100644
--- a/nanobot/agent/tools/file_state.py
+++ b/nanobot/agent/tools/file_state.py
@@ -130,6 +130,26 @@ class FileStates:
         self._state.clear()
 
 
+class FileStateStore:
+    """Lookup table for per-session file read/write state."""
+
+    __slots__ = ("_states_by_key",)
+
+    def __init__(self) -> None:
+        self._states_by_key: dict[str, FileStates] = {}
+
+    def for_session(self, session_key: str | None) -> FileStates:
+        key = session_key or "__default__"
+        states = self._states_by_key.get(key)
+        if states is None:
+            states = FileStates()
+            self._states_by_key[key] = states
+        return states
+
+    def clear(self) -> None:
+        self._states_by_key.clear()
+
+
 _current_file_states: ContextVar[FileStates | None] = ContextVar(
     "nanobot_file_states",
     default=None,