fix(memory): drain queued writes on shutdown

agent/memory_manager.py

@@ -30,7 +30,7 @@ import logging
 import re
 import inspect
 import threading
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import Future, ThreadPoolExecutor, wait
 from typing import Any, Callable, Dict, List, Optional
 
 from agent.memory_provider import MemoryProvider
@@ -368,6 +368,16 @@ class MemoryManager:
         # _submit_background() and the sync_all/queue_prefetch_all rationale.
         self._sync_executor: Optional[ThreadPoolExecutor] = None
         self._sync_executor_lock = threading.Lock()
+        # Futures are tracked by durability class so shutdown can give writes
+        # a bounded FIFO drain, then explicitly report anything abandoned.
+        self._background_futures: Dict[Future, str] = {}
+        self._shutting_down = False
+        self._shutdown_drain_state: Dict[str, Any] = {
+            "status": "not_started",
+            "abandoned_writes": 0,
+            "abandoned_prefetches": 0,
+            "active_tasks": 0,
+        }
 
     # -- Registration --------------------------------------------------------
 
@@ -539,7 +549,7 @@ class MemoryManager:
                         provider.name, e,
                     )
 
-        self._submit_background(_run)
+        self._submit_background(_run, kind="prefetch")
 
     # -- Sync ----------------------------------------------------------------
 
@@ -615,46 +625,57 @@ class MemoryManager:
 
     # -- Background dispatch -------------------------------------------------
 
-    def _submit_background(self, fn) -> None:
-        """Run ``fn`` on the manager's background worker.
-
-        The executor is created lazily and shared across calls. If the
-        executor can't be created or has already been shut down, ``fn``
-        runs inline as a last-resort fallback — losing the async benefit
-        but never losing the write itself. ``fn`` must do its own
-        per-provider error handling; this wrapper only guards executor
-        plumbing.
-        """
+    def _submit_background(self, fn, *, kind: str = "write") -> None:
+        """Queue ``fn`` on the serialized worker and track its durability class."""
         executor = self._get_sync_executor()
         if executor is None:
-            # Executor unavailable (shut down / creation failed) — run
-            # inline rather than drop the work. Slow, but correct.
+            if self._shutting_down:
+                logger.warning("Memory manager is shutting down; rejecting late %s task", kind)
+                return
+            # Creation failure outside shutdown: preserve the historical
+            # fail-safe behavior and run the operation inline.
             try:
                 fn()
             except Exception as e:  # pragma: no cover - fn guards internally
                 logger.debug("Inline memory background task failed: %s", e)
             return
         try:
-            executor.submit(fn)
+            # Make submit+tracking atomic with the shutdown snapshot. The
+            # callback is attached after releasing the lock because an already
+            # completed future invokes callbacks synchronously.
+            with self._sync_executor_lock:
+                if self._shutting_down:
+                    logger.warning("Memory manager is shutting down; rejecting late %s task", kind)
+                    return
+                future = executor.submit(fn)
+                self._background_futures[future] = kind
+            future.add_done_callback(self._forget_background_future)
         except RuntimeError:
-            # Executor was shut down between the get and the submit
-            # (teardown race). Fall back to inline.
+            if self._shutting_down:
+                logger.warning("Memory manager shut down during %s submission; task rejected", kind)
+                return
             try:
                 fn()
             except Exception as e:  # pragma: no cover - fn guards internally
                 logger.debug("Inline memory background task failed: %s", e)
 
+    def _forget_background_future(self, future: Future) -> None:
+        with self._sync_executor_lock:
+            self._background_futures.pop(future, None)
+
     def _get_sync_executor(self) -> Optional[ThreadPoolExecutor]:
         """Lazily create the single-worker background executor."""
+        if self._shutting_down:
+            return None
         if self._sync_executor is not None:
             return self._sync_executor
         with self._sync_executor_lock:
+            if self._shutting_down:
+                return None
             if self._sync_executor is None:
                 try:
                     # Daemon workers (see tools.daemon_pool): a provider wedged
-                    # on a network call must never block interpreter exit —
-                    # stdlib ThreadPoolExecutor's atexit hook would join it
-                    # unconditionally even after shutdown(wait=False).
+                    # on a network call must never block interpreter exit.
                     from tools.daemon_pool import DaemonThreadPoolExecutor
                     self._sync_executor = DaemonThreadPoolExecutor(
                         max_workers=1,
@@ -1069,51 +1090,66 @@ class MemoryManager:
                     provider.name, e,
                 )
 
-    def _drain_sync_executor(self) -> None:
-        """Shut down the background executor, waiting briefly for drain.
-
-        Bounded by ``_SYNC_DRAIN_TIMEOUT_S``: a wedged provider must never
-        hang process/session teardown. We stop accepting new work and
-        cancel anything still queued, then wait at most the drain timeout
-        for the currently-running task on a watcher thread. The worker is
-        daemon, so an over-running task dies with the interpreter.
-        """
+    @property
+    def shutdown_drain_state(self) -> Dict[str, Any]:
+        """Snapshot of the most recent bounded shutdown drain outcome."""
         with self._sync_executor_lock:
+            return dict(self._shutdown_drain_state)
+
+    def _drain_sync_executor(self) -> None:
+        """Give queued FIFO work a bounded chance, then abandon explicitly."""
+        with self._sync_executor_lock:
+            self._shutting_down = True
             executor = self._sync_executor
             self._sync_executor = None
+            tracked = dict(self._background_futures)
+            self._shutdown_drain_state = {
+                "status": "draining" if executor is not None else "drained",
+                "abandoned_writes": 0,
+                "abandoned_prefetches": 0,
+                "active_tasks": sum(not future.done() for future in tracked),
+            }
         if executor is None:
             return
-        try:
-            # Stop accepting new work and drop anything still queued, but
-            # do NOT block here — cancel_futures cancels not-yet-started
-            # tasks; the in-flight one keeps running on its daemon thread.
-            executor.shutdown(wait=False, cancel_futures=True)
-        except TypeError:
-            # Older Python without cancel_futures kwarg.
-            try:
-                executor.shutdown(wait=False)
-            except Exception as e:  # pragma: no cover
-                logger.debug("Memory sync executor shutdown failed: %s", e)
-            return
-        except Exception as e:  # pragma: no cover
-            logger.debug("Memory sync executor shutdown failed: %s", e)
-            return
-        # Give an in-flight sync a bounded chance to finish on a watcher
-        # thread so we don't block the caller past the drain timeout.
-        drainer = threading.Thread(
-            target=lambda: self._bounded_executor_wait(executor),
-            daemon=True,
-            name="mem-sync-drain",
-        )
-        drainer.start()
-        drainer.join(timeout=_SYNC_DRAIN_TIMEOUT_S)
 
-    @staticmethod
-    def _bounded_executor_wait(executor: ThreadPoolExecutor) -> None:
-        try:
-            executor.shutdown(wait=True)
-        except Exception as e:  # pragma: no cover
-            logger.debug("Memory sync executor drain wait failed: %s", e)
+        # shutdown(wait=False) closes submission without touching the FIFO.
+        # Waiting on the tracked futures lets the real single-worker executor
+        # run every queued write/boundary task in order up to the deadline.
+        executor.shutdown(wait=False, cancel_futures=False)
+        _, pending = wait(tuple(tracked), timeout=_SYNC_DRAIN_TIMEOUT_S)
+        if not pending:
+            with self._sync_executor_lock:
+                self._shutdown_drain_state.update(status="drained", active_tasks=0)
+            return
+
+        abandoned_writes = 0
+        abandoned_prefetches = 0
+        active_tasks = 0
+        for future in pending:
+            kind = tracked[future]
+            if future.cancel():
+                if kind == "prefetch":
+                    abandoned_prefetches += 1
+                else:
+                    abandoned_writes += 1
+            else:
+                active_tasks += 1
+
+        with self._sync_executor_lock:
+            self._shutdown_drain_state.update(
+                status="timed_out",
+                abandoned_writes=abandoned_writes,
+                abandoned_prefetches=abandoned_prefetches,
+                active_tasks=active_tasks,
+            )
+        logger.warning(
+            "Memory shutdown drain timed out after %.2fs; abandoning %d queued "
+            "memory write(s) and %d queued prefetch(es); %d active task(s) remain detached",
+            _SYNC_DRAIN_TIMEOUT_S,
+            abandoned_writes,
+            abandoned_prefetches,
+            active_tasks,
+        )
 
     def initialize_all(self, session_id: str, **kwargs) -> None:
         """Initialize all providers.

tests/agent/test_memory_async_sync.py

@@ -15,6 +15,8 @@ The fix dispatches provider work to a single-worker background executor.
 for session boundaries and deterministic tests. ``shutdown_all`` drains the
 executor with a bounded timeout so a wedged provider can't hang teardown.
 """
+import logging
+import threading
 import time
 
 import pytest
@@ -136,3 +138,81 @@ def test_writes_are_serialized_in_order():
         mgr.sync_all(f"turn-{i}", "resp", session_id="s1")
     assert mgr.flush_pending(timeout=10) is True
     assert order == [f"turn-{i}" for i in range(5)]
+
+
+def test_shutdown_drains_queued_writes_and_boundary_in_fifo_order():
+    """Shutdown must not cancel durable work merely because it is still queued."""
+    started = threading.Event()
+    release = threading.Event()
+    calls = []
+
+    class _BlockingProvider(_SlowProvider):
+        def sync_turn(self, user_content, assistant_content, *, session_id="", messages=None):
+            if user_content == "turn-0":
+                started.set()
+                assert release.wait(timeout=2)
+            calls.append(("sync", user_content))
+
+        def on_session_end(self, messages):
+            calls.append(("end", messages[0]["content"]))
+
+        def on_session_switch(self, new_session_id, **kwargs):
+            calls.append(("switch", new_session_id))
+
+    mgr = MemoryManager()
+    mgr.add_provider(_BlockingProvider(delay=0))
+    mgr.sync_all("turn-0", "response")
+    assert started.wait(timeout=1)
+    mgr.sync_all("turn-1", "response")
+    mgr.commit_session_boundary_async(
+        [{"role": "user", "content": "old-session"}],
+        new_session_id="new-session",
+    )
+
+    threading.Timer(0.05, release.set).start()
+    mgr.shutdown_all()
+
+    assert calls == [
+        ("sync", "turn-0"),
+        ("sync", "turn-1"),
+        ("end", "old-session"),
+        ("switch", "new-session"),
+    ]
+    assert mgr.shutdown_drain_state["status"] == "drained"
+    assert mgr.shutdown_drain_state["abandoned_writes"] == 0
+
+
+def test_shutdown_timeout_abandons_queued_write_with_state_and_log(monkeypatch, caplog):
+    """A wedged active write bounds shutdown and reports queued data loss."""
+    import agent.memory_manager as memory_manager_module
+
+    started = threading.Event()
+    release = threading.Event()
+    calls = []
+
+    class _WedgedProvider(_SlowProvider):
+        def sync_turn(self, user_content, assistant_content, *, session_id="", messages=None):
+            if user_content == "active":
+                started.set()
+                release.wait(timeout=2)
+            calls.append(user_content)
+
+    monkeypatch.setattr(memory_manager_module, "_SYNC_DRAIN_TIMEOUT_S", 0.1)
+    mgr = MemoryManager()
+    mgr.add_provider(_WedgedProvider(delay=0))
+    mgr.sync_all("active", "response")
+    assert started.wait(timeout=1)
+    mgr.sync_all("queued", "response")
+
+    with caplog.at_level(logging.WARNING, logger="agent.memory_manager"):
+        t0 = time.monotonic()
+        mgr.shutdown_all()
+        elapsed = time.monotonic() - t0
+
+    state = mgr.shutdown_drain_state
+    assert elapsed < 0.5
+    assert state["status"] == "timed_out"
+    assert state["abandoned_writes"] == 1
+    assert "queued" not in calls
+    assert "abandoning 1 queued memory write" in caplog.text
+    release.set()