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

fix(deriver): eliminate create_documents deadlock and stop silently burning batches on transient errors 

Two concurrent work units writing the same (workspace, observer, observed)
collection deadlocked on times_derived reinforcement UPDATEs issued in
batch order (DEV-1975, 682 events in 90 days). The deadlock was swallowed
per-document, the loop cascaded PendingRollbackErrors against the dead
session, the whole batch was lost, and the queue item was marked processed.

- serialize writers per collection with a transaction-scoped advisory lock
  (pg_advisory_xact_lock + SET LOCAL lock_timeout), skipped for insert-only
  batches; covers all three row-lock sites in one move
- hoist external-vector-store dup-candidate resolution ahead of the first
  DB statement so the lock's critical section contains no network calls
- abort the batch on SQLAlchemyError instead of continuing through an
  aborted transaction; per-document skip semantics kept for non-DB errors
- classify transient errors (new src/utils/retryable_errors.py) and retry
  them via a bounded in-process counter instead of marking items errored
This commit is contained in:
Aakash Kattelu 2026-08-18 11:04:55 -04:00
parent 9379c634ed
commit 8e0c3afd75
6 changed files with 1090 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

144
src/crud/document.py
View File

@ -1,3 +1,4 @@
import asyncio
import datetime
from collections.abc import Sequence
from dataclasses import dataclass, field
@ -5,9 +6,9 @@ from enum import Enum
from logging import getLogger
from typing import Any, cast
from sqlalchemy import delete, select, update
from sqlalchemy import delete, select, text, update
from sqlalchemy.engine import CursorResult
from sqlalchemy.exc import IntegrityError
from sqlalchemy.exc import DBAPIError, IntegrityError, SQLAlchemyError
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.sql import Select
from sqlalchemy.sql.functions import func
@ -210,6 +211,24 @@ def _uses_pgvector() -> bool:
)
# Semantic-dup candidate search parameters, shared by is_rejected_duplicate
# and the pre-lock candidate resolution in create_documents.
_SEMANTIC_DUP_MAX_DISTANCE = 0.05
_SEMANTIC_DUP_TOP_K = 1
def _semantic_dup_filters(doc: schemas.DocumentCreate) -> dict[str, Any] | None:
"""Merge scope for semantic dedup: never across levels, never across
sessions for explicit documents. None when the document has no valid
merge partner (session-less explicit)."""
filters: dict[str, Any] = {"level": doc.level}
if doc.level == "explicit":
if doc.session_name is None:
return None
filters["session_name"] = doc.session_name
return filters
async def query_external_vector_document_ids(
workspace_name: str,
observer: str,
@ -512,6 +531,32 @@ async def create_documents(
# Store (document_model, embedding) pairs - IDs aren't available until after commit
docs_with_embeddings: list[tuple[models.Document, list[float]]] = []
# Resolve external-vector-store dup candidates up front, before the first
# DB statement: the advisory lock's critical section below must contain no
# network round trips, and the lazy session has no open transaction yet.
# A None entry falls back to in-place resolution in is_rejected_duplicate
# (the pgvector path, or a document with no valid merge partner).
semantic_candidates: list[list[str] | None] = [None] * len(documents)
if deduplicate and not _uses_pgvector():
async def _resolve_candidates(index: int, doc: schemas.DocumentCreate) -> None:
filters = _semantic_dup_filters(doc)
if filters is None:
return
semantic_candidates[index] = await query_external_vector_document_ids(
workspace_name=workspace_name,
observer=observer,
observed=observed,
embedding=doc.embedding,
top_k=_SEMANTIC_DUP_TOP_K,
max_distance=_SEMANTIC_DUP_MAX_DISTANCE,
filters=filters,
)
await asyncio.gather(
*(_resolve_candidates(i, doc) for i, doc in enumerate(documents))
)
# exact-content dedup (independent of `deduplicate`): pre-fetch
# existing live documents whose normalized content matches anything in this
# batch, scoped to (workspace, observer, observed). The SQL normalization must
@ -557,6 +602,24 @@ async def create_documents(
existing_doc,
)
# Serialize writers per (workspace, observer, observed) collection: two
# batches reinforcing the same rows in different orders deadlock otherwise
# (DEV-1975), and work-unit keys don't prevent concurrent writers to one
# collection. Skipped when the only writes are INSERTs of new rows, which
# can't deadlock on document rows. This advisory lock must stay the
# OUTERMOST lock taken in this transaction — don't add row-locking reads
# earlier in this function. SET LOCAL persists for the transaction, so the
# timeout also bounds the FOR KEY SHARE wait at the final INSERT;
# lock_timeout raises 55P03, which the queue layer classifies retryable.
# hashtextextended needs PG 11+; \x1f because names can contain ':'.
if documents and (deduplicate or existing_by_key):
lock_key = f"honcho:documents:{workspace_name}\x1f{observer}\x1f{observed}"
await db.execute(text("SET LOCAL lock_timeout = '30s'"))
await db.execute(
text("SELECT pg_advisory_xact_lock(hashtextextended(:key, 0))"),
{"key": lock_key},
)
# Tracks dedup keys already accepted from this batch so exact
# duplicates within a single inference call collapse to one document.
seen_in_batch: set[tuple[str, str, str | None]] = set()
@ -565,7 +628,7 @@ async def create_documents(
exact_dup_in_batch_count = 0
semantic_dup_rejected_count = 0
semantic_dup_replaced_count = 0
for doc in documents:
for index, doc in enumerate(documents):
try:
# Session-purity invariant: an explicit document must always carry
# the session it was derived from. Refuse to write session-less
@ -601,6 +664,7 @@ async def create_documents(
# reinforcement (times_derived > 1, e.g. a future re-ingestion or
# collection-merge path). Mirrors the superior-replacement branch
# in is_rejected_duplicate.
# Row lock held until commit; ordered by the advisory lock above.
existing_match.times_derived = func.greatest(
models.Document.times_derived + 1,
doc.times_derived,
@ -614,7 +678,12 @@ async def create_documents(
# as a duplicate OR deletes an existing document that is a duplicate.
if deduplicate:
duplicate_result = await is_rejected_duplicate(
db, doc, workspace_name, observer=observer, observed=observed
db,
doc,
workspace_name,
observer=observer,
observed=observed,
candidate_document_ids=semantic_candidates[index],
)
if duplicate_result is SemanticRejectionResult.REPLACED_EXISTING:
# Existing doc was soft-deleted in favor of this one; the
@ -670,7 +739,14 @@ async def create_documents(
if doc.embedding:
docs_with_embeddings.append((new_doc, doc.embedding))
except SQLAlchemyError:
# The session/transaction is dead; continuing the loop would only
# cascade PendingRollbackErrors and lose the whole batch silently.
await db.rollback()
raise
except Exception as e:
# Genuinely per-document failures (bad content, metadata, token
# overflow) skip the document without poisoning the batch.
logger.error(
f"Error adding new document to {workspace_name}/{doc.session_name}/{observer}/{observed}: {e}"
)
@ -772,6 +848,11 @@ async def create_documents(
raise ValidationException(
"Failed to create documents due to integrity constraint violation"
) from e
except DBAPIError:
# Leave the session clean for callers that own it (e.g. a deadlock
# at the final commit); the queue layer classifies and retries.
await db.rollback()
raise
return CreateDocumentsResult(
created_documents=accepted_documents,
@ -1160,6 +1241,7 @@ async def is_rejected_duplicate(
*,
observer: str,
observed: str,
candidate_document_ids: list[str] | None = None,
) -> SemanticRejectionResult:
"""
Check if a document is a duplicate of an existing document.
@ -1185,27 +1267,41 @@ async def is_rejected_duplicate(
sessions for explicit-level documents (session-purity invariant: an
explicit document records what was derived from exactly one session, so
a near-duplicate from another session must not reinforce or replace it).
``candidate_document_ids`` carries pre-resolved external-vector-store
candidates (see create_documents), keeping vector-store round trips out
of the caller's transaction; the ``deleted_at IS NULL`` filter in
``fetch_documents_by_ids`` re-validates candidates that went stale since
resolution. None means resolve in place.
"""
filters: dict[str, Any] = {"level": doc.level}
if doc.level == "explicit":
if doc.session_name is None:
# create_documents refuses session-less explicit documents; if one
# reaches here anyway it has no valid merge partner.
return SemanticRejectionResult.NOT_DUPLICATE
filters["session_name"] = doc.session_name
filters = _semantic_dup_filters(doc)
if filters is None:
# create_documents refuses session-less explicit documents; if one
# reaches here anyway it has no valid merge partner.
return SemanticRejectionResult.NOT_DUPLICATE
# Step 1: Find potential duplicates using cosine similarity
similar_docs = await query_documents(
db=db,
workspace_name=workspace_name,
query=doc.content,
observer=observer,
observed=observed,
filters=filters,
max_distance=0.05,
top_k=1,
embedding=doc.embedding,
)
if candidate_document_ids is not None:
similar_docs: Sequence[models.Document] = await fetch_documents_by_ids(
db=db,
workspace_name=workspace_name,
observer=observer,
observed=observed,
document_ids=candidate_document_ids,
filters=filters,
)
else:
similar_docs = await query_documents(
db=db,
workspace_name=workspace_name,
query=doc.content,
observer=observer,
observed=observed,
filters=filters,
max_distance=_SEMANTIC_DUP_MAX_DISTANCE,
top_k=_SEMANTIC_DUP_TOP_K,
embedding=doc.embedding,
)
if not similar_docs:
return SemanticRejectionResult.NOT_DUPLICATE
@ -1233,6 +1329,8 @@ async def is_rejected_duplicate(
# another derivation rather than resetting times_derived to 1.
doc.times_derived = max(doc.times_derived, existing_doc.times_derived + 1)
# Soft-delete the existing document - reconciliation will clean up vectors and hard-delete
# Row lock held until commit; concurrent create_documents writers are
# serialized by the advisory lock taken there.
existing_doc.deleted_at = datetime.datetime.now(datetime.timezone.utc)
await db.flush()
return (
@ -1244,6 +1342,8 @@ async def is_rejected_duplicate(
# the increment atomic server-side -- concurrent workers reinforcing the same
# document must not lose updates -- while still honoring an incoming doc that
# already carries accumulated reinforcement (times_derived > 1).
# Row lock held until commit; concurrent create_documents writers are
# serialized by the advisory lock taken there.
existing_doc.times_derived = func.greatest(
models.Document.times_derived + 1,
doc.times_derived,

62
src/deriver/queue_manager.py
View File

@ -43,6 +43,7 @@ from src.reconciler import (
from src.schemas import ResolvedConfiguration
from src.telemetry import prometheus_metrics
from src.telemetry.sentry import initialize_sentry
from src.utils.retryable_errors import is_retryable_error
from src.utils.work_unit import parse_work_unit_key
from src.webhooks.events import (
QueueEmptyEvent,
@ -53,6 +54,11 @@ logger = getLogger(__name__)
load_dotenv(override=True)
# Total processing attempts per work unit for transient errors, counted
# per instance (N deriver instances give N x this many attempts).
MAX_RETRYABLE_ATTEMPTS = 3
RETRY_BACKOFF_SECONDS = 1.0
class WorkerOwnership(NamedTuple):
"""Represents the instance of a work unit that a worker is processing."""
@ -129,6 +135,10 @@ class QueueManager:
self.worker_ownership: dict[str, WorkerOwnership] = {}
self.queue_empty_flag: asyncio.Event = asyncio.Event()
# Transient-failure attempts per work-unit key. Entries are removed
# on success or terminal failure, never on the retry path itself.
self._retry_attempts: dict[str, int] = {}
# Current adaptive polling interval; grows while idle/erroring and
# resets to the base interval as soon as work is claimed.
self._current_poll_interval: float = (
@ -579,11 +589,19 @@ class QueueManager:
items: list[QueueItem],
work_unit_key: str,
context: str,
) -> None:
) -> bool:
"""
Handle processing errors by marking queue items as errored, logging, and forwarding to Sentry.
We only mark the first queue item as errored so we don't potentially throw away a batch. This allows us
to incrementally attempt to process the batch while still maintaining progress in a work unit.
Handle a processing error. Returns True when the caller should stop
processing and release the work unit for a later re-claim.
Transient errors (is_retryable_error) get up to MAX_RETRYABLE_ATTEMPTS
attempts per work unit: items stay unprocessed with no error recorded.
Reprocessing is safe because a retried batch re-derives the same
observations and exact dedup collapses them into reinforcement.
Terminal errors mark only the first queue item as errored so we don't
potentially throw away a batch. This allows us to incrementally attempt
to process the batch while still maintaining progress in a work unit.
Args:
error: The exception that occurred
@ -591,6 +609,21 @@ class QueueManager:
work_unit_key: The work unit key for the queue items
context: Context string describing what was being processed (e.g., "processing representation batch")
"""
if is_retryable_error(error):
attempts = self._retry_attempts.get(work_unit_key, 0) + 1
if attempts < MAX_RETRYABLE_ATTEMPTS:
self._retry_attempts[work_unit_key] = attempts
logger.warning(
"Transient error %s for work unit %s (attempt %d/%d); leaving items unprocessed for retry",
context,
work_unit_key,
attempts,
MAX_RETRYABLE_ATTEMPTS,
exc_info=error,
)
return True
self._retry_attempts.pop(work_unit_key, None)
error_msg = f"{error.__class__.__name__}: {str(error)}"
try:
if items:
@ -609,6 +642,7 @@ class QueueManager:
)
if settings.SENTRY.ENABLED:
sentry_sdk.capture_exception(error)
return False
async def process_work_unit(self, work_unit_key: str, worker_id: str) -> None:
"""Process all queue items for a specific work unit by routing to the correct handler."""
@ -672,14 +706,21 @@ class QueueManager:
await self.mark_queue_items_as_processed(
items_to_process, work_unit_key
)
self._retry_attempts.pop(work_unit_key, None)
queue_item_count += len(items_to_process)
except Exception as e:
await self._handle_processing_error(
if await self._handle_processing_error(
e,
items_to_process,
work_unit_key,
f"processing {work_unit.task_type} batch",
)
):
# Release the work unit (via the finally
# below) and let a later poll re-claim it.
await asyncio.sleep(
self._jitter(RETRY_BACKOFF_SECONDS)
)
break
else:
queue_item = await self.get_next_queue_item(
@ -696,14 +737,19 @@ class QueueManager:
await self.mark_queue_items_as_processed(
[queue_item], work_unit_key
)
self._retry_attempts.pop(work_unit_key, None)
queue_item_count += 1
except Exception as e:
await self._handle_processing_error(
if await self._handle_processing_error(
e,
[queue_item],
work_unit_key,
"processing queue item",
)
):
await asyncio.sleep(
self._jitter(RETRY_BACKOFF_SECONDS)
)
break
except Exception as e:
logger.error(

86
src/utils/retryable_errors.py Normal file
View File

@ -0,0 +1,86 @@
"""Classify exceptions as transient (safe to retry) or terminal.
Imports only exception taxonomies, so it is importable from anywhere and
unit-testable without a DB.
"""
import asyncio
from collections.abc import Iterator
import httpx
from sqlalchemy.exc import DBAPIError
__all__ = ["is_retryable_db_error", "is_retryable_error"]
_RETRYABLE_SQLSTATES = frozenset(
{
"40001", # serialization_failure
"40P01", # deadlock_detected
"55P03", # lock_not_available (lock_timeout / NOWAIT)
"57014", # query_canceled (statement_timeout)
"08000", # connection_exception family
"08001",
"08003",
"08004",
"08006",
}
)
# Provider/network transport failures. SDK wrappers (anthropic/openai
# APIConnectionError etc.) chain to these via __cause__.
_TRANSPORT_ERRORS = (
httpx.TransportError,
ConnectionError,
asyncio.TimeoutError,
TimeoutError,
)
def _iter_cause_chain(exc: BaseException) -> Iterator[BaseException]:
seen: set[int] = set()
current: BaseException | None = exc
while current is not None and id(current) not in seen:
seen.add(id(current))
yield current
current = current.__cause__
def _sqlstate(exc: DBAPIError) -> str | None:
"""Extract the SQLSTATE off ``DBAPIError.orig``, driver-agnostically."""
orig = getattr(exc, "orig", None)
for candidate in (orig, getattr(orig, "__cause__", None)):
code = getattr(candidate, "sqlstate", None)
if isinstance(code, str):
return code
return None
def is_retryable_db_error(exc: BaseException) -> bool:
"""True for transient DB failures: deadlock, serialization failure,
lock/statement timeout, or a lost connection.
Integrity (23xxx), data (22xxx), and programming (42xxx) errors are
deliberately terminal.
"""
for current in _iter_cause_chain(exc):
if not isinstance(current, DBAPIError):
continue
if current.connection_invalidated:
return True
if _sqlstate(current) in _RETRYABLE_SQLSTATES:
return True
return False
def is_retryable_error(exc: BaseException) -> bool:
"""Superset of ``is_retryable_db_error``: also transient network/provider
transport failures (timeouts, connection refused/reset).
Auth failures (401 from a rotated key) are deliberately terminal: they
never self-heal, so retrying only delays the burn.
"""
if is_retryable_db_error(exc):
return True
return any(
isinstance(current, _TRANSPORT_ERRORS) for current in _iter_cause_chain(exc)
)

479
tests/crud/test_document.py
View File

@ -1,10 +1,13 @@
import asyncio
import datetime
from typing import Any
from unittest.mock import AsyncMock, patch
import pytest
from nanoid import generate as generate_nanoid
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.exc import OperationalError
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
from src import crud, models, schemas
from src.crud.document import SemanticRejectionResult, is_rejected_duplicate
@ -1334,3 +1337,477 @@ class TestSessionPurityInvariant:
)
assert rejected is SemanticRejectionResult.NOT_DUPLICATE
mock_query.assert_not_awaited()
class TestCreateDocumentsConcurrency:
"""Concurrency regression tests for create_documents (DEV-1975).
Two deriver work units can write the same (workspace, observer, observed)
collection concurrently: work-unit keys include session but not observer,
so one peer active in two sessions produces two writers, and the dreamer
adds a third through a different key. Reinforcement UPDATEs issued in
batch order used to deadlock when two writers touched the same rows in
opposite orders; create_documents now serializes writers per collection
with a transaction-scoped advisory lock.
"""
N_DOCS: int = 20
N_ROUNDS: int = 5
async def _setup(
self,
db_session: AsyncSession,
test_workspace: models.Workspace,
test_peer: models.Peer,
) -> tuple[models.Peer, models.Session]:
"""Create an observed peer, session, and collection, committed so
they are visible to independent concurrent sessions."""
test_peer2 = models.Peer(
name=str(generate_nanoid()), workspace_name=test_workspace.name
)
test_session = models.Session(
name=str(generate_nanoid()), workspace_name=test_workspace.name
)
db_session.add_all([test_peer2, test_session])
await db_session.flush()
collection = models.Collection(
workspace_name=test_workspace.name,
observer=test_peer.name,
observed=test_peer2.name,
)
db_session.add(collection)
await db_session.commit()
return test_peer2, test_session
def _batch(self, session_name: str) -> list[schemas.DocumentCreate]:
return [
schemas.DocumentCreate(
content=f"user fact number {i}",
embedding=[0.1] * 1536,
session_name=session_name,
metadata=schemas.DocumentMetadata(
message_ids=[i],
message_created_at="2026-01-01T00:00:00Z",
),
)
for i in range(self.N_DOCS)
]
@staticmethod
def _chain(exc: BaseException) -> str:
parts: list[str] = []
seen: set[int] = set()
e: BaseException | None = exc
while e is not None and id(e) not in seen:
seen.add(id(e))
parts.append(f"{type(e).__name__}: {e}")
e = e.__cause__ or e.__context__
return " <- ".join(parts)
@pytest.mark.asyncio
async def test_concurrent_reinforcement_does_not_deadlock(
self,
db_engine: "AsyncEngine",
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
):
"""Concurrent same-collection batches locking rows in opposite orders
must serialize, not deadlock (DEV-1975)."""
test_workspace, test_peer = sample_data
test_peer2, test_session = await self._setup(
db_session, test_workspace, test_peer
)
# Seed the rows both writers will reinforce.
await crud.create_documents(
db_session,
self._batch(test_session.name),
workspace_name=test_workspace.name,
observer=test_peer.name,
observed=test_peer2.name,
)
session_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False)
for round_num in range(self.N_ROUNDS):
forward = self._batch(test_session.name)
backward = list(reversed(self._batch(test_session.name)))
async def _run(batch: list[schemas.DocumentCreate]) -> None:
async with session_factory() as db:
await crud.create_documents(
db,
batch,
workspace_name=test_workspace.name,
observer=test_peer.name,
observed=test_peer2.name,
)
results = await asyncio.gather(
_run(forward), _run(backward), return_exceptions=True
)
errors = [r for r in results if isinstance(r, BaseException)]
assert not errors, (
f"round {round_num}: concurrent create_documents failed: "
+ "; ".join(self._chain(e) for e in errors)
)
# Every round reinforced the same rows: 1 seed + 2 per round.
docs = (
(
await db_session.execute(
select(models.Document).where(
models.Document.workspace_name == test_workspace.name,
models.Document.observer == test_peer.name,
models.Document.observed == test_peer2.name,
models.Document.deleted_at.is_(None),
)
)
)
.scalars()
.all()
)
assert len(docs) == self.N_DOCS
assert all(d.times_derived == 1 + 2 * self.N_ROUNDS for d in docs)
class TestCreateDocumentsErrorHandling:
"""A dead transaction aborts the batch loudly; per-document failures
still skip just that document (DEV-1975)."""
async def _setup(
self,
db_session: AsyncSession,
test_workspace: models.Workspace,
test_peer: models.Peer,
) -> tuple[models.Peer, models.Session]:
test_peer2 = models.Peer(
name=str(generate_nanoid()), workspace_name=test_workspace.name
)
test_session = models.Session(
name=str(generate_nanoid()), workspace_name=test_workspace.name
)
db_session.add_all([test_peer2, test_session])
await db_session.flush()
collection = models.Collection(
workspace_name=test_workspace.name,
observer=test_peer.name,
observed=test_peer2.name,
)
db_session.add(collection)
await db_session.commit()
return test_peer2, test_session
def _doc(self, content: str, session_name: str) -> schemas.DocumentCreate:
return schemas.DocumentCreate(
content=content,
embedding=[0.1] * 1536,
session_name=session_name,
metadata=schemas.DocumentMetadata(
message_ids=[1],
message_created_at="2026-01-01T00:00:00Z",
),
)
@pytest.mark.asyncio
async def test_db_error_in_loop_aborts_batch(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
):
"""A DB error on an in-loop flush must raise out of create_documents
instead of cascading through a dead session, and commit nothing."""
test_workspace, test_peer = sample_data
test_peer2, test_session = await self._setup(
db_session, test_workspace, test_peer
)
# Plain strings: the rollback below expires ORM objects in the session.
workspace_name = test_workspace.name
observer = test_peer.name
observed = test_peer2.name
session_name = test_session.name
await crud.create_documents(
db_session,
[self._doc("existing fact", session_name)],
workspace_name=workspace_name,
observer=observer,
observed=observed,
)
class FakePGError(Exception):
sqlstate: str = "40P01"
deadlock = OperationalError("UPDATE documents", {}, FakePGError())
# Batch: an exact dup (triggers the reinforcement flush that "deadlocks")
# followed by a new document that must NOT be committed.
with (
patch.object(db_session, "flush", AsyncMock(side_effect=deadlock)),
pytest.raises(OperationalError),
):
await crud.create_documents(
db_session,
[
self._doc("existing fact", session_name),
self._doc("a brand new fact", session_name),
],
workspace_name=workspace_name,
observer=observer,
observed=observed,
)
docs = (
(
await db_session.execute(
select(models.Document).where(
models.Document.workspace_name == workspace_name,
models.Document.observer == observer,
models.Document.observed == observed,
)
)
)
.scalars()
.all()
)
assert [d.content for d in docs] == ["existing fact"]
assert docs[0].times_derived == 1
@pytest.mark.asyncio
async def test_per_document_error_still_skips_only_that_document(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
):
"""Non-DB per-document failures keep their skip semantics."""
test_workspace, test_peer = sample_data
test_peer2, test_session = await self._setup(
db_session, test_workspace, test_peer
)
from src.crud import document as document_module
real_dedup_key = document_module._dedup_key # pyright: ignore[reportPrivateUsage]
def flaky_dedup_key(
content: str, level: str, session_name: str | None
) -> tuple[str, str, str | None]:
if content == "poison":
raise ValueError("bad content")
return real_dedup_key(content, level, session_name)
with patch.object(document_module, "_dedup_key", flaky_dedup_key):
result = await crud.create_documents(
db_session,
[
self._doc("good fact one", test_session.name),
self._doc("poison", test_session.name),
self._doc("good fact two", test_session.name),
],
workspace_name=test_workspace.name,
observer=test_peer.name,
observed=test_peer2.name,
)
assert sorted(d.content for d in result.created_documents) == [
"good fact one",
"good fact two",
]
class TestAdvisoryLockBehavior:
"""The per-collection advisory lock in create_documents: held for the
transaction, keyed by (workspace, observer, observed), and skipped when
the batch can only INSERT."""
async def _setup(
self,
db_session: AsyncSession,
test_workspace: models.Workspace,
test_peer: models.Peer,
n_observed: int = 1,
) -> tuple[list[models.Peer], models.Session]:
observed_peers = [
models.Peer(name=str(generate_nanoid()), workspace_name=test_workspace.name)
for _ in range(n_observed)
]
test_session = models.Session(
name=str(generate_nanoid()), workspace_name=test_workspace.name
)
db_session.add_all([*observed_peers, test_session])
await db_session.flush()
for observed_peer in observed_peers:
db_session.add(
models.Collection(
workspace_name=test_workspace.name,
observer=test_peer.name,
observed=observed_peer.name,
)
)
await db_session.commit()
return observed_peers, test_session
def _doc(self, content: str, session_name: str) -> schemas.DocumentCreate:
return schemas.DocumentCreate(
content=content,
embedding=[0.1] * 1536,
session_name=session_name,
metadata=schemas.DocumentMetadata(
message_ids=[1],
message_created_at="2026-01-01T00:00:00Z",
),
)
@pytest.mark.asyncio
async def test_lock_blocks_same_key_not_other_keys(
self,
db_engine: AsyncEngine,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
):
"""While one writer holds a collection's lock, a second writer on the
same collection blocks and a writer on a different collection does not."""
test_workspace, test_peer = sample_data
(observed_a, observed_b), test_session = await self._setup(
db_session, test_workspace, test_peer, n_observed=2
)
session_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False)
holder_in_lock = asyncio.Event()
release_holder = asyncio.Event()
real_is_rejected_duplicate = crud.document.is_rejected_duplicate
async def pausing_is_rejected_duplicate(*args: Any, **kwargs: Any) -> Any:
holder_in_lock.set()
await release_holder.wait()
return await real_is_rejected_duplicate(*args, **kwargs)
async def _create(observed: str, content: str) -> None:
async with session_factory() as db:
await crud.create_documents(
db,
[self._doc(content, test_session.name)],
workspace_name=test_workspace.name,
observer=test_peer.name,
observed=observed,
deduplicate=True,
)
with patch(
"src.crud.document.is_rejected_duplicate",
side_effect=pausing_is_rejected_duplicate,
):
holder = asyncio.create_task(_create(observed_a.name, "holder fact"))
await asyncio.wait_for(holder_in_lock.wait(), timeout=10)
# Holder is parked inside its critical section; the patch is no longer
# active so contender/other run the real dedup path.
contender = asyncio.create_task(_create(observed_a.name, "contender fact"))
other_key = asyncio.create_task(_create(observed_b.name, "other fact"))
# A different collection's writer completes while the lock is held...
await asyncio.wait_for(other_key, timeout=10)
# ...but the same collection's writer stays blocked.
await asyncio.sleep(0.3)
assert not contender.done()
release_holder.set()
await asyncio.wait_for(asyncio.gather(holder, contender), timeout=10)
@pytest.mark.asyncio
async def test_lock_skipped_for_insert_only_batch(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
):
"""deduplicate=False with no exact-dup matches only INSERTs, so no
advisory lock is taken; a batch with an existing exact dup takes it."""
test_workspace, test_peer = sample_data
(observed_peer,), test_session = await self._setup(
db_session, test_workspace, test_peer
)
executed_sql: list[str] = []
real_execute = db_session.execute
async def spying_execute(statement: Any, *args: Any, **kwargs: Any) -> Any:
executed_sql.append(str(statement))
return await real_execute(statement, *args, **kwargs)
def lock_statements() -> list[str]:
return [s for s in executed_sql if "pg_advisory_xact_lock" in s]
with patch.object(db_session, "execute", side_effect=spying_execute):
await crud.create_documents(
db_session,
[self._doc("all new fact", test_session.name)],
workspace_name=test_workspace.name,
observer=test_peer.name,
observed=observed_peer.name,
deduplicate=False,
)
assert not lock_statements()
# Same content again: the exact-dup prefetch now matches, so the
# reinforcement UPDATE must run under the lock.
await crud.create_documents(
db_session,
[self._doc("all new fact", test_session.name)],
workspace_name=test_workspace.name,
observer=test_peer.name,
observed=observed_peer.name,
deduplicate=False,
)
assert len(lock_statements()) == 1
@pytest.mark.asyncio
async def test_external_candidates_resolved_before_lock(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
monkeypatch: pytest.MonkeyPatch,
):
"""On the external-vector-store path, every candidate resolution
happens before the advisory lock is acquired."""
from src.config import settings
test_workspace, test_peer = sample_data
(observed_peer,), test_session = await self._setup(
db_session, test_workspace, test_peer
)
monkeypatch.setattr(settings.VECTOR_STORE, "TYPE", "turbopuffer")
monkeypatch.setattr(settings.VECTOR_STORE, "MIGRATED", True)
events: list[str] = []
real_execute = db_session.execute
async def spying_execute(statement: Any, *args: Any, **kwargs: Any) -> Any:
if "pg_advisory_xact_lock" in str(statement):
events.append("lock")
return await real_execute(statement, *args, **kwargs)
async def fake_resolve(*_args: Any, **_kwargs: Any) -> list[str]:
events.append("resolve")
return []
with (
patch.object(db_session, "execute", side_effect=spying_execute),
patch(
"src.crud.document.query_external_vector_document_ids",
side_effect=fake_resolve,
),
):
result = await crud.create_documents(
db_session,
[
self._doc("fact one", test_session.name),
self._doc("fact two", test_session.name),
],
workspace_name=test_workspace.name,
observer=test_peer.name,
observed=observed_peer.name,
deduplicate=True,
)
assert len(result.created_documents) == 2
assert events == ["resolve", "resolve", "lock"]

233
tests/deriver/test_queue_processing.py
View File

@ -7,6 +7,7 @@ from unittest.mock import patch
import pytest
from nanoid import generate as generate_nanoid
from sqlalchemy import select
from sqlalchemy.exc import OperationalError
from sqlalchemy.ext.asyncio import AsyncSession
from src import models
@ -1874,3 +1875,235 @@ class TestPollingJitter:
qm.shutdown_event.set()
# A shutdown already signalled must short-circuit the (long) jitter sleep.
await asyncio.wait_for(qm._sleep_startup_jitter(), timeout=1.0) # pyright: ignore[reportPrivateUsage]
@pytest.mark.asyncio
class TestQueueRetry:
"""Bounded retry of transient errors in process_work_unit (DEV-1975).
A transient failure (deadlock, lost connection, provider transport) must
leave the batch's queue items unprocessed and release the work unit for
re-claim, up to MAX_RETRYABLE_ATTEMPTS per work unit; terminal failures
keep today's burn-one-item behavior.
"""
async def _seed_work_unit(
self,
db_session: AsyncSession,
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
create_queue_payload: Callable[..., Any],
n_messages: int = 1,
) -> tuple[QueueManager, str, str, list[models.QueueItem]]:
"""Seed a claimed representation work unit owned by a test worker."""
session, peers = sample_session_with_peers
peer = peers[0]
messages: list[models.Message] = []
for index in range(n_messages):
message = models.Message(
session_name=session.name,
workspace_name=session.workspace_name,
peer_name=peer.name,
content=f"Message {index}",
token_count=10,
seq_in_session=index + 1,
)
db_session.add(message)
messages.append(message)
await db_session.commit()
for message in messages:
await db_session.refresh(message)
queue_items: list[models.QueueItem] = []
work_unit_key = ""
for message in messages:
payload = create_queue_payload(
message=message,
task_type="representation",
observed=peer.name,
observer=peer.name,
)
work_unit_key = work_unit_key or construct_work_unit_key(
session.workspace_name, payload
)
queue_item = models.QueueItem(
session_id=session.id,
task_type="representation",
work_unit_key=work_unit_key,
payload=payload,
processed=False,
workspace_name=session.workspace_name,
message_id=message.id,
)
db_session.add(queue_item)
queue_items.append(queue_item)
await db_session.commit()
for queue_item in queue_items:
await db_session.refresh(queue_item)
qm = QueueManager()
worker_id = "test_worker"
claimed_units = await qm.claim_work_units(db_session, [work_unit_key])
qm.worker_ownership[worker_id] = WorkerOwnership(
work_unit_key=work_unit_key, aqs_id=claimed_units[work_unit_key]
)
await db_session.commit()
return qm, work_unit_key, worker_id, queue_items
@staticmethod
def _retryable_error() -> OperationalError:
class FakePGError(Exception):
sqlstate: str = "40P01"
return OperationalError("UPDATE documents", {}, FakePGError())
async def _fetch_items(
self, db_session: AsyncSession, work_unit_key: str
) -> list[models.QueueItem]:
db_session.expire_all()
return list(
(
await db_session.execute(
select(models.QueueItem)
.where(models.QueueItem.work_unit_key == work_unit_key)
.order_by(models.QueueItem.id)
)
)
.scalars()
.all()
)
async def _aqs_rows(self, db_session: AsyncSession, work_unit_key: str) -> int:
return len(
(
await db_session.execute(
select(models.ActiveQueueSession).where(
models.ActiveQueueSession.work_unit_key == work_unit_key
)
)
)
.scalars()
.all()
)
async def test_retryable_error_leaves_items_unprocessed(
self,
db_session: AsyncSession,
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
create_queue_payload: Callable[..., Any],
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A transient error stops the work unit after ONE batch fetch (no
tight loop), leaves items unprocessed with no error, and releases
the ActiveQueueSession row."""
monkeypatch.setattr("src.deriver.queue_manager.RETRY_BACKOFF_SECONDS", 0.0)
qm, work_unit_key, worker_id, _ = await self._seed_work_unit(
db_session, sample_session_with_peers, create_queue_payload, n_messages=2
)
initial_semaphore_value = qm.semaphore._value
batch_fetches = 0
original_get_batch = qm.get_queue_item_batch
async def counting_get_batch(*args: Any, **kwargs: Any) -> Any:
nonlocal batch_fetches
batch_fetches += 1
return await original_get_batch(*args, **kwargs)
with (
patch.object(qm, "get_queue_item_batch", side_effect=counting_get_batch),
patch(
"src.deriver.queue_manager.process_representation_batch",
side_effect=self._retryable_error(),
),
):
await qm.process_work_unit(work_unit_key, worker_id)
assert batch_fetches == 1
items = await self._fetch_items(db_session, work_unit_key)
assert all(not item.processed for item in items)
assert all(item.error is None for item in items)
assert await self._aqs_rows(db_session, work_unit_key) == 0
assert qm._retry_attempts[work_unit_key] == 1 # pyright: ignore[reportPrivateUsage]
assert qm.semaphore._value == initial_semaphore_value
async def test_retry_exhaustion_is_terminal(
self,
db_session: AsyncSession,
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
create_queue_payload: Callable[..., Any],
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""At the attempt cap a transient error burns the first item exactly
like today's terminal path and clears the counter."""
from src.deriver.queue_manager import MAX_RETRYABLE_ATTEMPTS
monkeypatch.setattr("src.deriver.queue_manager.RETRY_BACKOFF_SECONDS", 0.0)
qm, work_unit_key, worker_id, _ = await self._seed_work_unit(
db_session, sample_session_with_peers, create_queue_payload
)
qm._retry_attempts[work_unit_key] = MAX_RETRYABLE_ATTEMPTS - 1 # pyright: ignore[reportPrivateUsage]
with patch(
"src.deriver.queue_manager.process_representation_batch",
side_effect=self._retryable_error(),
):
await qm.process_work_unit(work_unit_key, worker_id)
items = await self._fetch_items(db_session, work_unit_key)
assert len(items) == 1
assert items[0].processed
assert items[0].error is not None
assert "OperationalError" in items[0].error
assert work_unit_key not in qm._retry_attempts # pyright: ignore[reportPrivateUsage]
async def test_non_retryable_error_burns_immediately(
self,
db_session: AsyncSession,
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
create_queue_payload: Callable[..., Any],
) -> None:
"""A non-retryable error keeps today's behavior verbatim: the first
item is marked errored on the first attempt."""
qm, work_unit_key, worker_id, _ = await self._seed_work_unit(
db_session, sample_session_with_peers, create_queue_payload
)
with patch(
"src.deriver.queue_manager.process_representation_batch",
side_effect=ValueError("bad batch"),
):
await qm.process_work_unit(work_unit_key, worker_id)
items = await self._fetch_items(db_session, work_unit_key)
assert len(items) == 1
assert items[0].processed
assert items[0].error is not None
assert "ValueError" in items[0].error
assert work_unit_key not in qm._retry_attempts # pyright: ignore[reportPrivateUsage]
async def test_counter_cleared_after_success(
self,
db_session: AsyncSession,
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
create_queue_payload: Callable[..., Any],
) -> None:
"""A success wipes the accumulated attempt count for the work unit."""
qm, work_unit_key, worker_id, _ = await self._seed_work_unit(
db_session, sample_session_with_peers, create_queue_payload
)
qm._retry_attempts[work_unit_key] = 1 # pyright: ignore[reportPrivateUsage]
async def noop_batch(*_args: Any, **_kwargs: Any) -> None:
return None
with patch(
"src.deriver.queue_manager.process_representation_batch",
side_effect=noop_batch,
):
await qm.process_work_unit(work_unit_key, worker_id)
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 work_unit_key not in qm._retry_attempts # pyright: ignore[reportPrivateUsage]

117
tests/utils/test_retryable_errors.py Normal file
View File

@ -0,0 +1,117 @@
"""DB-free unit tests for src/utils/retryable_errors.py."""
import asyncio
from typing import cast
import httpx
import pytest
from sqlalchemy.exc import DBAPIError, OperationalError
from src.utils.retryable_errors import is_retryable_db_error, is_retryable_error
class FakePGError(Exception):
"""Stands in for a driver exception carrying a SQLSTATE."""
sqlstate: str | None
def __init__(self, sqlstate: str | None) -> None:
super().__init__(f"fake pg error ({sqlstate})")
self.sqlstate = sqlstate
def _dbapi_error(
sqlstate: str | None,
*,
orig: BaseException | None = None,
connection_invalidated: bool = False,
) -> DBAPIError:
if orig is None and sqlstate is not None:
orig = FakePGError(sqlstate)
return OperationalError(
"SELECT 1",
{},
cast(BaseException, orig),
connection_invalidated=connection_invalidated,
)
@pytest.mark.parametrize(
("sqlstate", "expected"),
[
("40P01", True), # deadlock_detected
("40001", True), # serialization_failure
("55P03", True), # lock_not_available
("57014", True), # query_canceled
("08006", True), # connection_failure
("23505", False), # unique_violation
("42P01", False), # undefined_table
("22P02", False), # invalid_text_representation
],
)
def test_sqlstate_classification(sqlstate: str, expected: bool):
exc = _dbapi_error(sqlstate)
assert is_retryable_db_error(exc) is expected
assert is_retryable_error(exc) is expected
def test_orig_none_is_terminal():
assert not is_retryable_db_error(_dbapi_error(None))
def test_sqlstate_on_orig_cause():
"""SQLSTATE found by walking orig.__cause__ when orig itself has none."""
wrapper = Exception("driver wrapper")
wrapper.__cause__ = FakePGError("40P01")
assert is_retryable_db_error(_dbapi_error(None, orig=wrapper))
def test_connection_invalidated_is_retryable():
exc = _dbapi_error(None, connection_invalidated=True)
assert is_retryable_db_error(exc)
def test_dbapi_error_nested_in_cause_chain():
outer = RuntimeError("save failed")
outer.__cause__ = _dbapi_error("40P01")
assert is_retryable_db_error(outer)
assert is_retryable_error(outer)
def test_non_db_exceptions_are_not_db_retryable():
assert not is_retryable_db_error(ValueError("bad input"))
assert not is_retryable_db_error(httpx.ConnectTimeout("timed out"))
@pytest.mark.parametrize(
("exc", "expected"),
[
(httpx.ConnectTimeout("timed out"), True),
(httpx.ReadTimeout("timed out"), True),
(httpx.ConnectError("connection refused"), True),
(ConnectionResetError("reset"), True),
(asyncio.TimeoutError(), True),
(TimeoutError(), True),
(ValueError("bad input"), False),
(httpx.HTTPStatusError("401", request=None, response=None), False), # pyright: ignore[reportArgumentType]
],
)
def test_transport_classification(exc: BaseException, expected: bool):
assert is_retryable_error(exc) is expected
assert not is_retryable_db_error(exc)
def test_transport_error_nested_in_cause_chain():
"""SDK wrappers (e.g. APIConnectionError) chain to httpx via __cause__."""
wrapper = RuntimeError("provider call failed")
wrapper.__cause__ = httpx.ConnectError("connection refused")
assert is_retryable_error(wrapper)
assert not is_retryable_db_error(wrapper)
def test_cause_cycle_terminates():
a = RuntimeError("a")
b = RuntimeError("b")
a.__cause__ = b
b.__cause__ = a
assert not is_retryable_error(a)