refactor(tools): move file state lookup out of loop

Made-with: Cursor
This commit is contained in:
Xubin Ren 2026-05-01 11:13:59 +00:00 committed by Xubin Ren
parent fae38319ca
commit 39c38b593f
2 changed files with 23 additions and 11 deletions

View File

@ -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,

View File

@ -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,