diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 2532471687c..38e45a5f631 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -76,6 +76,7 @@ from nanobot.session.history_visibility import HIDDEN_HISTORY_META from nanobot.session.keys import UNIFIED_SESSION_KEY, remember_last_channel from nanobot.session.manager import ( + SESSION_CACHE_MAX_SIZE, Session, SessionManager, replay_max_messages_for_context, @@ -376,11 +377,15 @@ def __init__( self.context = ContextBuilder(workspace, timezone=timezone, disabled_skills=disabled_skills) self.sessions = session_manager or SessionManager(workspace) - self.sessions.set_file_cap_archiver(self.context.memory.raw_archive) - self.tools = tool_registry if tool_registry is not None else 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_state_store = FileStateStore() + self._file_state_store = FileStateStore(max_sessions=SESSION_CACHE_MAX_SIZE) + # SessionManager owns every durable deletion entrypoint, including the + # WebUI and fork rollback paths. Observe that boundary once instead of + # duplicating cleanup in each consumer. + self.sessions.set_delete_observer(self._file_state_store.discard) + self.sessions.set_file_cap_archiver(self.context.memory.raw_archive) + self.tools = tool_registry if tool_registry is not None else ToolRegistry() self._exec_session_manager = ExecSessionManager() self.runner = AgentRunner() self.subagents = SubagentManager( @@ -818,8 +823,13 @@ async def discard_session(self, key: str) -> None: self.sessions.invalidate(key) await self._cancel_active_tasks(key) finally: + self.discard_session_file_state(key) self._discarding_sessions.discard(key) + def discard_session_file_state(self, key: str) -> None: + """Forget ephemeral file-read state for a reset or removed session.""" + self._file_state_store.discard(key) + def _effective_session_key(self, msg: InboundMessage) -> str: """Return the session key used for task routing and mid-turn injections.""" if self._unified_session and not msg.session_key_override: diff --git a/nanobot/agent/tools/file_state.py b/nanobot/agent/tools/file_state.py index 3dd4667d5ec..81198191381 100644 --- a/nanobot/agent/tools/file_state.py +++ b/nanobot/agent/tools/file_state.py @@ -4,6 +4,7 @@ import hashlib import os +from collections import OrderedDict from contextvars import ContextVar, Token from dataclasses import dataclass from pathlib import Path @@ -135,21 +136,30 @@ def clear(self) -> None: class FileStateStore: - """Lookup table for per-session file read/write state.""" + """Bounded lookup table for per-session file read/write state.""" - __slots__ = ("_states_by_key",) + __slots__ = ("_max_sessions", "_states_by_key") - def __init__(self) -> None: - self._states_by_key: dict[str, FileStates] = {} + def __init__(self, *, max_sessions: int = 128) -> None: + if max_sessions <= 0: + raise ValueError("max_sessions must be positive") + self._max_sessions = max_sessions + self._states_by_key: OrderedDict[str, FileStates] = OrderedDict() def for_session(self, session_key: str | None) -> FileStates: key = session_key or "__default__" - states = self._states_by_key.get(key) + states = self._states_by_key.pop(key, None) if states is None: states = FileStates() - self._states_by_key[key] = states + self._states_by_key[key] = states + while len(self._states_by_key) > self._max_sessions: + self._states_by_key.popitem(last=False) return states + def discard(self, session_key: str | None) -> None: + """Forget file state when a session is reset or removed.""" + self._states_by_key.pop(session_key or "__default__", None) + def clear(self) -> None: self._states_by_key.clear() diff --git a/nanobot/command/builtin.py b/nanobot/command/builtin.py index 7ef7c6368a7..c5fe5910572 100644 --- a/nanobot/command/builtin.py +++ b/nanobot/command/builtin.py @@ -302,6 +302,7 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage: """Stop active task and start a fresh session.""" loop = ctx.loop await loop._cancel_active_tasks(ctx.key) # pyright: ignore[reportPrivateUsage] + loop.discard_session_file_state(ctx.key) session = ctx.session or loop.sessions.get_or_create(ctx.key) snapshot = session.messages[session.last_consolidated:] runtime = None diff --git a/nanobot/sdk/clients.py b/nanobot/sdk/clients.py index a1d189af34b..afd97576460 100644 --- a/nanobot/sdk/clients.py +++ b/nanobot/sdk/clients.py @@ -138,6 +138,7 @@ async def restore( def clear(self, session_key: str) -> SessionSnapshot: """Clear one session and persist the empty session.""" + self._loop.discard_session_file_state(session_key) session = self._loop.sessions.get_or_create(session_key) session.clear() self._loop.sessions.save(session) diff --git a/nanobot/session/manager.py b/nanobot/session/manager.py index fa46f8e3b55..3a8e5f95d2b 100644 --- a/nanobot/session/manager.py +++ b/nanobot/session/manager.py @@ -1523,6 +1523,7 @@ def __init__( self._overflow_cache: WeakValueDictionary[str, Session] = WeakValueDictionary() self._max_cached_sessions = SESSION_CACHE_MAX_SIZE self._file_cap_archiver: Callable[..., None] | None = None + self._delete_observer: Callable[[str], None] | None = None def _remember(self, session: Session) -> None: """Keep recent sessions strongly cached without duplicating live objects.""" @@ -1552,6 +1553,10 @@ def set_file_cap_archiver(self, archiver: Callable[..., None]) -> None: """Archive unconsolidated overflow whenever a session is persisted.""" self._file_cap_archiver = archiver + def set_delete_observer(self, observer: Callable[[str], None]) -> None: + """Observe explicit session deletion for process-local state cleanup.""" + self._delete_observer = observer + @staticmethod def safe_key(key: str) -> str: """Public helper used by HTTP handlers to map an arbitrary key to a stable filename stem.""" @@ -1684,7 +1689,10 @@ def invalidate(self, key: str) -> None: def delete_session(self, key: str) -> bool: """Delete a persisted session and invalidate its cache entry.""" self.invalidate(key) - return self._store.delete(key) + deleted = self._store.delete(key) + if self._delete_observer is not None: + self._delete_observer(key) + return deleted def restore_sessions_to_workspace(self) -> SessionRestoreResult: """Restore session files to the pre-relocation path for an explicit rollback.""" diff --git a/tests/agent/test_loop_session_policy.py b/tests/agent/test_loop_session_policy.py index 3a2a065f505..adee6817bab 100644 --- a/tests/agent/test_loop_session_policy.py +++ b/tests/agent/test_loop_session_policy.py @@ -131,6 +131,7 @@ async def wait_for_discard(key: str) -> None: terminate_exec_sessions, ) key = "websocket:transient-cancelled" + previous_file_state = loop._file_state_store.for_session(key) loop.sessions.get_or_create_transient( key, disabled_tools={"create_goal", "update_goal", "spawn", "cron"}, @@ -157,6 +158,7 @@ async def wait_for_discard(key: str) -> None: await asyncio.wait_for(active_task, timeout=2) await asyncio.wait_for(wait_for_discard(key), timeout=2) assert loop.sessions.get_cached(key) is None + assert loop._file_state_store.for_session(key) is not previous_file_state terminate_exec_sessions.assert_awaited_once_with(key) loop.stop() diff --git a/tests/agent/test_session_delete.py b/tests/agent/test_session_delete.py index 97955fc1b93..74b7b0a4663 100644 --- a/tests/agent/test_session_delete.py +++ b/tests/agent/test_session_delete.py @@ -34,6 +34,17 @@ def test_delete_session_returns_false_when_missing(tmp_path: Path) -> None: assert sm.delete_session("nope:none") is False +def test_delete_session_notifies_process_local_state_observer(tmp_path: Path) -> None: + sm = _seed(tmp_path, "websocket:abc") + deleted_keys: list[str] = [] + sm.set_delete_observer(deleted_keys.append) + + assert sm.delete_session("websocket:abc") is True + assert sm.delete_session("websocket:missing") is False + + assert deleted_keys == ["websocket:abc", "websocket:missing"] + + def test_read_session_file_returns_metadata_and_messages(tmp_path: Path) -> None: sm = _seed(tmp_path, "telegram:abc") data = sm.read_session_file("telegram:abc") diff --git a/tests/agent/test_unified_session.py b/tests/agent/test_unified_session.py index 7c0656ec299..1acede1c6d5 100644 --- a/tests/agent/test_unified_session.py +++ b/tests/agent/test_unified_session.py @@ -20,6 +20,7 @@ import pytest from nanobot.agent.loop import AgentLoop +from nanobot.agent.tools.file_state import FileStateStore from nanobot.bus.events import InboundMessage from nanobot.bus.queue import MessageBus from nanobot.command.builtin import cmd_new, register_builtin_commands @@ -250,10 +251,16 @@ async def test_cmd_new_clears_unified_session(self, tmp_path: Path): # asyncio.create_task(). Mirror that exactly so the coroutine is consumed # and no RuntimeWarning is emitted. admitted_runtime = MagicMock(name="admitted_runtime") + file_state_store = FileStateStore() + previous_file_state = file_state_store.for_session("unified:default") + tracked_file = tmp_path / "tracked.txt" + tracked_file.write_text("tracked", encoding="utf-8") + previous_file_state.record_read(tracked_file) loop = SimpleNamespace( sessions=sessions, consolidator=SimpleNamespace(archive=AsyncMock(return_value=True)), _cancel_active_tasks=AsyncMock(return_value=0), + discard_session_file_state=file_state_store.discard, llm_runtime=MagicMock(return_value=MagicMock()), schedule_background=lambda coro: asyncio.ensure_future(coro), ) @@ -278,6 +285,9 @@ async def test_cmd_new_clears_unified_session(self, tmp_path: Path): sessions.invalidate("unified:default") reloaded = sessions.get_or_create("unified:default") assert reloaded.messages == [] + reset_file_state = file_state_store.for_session("unified:default") + assert reset_file_state is not previous_file_state + assert reset_file_state.is_unchanged(tracked_file) is False loop.consolidator.archive.assert_called_once_with( expected_snapshot, runtime=admitted_runtime, @@ -302,6 +312,7 @@ async def test_cmd_new_in_unified_mode_does_not_affect_other_sessions(self, tmp_ sessions=sessions, consolidator=SimpleNamespace(archive=AsyncMock(return_value=True)), _cancel_active_tasks=AsyncMock(return_value=0), + discard_session_file_state=MagicMock(), runtime_for_session=MagicMock(return_value=MagicMock()), schedule_background=lambda coro: asyncio.ensure_future(coro), ) diff --git a/tests/test_nanobot_facade.py b/tests/test_nanobot_facade.py index dfddf7150fa..d3ef4d23397 100644 --- a/tests/test_nanobot_facade.py +++ b/tests/test_nanobot_facade.py @@ -1500,10 +1500,15 @@ async def test_session_helpers_get_list_export_clear_delete_flush(tmp_path): exported.messages[0]["content"] = "mutated copy" assert bot.sessions.get("sdk:first").messages[0]["content"] == "hello" + state_before_clear = bot._loop._file_state_store.for_session("sdk:first") cleared = bot.sessions.clear("sdk:first") assert cleared.messages == [] + state_after_clear = bot._loop._file_state_store.for_session("sdk:first") + assert state_after_clear is not state_before_clear assert bot.sessions.flush() >= 1 + state_before_delete = state_after_clear assert bot.sessions.delete("sdk:first") is True + assert bot._loop._file_state_store.for_session("sdk:first") is not state_before_delete assert bot.sessions.get("sdk:first") is None diff --git a/tests/tools/test_file_state_store.py b/tests/tools/test_file_state_store.py new file mode 100644 index 00000000000..2d241ef51c5 --- /dev/null +++ b/tests/tools/test_file_state_store.py @@ -0,0 +1,29 @@ +import pytest + +from nanobot.agent.tools.file_state import FileStateStore + + +def test_file_state_store_evicts_least_recently_used_session() -> None: + store = FileStateStore(max_sessions=2) + first = store.for_session("first") + second = store.for_session("second") + + assert store.for_session("first") is first + store.for_session("third") + + assert store.for_session("first") is first + assert store.for_session("second") is not second + + +def test_file_state_store_discards_reset_session() -> None: + store = FileStateStore() + previous = store.for_session("websocket:chat") + + store.discard("websocket:chat") + + assert store.for_session("websocket:chat") is not previous + + +def test_file_state_store_requires_positive_capacity() -> None: + with pytest.raises(ValueError, match="max_sessions must be positive"): + FileStateStore(max_sessions=0)