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, pending_ask_user_id,
) )
from nanobot.agent.tools.cron import CronTool 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.filesystem import EditFileTool, ListDirTool, ReadFileTool, WriteFileTool
from nanobot.agent.tools.message import MessageTool from nanobot.agent.tools.message import MessageTool
from nanobot.agent.tools.notebook import NotebookEditTool from nanobot.agent.tools.notebook import NotebookEditTool
@ -250,7 +250,7 @@ class AgentLoop:
self.tools = ToolRegistry() self.tools = ToolRegistry()
# One file-read/write tracker per logical session. The tool registry is # One file-read/write tracker per logical session. The tool registry is
# shared by this loop, so tools resolve the active state via contextvars. # 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.runner = AgentRunner(provider)
self.subagents = SubagentManager( self.subagents = SubagentManager(
provider=provider, provider=provider,
@ -312,14 +312,6 @@ class AgentLoop:
self.commands = CommandRouter() self.commands = CommandRouter()
register_builtin_commands(self.commands) 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: def _sync_subagent_runtime_limits(self) -> None:
"""Keep subagent runtime limits aligned with mutable loop settings.""" """Keep subagent runtime limits aligned with mutable loop settings."""
self.subagents.max_iterations = self.max_iterations self.subagents.max_iterations = self.max_iterations
@ -633,7 +625,7 @@ class AgentLoop:
return items return items
active_session_key = session.key if session else session_key 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: try:
result = await self.runner.run(AgentRunSpec( result = await self.runner.run(AgentRunSpec(
initial_messages=initial_messages, initial_messages=initial_messages,

View File

@ -130,6 +130,26 @@ class FileStates:
self._state.clear() 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( _current_file_states: ContextVar[FileStates | None] = ContextVar(
"nanobot_file_states", "nanobot_file_states",
default=None, default=None,