From 5cc5c58e019ef91802a5a9f7ba12b69ad281f685 Mon Sep 17 00:00:00 2001 From: Enough1122 Date: Tue, 28 Jul 2026 13:40:13 +0100 Subject: [PATCH] fix(gateway): flush pending memory writes before session teardown (#73297) The gateway's /reset cleanup path called shutdown_memory_provider without first draining the memory manager's serialized background write worker. shutdown_all only gives that worker a bounded (~5s) drain and abandons whatever is still queued past it, so a /reset could silently drop writes the session had already handed off -- the next session then loaded stale MEMORY.md. The CLI exit path already drains via MemoryManager.flush_pending before shutdown; this PR pins the same contract on the gateway cleanup path. Cleanup now calls agent._memory_manager.flush_pending(timeout=10) before the existing shutdown_memory_provider step. The flush is best-effort: a flush failure must never block teardown, so it is wrapped in try/except and the existing shutdown path remains the fallback. Closes #73297 --- gateway/run.py | 18 ++ .../test_73297_memory_flush_on_reset.py | 163 ++++++++++++++++++ 2 files changed, 181 insertions(+) create mode 100644 tests/gateway/test_73297_memory_flush_on_reset.py diff --git a/gateway/run.py b/gateway/run.py index f5f902bc9c9..e7e641601bc 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -7094,6 +7094,24 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew return try: if hasattr(agent, "shutdown_memory_provider"): + # Drain queued memory writes BEFORE tearing the provider down. + # The memory manager persists per-turn sync and end-of-session + # extraction on a single serialized background worker. + # shutdown_memory_provider() -> shutdown_all() only gives that + # worker a ~5s bounded drain and abandons (cancels) anything + # still queued past it, so a /reset — or any gateway session + # rotation that reaches this cleanup path — could silently drop + # writes the session had already handed off. The next session + # then loads stale memory (#73297). Give pending work a bounded + # head start through the manager's own barrier first, mirroring + # the CLI exit path (cli.py). Best-effort: a flush failure must + # never block teardown. + _mm = getattr(agent, "_memory_manager", None) + if _mm is not None and hasattr(_mm, "flush_pending"): + try: + _mm.flush_pending(timeout=10) + except Exception: + pass # Pass the agent's own conversation transcript so memory # providers' ``on_session_end`` hooks see the real messages # instead of the empty default (#15165). ``_session_messages`` diff --git a/tests/gateway/test_73297_memory_flush_on_reset.py b/tests/gateway/test_73297_memory_flush_on_reset.py new file mode 100644 index 00000000000..ddd9de20f90 --- /dev/null +++ b/tests/gateway/test_73297_memory_flush_on_reset.py @@ -0,0 +1,163 @@ +"""Regression tests for #73297: memory rollback after /reset. + +The gateway's ``_cleanup_agent_resources`` (the cleanup chokepoint that +``/reset`` and every other gateway session rotation invokes to tear down the +cached agent) used to call ``shutdown_memory_provider()`` WITHOUT first +draining the memory manager's serialized background write worker. +``shutdown_memory_provider`` -> ``shutdown_all`` only gives that worker a +bounded (~5s) drain and abandons whatever is still queued past it, so a +/reset could silently drop writes the session had already handed off — the +next session then loaded stale memory. + +The CLI exit path already drains via ``MemoryManager.flush_pending`` before +shutdown (cli.py); these tests pin the same contract on the gateway path. + +The fix: in ``_cleanup_agent_resources``, call +``agent._memory_manager.flush_pending(timeout=10)`` BEFORE +``shutdown_memory_provider``. +""" +import time +from types import SimpleNamespace +from unittest.mock import MagicMock, call + +import pytest + +from agent.memory_manager import MemoryManager +from agent.memory_provider import MemoryProvider +import agent.memory_manager as _mm_module +from gateway.run import GatewayRunner + + +# How long the "gate" write occupies the single background worker. Chosen so a +# queued write reliably stays PENDING behind it for the duration of the test +# (the worker is FIFO, single-threaded). Kept small for test speed. +_GATE_DELAY_S = 0.6 + + +class _RecordingProvider(MemoryProvider): + """Provider that records completed writes to ``self.recorded`` (the "disk"). + + A write whose assistant payload is the sentinel ``__GATE__`` sleeps for + ``gate_delay`` first — it occupies the manager's single background worker + so a subsequent real write sits PENDING behind it until the gate clears. + """ + + _name = "recording" + + def __init__(self, gate_delay: float = _GATE_DELAY_S): + self._gate_delay = gate_delay + self.recorded: list[str] = [] + + @property + def name(self) -> str: + return self._name + + def is_available(self) -> bool: + return True + + def initialize(self, session_id: str = "", **kwargs) -> None: + pass + + def system_prompt_block(self) -> str: + return "" + + def prefetch(self, query, *, session_id: str = "") -> str: + return "" + + def get_tool_schemas(self): + return [] + + def handle_tool_call(self, tool_name, args, **kwargs) -> str: + return "" + + def sync_turn(self, user_content, assistant_content, *, session_id: str = "", messages=None) -> None: + if assistant_content == "__GATE__": + time.sleep(self._gate_delay) + return + self.recorded.append(assistant_content) + + +def _make_agent(mgr: MemoryManager) -> SimpleNamespace: + """Build a minimal agent wired to a real MemoryManager. + + ``shutdown_memory_provider`` mirrors ``AIAgent.shutdown_memory_provider`` + (run_agent.py): end-of-session notification then ``shutdown_all``. The + memory manager is the live one carrying the queued writes, so the gateway + cleanup path exercises the real flush -> shutdown ordering. + """ + + def _shutdown_memory_provider(messages=None): + mgr.on_session_end(messages or []) + mgr.shutdown_all() + + return SimpleNamespace( + _memory_manager=mgr, + _session_messages=[{"role": "user", "content": "earlier turn"}], + shutdown_memory_provider=_shutdown_memory_provider, + ) + + +def test_cleanup_flushes_pending_writes_before_shutdown(monkeypatch): + """#73297 invariant: ``disk_state_after_reset >= last_known_state_before_reset``. + + A real write ("fact-1") is queued on the memory manager's background + worker behind a slow gate write, so it is PENDING when /reset cleanup + runs. ``shutdown_all``'s drain is shortened (via ``_SYNC_DRAIN_TIMEOUT_S``) + to model the production condition where the bounded drain abandons queued + work — the condition the issue reports. Without the pre-shutdown + ``flush_pending``, "fact-1" is abandoned and never reaches the provider + (the on-disk store). With the fix, ``flush_pending`` drains the queue + BEFORE the unreliable shutdown drain, so the write lands. + """ + # Shorten shutdown_all's drain so it abandons the still-pending write, + # modelling the bounded-drain abandonment at the heart of #73297. The + # fix's flush_pending uses its OWN (longer) barrier, independent of this + # timeout, so it still drains the queue. + monkeypatch.setattr(_mm_module, "_SYNC_DRAIN_TIMEOUT_S", 0.1) + + mgr = MemoryManager() + provider = _RecordingProvider() + mgr.add_provider(provider) + mgr.initialize_all("sess-old") + + # Occupy the single worker with a slow gate write, then queue the real + # write behind it. "fact-1" is now PENDING (FIFO, worker busy with gate). + mgr.sync_all("user msg a", "__GATE__", session_id="sess-old") + mgr.sync_all("user msg b", "fact-1", session_id="sess-old") + + last_known_state_before_reset = ["fact-1"] + assert provider.recorded == [], "precondition: write must still be pending" + + agent = _make_agent(mgr) + + # Run the gateway cleanup path the /reset handler runs. No explicit flush + # here — the contract is that cleanup itself flushes before shutdown. + GatewayRunner._cleanup_agent_resources(object(), agent) + + disk_state_after_reset = provider.recorded + assert set(last_known_state_before_reset) <= set(disk_state_after_reset), ( + f"#73297 regression: pending memory write lost across reset. " + f"expected >= {last_known_state_before_reset}, got {disk_state_after_reset}" + ) + + +def test_cleanup_calls_flush_pending_before_shutdown_memory_provider(): + """Structural contract: the gateway cleanup drains the memory manager + BEFORE shutting the provider down (ordering matters — shutdown's drain is + the unreliable one the fix works around).""" + agent = MagicMock() + agent._session_messages = [{"role": "user", "content": "hi"}] + agent._memory_manager = MagicMock() + + # Track call order across both objects via a parent mock. + parent = MagicMock() + parent.attach_mock(agent._memory_manager.flush_pending, "flush_pending") + parent.attach_mock(agent.shutdown_memory_provider, "shutdown_memory_provider") + + GatewayRunner._cleanup_agent_resources(object(), agent) + + agent._memory_manager.flush_pending.assert_called_once_with(timeout=10) + parent.assert_has_calls( + [call.flush_pending(timeout=10), call.shutdown_memory_provider([{"role": "user", "content": "hi"}])], + any_order=False, + )