diff --git a/plugins/memory/hindsight/__init__.py b/plugins/memory/hindsight/__init__.py index 03632b98da71..801bfc8d6adf 100644 --- a/plugins/memory/hindsight/__init__.py +++ b/plugins/memory/hindsight/__init__.py @@ -2292,20 +2292,9 @@ def on_session_switch( # everything before mutating self._* so metadata + tags + doc_id # all reference the old session consistently. if self._session_turns: - old_turns = list(self._session_turns) old_session_id = self._session_id old_parent_session_id = self._parent_session_id old_turn_index = self._turn_index - old_metadata = self._build_metadata( - message_count=len(old_turns) * 2, - turn_index=old_turn_index, - ) - old_lineage_tags: list[str] = [] - if old_session_id: - old_lineage_tags.append(f"session:{old_session_id}") - if old_parent_session_id: - old_lineage_tags.append(f"parent:{old_parent_session_id}") - old_content = "[" + ",".join(old_turns) + "]" # Resolve doc_id + update_mode against the OLD session BEFORE # we rotate _session_id, so the flush lands in the old # session's document either way (legacy: per-process unique; @@ -2313,44 +2302,68 @@ def on_session_switch( old_document_id, old_update_mode = self._resolve_retain_target( self._document_id ) - - def _flush(): - try: - item = self._build_retain_kwargs( - old_content, - context=self._retain_context, - metadata=old_metadata, - tags=old_lineage_tags or None, - ) - item.pop("bank_id", None) - item.pop("retain_async", None) - if old_update_mode is not None: - item["update_mode"] = old_update_mode - logger.debug( - "Hindsight flush-on-switch: bank=%s, doc=%s, mode=%s, num_turns=%d", - self._bank_id, old_document_id, old_update_mode, len(old_turns), - ) - self._run_hindsight_operation( - lambda client: client.aretain_batch( - bank_id=self._bank_id, - items=[item], - document_id=old_document_id, - retain_async=self._retain_async, + # Only flush turns a completed retain hasn't already persisted. + # The watermark tracks that for both APIs: on append-capable + # servers ship just the delta since the last boundary (the + # server appends); on overwrite APIs re-send the whole session + # because each retain replaces the document. Skip entirely when + # nothing new is buffered so we never re-ship / duplicate + # already-retained turns at switch time. + turns_since_retain = self._session_turns[self._last_retained_turn_count:] + if turns_since_retain: + old_turns = ( + turns_since_retain + if old_update_mode == "append" + else list(self._session_turns) + ) + old_metadata = self._build_metadata( + message_count=len(old_turns) * 2, + turn_index=old_turn_index, + ) + old_lineage_tags: list[str] = [] + if old_session_id: + old_lineage_tags.append(f"session:{old_session_id}") + if old_parent_session_id: + old_lineage_tags.append(f"parent:{old_parent_session_id}") + old_content = "[" + ",".join(old_turns) + "]" + + def _flush(): + try: + item = self._build_retain_kwargs( + old_content, + context=self._retain_context, + metadata=old_metadata, + tags=old_lineage_tags or None, ) - ) - except Exception as e: - logger.warning("Hindsight flush-on-switch failed: %s", e, exc_info=True) - - # Route the flush through the same writer queue sync_turn - # uses. That serializes it behind any still-queued retains - # from the old session (FIFO by document_id), avoids racing - # two threads on aretain_batch against the same document, and - # keeps shutdown's drain semantics intact. Skip enqueue if - # shutdown has already fired — the writer is draining/gone. - if not self._shutting_down.is_set(): - self._ensure_writer() - self._register_atexit() - self._retain_queue.put(_flush) + item.pop("bank_id", None) + item.pop("retain_async", None) + if old_update_mode is not None: + item["update_mode"] = old_update_mode + logger.debug( + "Hindsight flush-on-switch: bank=%s, doc=%s, mode=%s, num_turns=%d", + self._bank_id, old_document_id, old_update_mode, len(old_turns), + ) + self._run_hindsight_operation( + lambda client: client.aretain_batch( + bank_id=self._bank_id, + items=[item], + document_id=old_document_id, + retain_async=self._retain_async, + ) + ) + except Exception as e: + logger.warning("Hindsight flush-on-switch failed: %s", e, exc_info=True) + + # Route the flush through the same writer queue sync_turn + # uses. That serializes it behind any still-queued retains + # from the old session (FIFO by document_id), avoids racing + # two threads on aretain_batch against the same document, + # and keeps shutdown's drain semantics intact. Skip enqueue + # if shutdown has already fired — the writer is draining/gone. + if not self._shutting_down.is_set(): + self._ensure_writer() + self._register_atexit() + self._retain_queue.put(_flush) # 2. Drain any in-flight prefetch from the old session and drop # its cached result so the new session doesn't see stale recall. diff --git a/tests/agent/test_memory_session_switch.py b/tests/agent/test_memory_session_switch.py index 2156ee475c03..8c2b5099098c 100644 --- a/tests/agent/test_memory_session_switch.py +++ b/tests/agent/test_memory_session_switch.py @@ -166,6 +166,11 @@ def _make_hindsight_provider(): provider._session_turns = ["turn-1", "turn-2"] provider._turn_counter = 2 provider._turn_index = 2 + # Retain watermark — set unconditionally by __init__ (and updated on + # every retain) so the buffer-flush path knows which turns have already + # been persisted. Seeded here gate-consistent with a provider that has + # not yet retained anything. + provider._last_retained_turn_count = 0 # Attrs read by _build_metadata / _build_retain_kwargs when the # buffer-flush path on session switch fires. Empty strings keep the # metadata minimal but well-formed. diff --git a/tests/plugins/memory/test_hindsight_provider.py b/tests/plugins/memory/test_hindsight_provider.py index d7114e00394e..e899e392a21e 100644 --- a/tests/plugins/memory/test_hindsight_provider.py +++ b/tests/plugins/memory/test_hindsight_provider.py @@ -1211,6 +1211,59 @@ def test_session_switch_flush_picks_capability_against_old_session( assert kw["items"][0]["update_mode"] == "append" +# --------------------------------------------------------------------------- + +class TestSessionSwitchBufferFlushNoReship: + """The session-switch flush must also respect the append watermark: + once a boundary retain persisted turns 1..N, a later switch must only + flush the partial block since the boundary — re-shipping the whole + buffer would duplicate already-retained turns in the document (same + data-loss/dup class as shutdown).""" + + @staticmethod + def _force_append(monkeypatch): + from plugins.memory.hindsight import ( + _append_capability_cache, + _append_capability_lock, + ) + with _append_capability_lock: + _append_capability_cache.clear() + monkeypatch.setattr( + "plugins.memory.hindsight._fetch_hindsight_api_version", + lambda *a, **kw: "0.5.6", + ) + + def test_switch_flush_does_not_reship_retained_turns( + self, provider_with_config, monkeypatch + ): + self._force_append(monkeypatch) + p = provider_with_config(retain_every_n_turns=3, retain_async=False) + client = p._client + + # Turn 3 hits the boundary -> retain dispatched, watermark advances. + p.sync_turn("t1-user", "t1-asst") + p.sync_turn("t2-user", "t2-asst") + p.sync_turn("t3-user", "t3-asst") + p._retain_queue.join() + assert client.aretain_batch.call_count == 1 + + # One more buffered turn, then a session switch. + p.sync_turn("t4-user", "t4-asst") + p.on_session_switch( + "new-sid", parent_session_id="test-session", reset=True + ) + p._retain_queue.join() + + # Exactly one extra aretain_batch, carrying ONLY the partial block. + assert client.aretain_batch.call_count == 2 + last_kw = client.aretain_batch.call_args_list[-1].kwargs + flat = json.dumps(last_kw["items"][0]["content"]) + assert "t4-user" in flat + assert "t1-user" not in flat + assert "t2-user" not in flat + assert "t3-user" not in flat + + # --------------------------------------------------------------------------- # System prompt tests # ---------------------------------------------------------------------------