CLI /new session memory extraction blocking fix

Fixes on record

Merged to main. A stamp counts the issues a fix is linked to.

• #61139 fix(cli): run /new session-boundary memory extraction off the command path (#16454 salvage)

  merged to main

  files: agent/memory_manager.py, cli.py, tests/agent/test_memory_boundary_commit.py, tests/cli/test_cli_new_session.py

  Show the diff

  • agent/memory_manager.py

    @@ -783,6 +783,55 @@ def on_session_end(self, messages: List[Dict[str, Any]]) -> None:
                         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.
    +        """
    +        if not self._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,

  • 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
    @@ -6956,17 +6969,69 @@ def _discard_session_if_empty(self, session_id: Optional[str]) -> bool:
                 )
                 return False
     
    +    def _launch_session_boundary_memory_flush(
    +        self,
    +        history_snapshot: list,
    +        *,
    +        session_id: Optional[str] = None,
    +    ) -> Optional[list]:
    +        """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. The returned snapshot is handed by ``new_session()`` to
    +        ``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).
    +
    +        Returns the history snapshot to queue, or ``None`` when there is
    +        nothing to extract (no agent / empty history / no memory manager).
    +        """
    +        agent = getattr(self, "agent", None)
    +        if not agent or not history_snapshot:
    +            return None
    +
    +        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,
    +                )
    +
    +        # No provider extraction to queue when no memory manager is
    +        # configured — new_session() falls back to the inline switch path.
    +        if getattr(agent, "_memory_manager", None) is None:
    +            return None
    +        return history_snapshot
    +
         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
    +        _boundary_snapshot = None
             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)
    +            # Deliver the context-engine boundary synchronously and get back
    +            # the history snapshot for the deferred provider extraction —
    +            # queued below (after rotation) so /new never blocks on the
    +            # LLM-bound extraction call.
    +            _boundary_snapshot = 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
    @@ -7054,15 +7119,29 @@ def new_session(self, silent=False, title=None):
                 # 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",
    -                    )
    +                    if _boundary_snapshot:
    +                        _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")
    @@ -7074,7 +7153,6 @@ def new_session(self, silent=False, title=None):
                     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.
     

  • 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

  • tests/cli/test_cli_new_session.py

    @@ -175,6 +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_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()
    +
    +
    +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
    +
    +    agent = MagicMock()
    +    mm = MagicMock()
    +    mm.flush_pending.return_value = True
    +    agent._memory_manager = mm
    +    agent._session_messages = []
    +
    +    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
    +
    +    mm.flush_pending.assert_called_once_with(timeout=10)
    +
    +
     def test_new_command_rotates_hermes_session_id_env_and_context(tmp_path):
         from gateway.session_context import _VAR_MAP, get_session_env