diff --git a/cli.py b/cli.py index 6091e3907bd..7de186e4a70 100644 --- a/cli.py +++ b/cli.py @@ -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 diff --git a/run_agent.py b/run_agent.py index 3cd458cd505..216425e486a 100644 --- a/run_agent.py +++ b/run_agent.py @@ -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: diff --git a/tests/cli/test_cli_new_session.py b/tests/cli/test_cli_new_session.py index c56ab63cf24..4bfb1d1ba35 100644 --- a/tests/cli/test_cli_new_session.py +++ b/tests/cli/test_cli_new_session.py @@ -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 diff --git a/tests/run_agent/test_commit_memory_session_context_engine.py b/tests/run_agent/test_commit_memory_session_context_engine.py index 307814891a2..956cf5117be 100644 --- a/tests/run_agent/test_commit_memory_session_context_engine.py +++ b/tests/run_agent/test_commit_memory_session_context_engine.py @@ -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()