diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index 56320442e9..86545d5532 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/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, diff --git a/hindsight-api-slim/tests/test_memories_extension.py b/hindsight-api-slim/tests/test_memories_extension.py index 4dcf5aebb1..271de85a25 100644 --- a/hindsight-api-slim/tests/test_memories_extension.py +++ b/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) diff --git a/hindsight-api-slim/tests/test_retain_outbox_session.py b/hindsight-api-slim/tests/test_retain_outbox_session.py new file mode 100644 index 0000000000..12baeb9a8d --- /dev/null +++ b/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" diff --git a/hindsight-api-slim/tests/test_webhooks.py b/hindsight-api-slim/tests/test_webhooks.py index af423ac768..e807aa17a6 100644 --- a/hindsight-api-slim/tests/test_webhooks.py +++ b/hindsight-api-slim/tests/test_webhooks.py @@ -11,6 +11,7 @@ import hmac import json import uuid +from collections.abc import AsyncIterator, Iterator from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch @@ -21,6 +22,7 @@ from hindsight_api import LLMConfig from hindsight_api.api import create_app from hindsight_api.engine.memory_engine import MemoryEngine +from hindsight_api.engine.query_analyzer import QueryAnalyzer from hindsight_api.extensions import OperationValidationError from hindsight_api.webhooks.manager import MAX_ATTEMPTS, RETRY_DELAYS, WebhookManager from hindsight_api.webhooks.models import ( @@ -39,6 +41,46 @@ # --------------------------------------------------------------------------- +@pytest.fixture(params=["postgres", "buffered-store"]) +def retain_count_store(request: pytest.FixtureRequest) -> Iterator[None]: + """Exercise counts against both immediate SQL writes and a commit-buffered store.""" + from hindsight_api.engine.memories import get_memories, set_memories + from tests.test_memories_extension import InMemoryMemories + + original_store = get_memories() + if request.param == "buffered-store": + set_memories(InMemoryMemories({})) + try: + yield + finally: + set_memories(original_store) + + +@pytest_asyncio.fixture +async def retain_count_memory(pg0_db_url: str, query_analyzer: QueryAnalyzer) -> AsyncIterator[MemoryEngine]: + """Real retain/outbox persistence with deterministic, torch-free embeddings.""" + from hindsight_api.engine.task_backend import SyncTaskBackend + from tests.test_llm_reasoning_effort_env import DummyCrossEncoder + from tests.test_retain_same_document_concurrency import _StubEmbeddings + + memory = MemoryEngine( + db_url=pg0_db_url, + memory_llm_provider="mock", + memory_llm_api_key="", + memory_llm_model="mock", + embeddings=_StubEmbeddings(), + cross_encoder=DummyCrossEncoder(), + query_analyzer=query_analyzer, + run_migrations=False, + task_backend=SyncTaskBackend(), + ) + try: + await memory.initialize() + yield memory + finally: + await memory.close() + + def _make_event(bank_id: str = "bank-1") -> WebhookEvent: return WebhookEvent( event=WebhookEventType.CONSOLIDATION_COMPLETED, @@ -1685,12 +1727,16 @@ async def _extract_no_facts(*args, **kwargs): await memory.delete_bank(bank_id, request_context=request_context) @pytest.mark.asyncio - async def test_retain_completed_payload_carries_memory_unit_count(self, memory: MemoryEngine, request_context): + @pytest.mark.parametrize("use_factory", [False, True]) + async def test_retain_completed_payload_carries_memory_unit_count( + self, retain_count_memory: MemoryEngine, request_context, retain_count_store, use_factory: bool + ): """The event reports how many memory units the document owns afterwards. Without it a receiver cannot tell a document that produced memories from one that produced none — the event body is otherwise identical (#3040). """ + memory = retain_count_memory bank_id = f"wh-count-{uuid.uuid4().hex[:8]}" webhook_id = uuid.uuid4() original_manager = memory._webhook_manager @@ -1710,6 +1756,8 @@ async def test_retain_completed_payload_carries_memory_unit_count(self, memory: ) contents = [{"content": "Alice works at Google", "document_id": "doc-counted"}] + if use_factory: + contents.append({"content": "Bob works at Microsoft", "document_id": "doc-counted-other"}) callback = memory._build_retain_outbox_callback( bank_id=bank_id, contents=contents, operation_id="op-counted" ) @@ -1718,19 +1766,28 @@ async def test_retain_completed_payload_carries_memory_unit_count(self, memory: bank_id=bank_id, contents=contents, request_context=request_context, - outbox_callback=callback, + outbox_callback=None if use_factory else callback, + outbox_callback_factory=( + memory._build_retain_outbox_callback_factory(bank_id, "op-counted") if use_factory else None + ), ) - stored_units = ( - await memory.list_memory_units( - bank_id, document_id="doc-counted", limit=1000, request_context=request_context - ) - )["total"] - assert stored_units > 0, "fixture precondition: the mock LLM must extract facts here" - payloads = await self._retain_delivery_payloads(memory._pool, bank_id) - assert len(payloads) == 1 - assert payloads[0]["data"]["memory_unit_count"] == stored_units + assert len(payloads) == len(contents) + assert {payload["data"]["document_id"] for payload in payloads} == { + content["document_id"] for content in contents + } + for payload in payloads: + stored_units = ( + await memory.list_memory_units( + bank_id, + document_id=payload["data"]["document_id"], + limit=1000, + request_context=request_context, + ) + )["total"] + assert stored_units > 0, "fixture precondition: the mock LLM must extract facts here" + assert payload["data"]["memory_unit_count"] == stored_units finally: memory._webhook_manager = original_manager async with memory._pool.acquire() as conn: @@ -1743,7 +1800,7 @@ async def test_retain_completed_payload_carries_memory_unit_count(self, memory: @pytest.mark.asyncio async def test_retain_completed_payload_reports_zero_for_zero_fact_document( - self, memory: MemoryEngine, request_context, monkeypatch + self, retain_count_memory: MemoryEngine, request_context, monkeypatch, retain_count_store ): """A document that extracted nothing must report ``memory_unit_count: 0``. @@ -1756,6 +1813,7 @@ async def test_retain_completed_payload_reports_zero_for_zero_fact_document( from hindsight_api.engine.retain.types import ExtractionResult bank_id = f"wh-zerocount-{uuid.uuid4().hex[:8]}" + memory = retain_count_memory webhook_id = uuid.uuid4() original_manager = memory._webhook_manager @@ -1805,7 +1863,7 @@ async def _extract_no_facts(*args, **kwargs): @pytest.mark.asyncio async def test_retain_completed_payload_counts_document_not_units_created( - self, memory: MemoryEngine, request_context + self, retain_count_memory: MemoryEngine, request_context, retain_count_store ): """Re-retaining unchanged content must not look like a zero-fact document. @@ -1813,6 +1871,7 @@ async def test_retain_completed_payload_counts_document_not_units_created( units while the document keeps every memory it already had. Reporting units *created* would raise a false alarm on every idempotent re-submit. """ + memory = retain_count_memory bank_id = f"wh-delta-count-{uuid.uuid4().hex[:8]}" webhook_id = uuid.uuid4() original_manager = memory._webhook_manager