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
This commit is contained in:
KDB
2026-07-18 17:37:07 +08:00
committed by Xubin Ren
parent d4f5abe004
commit d35f99abfc
4 changed files with 137 additions and 9 deletions
+2 -2
View File
@@ -66,7 +66,7 @@ class SessionClient:
def get(self, session_key: str) -> SessionSnapshot | None: def get(self, session_key: str) -> SessionSnapshot | None:
"""Return a display-safe snapshot without creating a new session on disk.""" """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: if cached is not None:
return snapshot_from_session(cached) return snapshot_from_session(cached)
payload = self._loop.sessions.read_session_file(session_key) payload = self._loop.sessions.read_session_file(session_key)
@@ -90,7 +90,7 @@ class SessionClient:
def export(self, session_key: str) -> SessionSnapshot | None: def export(self, session_key: str) -> SessionSnapshot | None:
"""Return a trusted full snapshot, including model-only runtime context.""" """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: if cached is not None:
return snapshot_from_session(cached, include_runtime_context=True) return snapshot_from_session(cached, include_runtime_context=True)
payload = self._loop.sessions.read_session_file(session_key) payload = self._loop.sessions.read_session_file(session_key)
+37 -7
View File
@@ -5,12 +5,14 @@ import json
import os import os
import re import re
import shutil import shutil
from collections import OrderedDict
from contextlib import suppress from contextlib import suppress
from copy import deepcopy from copy import deepcopy
from dataclasses import dataclass, field from dataclasses import dataclass, field
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
from weakref import WeakValueDictionary
from loguru import logger from loguru import logger
@@ -31,6 +33,7 @@ from nanobot.utils.helpers import (
from nanobot.utils.subagent_channel_display import scrub_subagent_announce_body from nanobot.utils.subagent_channel_display import scrub_subagent_announce_body
FILE_MAX_MESSAGES = 2000 FILE_MAX_MESSAGES = 2000
SESSION_CACHE_MAX_SIZE = 128
MIN_REPLAY_MAX_MESSAGES = 120 MIN_REPLAY_MAX_MESSAGES = 120
REPLAY_TOKENS_PER_MESSAGE = 100 REPLAY_TOKENS_PER_MESSAGE = 100
_MESSAGE_TIME_PREFIX_RE = re.compile(r"^\[Message Time: [^\]]+\]\n?") _MESSAGE_TIME_PREFIX_RE = re.compile(r"^\[Message Time: [^\]]+\]\n?")
@@ -418,7 +421,30 @@ class SessionManager:
self.workspace = workspace self.workspace = workspace
self.sessions_dir = ensure_dir(self.workspace / "sessions") self.sessions_dir = ensure_dir(self.workspace / "sessions")
self.legacy_sessions_dir = get_legacy_sessions_dir() 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 @staticmethod
def safe_key(key: str) -> str: def safe_key(key: str) -> str:
@@ -482,14 +508,15 @@ class SessionManager:
Returns: Returns:
The session. The session.
""" """
if key in self._cache: session = self._cached(key)
return self._cache[key] if session is not None:
return session
session = self._load(key) session = self._load(key)
if session is None: if session is None:
session = Session(key=key) session = Session(key=key)
self._cache[key] = session self._remember(session)
return session return session
def _load(self, key: str) -> Session | None: def _load(self, key: str) -> Session | None:
@@ -673,7 +700,7 @@ class SessionManager:
tmp_path.unlink(missing_ok=True) tmp_path.unlink(missing_ok=True)
raise raise
self._cache[session.key] = session self._remember(session)
def flush_all(self) -> int: def flush_all(self) -> int:
"""Re-save every cached session with fsync for durable shutdown. """Re-save every cached session with fsync for durable shutdown.
@@ -683,7 +710,9 @@ class SessionManager:
flushed. flushed.
""" """
flushed = 0 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: try:
self.save(session, fsync=True) self.save(session, fsync=True)
flushed += 1 flushed += 1
@@ -694,6 +723,7 @@ class SessionManager:
def invalidate(self, key: str) -> None: def invalidate(self, key: str) -> None:
"""Remove a session from the in-memory cache.""" """Remove a session from the in-memory cache."""
self._cache.pop(key, None) self._cache.pop(key, None)
self._overflow_cache.pop(key, None)
def delete_session(self, key: str) -> bool: def delete_session(self, key: str) -> bool:
"""Remove a session from disk (both workspace and legacy locations) and cache. """Remove a session from disk (both workspace and legacy locations) and cache.
@@ -733,7 +763,7 @@ class SessionManager:
""" """
if before_user_index < 0: if before_user_index < 0:
return None 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: if source is None:
return None return None
+75
View File
@@ -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)}
+23
View File
@@ -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 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 @pytest.mark.asyncio
async def test_session_export_and_restore_preserve_runtime_context(tmp_path): async def test_session_export_and_restore_preserve_runtime_context(tmp_path):
config_path = _write_config(tmp_path) config_path = _write_config(tmp_path)