fix: run CLI new-session memory flush off-thread

This commit is contained in:
Tosko4 2026-07-05 10:59:25 +02:00 committed by Tosko4
parent d577408f3f
commit 3e82a861e6
4 changed files with 103 additions and 5 deletions

40
cli.py
View file

@ -4127,6 +4127,7 @@ class HermesCLI(CLIAgentSetupMixin, CLICommandsMixin):
# Background task tracking: {task_id: threading.Thread}
self._background_tasks: Dict[str, threading.Thread] = {}
self._background_task_counter = 0
self._last_session_boundary_thread: Optional[threading.Thread] = None
def _claim_active_session(self, surface: str = "cli", *, stderr: bool = False) -> bool:
"""Claim a global active-session slot for this CLI process."""
@ -6925,17 +6926,50 @@ class HermesCLI(CLIAgentSetupMixin, CLICommandsMixin):
)
return False
def _launch_session_boundary_memory_flush(
self,
history_snapshot: list,
*,
session_id: Optional[str] = None,
) -> None:
"""Kick old-session memory extraction off-thread so /new is responsive."""
self._last_session_boundary_thread = None
agent = getattr(self, "agent", None)
if not agent or not history_snapshot:
return
def _run_boundary_flush() -> None:
if hasattr(agent, "commit_memory_session"):
try:
agent.commit_memory_session(history_snapshot, session_id=session_id)
except Exception:
pass
thread = threading.Thread(
target=_run_boundary_flush,
daemon=True,
name=f"session-boundary-flush-{(session_id or 'unknown')[:24]}",
)
self._last_session_boundary_thread = thread
thread.start()
def new_session(self, silent=False, title=None):
"""Start a fresh session with a new session ID and cleared agent state."""
old_session_id = self.session_id
if self.agent and self.conversation_history:
# Trigger memory extraction on the old session before session_id rotates.
self.agent.commit_memory_session(self.conversation_history)
# Trigger memory extraction on the old session without blocking
# session rotation. Snapshot the history and old session id first:
# the background thread may run after ``self.agent.session_id`` is
# updated to the new session below.
self._launch_session_boundary_memory_flush(
list(self.conversation_history),
session_id=old_session_id,
)
self._notify_session_boundary("on_session_finalize")
elif self.agent:
# First session or empty history — still finalize the old session
self._notify_session_boundary("on_session_finalize")
old_session_id = self.session_id
if self._session_db and old_session_id:
# Flush any un-persisted messages from the current turn to the
# old session *before* rotating. /new can be called mid-turn

View file

@ -3307,11 +3307,18 @@ class AIAgent:
except Exception:
pass
def commit_memory_session(self, messages: list = None) -> None:
def commit_memory_session(
self,
messages: list = None,
session_id: Optional[str] = None,
) -> None:
"""Trigger end-of-session extraction without tearing providers down.
Called when session_id rotates (e.g. /new, context compression);
providers keep their state and continue running under the old
session_id they just flush pending extraction now."""
ended_session_id = (
session_id if session_id is not None else (self.session_id or "")
)
if self._memory_manager:
try:
self._memory_manager.on_session_end(messages or [])
@ -3326,7 +3333,7 @@ class AIAgent:
if hasattr(self, "context_compressor") and self.context_compressor:
try:
self.context_compressor.on_session_end(
self.session_id or "",
ended_session_id,
messages or [],
)
except Exception:

View file

@ -156,6 +156,10 @@ def test_new_command_creates_real_fresh_session_and_resets_agent_state(tmp_path)
cli.process_command("/new")
boundary_thread = getattr(cli, "_last_session_boundary_thread", None)
if boundary_thread is not None:
boundary_thread.join(timeout=1)
assert cli.session_id != old_session_id
old_session = cli._session_db.get_session(old_session_id)
@ -175,6 +179,46 @@ def test_new_command_creates_real_fresh_session_and_resets_agent_state(tmp_path)
cli.agent._invalidate_system_prompt.assert_called_once()
def test_new_session_dispatches_memory_flush_on_background_thread(tmp_path):
cli = _prepare_cli_with_active_session(tmp_path)
old_session_id = cli.session_id
import cli as _cli_mod
started = {}
class _DummyThread:
def __init__(self, *, target=None, args=(), kwargs=None, daemon=None, name=None):
started["target"] = target
started["args"] = args
started["kwargs"] = kwargs or {}
started["daemon"] = daemon
started["name"] = name
started["started"] = False
def start(self):
started["started"] = True
def join(self, timeout=None):
return None
with patch.object(_cli_mod.threading, "Thread", _DummyThread):
cli.process_command("/new")
assert started["started"] is True
assert callable(started["target"])
assert started["daemon"] is True
assert started["name"] == f"session-boundary-flush-{old_session_id[:24]}"
cli.agent.commit_memory_session.assert_not_called()
started["target"]()
cli.agent.commit_memory_session.assert_called_once_with(
[{"role": "user", "content": "hello"}],
session_id=old_session_id,
)
def test_new_command_rotates_hermes_session_id_env_and_context(tmp_path):
from gateway.session_context import _VAR_MAP, get_session_env

View file

@ -56,6 +56,19 @@ def test_commit_memory_session_with_no_messages_passes_empty_list():
ctx.on_session_end.assert_called_once_with("sess-7", [])
def test_commit_memory_session_can_finalize_explicit_old_session_id():
"""Async callers may run after agent.session_id already changed."""
mm = MagicMock()
ctx = MagicMock()
agent = _make_minimal_agent(mm, ctx, session_id="new-session")
msgs = [{"role": "user", "content": "old turn"}]
agent.commit_memory_session(msgs, session_id="old-session")
mm.on_session_end.assert_called_once_with(msgs)
ctx.on_session_end.assert_called_once_with("old-session", msgs)
def test_commit_memory_session_no_memory_manager_still_notifies_context_engine():
"""If only the context engine is configured, it still gets the hook."""
ctx = MagicMock()