From d35f99abfc2840cdfceefd3b0d6c1d5b37b456b8 Mon Sep 17 00:00:00 2001 From: KDB <937925477@qq.com> Date: Thu, 16 Jul 2026 21:52:22 +0800 Subject: [PATCH] fix(session): bound the in-memory session cache Keep only 128 recently used sessions strongly cached while retaining weak references to evicted sessions still owned by active callers. This bounds idle memory growth without allowing duplicate live Session objects or skipping shutdown flushes. Add LRU, lifecycle, SDK, and flush regression coverage. Refs #4786 --- nanobot/sdk/clients.py | 4 +- nanobot/session/manager.py | 44 ++++++++++++++--- tests/session/test_session_cache.py | 75 +++++++++++++++++++++++++++++ tests/test_nanobot_facade.py | 23 +++++++++ 4 files changed, 137 insertions(+), 9 deletions(-) create mode 100644 tests/session/test_session_cache.py 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)