From 1be353bf9c1eae871f55597810e2c65bbd6472d1 Mon Sep 17 00:00:00 2001 From: And Date: Mon, 13 Jul 2026 02:30:05 +0300 Subject: [PATCH] fix(memory/hindsight): gate prefetch on server-side retain completion, not just queue drain Address PR #62871 review: with the default retain_async=True, aretain_batch returns when the write is accepted, not when it's durable/recall-visible, so draining the local writer queue (task_done) is not a read-after-write signal. The next-turn prefetch could still recall before the just-completed turn was observable on the server. - Track the async operation_id/operation_ids returned by aretain_batch - _wait_for_retains_drained now applies two ordered, budget-bounded barriers: (1) local writer queue drains, then (2) tracked server-side async ops report completion via operations.get_operation_status (an explicit read-after-write condition). NotFound (completed+evicted) counts as done; transient errors keep waiting until the deadline - Completed ops are removed from the pending set so later prefetches don't re-poll them; the whole wait stays off the reply path - Add TestPrefetchServerRetainVisibility: op-id tracking (single/multiple), no-op tracking when retain_async=False, prefetch waiting for server completion before recall, timeout fallback on a wedged op, and NotFound / transient-error status handling --- plugins/memory/hindsight/__init__.py | 176 +++++++++++++++--- .../plugins/memory/test_hindsight_provider.py | 131 +++++++++++++ 2 files changed, 282 insertions(+), 25 deletions(-) diff --git a/plugins/memory/hindsight/__init__.py b/plugins/memory/hindsight/__init__.py index b64992a425..c12238518c 100644 --- a/plugins/memory/hindsight/__init__.py +++ b/plugins/memory/hindsight/__init__.py @@ -727,6 +727,16 @@ class HindsightMemoryProvider(MemoryProvider): self._writer_thread: threading.Thread | None = None self._shutting_down = threading.Event() self._atexit_registered = False + # Server-side async retain operations still in flight. With + # retain_async=True, aretain_batch returns as soon as the write is + # *accepted*, not when it's durable/recall-visible, so the returned + # operation_id(s) stay "pending" until the server finishes. The + # background prefetch gates on these via get_operation_status so recall + # observes the just-completed turn (draining the local queue alone is + # not a read-after-write signal for async retains). + self._pending_retain_ops: set[str] = set() + self._pending_retain_ops_lock = threading.Lock() + self._retain_ops_bank_id = "" # Legacy alias — older tests/callers reference _sync_thread directly. # Points at _writer_thread once the writer is running. self._sync_thread = None @@ -745,10 +755,13 @@ class HindsightMemoryProvider(MemoryProvider): self._retain_async = True # Async retain never blocks the reply (writes drain on the single # writer thread). But the next turn's warm prefetch runs on its own - # thread and could read BEFORE the just-enqueued retain write lands, - # dropping the latest turn from recall. When True, the background - # prefetch waits (bounded) for pending retains to drain first, closing - # that race without putting any write on the reply path. + # thread and could read BEFORE the just-completed retain is + # recall-visible on the server, dropping the latest turn from recall. + # When True, the background prefetch first waits (bounded) for the + # local writer queue to drain AND for the server-side async retain + # operation(s) to report completion, an explicit read-after-write + # signal — closing that race without putting any write on the reply + # path. self._prefetch_waits_for_retain = True self._prefetch_retain_drain_timeout = 10.0 self._retain_context = "conversation between Hermes Agent and the User" @@ -1067,8 +1080,8 @@ class HindsightMemoryProvider(MemoryProvider): {"key": "auto_retain", "description": "Automatically retain conversation turns", "default": True}, {"key": "retain_every_n_turns", "description": "Retain every N turns (1 = every turn)", "default": 1}, {"key": "retain_async","description": "Process retain asynchronously on the Hindsight server", "default": True}, - {"key": "prefetch_waits_for_retain", "description": "Have the background next-turn prefetch wait for pending retain writes to drain first, so recall includes the just-completed turn (runs off the reply path, adds no response latency)", "default": True}, - {"key": "prefetch_retain_drain_timeout", "description": "Max seconds the background prefetch waits for pending retains to drain before recalling anyway", "default": 10.0}, + {"key": "prefetch_waits_for_retain", "description": "Have the background next-turn prefetch wait for the just-completed retain to become recall-visible on the server (local queue drain + async operation completion) before recalling, so recall includes the just-completed turn (runs off the reply path, adds no response latency)", "default": True}, + {"key": "prefetch_retain_drain_timeout", "description": "Max seconds the background prefetch waits for the retain to become recall-visible (queue drain + server-side completion) before recalling anyway", "default": 10.0}, {"key": "retain_context", "description": "Context label for retained memories", "default": "conversation between Hermes Agent and the User"}, {"key": "recall_max_tokens", "description": "Maximum tokens for recall results", "default": 4096}, {"key": "recall_max_input_chars", "description": "Maximum input query length for auto-recall", "default": 800}, @@ -1173,31 +1186,136 @@ class HindsightMemoryProvider(MemoryProvider): self._sync_thread = thread thread.start() - def _wait_for_retains_drained(self, timeout: float) -> bool: - """Block up to *timeout* seconds for queued retains to finish. + def _track_retain_ops(self, retain_response, bank_id: str) -> None: + """Record server-side async operation id(s) from an aretain_batch reply. - Used by the background prefetch so the next turn's recall observes - the just-completed turn's write instead of racing ahead of it. Runs - only on the background prefetch thread — never on the reply path. - - Returns True if the queue drained within the budget, False on timeout. - Polls ``unfinished_tasks`` rather than ``queue.join()`` so the wait is - bounded even if a write is wedged (the writer is best-effort). + Async retains return ``operation_id`` / ``operation_ids`` that stay + ``pending`` on the server until the write is durable and recall-visible. + The bank_id is captured alongside so completion can be polled with the + same bank the write targeted. """ - if timeout <= 0: - return self._retain_queue.unfinished_tasks == 0 - deadline = time.monotonic() + timeout + ids: list[str] = [] + single = getattr(retain_response, "operation_id", None) + if single: + ids.append(str(single)) + multiple = getattr(retain_response, "operation_ids", None) + if multiple: + ids.extend(str(op) for op in multiple if op) + if not ids: + # Server didn't hand back an op id (older API, or it completed + # synchronously). Nothing to poll — local queue drain is the only + # available signal in that case. + return + self._retain_ops_bank_id = bank_id + with self._pending_retain_ops_lock: + self._pending_retain_ops.update(ids) + + def _is_retain_op_complete(self, bank_id: str, op_id: str) -> bool: + """Return True when a server-side async retain op is done (or gone). + + ``get_operation_status`` returns ``completed``/``failed`` for a known + op; completed ops are evicted server-side, so a NotFound (404) also + means "no longer pending" and is treated as done. Transient errors + return False so the caller keeps waiting until its deadline. + """ + from hindsight_client_api.exceptions import NotFoundException + + try: + resp = self._run_hindsight_operation( + lambda client: client.operations.get_operation_status( + bank_id=bank_id, operation_id=op_id + ) + ) + except NotFoundException: + return True + except Exception as exc: + logger.debug("Prefetch: operation status check failed for %s: %s", op_id, exc) + return False + status = str(getattr(resp, "status", "") or "").lower() + return status in {"completed", "failed"} + + def _wait_for_retains_drained(self, timeout: float) -> bool: + """Block up to *timeout* seconds for the just-completed turn's retain to + become recall-visible on the server. + + Used by the background prefetch so the next turn's recall observes the + just-completed turn's write instead of racing ahead of it. Runs only on + the background prefetch thread — never on the reply path. + + Two ordered barriers, both bounded by the shared *timeout* budget: + + 1. Local writer queue drains (the retain call has been *dispatched* to + the server). Polls ``unfinished_tasks`` rather than ``queue.join()`` + so a wedged write can't hang the prefetch. + 2. Server-side async operations complete. With ``retain_async=True`` the + dispatched call returns on *acceptance*, not durability, so draining + the local queue alone is NOT a read-after-write signal. We poll + ``get_operation_status`` for the tracked op id(s) until the server + reports completion (an explicit read-after-write condition). + + Returns True if both barriers cleared within the budget, False on + timeout/shutdown. + """ + deadline = None if timeout <= 0 else time.monotonic() + timeout + + def _expired() -> bool: + return deadline is not None and time.monotonic() >= deadline + + # Barrier 1: local queue drain (retain dispatched to the server). while self._retain_queue.unfinished_tasks > 0: if self._shutting_down.is_set(): return False - if time.monotonic() >= deadline: + if _expired(): logger.debug( "Prefetch: retain drain timed out after %.1fs (%d pending)", timeout, self._retain_queue.unfinished_tasks, ) return False time.sleep(0.05) - return True + + # Barrier 2: server-side async retain completion (read-after-write). + return self._wait_for_server_retain_ops(deadline, timeout) + + def _wait_for_server_retain_ops(self, deadline: float | None, timeout: float) -> bool: + """Poll tracked async retain ops until complete or the deadline passes. + + *deadline* is a ``time.monotonic()`` value (None = no bound). Completed + ops are removed from the pending set as they finish so a later prefetch + doesn't re-poll them. + """ + while True: + with self._pending_retain_ops_lock: + bank_id = getattr(self, "_retain_ops_bank_id", "") or self._bank_id + pending = list(self._pending_retain_ops) + if not pending: + return True + if self._shutting_down.is_set(): + return False + + done: set[str] = set() + for op_id in pending: + if self._shutting_down.is_set(): + return False + if deadline is not None and time.monotonic() >= deadline: + logger.debug( + "Prefetch: server retain visibility timed out after %.1fs " + "(%d op(s) still pending)", + timeout, len(pending) - len(done), + ) + with self._pending_retain_ops_lock: + self._pending_retain_ops.difference_update(done) + return False + if self._is_retain_op_complete(bank_id, op_id): + done.add(op_id) + + with self._pending_retain_ops_lock: + self._pending_retain_ops.difference_update(done) + still_pending = bool(self._pending_retain_ops) + if not still_pending: + return True + if deadline is not None and time.monotonic() >= deadline: + return False + time.sleep(0.05) def _writer_loop(self) -> None: """Drain the retain queue serially. Exits on sentinel. @@ -1589,10 +1707,13 @@ class HindsightMemoryProvider(MemoryProvider): query = query[:self._recall_max_input_chars] def _run(): - # Ensure the just-completed turn's retain has landed before we - # recall, so the warmed context for the next turn includes it. - # This runs on the background prefetch thread, never the reply - # path, so it adds no latency to the user's response. + # Ensure the just-completed turn's retain is recall-visible on the + # server before we recall, so the warmed context for the next turn + # includes it. This waits for the local writer queue to drain AND + # for the server-side async retain op(s) to complete (an explicit + # read-after-write signal), because async retain returns on + # acceptance rather than durability. Runs on the background prefetch + # thread, never the reply path, so it adds no response latency. if self._prefetch_waits_for_retain: self._wait_for_retains_drained(self._prefetch_retain_drain_timeout) try: @@ -1776,7 +1897,7 @@ class HindsightMemoryProvider(MemoryProvider): item["update_mode"] = update_mode logger.debug("Hindsight retain: bank=%s, doc=%s, mode=%s, async=%s, content_len=%d, num_turns=%d", bank_id, document_id, update_mode, retain_async_flag, len(content), num_turns) - self._run_hindsight_operation( + resp = self._run_hindsight_operation( lambda client: client.aretain_batch( bank_id=bank_id, items=[item], @@ -1784,6 +1905,11 @@ class HindsightMemoryProvider(MemoryProvider): retain_async=retain_async_flag, ) ) + # For async retains the write is only *accepted* here; track the + # returned operation id(s) so the next-turn prefetch can wait for + # true server-side completion (read-after-write) before recalling. + if retain_async_flag: + self._track_retain_ops(resp, bank_id) logger.debug("Hindsight retain succeeded") self._ensure_writer() diff --git a/tests/plugins/memory/test_hindsight_provider.py b/tests/plugins/memory/test_hindsight_provider.py index 092061830e..fef2bf3e7f 100644 --- a/tests/plugins/memory/test_hindsight_provider.py +++ b/tests/plugins/memory/test_hindsight_provider.py @@ -543,6 +543,137 @@ class TestPrefetch: assert p._prefetch_waits_for_retain is False +class TestPrefetchServerRetainVisibility: + """PR #62871 review follow-up: draining the local writer queue is not a + read-after-write signal for async retains. With ``retain_async=True`` the + server accepts the write and returns an ``operation_id`` that stays + ``pending`` until the write is durable/recall-visible. The background + prefetch must gate on server-side operation completion, not just the local + queue, before recalling. + """ + + def _client_with_ops(self, statuses): + """Mock client whose aretain_batch returns an async operation_id and + whose operations.get_operation_status yields *statuses* in order + (last value repeats).""" + client = _make_mock_client() + client.aretain_batch = AsyncMock( + return_value=SimpleNamespace(operation_id="op-1", operation_ids=None) + ) + seq = list(statuses) + + async def _status(**kwargs): + value = seq.pop(0) if len(seq) > 1 else seq[0] + return SimpleNamespace(status=value) + + client.operations = MagicMock() + client.operations.get_operation_status = AsyncMock(side_effect=_status) + return client + + def test_tracks_async_operation_id_from_retain(self, provider): + provider._client.aretain_batch = AsyncMock( + return_value=SimpleNamespace(operation_id="op-async-1", operation_ids=None) + ) + provider.sync_turn("hello", "world") + provider._retain_queue.join() + assert "op-async-1" in provider._pending_retain_ops + + def test_tracks_multiple_operation_ids(self, provider): + provider._client.aretain_batch = AsyncMock( + return_value=SimpleNamespace( + operation_id=None, operation_ids=["op-a", "op-b"] + ) + ) + provider.sync_turn("hello", "world") + provider._retain_queue.join() + assert {"op-a", "op-b"} <= provider._pending_retain_ops + + def test_sync_retain_tracks_no_ops(self, provider_with_config): + p = provider_with_config(retain_async=False) + p._client = _make_mock_client() + p._client.aretain_batch = AsyncMock( + return_value=SimpleNamespace(operation_id="op-x", operation_ids=None) + ) + p.sync_turn("hello", "world") + p._retain_queue.join() + # retain_async=False → no server-side op to wait on. + assert p._pending_retain_ops == set() + + def test_prefetch_waits_for_server_completion_before_recall(self, provider): + """Recall must not run until the tracked async op reports completed.""" + order = [] + + async def _recall(**kwargs): + order.append("recall") + return SimpleNamespace(results=[SimpleNamespace(text="m")]) + + provider._client = self._client_with_ops(["pending", "pending", "completed"]) + provider._client.arecall = AsyncMock(side_effect=_recall) + + provider.sync_turn("hello", "world") + provider._retain_queue.join() + assert "op-1" in provider._pending_retain_ops + + provider.queue_prefetch("next turn query") + if provider._prefetch_thread: + provider._prefetch_thread.join(timeout=5.0) + + # Recall ran, the op was polled to completion, and the pending set + # was cleared (so a later prefetch won't re-poll it). + assert order == ["recall"] + assert provider._client.operations.get_operation_status.await_count >= 3 + assert provider._pending_retain_ops == set() + + def test_prefetch_proceeds_after_server_wait_timeout(self, provider_with_config): + """A wedged/never-completing async op must not hang prefetch forever; + it recalls anyway once the drain budget is exhausted.""" + p = provider_with_config(prefetch_retain_drain_timeout=0.3) + order = [] + + async def _recall(**kwargs): + order.append("recall") + return SimpleNamespace(results=[SimpleNamespace(text="m")]) + + p._client = self._client_with_ops(["pending"]) # never completes + p._client.arecall = AsyncMock(side_effect=_recall) + + p.sync_turn("hello", "world") + p._retain_queue.join() + + start = time.monotonic() + p.queue_prefetch("next turn query") + if p._prefetch_thread: + p._prefetch_thread.join(timeout=5.0) + elapsed = time.monotonic() - start + + assert order == ["recall"], "prefetch should recall after the timeout" + assert elapsed < 3.0, "prefetch must not block well past the drain budget" + + def test_operation_notfound_treated_as_complete(self, provider): + """A NotFound (completed+evicted) op is treated as done, not pending.""" + from hindsight_client_api.exceptions import NotFoundException + + client = _make_mock_client() + client.operations = MagicMock() + client.operations.get_operation_status = AsyncMock( + side_effect=NotFoundException(status=404, reason="gone") + ) + provider._client = client + + assert provider._is_retain_op_complete("bank", "op-gone") is True + + def test_transient_status_error_keeps_waiting(self, provider): + """A transient status-check error means 'unknown', so keep waiting.""" + client = _make_mock_client() + client.operations = MagicMock() + client.operations.get_operation_status = AsyncMock( + side_effect=RuntimeError("temporary blip") + ) + provider._client = client + + assert provider._is_retain_op_complete("bank", "op-1") is False + + # --------------------------------------------------------------------------- # sync_turn tests # ---------------------------------------------------------------------------