ouyangbo/hermes-agent
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

fix(memory/hindsight): gate prefetch on server-side retain completion, not just queue drain

Browse Source 
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
This commit is contained in:
And
2026-07-13 02:30:05 +03:00
committed by kshitij
parent 94b10eccf5
commit 1be353bf9c
2 changed files with 282 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
plugins/memory/hindsight/__init__.py
+151 -25
View File
@@ -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()
tests/plugins/memory/test_hindsight_provider.py
+131
View File
@@ -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
# ---------------------------------------------------------------------------