diff --git a/plugins/memory/honcho/session.py b/plugins/memory/honcho/session.py index 9ba60d2b5d..bdc2dbd8ef 100644 --- a/plugins/memory/honcho/session.py +++ b/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 ----- diff --git a/tests/honcho_plugin/test_async_memory.py b/tests/honcho_plugin/test_async_memory.py index aaddf249e3..15ebcefe73 100644 --- a/tests/honcho_plugin/test_async_memory.py +++ b/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()