fix(honcho): give the writer join and its drain the shutdown deadline

stop_async_writer bounded only the join and then drained the queue with no deadline, so a shutdown whose budget was already spent could still start uploads, and honcho-ai's add_messages has no per-call timeout. The join and the drain now share the shutdown deadline, no upload starts once it has passed, the writer skips its 2s retry once shutdown began, and shutdown logs one warning with the count left unsynced. The docstrings now say what the budget can do: stop new uploads and bound lock waits, while an upload already in flight runs to the client's HTTP timeout.

plugins/memory/honcho/session.py

@@ -405,8 +405,9 @@ class HonchoSessionManager(SessionAuthMixin, SessionPeersMixin, SessionContextMi
             return self._flush_session_locked(session)
 
     def _flush_session_before(self, session: HonchoSession, deadline: float | None) -> bool:
-        """_flush_session bounded by ``deadline``. False, with nothing sent, when the budget is spent or another
-        flush of this session holds the lock past it."""
+        """_flush_session that starts only while ``deadline`` has time left and waits for the session's lock no
+        longer than that; False, with nothing sent, otherwise. An upload that has started is not interrupted: the
+        SDK has no per-call timeout, so it runs to the client's HTTP timeout."""
         if deadline is None:
             self._flush_session(session)
             return True
@@ -414,7 +415,7 @@ class HonchoSessionManager(SessionAuthMixin, SessionPeersMixin, SessionContextMi
         if remaining <= 0 or not session._flush_lock.acquire(timeout=remaining):
             return False
         try:
-            self._flush_session_locked(session)
+            self._flush_session(session)  # re-enters the RLock already held above
         finally:
             session._flush_lock.release()
         return True
@@ -479,6 +480,10 @@ class HonchoSessionManager(SessionAuthMixin, SessionPeersMixin, SessionContextMi
                 if item is _ASYNC_SHUTDOWN:
                     break
                 if not self._try_flush(item, logging.WARNING, "Honcho async write failed, retrying once"):
+                    if self._shutting_down:
+                        # The shutdown flush is already attempting this session within its own budget.
+                        logger.error("Honcho async write failed while shutting down, so the batch stays unsynced")
+                        continue
                     time.sleep(2)
                     self._try_flush(item, logging.ERROR, "Honcho async write retry failed, dropping batch")
             except queue.Empty:
@@ -537,8 +542,15 @@ class HonchoSessionManager(SessionAuthMixin, SessionPeersMixin, SessionContextMi
 
     def flush_all(self, timeout: float | None = None) -> None:
         """Flush unsynced messages for all cached sessions, then drain the async queue inline. ``timeout`` bounds
-        the whole pass: a session it cannot reach in time keeps its messages and is counted in one warning."""
+        when a flush may start and how long it waits for a session's lock, not an upload already in flight, which
+        runs to the client's HTTP timeout. A session skipped keeps its messages and is counted in one warning."""
         deadline = None if timeout is None else time.monotonic() + timeout
+        skipped = self._flush_cached_before(deadline)
+        skipped.extend(self._drain_async_queue(deadline))
+        self._warn_unsynced(skipped, timeout)
+
+    def _flush_cached_before(self, deadline: float | None) -> list[HonchoSession]:
+        """Flush every cached and retry-listed session that ``deadline`` allows; returns the ones it did not."""
         with self._cache_lock:
             sessions = list(self._cache.values())
             sessions += [s for s in self._retry_sessions if not any(s is c for c in sessions)]
@@ -551,12 +563,14 @@ class HonchoSessionManager(SessionAuthMixin, SessionPeersMixin, SessionContextMi
                 logger.error("Honcho flush_all error for %s: %s", session.key, e)
         with self._cache_lock:
             self._retry_sessions = [s for s in self._retry_sessions if self._has_unsynced(s)]
-        skipped.extend(self._drain_async_queue(deadline))
+        return skipped
+
+    def _warn_unsynced(self, skipped: list[HonchoSession], timeout: float | None) -> None:
         left = [s for s in {id(s): s for s in skipped}.values() if self._has_unsynced(s)]
         if left:
             unsynced = sum(1 for s in left for m in list(s.messages) if not m.get("_synced"))
             logger.warning("Honcho flush ran out of time after %.1fs with %d message(s) in %d session(s) still unsynced",
-                           timeout, unsynced, len(left))
+                           timeout or 0.0, unsynced, len(left))
 
     def _drain_async_queue(self, deadline: float | None = None) -> list[HonchoSession]:
         """Flush every queued session inline. Returns the sessions ``deadline`` left unflushed."""
@@ -585,22 +599,32 @@ class HonchoSessionManager(SessionAuthMixin, SessionPeersMixin, SessionContextMi
             self._async_thread.start()
 
     def stop_async_writer(self, timeout: float = 10.0) -> None:
-        """Join the async writer, then drain whatever was queued before the join. saveMessages: false never
-        enqueues, so the drain is a no-op there and the exit stays clean."""
+        """Join the async writer, then drain whatever was queued before the join, both within ``timeout``.
+        saveMessages: false never enqueues, so the drain is a no-op there and the exit stays clean."""
+        self._warn_unsynced(self._stop_async_writer_before(time.monotonic() + timeout), timeout)
+
+    def _stop_async_writer_before(self, deadline: float) -> list[HonchoSession]:
+        """Join the writer for what is left of ``deadline``, then drain the queue under the same deadline. A writer
+        still inside an upload keeps running to the client's HTTP timeout; it holds the shared client, so it is
+        joined, never abandoned."""
         with self._async_thread_lock:
             self._shutting_down = True
         if self._async_queue is not None and self._async_thread is not None and self._async_thread.is_alive():
             self._async_queue.put(_ASYNC_SHUTDOWN)
-            self._async_thread.join(timeout=timeout)
-        self._drain_async_queue()
+            self._async_thread.join(timeout=max(0.0, deadline - time.monotonic()))
+        return self._drain_async_queue(deadline)
 
     def shutdown(self, timeout: float = 10.0) -> None:
-        """Flush everything, then stop the async writer thread, both within ``timeout``."""
+        """Flush everything, then stop the async writer thread, within ``timeout``. The budget stops new uploads
+        from starting and bounds the lock waits and the join; an upload already in flight runs to the client's
+        HTTP timeout. Whatever stayed unsynced is counted in one warning."""
         self._shutting_down = True
         if self._async_queue is not None:
             deadline = time.monotonic() + timeout
-            self.flush_all(timeout=timeout)
-            self.stop_async_writer(timeout=max(0.0, deadline - time.monotonic()))
+            skipped = self._flush_cached_before(deadline)
+            skipped.extend(self._drain_async_queue(deadline))
+            skipped.extend(self._stop_async_writer_before(deadline))
+            self._warn_unsynced(skipped, timeout)
 
     # ----- Prefetch cache -----
 

tests/honcho_plugin/test_async_memory.py

@@ -10,7 +10,9 @@ Covers:
 """
 
 import json
+import logging
 import threading
+import time
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -380,11 +382,51 @@ class TestStopAsyncWriterDrain:
     def test_shutdown_gives_the_writer_join_what_the_flush_left_of_the_timeout(self, make_manager, monkeypatch):
         mgr = make_manager("async")
         seen = {}
-        monkeypatch.setattr(mgr, "stop_async_writer", lambda timeout=10.0: seen.setdefault("timeout", timeout))
+        monkeypatch.setattr(mgr, "_stop_async_writer_before",
+                            lambda deadline: seen.setdefault("remaining", deadline - time.monotonic()) and [])
 
         mgr.shutdown(timeout=2.5)
 
-        assert 2.0 < seen["timeout"] <= 2.5
+        assert 2.0 < seen["remaining"] <= 2.5
+
+    def _pending_session(self, mgr, uploads):
+        session = _make_session(key="pending")
+        session.add_message("user", "pending")
+        mgr._cache["pending"] = session
+        mgr._async_queue.put(session)
+        mgr._flush_session = lambda s: uploads.append(s.key) or True
+        mgr._flush_session_locked = lambda s: uploads.append(s.key) or True
+        return session
+
+    def test_shutdown_with_the_budget_spent_starts_no_upload_and_warns_once(self, make_manager, caplog):
+        """The SDK has no per-call timeout, so the budget can only stop uploads from starting. With no time left,
+        shutdown must not open one and must say what stayed behind."""
+        mgr = make_manager("async")
+        uploads = []
+        session = self._pending_session(mgr, uploads)
+
+        started = time.monotonic()
+        with caplog.at_level(logging.WARNING, logger="plugins.memory.honcho"):
+            mgr.shutdown(timeout=0)
+
+        assert uploads == []
+        assert time.monotonic() - started < 1.0
+        assert mgr._async_queue.empty()
+        assert session.messages[0].get("_synced") is None
+        assert caplog.text.count("still unsynced") == 1
+        assert "1 message(s) in 1 session(s) still unsynced" in caplog.text
+
+    def test_stop_async_writer_drains_only_within_its_timeout(self, make_manager, caplog):
+        mgr = make_manager("async")
+        uploads = []
+        self._pending_session(mgr, uploads)
+
+        with caplog.at_level(logging.WARNING, logger="plugins.memory.honcho"):
+            mgr.stop_async_writer(timeout=0)
+
+        assert uploads == []
+        assert mgr._async_queue.empty()
+        assert "1 message(s) in 1 session(s) still unsynced" in caplog.text
 
 
 class TestAsyncWriterRetry:
@@ -413,6 +455,32 @@ class TestAsyncWriterRetry:
         mgr.shutdown()
         assert call_count[0] == 2
 
+    def test_does_not_retry_once_shutdown_began(self, make_manager):
+        """The shutdown flush already attempts the session within its budget; a 2s sleep and a second upload from
+        the writer would run past it."""
+        mgr = make_manager(write_frequency="async")
+        mgr._ensure_async_writer()
+        sess = _make_session()
+        sess.add_message("user", "msg")
+        calls = []
+        failed = threading.Event()
+
+        def failing_flush(session):
+            calls.append(session)
+            failed.set()
+            return False
+
+        mgr._flush_session = failing_flush
+        mgr._shutting_down = True
+        mgr._async_queue.put(sess)
+        assert failed.wait(timeout=5), "async writer never picked up the batch"
+
+        started = time.monotonic()
+        mgr.stop_async_writer(timeout=5)
+
+        assert time.monotonic() - started < 1.5
+        assert len(calls) == 1
+
     def test_drops_after_two_failures(self, make_manager):
         mgr = make_manager(write_frequency="async")
         mgr._ensure_async_writer()