RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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
This commit is contained in:
Aakash Kattelu 2026-08-24 11:14:53 -04:00
parent 301564bb6b
commit 772197417c
4 changed files with 242 additions and 62 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

136
src/crud/document.py
View File

@ -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

54
src/deriver/queue_manager.py
View File

@ -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)

113
tests/crud/test_document.py
View File

@ -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."""

1
tests/deriver/test_queue_processing.py
View File

@ -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(