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:
Kshitij Kapoor 2026-07-09 02:48:16 +05:30 committed by kshitij
parent 2a293319b1
commit d8bc4f242f
4 changed files with 314 additions and 57 deletions

View file

@ -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):