diff --git a/nanobot/sdk/clients.py b/nanobot/sdk/clients.py index 13bff038..be4bea86 100644 --- a/nanobot/sdk/clients.py +++ b/nanobot/sdk/clients.py @@ -66,7 +66,7 @@ class SessionClient: def get(self, session_key: str) -> SessionSnapshot | None: """Return a display-safe snapshot without creating a new session on disk.""" - cached = self._loop.sessions._cache.get(session_key) + cached = self._loop.sessions._cached(session_key) if cached is not None: return snapshot_from_session(cached) payload = self._loop.sessions.read_session_file(session_key) @@ -90,7 +90,7 @@ class SessionClient: def export(self, session_key: str) -> SessionSnapshot | None: """Return a trusted full snapshot, including model-only runtime context.""" - cached = self._loop.sessions._cache.get(session_key) + cached = self._loop.sessions._cached(session_key) if cached is not None: return snapshot_from_session(cached, include_runtime_context=True) payload = self._loop.sessions.read_session_file(session_key) diff --git a/nanobot/session/manager.py b/nanobot/session/manager.py index 87fde59c..dedf25cb 100644 --- a/nanobot/session/manager.py +++ b/nanobot/session/manager.py @@ -5,12 +5,14 @@ import json import os import re import shutil +from collections import OrderedDict from contextlib import suppress from copy import deepcopy from dataclasses import dataclass, field from datetime import datetime from pathlib import Path from typing import Any +from weakref import WeakValueDictionary from loguru import logger @@ -31,6 +33,7 @@ from nanobot.utils.helpers import ( from nanobot.utils.subagent_channel_display import scrub_subagent_announce_body FILE_MAX_MESSAGES = 2000 +SESSION_CACHE_MAX_SIZE = 128 MIN_REPLAY_MAX_MESSAGES = 120 REPLAY_TOKENS_PER_MESSAGE = 100 _MESSAGE_TIME_PREFIX_RE = re.compile(r"^\[Message Time: [^\]]+\]\n?") @@ -418,7 +421,30 @@ class SessionManager: self.workspace = workspace self.sessions_dir = ensure_dir(self.workspace / "sessions") self.legacy_sessions_dir = get_legacy_sessions_dir() - self._cache: dict[str, Session] = {} + self._cache: OrderedDict[str, Session] = OrderedDict() + # Preserve identity for sessions held by active callers without retaining idle ones. + self._overflow_cache: WeakValueDictionary[str, Session] = WeakValueDictionary() + self._max_cached_sessions = SESSION_CACHE_MAX_SIZE + + def _remember(self, session: Session) -> None: + """Keep recent sessions strongly cached without duplicating live objects.""" + self._overflow_cache.pop(session.key, None) + self._cache[session.key] = session + self._cache.move_to_end(session.key) + while len(self._cache) > self._max_cached_sessions: + key, evicted = self._cache.popitem(last=False) + self._overflow_cache[key] = evicted + + def _cached(self, key: str) -> Session | None: + session = self._cache.get(key) + if session is not None: + self._cache.move_to_end(key) + return session + + session = self._overflow_cache.get(key) + if session is not None: + self._remember(session) + return session @staticmethod def safe_key(key: str) -> str: @@ -482,14 +508,15 @@ class SessionManager: Returns: The session. """ - if key in self._cache: - return self._cache[key] + session = self._cached(key) + if session is not None: + return session session = self._load(key) if session is None: session = Session(key=key) - self._cache[key] = session + self._remember(session) return session def _load(self, key: str) -> Session | None: @@ -673,7 +700,7 @@ class SessionManager: tmp_path.unlink(missing_ok=True) raise - self._cache[session.key] = session + self._remember(session) def flush_all(self) -> int: """Re-save every cached session with fsync for durable shutdown. @@ -683,7 +710,9 @@ class SessionManager: flushed. """ flushed = 0 - for key, session in list(self._cache.items()): + cached = dict(self._overflow_cache.items()) + cached.update(self._cache) + for key, session in cached.items(): try: self.save(session, fsync=True) flushed += 1 @@ -694,6 +723,7 @@ class SessionManager: def invalidate(self, key: str) -> None: """Remove a session from the in-memory cache.""" self._cache.pop(key, None) + self._overflow_cache.pop(key, None) def delete_session(self, key: str) -> bool: """Remove a session from disk (both workspace and legacy locations) and cache. @@ -733,7 +763,7 @@ class SessionManager: """ if before_user_index < 0: return None - source = self._cache.get(source_key) or self._load(source_key) + source = self._cached(source_key) or self._load(source_key) if source is None: return None diff --git a/tests/session/test_session_cache.py b/tests/session/test_session_cache.py new file mode 100644 index 00000000..cf1445b4 --- /dev/null +++ b/tests/session/test_session_cache.py @@ -0,0 +1,75 @@ +import gc +import weakref + +from nanobot.session.manager import SESSION_CACHE_MAX_SIZE, SessionManager + + +def _bounded_manager(tmp_path, limit: int) -> SessionManager: + manager = SessionManager(tmp_path) + manager._max_cached_sessions = limit + return manager + + +def test_default_session_cache_is_bounded(tmp_path) -> None: + manager = SessionManager(tmp_path) + + for index in range(SESSION_CACHE_MAX_SIZE + 1): + manager.get_or_create(f"test:{index}") + + assert len(manager._cache) == SESSION_CACHE_MAX_SIZE + + +def test_session_cache_releases_inactive_lru_entries(tmp_path) -> None: + manager = _bounded_manager(tmp_path, 1) + first = manager.get_or_create("test:first") + first.add_message("user", "persist me") + manager.save(first) + first_ref = weakref.ref(first) + + second = manager.get_or_create("test:second") + manager.save(second) + del first + gc.collect() + + assert len(manager._cache) == 1 + assert first_ref() is None + assert manager.get_or_create("test:first").messages[0]["content"] == "persist me" + + +def test_session_cache_keeps_identity_for_evicted_active_sessions(tmp_path) -> None: + manager = _bounded_manager(tmp_path, 1) + active = manager.get_or_create("test:active") + manager.save(active) + + manager.save(manager.get_or_create("test:other")) + + assert manager.get_or_create("test:active") is active + + +def test_session_cache_refreshes_lru_order_on_access(tmp_path) -> None: + manager = _bounded_manager(tmp_path, 2) + manager.save(manager.get_or_create("test:first")) + manager.save(manager.get_or_create("test:second")) + + manager.get_or_create("test:first") + manager.save(manager.get_or_create("test:third")) + + assert list(manager._cache) == ["test:first", "test:third"] + + +def test_flush_all_includes_live_sessions_outside_strong_cache(tmp_path, monkeypatch) -> None: + manager = _bounded_manager(tmp_path, 1) + active = manager.get_or_create("test:active") + manager.save(active) + manager.save(manager.get_or_create("test:other")) + saved: list[tuple[str, bool]] = [] + original_save = manager.save + + def recording_save(session, *, fsync=False): + saved.append((session.key, fsync)) + original_save(session, fsync=fsync) + + monkeypatch.setattr(manager, "save", recording_save) + + assert manager.flush_all() == 2 + assert set(saved) == {("test:active", True), ("test:other", True)} diff --git a/tests/test_nanobot_facade.py b/tests/test_nanobot_facade.py index b1462a39..f1e8c75c 100644 --- a/tests/test_nanobot_facade.py +++ b/tests/test_nanobot_facade.py @@ -1241,6 +1241,29 @@ async def test_session_helpers_get_list_export_clear_delete_flush(tmp_path): assert bot.sessions.get("sdk:first") is None +def test_session_helpers_read_live_session_outside_strong_cache(tmp_path): + config_path = _write_config(tmp_path) + bot = Nanobot.from_config(config_path, workspace=tmp_path) + sessions = bot._loop.sessions + sessions._max_cached_sessions = 1 + active = sessions.get_or_create("sdk:active") + active.add_message("user", "persisted") + sessions.save(active) + + sessions.save(sessions.get_or_create("sdk:other")) + active.add_message("assistant", "not saved yet") + + visible = bot.sessions.get("sdk:active") + exported = bot.sessions.export("sdk:active") + assert visible is not None + assert exported is not None + assert [message["content"] for message in visible.messages] == [ + "persisted", + "not saved yet", + ] + assert exported.messages == visible.messages + + @pytest.mark.asyncio async def test_session_export_and_restore_preserve_runtime_context(tmp_path): config_path = _write_config(tmp_path)