diff --git a/agent/memory_manager.py b/agent/memory_manager.py index c8b80a1514e..342656b4292 100644 --- a/agent/memory_manager.py +++ b/agent/memory_manager.py @@ -783,6 +783,56 @@ class MemoryManager: exc_info=True, ) + def commit_session_boundary_async( + self, + messages: List[Dict[str, Any]], + *, + new_session_id: str, + parent_session_id: str = "", + reason: str = "new_session", + ) -> None: + """Queue old-session extraction + provider rebinding as ONE serialized task. + + Session rotation (/new) must deliver ``on_session_end`` (end-of-session + extraction — an LLM-bound call that can take seconds) strictly BEFORE + ``on_session_switch`` (which rebinds provider-internal ``_session_id`` / + turn buffers to the new session). Running extraction inline blocked the + /new command for the whole LLM round-trip (#16454); running it on an + ad-hoc thread raced the inline switch — providers key off internal + state, so a late ``on_session_end`` ran against post-switch bindings + (transcript misattributed to the new session id, double-ingest of the + old turn buffer, new-session buffers cleared). + + Submitting BOTH hooks as one task on the manager's single background + worker gives both properties at a single chokepoint: the caller returns + immediately, and the worker's FIFO order serializes end→switch against + every other provider write (per-turn ``sync_all``, prefetches), which + already share the same worker. If the executor is unavailable, + ``_submit_background`` degrades to inline execution — the pre-#16454 + synchronous behavior, slow but correct. + """ + providers = list(self._providers) + if not providers: + return + snapshot = list(messages or []) + + def _run() -> None: + try: + self.on_session_end(snapshot) + except Exception as e: # pragma: no cover - on_session_end guards per-provider + logger.warning("Session-boundary extraction failed: %s", e) + try: + self.on_session_switch( + new_session_id, + parent_session_id=parent_session_id, + reset=True, + reason=reason, + ) + except Exception as e: # pragma: no cover - on_session_switch guards per-provider + logger.warning("Session-boundary switch failed: %s", e) + + self._submit_background(_run) + def on_session_switch( self, new_session_id: str, diff --git a/cli.py b/cli.py index 0fa4858236d..c394f5545f6 100644 --- a/cli.py +++ b/cli.py @@ -1110,6 +1110,19 @@ def _run_cleanup(*, notify_session_finalize: bool = True): ) try: if _active_agent_ref and hasattr(_active_agent_ref, 'shutdown_memory_provider'): + # A /new shortly before exit leaves its end→switch boundary task + # (old-session extraction, LLM-bound) queued on the memory + # manager's serialized worker. shutdown_all()'s drain only waits + # ~5s and cancels queued tasks, so give pending work a bounded + # head start via the manager's own barrier — otherwise a + # "/new then quit" silently drops the old session's extraction. + # The 30s exit watchdog remains the hard backstop. + _mm = getattr(_active_agent_ref, '_memory_manager', None) + if _mm is not None and hasattr(_mm, 'flush_pending'): + try: + _mm.flush_pending(timeout=10) + except Exception: + pass # Forward the agent's own transcript so memory providers' # ``on_session_end`` hooks see the real conversation instead of # an empty list (#15165). ``_session_messages`` is set on @@ -4131,7 +4144,9 @@ 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 + # History snapshot staged by _launch_session_boundary_memory_flush for + # the /new end→switch boundary task (see new_session()). + self._session_boundary_snapshot: Optional[list] = None def _claim_active_session(self, surface: str = "cli", *, stderr: bool = False) -> bool: """Claim a global active-session slot for this CLI process.""" @@ -6963,26 +6978,38 @@ class HermesCLI(CLIAgentSetupMixin, CLICommandsMixin): *, session_id: Optional[str] = None, ) -> None: - """Kick old-session memory extraction off-thread so /new is responsive.""" - self._last_session_boundary_thread = None + """Stage old-session memory extraction so /new stays responsive. + + The context-engine ``on_session_end`` boundary is delivered + synchronously here: it is cheap (local state clear, no LLM call) and + ordering-sensitive — it must land before ``reset_session_state()`` + rebinds the engine to the new session. + + The memory-provider half (LLM-bound extraction, seconds) is NOT run + here. It is queued later in ``new_session()`` via + ``MemoryManager.commit_session_boundary_async`` as a single + end→switch task on the manager's serialized background worker, so + extraction can never race the provider rebinding (providers key off + internal ``_session_id`` state — a late ``on_session_end`` after + ``on_session_switch`` would misattribute the old transcript to the + new session). This method just snapshots the history for that call. + """ + self._session_boundary_snapshot = 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 + engine = getattr(agent, "context_compressor", None) + if engine is not None and hasattr(engine, "on_session_end"): + try: + engine.on_session_end(session_id or "", history_snapshot) + except Exception: + logger.debug( + "Context engine on_session_end failed at /new boundary", + exc_info=True, + ) - 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() + self._session_boundary_snapshot = list(history_snapshot) def new_session(self, silent=False, title=None): """Start a fresh session with a new session ID and cleared agent state.""" @@ -7088,15 +7115,33 @@ class HermesCLI(CLIAgentSetupMixin, CLICommandsMixin): # per-session state (_session_turns, _turn_counter, _document_id). # Fires BEFORE the plugin on_session_reset hook (shell hooks only # see the new id; Python providers see the transition). See #6672. + # + # When the old session has history, end-of-session extraction + # (LLM-bound, seconds) and this switch are queued as ONE task on + # the memory manager's serialized worker — end strictly before + # switch, without blocking /new (#16454). With no history there + # is nothing to extract; switch inline as before. try: _mm = getattr(self.agent, "_memory_manager", None) if _mm is not None: - _mm.on_session_switch( - self.session_id, - parent_session_id=old_session_id or "", - reset=True, - reason="new_session", + _boundary_snapshot = getattr( + self, "_session_boundary_snapshot", None ) + if _boundary_snapshot: + self._session_boundary_snapshot = None + _mm.commit_session_boundary_async( + _boundary_snapshot, + new_session_id=self.session_id, + parent_session_id=old_session_id or "", + reason="new_session", + ) + else: + _mm.on_session_switch( + self.session_id, + parent_session_id=old_session_id or "", + reset=True, + reason="new_session", + ) except Exception: pass self._notify_session_boundary("on_session_reset") @@ -7108,7 +7153,6 @@ class HermesCLI(CLIAgentSetupMixin, CLICommandsMixin): print("(^_^)v New session started!") - def _consume_pending_resume_selection(self, text: str) -> bool: """Resolve a bare numeric reply that follows a bare ``/resume`` prompt. diff --git a/tests/agent/test_memory_boundary_commit.py b/tests/agent/test_memory_boundary_commit.py new file mode 100644 index 00000000000..521e199922c --- /dev/null +++ b/tests/agent/test_memory_boundary_commit.py @@ -0,0 +1,124 @@ +"""Tests for MemoryManager.commit_session_boundary_async. + +The /new session boundary must deliver on_session_end (old-session +extraction) strictly BEFORE on_session_switch (provider rebinding to the +new session), without blocking the caller. Both hooks run as one task on +the manager's single serialized background worker. +""" + +from __future__ import annotations + +import threading +import time +from typing import Any, Dict, List + +from agent.memory_manager import MemoryManager +from agent.memory_provider import MemoryProvider + + +class _RecordingProvider(MemoryProvider): + """Provider that records hook invocations with thread identity.""" + + def __init__(self, end_delay: float = 0.0): + self.calls: List[tuple] = [] + self._end_delay = end_delay + self._caller_thread_ids: List[int] = [] + + # Required ABC surface (minimal no-ops) + @property + def name(self) -> str: + return "recorder" + + def is_available(self) -> bool: + return True + + def get_tool_schemas(self) -> List[Dict[str, Any]]: + return [] + + def initialize(self, agent: Any = None, **kwargs) -> bool: # type: ignore[override] + return True + + def build_system_prompt(self) -> str: # type: ignore[override] + return "" + + def sync_turn(self, user_content: str, assistant_content: str, **kwargs) -> None: # type: ignore[override] + self.calls.append(("sync_turn", kwargs.get("session_id", ""))) + + def on_session_end(self, messages: List[Dict[str, Any]]) -> None: + if self._end_delay: + time.sleep(self._end_delay) + self._caller_thread_ids.append(threading.get_ident()) + self.calls.append(("end", list(messages))) + + def on_session_switch(self, new_session_id: str, **kwargs) -> None: + self.calls.append(("switch", new_session_id, kwargs.get("reset"))) + + +def _make_manager(provider: _RecordingProvider) -> MemoryManager: + mm = MemoryManager() + mm._providers.append(provider) # bypass add_provider validation for the stub + return mm + + +def test_boundary_commit_delivers_end_strictly_before_switch(): + """Even with a slow (LLM-like) extraction, switch waits for end.""" + provider = _RecordingProvider(end_delay=0.15) + mm = _make_manager(provider) + + msgs = [{"role": "user", "content": "old turn"}] + t0 = time.monotonic() + mm.commit_session_boundary_async( + msgs, new_session_id="new-sid", parent_session_id="old-sid" + ) + # Caller returns immediately — the slow extraction must not block /new. + assert time.monotonic() - t0 < 0.1 + + assert mm.flush_pending(timeout=5) + + kinds = [c[0] for c in provider.calls] + assert kinds == ["end", "switch"], f"ordering violated: {provider.calls}" + assert provider.calls[0] == ("end", msgs) + assert provider.calls[1] == ("switch", "new-sid", True) + # And it genuinely ran off the caller's thread. + assert provider._caller_thread_ids[0] != threading.get_ident() + + +def test_boundary_commit_serializes_against_turn_syncs(): + """The boundary task shares the single worker with sync_all — FIFO order + means a queued boundary can't interleave into a later turn's sync.""" + provider = _RecordingProvider(end_delay=0.05) + mm = _make_manager(provider) + + mm.commit_session_boundary_async( + [{"role": "user", "content": "old"}], + new_session_id="new-sid", + ) + mm.sync_all("next-session user msg", "assistant reply", session_id="new-sid") + + assert mm.flush_pending(timeout=5) + + kinds = [c[0] for c in provider.calls] + assert kinds == ["end", "switch", "sync_turn"], f"unexpected order: {provider.calls}" + + +def test_boundary_commit_switch_still_fires_when_end_raises(): + """A failing provider extraction must not strand providers on the old sid.""" + + class _ExplodingEndProvider(_RecordingProvider): + def on_session_end(self, messages): # type: ignore[override] + raise RuntimeError("provider extraction blew up") + + provider = _ExplodingEndProvider() + mm = _make_manager(provider) + + mm.commit_session_boundary_async([{"role": "user", "content": "x"}], new_session_id="new-sid") + assert mm.flush_pending(timeout=5) + + assert ("switch", "new-sid", True) in provider.calls + + +def test_boundary_commit_noop_without_providers(): + mm = MemoryManager() + # Must not create the executor or raise. + mm.commit_session_boundary_async([{"role": "user", "content": "x"}], new_session_id="s") + assert mm._sync_executor is None diff --git a/tests/cli/test_cli_new_session.py b/tests/cli/test_cli_new_session.py index 4bfb1d1ba35..c9f5cfea1ea 100644 --- a/tests/cli/test_cli_new_session.py +++ b/tests/cli/test_cli_new_session.py @@ -156,10 +156,6 @@ 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) @@ -179,44 +175,87 @@ 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): +def test_new_session_queues_boundary_commit_with_snapshot(tmp_path): + """/new hands the OLD session's history + ids to the memory manager's + serialized boundary task instead of blocking on extraction inline.""" cli = _prepare_cli_with_active_session(tmp_path) old_session_id = cli.session_id + mm = MagicMock() + cli.agent._memory_manager = mm + + cli.process_command("/new") + + mm.commit_session_boundary_async.assert_called_once() + args, kwargs = mm.commit_session_boundary_async.call_args + assert args[0] == [{"role": "user", "content": "hello"}] + assert kwargs["new_session_id"] == cli.session_id + assert kwargs["parent_session_id"] == old_session_id + assert kwargs["reason"] == "new_session" + # The queued path replaces the inline switch — not both. + mm.on_session_switch.assert_not_called() + # Snapshot is consumed (a later /new without history must not re-send it). + assert cli._session_boundary_snapshot is None + + +def test_new_session_without_history_switches_inline(tmp_path): + """No old-session history → nothing to extract → plain inline switch.""" + cli = _prepare_cli_with_active_session(tmp_path) + cli.conversation_history = [] + + mm = MagicMock() + cli.agent._memory_manager = mm + + cli.process_command("/new") + + mm.commit_session_boundary_async.assert_not_called() + mm.on_session_switch.assert_called_once() + _, kwargs = mm.on_session_switch.call_args + assert kwargs["reset"] is True + + +def test_new_session_delivers_context_engine_boundary_synchronously(tmp_path): + """The context-engine on_session_end must fire during /new itself. + + It is cheap local state work and ordering-sensitive: it must land before + reset_session_state() rebinds the engine to the new session. The LLM-bound + provider extraction is what gets deferred, not this.""" + cli = _prepare_cli_with_active_session(tmp_path) + old_session_id = cli.session_id + + engine_calls = [] + cli.agent.context_compressor.on_session_end = ( + lambda sid, msgs: engine_calls.append((sid, list(msgs))) + ) + + cli.process_command("/new") + + assert engine_calls == [(old_session_id, [{"role": "user", "content": "hello"}])] + + +def test_run_cleanup_flushes_pending_memory_manager_work(tmp_path): + """A '/new then quit' must not drop the queued old-session extraction. + + _run_cleanup gives the manager's serialized worker a bounded drain via + flush_pending() before shutdown_all()'s short-fuse drain runs.""" import cli as _cli_mod - started = {} + agent = MagicMock() + mm = MagicMock() + mm.flush_pending.return_value = True + agent._memory_manager = mm + agent._session_messages = [] - 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 + old_ref = _cli_mod._active_agent_ref + _cli_mod._active_agent_ref = agent + _cli_mod._cleanup_done = False + try: + _cli_mod._run_cleanup(notify_session_finalize=False) + finally: + _cli_mod._cleanup_done = True + _cli_mod._active_agent_ref = old_ref - 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, - ) + mm.flush_pending.assert_called_once_with(timeout=10) def test_new_command_rotates_hermes_session_id_env_and_context(tmp_path):