mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-15 14:22:43 +00:00
fix: serialize /new end→switch boundary on the memory manager worker
Deep review of the cherry-picked #16454 found the ad-hoc flush thread raced new_session()'s inline on_session_switch(reset=True): memory providers key off internal _session_id state (MemoryManager.on_session_end takes no session id), so a late off-thread extraction ran against post-rotation bindings — misattributing the old transcript to the new session id, double-ingesting the old turn buffer (supermemory), or double-committing (openviking already async-finalizes in on_session_switch). Redesign: new MemoryManager.commit_session_boundary_async queues on_session_end + on_session_switch as ONE task on the manager's existing single-worker background executor (the same worker sync_all already uses). This preserves the strict end→switch ordering providers depend on, serializes against per-turn syncs FIFO, keeps /new non-blocking, and degrades to inline (pre-#16454 behavior) when the executor is unavailable. No ad-hoc threads; no per-provider changes needed. The context-engine on_session_end half stays synchronous in _launch_session_boundary_memory_flush (cheap, must land before reset_session_state rebinds the engine). Exit durability: _run_cleanup calls the manager's existing flush_pending(timeout=10) barrier before shutdown, so '/new then quit' doesn't drop the queued extraction (shutdown_all's own drain is ~5s and cancels queued tasks). Bounded well inside the 30s exit watchdog. Tests: ordering invariant with slow (LLM-like) extraction, FIFO serialization vs sync_all, switch-fires-even-if-end-raises, no-provider no-op, CLI snapshot handoff + inline-switch fallback, sync engine boundary, cleanup flush_pending.
This commit is contained in:
parent
2a293319b1
commit
d8bc4f242f
4 changed files with 314 additions and 57 deletions
|
|
@ -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,
|
||||
|
|
|
|||
88
cli.py
88
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.
|
||||
|
||||
|
|
|
|||
124
tests/agent/test_memory_boundary_commit.py
Normal file
124
tests/agent/test_memory_boundary_commit.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue