fix: defer retain completion outbox until store commit

hindsight-api-slim/hindsight_api/engine/memory_engine.py

@@ -5773,6 +5773,31 @@ async def _run_retain_execution(
             _session_config = await self._resolve_retain_config(bank_id, request_context, strategy)
             retain_session = await _store.begin_retain(bank_id=bank_id, config=_session_config)
 
+        pending_outbox_callbacks: list[RetainOutboxCallback] = []
+        if retain_session is not None:
+            # A session may buffer all memories until commit. Running the outbox in the
+            # pipeline's SQL transaction used to publish a pre-commit document count (#4189).
+            # Record only callbacks the pipeline actually reaches (including per-document
+            # factories), and run them after the owning session has committed successfully.
+            # Stores without sessions and the SQL path keep their existing transaction boundary.
+            def defer_outbox(callback: RetainOutboxCallback | None) -> RetainOutboxCallback | None:
+                if callback is None:
+                    return None
+
+                async def enqueue(_conn: asyncpg.Connection) -> None:
+                    pending_outbox_callbacks.append(callback)
+
+                return enqueue
+
+            outbox_callback = defer_outbox(outbox_callback)
+            if outbox_callback_factory is not None:
+                original_factory = outbox_callback_factory
+
+                def deferred_factory(callback_contents: list[RetainContentDict]) -> RetainOutboxCallback | None:
+                    return defer_outbox(original_factory(callback_contents))
+
+                outbox_callback_factory = deferred_factory
+
         # A store that owns persistence does NOT sub-batch. Splitting exists to bound what one
         # unit of work holds and to give the sub-batches something to run concurrently over — and
         # neither survives the session: the session buffers until commit either way, so slicing no
@@ -6103,6 +6128,31 @@ async def _run_sub(idx: int, contents_, origins_, offset_, is_last_, body_, body
             # Progress for this path is emitted by the streaming pipeline as
             # "storing N/total chunks" via progress_callback (see _retain_batch_async_internal).
 
+        if pending_outbox_callbacks:
+            # Do not hold a SQL connection while committing an external store. The outbox
+            # still lives in SQL, but cannot share a transaction with that store's commit.
+            #
+            # The memories are already committed in the store by the time this runs, so a
+            # failure here cannot be undone by failing the retain — it would only mark a
+            # successful retain as failed and invite the caller to re-submit a document that
+            # is already stored. The event is the lossy side of a boundary that is not
+            # transactional either way (see the deferral note above): log it loudly and let
+            # the retain report the truth, which is that it succeeded.
+            try:
+                backend = await self._get_backend()
+                async with acquire_with_retry(backend) as conn:
+                    async with conn.transaction():
+                        for callback in pending_outbox_callbacks:
+                            await callback(conn)
+            except Exception:
+                logger.error(
+                    "[BATCH_RETAIN] bank=%s operation=%s retained successfully but the retain.completed "
+                    "outbox write failed; the completion event is lost for this operation",
+                    bank_id,
+                    operation_id,
+                    exc_info=True,
+                )
+
         return _RetainExecutionResult(
             unit_ids=result,
             usage=total_usage,

hindsight-api-slim/tests/test_memories_extension.py

@@ -420,6 +420,8 @@ async def delete_stale_observations(self, *, conn, ops, fq_table, bank_id, fact_
     async def list_memory_units(self, *, conn, ops, fq_table, bank_id, limit=100, offset=0, **kwargs):
         self.calls.append("list_memory_units")
         ordered = list(self.rows.values())
+        if kwargs.get("document_id") is not None:
+            ordered = [row for row in ordered if row.document_id == kwargs["document_id"]]
         return {"items": ordered[offset : offset + limit], "total": len(ordered), "limit": limit, "offset": offset}
 
     async def get_memory_unit(self, *, conn, ops, fq_table, bank_id, unit_id):
@@ -759,19 +761,40 @@ async def commit(self) -> RetainResult:
         self._store.calls.append("session.commit")
         unit_ids: dict[str, list[str]] = {}
         for part in self._parts:
-            doc = self._store.documents.setdefault(part.document_id, {"chunks": [], "text": ""})
+            # Session parts carry FactRecord objects and the store's document schema,
+            # not the SQL pipeline's processed facts or an alternate chunks/text shape.
+            if part.document_id not in self._store.documents:
+                await self._store.put_document(
+                    bank_id=self._bank_id,
+                    document_id=part.document_id,
+                    content_hash=part.content_hash,
+                    original_text=part.document_body or "",
+                    chunk_texts=[],
+                    tags=part.tags,
+                    metadata=part.metadata,
+                )
+            doc = self._store.documents[part.document_id]
             if part.document_body is not None:
-                doc["text"] = part.document_body
+                doc["original_text"] = part.document_body
+            doc["content_hash"] = part.content_hash
             if part.chunk_texts:
                 # `chunk_offset` is per document, so a part is placed at its offset rather than
                 # appended — two parts of one document can arrive in either order.
                 needed = part.chunk_offset + len(part.chunk_texts)
-                if len(doc["chunks"]) < needed:
-                    doc["chunks"].extend([""] * (needed - len(doc["chunks"])))
-                doc["chunks"][part.chunk_offset : needed] = list(part.chunk_texts)
+                if len(doc["chunk_texts"]) < needed:
+                    doc["chunk_texts"].extend([""] * (needed - len(doc["chunk_texts"])))
+                doc["chunk_texts"][part.chunk_offset : needed] = list(part.chunk_texts)
             if part.facts:
-                ids = self._store.allocate_unit_ids(len(part.facts))
-                await self._store.index_facts(self._bank_id, ids, part.facts, part.document_id)
+                ids = [fact.unit_id for fact in part.facts]
+                for fact in part.facts:
+                    self._store.rows[fact.unit_id] = StoredMemory(
+                        unit_id=fact.unit_id,
+                        text=fact.text,
+                        fact_type=fact.fact_type,
+                        document_id=part.document_id,
+                        tags=list(fact.tags),
+                        created_at=fact.created_at or datetime.now(timezone.utc),
+                    )
                 unit_ids.setdefault(part.document_id, []).extend(ids)
         self._parts.clear()
         return RetainResult(unit_ids=unit_ids)

hindsight-api-slim/tests/test_retain_outbox_session.py

@@ -0,0 +1,124 @@
+"""Completion events must observe committed store-owned memories."""
+
+import logging
+import time
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from hindsight_api import RequestContext
+from hindsight_api.engine.memory_engine import MemoryEngine
+from hindsight_api.engine.response_models import TokenUsage
+from hindsight_api.engine.retain.types import RetainBatchResult, RetainContentDict
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("use_factory", [False, True])
+@pytest.mark.parametrize("previous_count,committed_count", [(0, 4), (4, 4), (0, 0)])
+@pytest.mark.parametrize("failure", [None, "commit", "retain", "outbox"])
+async def test_completion_counts_committed_document(
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+    use_factory: bool,
+    previous_count: int,
+    committed_count: int,
+    failure: str | None,
+) -> None:
+    engine = MemoryEngine.__new__(MemoryEngine)
+    fire_event = AsyncMock(side_effect=RuntimeError("outbox failed") if failure == "outbox" else None)
+    engine._webhook_manager = MagicMock(fire_event_with_conn=fire_event)
+    engine._resolve_retain_config = AsyncMock()
+    # Tokenization is unrelated to completion ordering; keep this unit test offline.
+    monkeypatch.setattr("hindsight_api.engine.memory_engine.count_tokens", lambda text: 1)
+    conn = MagicMock()
+    conn.transaction.return_value.__aenter__ = AsyncMock()
+    conn.transaction.return_value.__aexit__ = AsyncMock(return_value=False)
+    steps: list[str] = []
+    visible_count = previous_count
+
+    async def commit() -> None:
+        nonlocal visible_count
+        steps.append("commit")
+        if failure == "commit":
+            raise RuntimeError("commit failed")
+        visible_count = committed_count
+
+    async def count(**kwargs) -> dict[str, int]:
+        assert kwargs["bank_id"] == "test-bank"
+        assert kwargs["document_ids"] == ["test-document"]
+        steps.append("count")
+        return {"test-document": visible_count}
+
+    session = MagicMock(commit=AsyncMock(side_effect=commit))
+    store = MagicMock(
+        store_owned_for=MagicMock(return_value=True),
+        begin_retain=AsyncMock(return_value=session),
+        document_memory_counts=AsyncMock(side_effect=count),
+    )
+    monkeypatch.setattr("hindsight_api.engine.memories.get_memories", lambda: store)
+
+    @asynccontextmanager
+    async def acquire(_backend) -> AsyncIterator[MagicMock]:
+        yield conn
+
+    monkeypatch.setattr("hindsight_api.engine.memory_engine.acquire_with_retry", acquire)
+    engine._get_backend = AsyncMock()
+    contents: list[RetainContentDict] = [{"content": "Alice works at Google", "document_id": "test-document"}]
+
+    async def retain(**kwargs) -> RetainBatchResult:
+        assert kwargs["retain_session"] is session
+        callback = kwargs["outbox_callback_factory"](contents) if use_factory else kwargs["outbox_callback"]
+        assert callback is not None
+        await callback(conn)
+        if failure == "retain":
+            raise RuntimeError("retain failed")
+        # Unchanged retains create no units, but must still report the existing total.
+        return RetainBatchResult([[]], TokenUsage(), 0)
+
+    engine._retain_batch_async_internal = AsyncMock(side_effect=retain)
+    callback = engine._build_retain_outbox_callback("test-bank", contents, "test-operation", schema="test-schema")
+    factory = engine._build_retain_outbox_callback_factory("test-bank", "test-operation", schema="test-schema")
+    execution = engine._run_retain_execution(
+        bank_id="test-bank",
+        contents=contents,
+        request_context=RequestContext(),
+        document_id=None,
+        fact_type_override=None,
+        document_tags=None,
+        operation_id="test-operation",
+        strategy=None,
+        outbox_callback=None if use_factory else callback,
+        outbox_callback_factory=factory if use_factory else None,
+        start_time=time.time(),
+    )
+
+    if failure == "outbox":
+        # The store has already committed by the time the outbox runs, so a failed
+        # delivery-row write cannot be undone by failing the retain — that would only
+        # report a stored document as lost and invite a duplicate re-submit. The retain
+        # succeeds and the dropped event is logged.
+        with caplog.at_level(logging.ERROR, logger="hindsight_api.engine.memory_engine"):
+            await execution
+        assert steps == ["commit", "count"]
+        assert "outbox write failed" in caplog.text
+        assert "test-operation" in caplog.text
+        return
+
+    if failure:
+        with pytest.raises(RuntimeError, match=f"{failure} failed"):
+            await execution
+        # The partial-retain cleanup still commits, but no success event may escape.
+        assert steps == ["commit"]
+        engine._webhook_manager.fire_event_with_conn.assert_not_awaited()
+        return
+
+    await execution
+    assert steps == ["commit", "count"]
+    engine._webhook_manager.fire_event_with_conn.assert_awaited_once()
+    call = engine._webhook_manager.fire_event_with_conn.await_args
+    assert call.args[0].data.memory_unit_count == committed_count
+    assert call.args[0].bank_id == "test-bank"
+    assert call.args[0].operation_id == "test-operation"
+    assert call.kwargs["schema"] == "test-schema"