From 772197417cddd27a32219a4117341285e766ed54 Mon Sep 17 00:00:00 2001 From: Aakash Kattelu Date: Mon, 24 Aug 2026 11:14:53 -0400 Subject: [PATCH] fix(deriver): harden retry metadata cleanup and stale reinforce fallback - Strip _retry_attempts from payloads in the same transaction as mark_queue_items_as_processed / mark_queue_item_as_errored - Clear shared retry metadata only after a successful terminal mark - On reinforce, if the locked target is gone or soft-deleted, insert the incoming document instead of dropping it - Skip pgvector semantic lookup when embedding is empty so query_documents cannot embed under an open session --- src/crud/document.py | 136 +++++++++++++++---------- src/deriver/queue_manager.py | 54 ++++++++-- tests/crud/test_document.py | 113 ++++++++++++++++++++ tests/deriver/test_queue_processing.py | 1 + 4 files changed, 242 insertions(+), 62 deletions(-) diff --git a/src/crud/document.py b/src/crud/document.py index c2fab22c..c3f06757 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -494,6 +494,9 @@ class _DocumentRowOp: kind: Literal["reinforce", "replace"] document_id: str incoming_times_derived: int = 1 + # When a reinforce skipped insert and the locked target is gone/deleted, + # insert this document instead of dropping it. + fallback_document: schemas.DocumentCreate | None = None @dataclass @@ -667,7 +670,12 @@ async def create_documents( current_td + 1, doc.times_derived ) row_ops.append( - _DocumentRowOp("reinforce", existing_match.id, doc.times_derived) + _DocumentRowOp( + "reinforce", + existing_match.id, + doc.times_derived, + fallback_document=doc, + ) ) exact_dup_existing_count += 1 continue @@ -703,54 +711,21 @@ async def create_documents( current_td + 1, doc.times_derived ) row_ops.append( - _DocumentRowOp("reinforce", existing_dup.id, doc.times_derived) + _DocumentRowOp( + "reinforce", + existing_dup.id, + doc.times_derived, + fallback_document=doc, + ) ) semantic_dup_rejected_count += 1 continue - metadata_dict = doc.metadata.model_dump(exclude_none=True) - - # Determine if we need to persist embeddings to postgres - # True when: TYPE=pgvector OR still migrating (dual-write to both stores) - store_embeddings_in_postgres = ( - settings.VECTOR_STORE.TYPE == "pgvector" - or not settings.VECTOR_STORE.MIGRATED + new_doc = _document_model_from_create( + doc, workspace_name=workspace_name, observer=observer, observed=observed ) - - if store_embeddings_in_postgres and doc.embedding: - new_doc = models.Document( - workspace_name=workspace_name, - observer=observer, - observed=observed, - content=doc.content, - level=doc.level, - times_derived=doc.times_derived, - internal_metadata=metadata_dict, - session_name=doc.session_name, - embedding=doc.embedding, - # Tree linkage column - source_ids=doc.source_ids, - ) - else: - new_doc = models.Document( - workspace_name=workspace_name, - observer=observer, - observed=observed, - content=doc.content, - level=doc.level, - times_derived=doc.times_derived, - internal_metadata=metadata_dict, - session_name=doc.session_name, - # Tree linkage column - source_ids=doc.source_ids, - ) - - if doc.embedding: - new_doc.sync_state = "pending" honcho_documents.append(new_doc) accepted_documents.append(doc) - - # Track embedding for vector store (ID will be available after commit) if doc.embedding: docs_with_embeddings.append((new_doc, doc.embedding)) @@ -771,13 +746,24 @@ async def create_documents( continue try: - await _apply_document_row_updates( + fallback_docs = await _apply_document_row_updates( db, row_ops, workspace_name=workspace_name, observer=observer, observed=observed, ) + for fallback_doc in fallback_docs: + new_doc = _document_model_from_create( + fallback_doc, + workspace_name=workspace_name, + observer=observer, + observed=observed, + ) + honcho_documents.append(new_doc) + accepted_documents.append(fallback_doc) + if fallback_doc.embedding: + docs_with_embeddings.append((new_doc, fallback_doc.embedding)) db.add_all(honcho_documents) # NOTE # If the process crashes after this commit but before vector upsert completes, @@ -1253,6 +1239,47 @@ async def create_observations( return honcho_documents +def _document_model_from_create( + doc: schemas.DocumentCreate, + *, + workspace_name: str, + observer: str, + observed: str, +) -> models.Document: + metadata_dict = doc.metadata.model_dump(exclude_none=True) + store_embeddings_in_postgres = ( + settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED + ) + if store_embeddings_in_postgres and doc.embedding: + new_doc = models.Document( + workspace_name=workspace_name, + observer=observer, + observed=observed, + content=doc.content, + level=doc.level, + times_derived=doc.times_derived, + internal_metadata=metadata_dict, + session_name=doc.session_name, + embedding=doc.embedding, + source_ids=doc.source_ids, + ) + else: + new_doc = models.Document( + workspace_name=workspace_name, + observer=observer, + observed=observed, + content=doc.content, + level=doc.level, + times_derived=doc.times_derived, + internal_metadata=metadata_dict, + session_name=doc.session_name, + source_ids=doc.source_ids, + ) + if doc.embedding: + new_doc.sync_state = "pending" + return new_doc + + async def _apply_document_row_updates( db: AsyncSession, ops: list[_DocumentRowOp], @@ -1260,10 +1287,10 @@ async def _apply_document_row_updates( workspace_name: str, observer: str, observed: str, -) -> None: - """Lock target rows by id, then apply reinforcements and replacements.""" +) -> list[schemas.DocumentCreate]: + """Lock target rows by id, apply ops, return fallbacks for vanished targets.""" if not ops: - return + return [] ids = sorted({op.document_id for op in ops}) result = await db.execute( select(models.Document) @@ -1279,17 +1306,21 @@ async def _apply_document_row_updates( ) locked = {doc.id: doc for doc in result.scalars()} now = datetime.datetime.now(datetime.timezone.utc) + fallbacks: list[schemas.DocumentCreate] = [] for op in ops: row = locked.get(op.document_id) - if row is None: - continue if op.kind == "replace": - row.deleted_at = now + if row is not None and row.deleted_at is None: + row.deleted_at = now continue - if row.deleted_at is not None: + # reinforce + if row is None or row.deleted_at is not None: + if op.fallback_document is not None: + fallbacks.append(op.fallback_document) continue row.times_derived = max(row.times_derived + 1, op.incoming_times_derived) await db.flush() + return fallbacks class SemanticRejectionResult(Enum): @@ -1322,6 +1353,9 @@ async def _semantic_dup_decision( filters=filters, ) elif _uses_pgvector(): + if not doc.embedding: + # Match external-store path: never embed under an open session. + return SemanticRejectionResult.NOT_DUPLICATE, None similar_docs = await query_documents( db=db, workspace_name=workspace_name, @@ -1331,7 +1365,7 @@ async def _semantic_dup_decision( filters=filters, max_distance=_SEMANTIC_DUP_MAX_DISTANCE, top_k=_SEMANTIC_DUP_TOP_K, - embedding=doc.embedding or None, + embedding=doc.embedding, ) else: return SemanticRejectionResult.NOT_DUPLICATE, None diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index ade79c80..3f828550 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -623,13 +623,15 @@ class QueueManager: ) return True - await self._clear_work_unit_retry_attempts(work_unit_key) error_msg = f"{error.__class__.__name__}: {str(error)}" try: if items: + # Clear retry metadata only after the terminal mark commits so a + # failed mark leaves the shared budget intact for the next claim. await self.mark_queue_item_as_errored( items[0], work_unit_key, error_msg ) + await self._clear_work_unit_retry_attempts(work_unit_key) except Exception as mark_error: logger.error( f"Failed to mark queue items as errored for work unit {work_unit_key}: {mark_error}", @@ -1159,8 +1161,16 @@ class QueueManager: ) await db.commit() + @staticmethod + def _payload_without_retry_attempts( + payload: dict[str, Any] | None, + ) -> dict[str, Any]: + cleaned = dict(payload or {}) + cleaned.pop(_RETRY_ATTEMPTS_PAYLOAD_KEY, None) + return cleaned + async def _clear_work_unit_retry_attempts(self, work_unit_key: str) -> None: - """Drop the shared attempt count after success or terminal failure.""" + """Drop the shared attempt count from remaining unprocessed items.""" async with tracked_db("clear_work_unit_retry_attempts") as db: result = await db.execute( select(models.QueueItem) @@ -1174,12 +1184,10 @@ class QueueManager: payload = item.payload or {} if _RETRY_ATTEMPTS_PAYLOAD_KEY not in payload: continue - new_payload = dict(payload) - new_payload.pop(_RETRY_ATTEMPTS_PAYLOAD_KEY, None) await db.execute( update(models.QueueItem) .where(models.QueueItem.id == item.id) - .values(payload=new_payload) + .values(payload=self._payload_without_retry_attempts(payload)) ) await db.commit() @@ -1191,12 +1199,23 @@ class QueueManager: async with tracked_db("process_queue_item_batch") as db: work_unit = parse_work_unit_key(work_unit_key) item_ids = [item.id for item in items] - await db.execute( - update(models.QueueItem) + result = await db.execute( + select(models.QueueItem) .where(models.QueueItem.id.in_(item_ids)) .where(models.QueueItem.work_unit_key == work_unit_key) - .values(processed=True) + .with_for_update() ) + for queue_item in result.scalars(): + await db.execute( + update(models.QueueItem) + .where(models.QueueItem.id == queue_item.id) + .values( + processed=True, + payload=self._payload_without_retry_attempts( + queue_item.payload + ), + ) + ) await db.execute( update(models.ActiveQueueSession) .where(models.ActiveQueueSession.work_unit_key == work_unit_key) @@ -1222,11 +1241,24 @@ class QueueManager: if not item: return async with tracked_db("mark_queue_item_as_errored") as db: - await db.execute( - update(models.QueueItem) + result = await db.execute( + select(models.QueueItem) .where(models.QueueItem.id == item.id) .where(models.QueueItem.work_unit_key == work_unit_key) - .values(processed=True, error=error[:65535]) # Truncate to TEXT limit + .with_for_update() + ) + queue_item = result.scalar_one_or_none() + if queue_item is None: + await db.commit() + return + await db.execute( + update(models.QueueItem) + .where(models.QueueItem.id == queue_item.id) + .values( + processed=True, + error=error[:65535], # Truncate to TEXT limit + payload=self._payload_without_retry_attempts(queue_item.payload), + ) ) await db.execute( update(models.ActiveQueueSession) diff --git a/tests/crud/test_document.py b/tests/crud/test_document.py index b9699fe4..a5fa8ab9 100644 --- a/tests/crud/test_document.py +++ b/tests/crud/test_document.py @@ -1653,6 +1653,119 @@ class TestCreateDocumentsErrorHandling: "good fact two", ] + @pytest.mark.asyncio + async def test_empty_embedding_skips_semantic_without_embed( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + monkeypatch: pytest.MonkeyPatch, + ): + """Empty embeddings must not trigger embed() under an open session.""" + from src.config import settings + + test_workspace, test_peer = sample_data + test_peer2, test_session = await self._setup( + db_session, test_workspace, test_peer + ) + monkeypatch.setattr(settings.VECTOR_STORE, "TYPE", "pgvector") + monkeypatch.setattr(settings.VECTOR_STORE, "MIGRATED", True) + + empty = self._doc("fact without vector", test_session.name) + empty.embedding = [] + + with patch( + "src.crud.document.embedding_client.embed", + new_callable=AsyncMock, + ) as mock_embed: + result = await crud.create_documents( + db_session, + [empty], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + deduplicate=True, + ) + + assert len(result.created_documents) == 1 + mock_embed.assert_not_awaited() + + @pytest.mark.asyncio + async def test_stale_reinforce_target_falls_back_to_insert( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """If a reinforce target vanishes under lock, insert the incoming doc.""" + from src.crud import document as document_module + + test_workspace, test_peer = sample_data + test_peer2, test_session = await self._setup( + db_session, test_workspace, test_peer + ) + workspace_name = test_workspace.name + observer = test_peer.name + observed = test_peer2.name + session_name = test_session.name + + seeded = await crud.create_documents( + db_session, + [self._doc("shared fact", session_name)], + workspace_name=workspace_name, + observer=observer, + observed=observed, + ) + assert len(seeded.created_documents) == 1 + + existing = ( + await db_session.execute( + select(models.Document).where( + models.Document.workspace_name == workspace_name, + models.Document.observer == observer, + models.Document.observed == observed, + models.Document.deleted_at.is_(None), + ) + ) + ).scalar_one() + + real_apply = document_module._apply_document_row_updates # pyright: ignore[reportPrivateUsage] + + async def delete_then_apply(*args: Any, **kwargs: Any) -> Any: + existing.deleted_at = datetime.datetime.now(datetime.timezone.utc) + await db_session.flush() + return await real_apply(*args, **kwargs) + + with patch.object( + document_module, + "_apply_document_row_updates", + side_effect=delete_then_apply, + ): + result = await crud.create_documents( + db_session, + [self._doc("shared fact", session_name)], + workspace_name=workspace_name, + observer=observer, + observed=observed, + ) + + assert len(result.created_documents) == 1 + live = ( + ( + await db_session.execute( + select(models.Document).where( + models.Document.workspace_name == workspace_name, + models.Document.observer == observer, + models.Document.observed == observed, + models.Document.deleted_at.is_(None), + ) + ) + ) + .scalars() + .all() + ) + assert len(live) == 1 + assert live[0].id != existing.id + assert live[0].content == "shared fact" + class TestExternalCandidateHoist: """External-store dup candidates resolve before the first DB statement.""" diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index 933299ae..76e3e925 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -2118,6 +2118,7 @@ class TestQueueRetry: items = await self._fetch_items(db_session, work_unit_key) assert all(item.processed for item in items) assert all(item.error is None for item in items) + assert all("_retry_attempts" not in (item.payload or {}) for item in items) assert await self._retry_attempts_on_items(db_session, work_unit_key) is None async def test_retry_budget_survives_reclaim_by_another_manager(