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

fix(deriver): fix create_documents deadlock (#1033) 

* 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

* fix(deriver): replace create_documents advisory lock with id-ordered row locks

Advisory locks are database-scoped and would serialize every writer to a
collection, including across Groudon tenants that share names. Collect
reinforcement and replace ops during the loop, lock target rows with
SELECT ... ORDER BY id FOR UPDATE, then apply. populate_existing reloads
times_derived so a prefetched identity-map row cannot lose a concurrent
increment.

* fix(deriver): harden create_documents candidate hoist and test isolation

Skip empty embeddings on the external-store path, isolate per-document
resolve failures, and keep replacement times_derived in the in-batch
ledger. Patch get_external_vector_store in the hoist test and cover
in-loop SQLAlchemyError abort.

* fix(deriver): address CodeRabbit findings on create_documents deadlock fix

- Distinguish external resolve failure ([] skip) from pgvector fallback (None)
  so _semantic_dup_decision never re-enters external I/O under an open session
- Bound external candidate hoist concurrency with a semaphore
- Map in-loop IntegrityError to ValidationException for a uniform contract
- Persist transient retry attempts on the oldest unprocessed queue item so
  every deriver instance shares one MAX_RETRYABLE_ATTEMPTS budget
- Cover resolve-failure skip and multi-manager reclaim of the retry budget

* 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

* fix(deriver): address review on deadlock retry and row-lock apply

Strip _retry_attempts before payload validation so non-representation
tasks are not burned as extra_forbidden. Re-raise retryable observer
save errors after telemetry so the queue actually retries. Skip
same-batch reinforce fallbacks after a replace. Revert unordered
FOR UPDATE on mark processed/errored and drop post-commit retry
cleanup from the success path.

* fix: add test and simplify queue query

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
Aakash Kattelu 2026-09-02 11:07:50 -04:00 committed by GitHub
parent a026bebdef
commit ccdb8ba113
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 1765 additions and 166 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
CLAUDE.md
View File

@ -118,6 +118,16 @@ cd sdks/typescript && bun run tsc --noEmit
- **Never hold a DB session during external calls** (LLM, embedding, HTTP). If a function needs both a DB session and an external call result, compute the external result first and pass it as a parameter. This avoids tying up DB connections during slow network I/O. Use `tracked_db` for short-lived, DB-only operations; pass a shared session when multiple DB-only calls can reuse one connection.
- **Never write through a read-only session** (`tracked_db(..., read_only=True)`, `get_read_db`, `ReadSessionLocal`). These run in AUTOCOMMIT mode with no transaction: writes are NOT blocked by the database — they silently commit immediately, and `begin_nested()` savepoints break. There is no runtime guard; this is enforced by convention only. Use `read_only=True` strictly for SELECT-only windows; anything that mutates (including get-or-create paths) must use a regular write session.
#### Multi-row locking and deadlocks
Tables written concurrently by more than one worker — `documents` (deriver, dreamer, scope backfill/removal, reconciler) and `queue` (every deriver replica) — deadlock when two writers touch an overlapping row set in different orders. Rules:
- **A multi-row `SELECT ... FOR UPDATE` MUST carry an explicit `ORDER BY <pk>`.** Without it Postgres locks in scan order, which differs per plan, so two writers with overlapping sets can cycle. `_apply_document_row_updates` in `src/crud/document.py` is the reference implementation.
- **`WHERE id IN (...)` does NOT impose an order**, so sorting the Python list is a no-op — the list order is discarded and the planner picks `Bitmap Heap Scan` (ctid order), `Index Scan` (id order), or `Seq Scan` per invocation. Deterministic ordering requires either a preceding `SELECT ... ORDER BY id FOR UPDATE` or `WHERE id IN (SELECT id ... ORDER BY id FOR UPDATE)`.
- **`Document.id` is a random nanoid** (`models.py`), so id order is uncorrelated with physical order — an unordered predicate `UPDATE`/`DELETE` is roughly a coin flip against an id-ordered locker per row pair, not a rare edge case. (`QueueItem.id` is an integer identity, so there id order is also chronological.)
- **Prefer no lock at all.** A single `UPDATE ... WHERE <predicate>` acquires row locks as it writes and has no separate lock phase to get wrong. Reach for `FOR UPDATE` only when a value must be read, computed in Python, and written back — that read-modify-write is the only reason `_apply_document_row_updates` locks (it replaced a server-side `func.greatest()`), and `populate_existing=True` is required with it so the identity map doesn't serve a stale pre-lock value. Server-side expressions (`func.greatest`, the JSONB `-` operator) avoid the lock entirely; see `_clear_work_unit_retry_attempts` in `src/deriver/queue_manager.py`.
- `FOR UPDATE SKIP LOCKED` (the reconciler's claim pattern) never waits, so it cannot be a deadlock partner — but holding those locks across an external call still stalls other writers. See the "never hold a DB session during external calls" rule above.
#### Auth scoping
- **`allow_member_read=True` (in `require_auth(...)`) is read-only — NEVER set it on a route that mutates state.** It lets a peer-scoped key reach a session route when its peer is an active member of the session, so on a mutating route it would hand any session member write access (message injection, config mutation, deletion). HTTP method is not a reliable read/write signal here (some read routes use POST for a richer body), so this is enforced by an explicit allowlist in `tests/routes/test_auth_route_policy.py` — adding the flag to a new route fails that test until you consciously add the route to `EXPECTED_MEMBER_READ_ROUTES`, and you must never add a mutating method there.

456
src/crud/document.py
View File

@ -1,13 +1,14 @@
import asyncio
import datetime
from collections.abc import Sequence
from dataclasses import dataclass, field
from enum import Enum
from logging import getLogger
from typing import Any, cast
from typing import Any, Literal, cast
from sqlalchemy import delete, select, 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:
)
# Shared by is_rejected_duplicate and create_documents candidate resolution.
_SEMANTIC_DUP_MAX_DISTANCE = 0.05
_SEMANTIC_DUP_TOP_K = 1
_SEMANTIC_CANDIDATE_CONCURRENCY = 8
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,
@ -473,6 +492,16 @@ def _dedup_key(
)
@dataclass(frozen=True, slots=True)
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
class CreateDocumentsResult:
created_documents: list[schemas.DocumentCreate] = field(default_factory=list)
@ -515,6 +544,43 @@ 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-store dup candidates before the first DB statement.
# None = pgvector in-place fallback; [] = skip semantic (no external I/O under db).
semantic_candidates: list[list[str] | None] = [None] * len(documents)
if deduplicate and not _uses_pgvector():
resolve_sem = asyncio.Semaphore(_SEMANTIC_CANDIDATE_CONCURRENCY)
async def _resolve_candidates(index: int, doc: schemas.DocumentCreate) -> None:
filters = _semantic_dup_filters(doc)
if filters is None or not doc.embedding:
semantic_candidates[index] = []
return
async with resolve_sem:
try:
ids = 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,
)
except Exception:
logger.exception(
"External semantic-candidate resolve failed for %s/%s/%s",
workspace_name,
observer,
observed,
)
semantic_candidates[index] = []
return
semantic_candidates[index] = ids or []
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
@ -563,12 +629,14 @@ async def create_documents(
# 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()
row_ops: list[_DocumentRowOp] = []
pending_times_derived: dict[str, int] = {}
exact_dup_existing_count = 0
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
@ -598,88 +666,107 @@ async def create_documents(
# the re-derivation as reinforcement on the existing row.
existing_match = existing_by_key.get(dedup_key)
if existing_match is not None:
# Reinforce the existing row. greatest(...) keeps the bump atomic
# server-side (concurrent workers can't lose an increment) while
# still honoring an incoming doc that already carries accumulated
# reinforcement (times_derived > 1, e.g. a future re-ingestion or
# collection-merge path). Mirrors the superior-replacement branch
# in is_rejected_duplicate.
existing_match.times_derived = func.greatest(
models.Document.times_derived + 1,
doc.times_derived,
current_td = pending_times_derived.get(
existing_match.id, existing_match.times_derived
)
pending_times_derived[existing_match.id] = max(
current_td + 1, doc.times_derived
)
row_ops.append(
_DocumentRowOp(
"reinforce",
existing_match.id,
doc.times_derived,
fallback_document=doc,
)
)
await db.flush()
exact_dup_existing_count += 1
continue
# for each document, if deduplicate is True, perform a process
# that checks against existing documents and either rejects this document
# 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
duplicate_result, existing_dup = await _semantic_dup_decision(
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
# new doc still gets inserted below.
if (
duplicate_result is SemanticRejectionResult.REPLACED_EXISTING
and existing_dup is not None
):
current_td = pending_times_derived.get(
existing_dup.id, existing_dup.times_derived
)
doc.times_derived = max(doc.times_derived, current_td + 1)
pending_times_derived[existing_dup.id] = doc.times_derived
row_ops.append(_DocumentRowOp("replace", existing_dup.id))
semantic_dup_replaced_count += 1
elif duplicate_result is SemanticRejectionResult.REJECTED:
elif (
duplicate_result is SemanticRejectionResult.REJECTED
and existing_dup is not None
):
current_td = pending_times_derived.get(
existing_dup.id, existing_dup.times_derived
)
pending_times_derived[existing_dup.id] = max(
current_td + 1, doc.times_derived
)
row_ops.append(
_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))
except IntegrityError as e:
await db.rollback()
raise ValidationException(
"Failed to create documents due to integrity constraint violation"
) from e
except SQLAlchemyError:
# Dead transaction: continuing would cascade PendingRollbackErrors.
await db.rollback()
raise
except Exception as e:
# Per-document failures (bad content, metadata, token overflow).
logger.error(
f"Error adding new document to {workspace_name}/{doc.session_name}/{observer}/{observed}: {e}"
)
continue
try:
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,
@ -775,6 +862,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,
@ -1152,12 +1244,163 @@ 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],
*,
workspace_name: str,
observer: str,
observed: str,
) -> list[schemas.DocumentCreate]:
"""Lock target rows by id, apply ops, return fallbacks for vanished targets."""
if not ops:
return []
# Deadlock fix: lock in id order (IN-clause order is ignored).
ids = sorted({op.document_id for op in ops})
result = await db.execute(
select(models.Document)
.where(
models.Document.id.in_(ids),
models.Document.workspace_name == workspace_name,
models.Document.observer == observer,
models.Document.observed == observed,
)
.order_by(models.Document.id)
.with_for_update()
# Reload identity-map rows so the Python max() sees concurrent increments.
.execution_options(populate_existing=True)
)
locked = {doc.id: doc for doc in result.scalars()}
now = datetime.datetime.now(datetime.UTC)
fallbacks: list[schemas.DocumentCreate] = []
stale_at_lock = {
op.document_id
for op in ops
if (locked_row := locked.get(op.document_id)) is None
or locked_row.deleted_at is not None
}
for op in ops:
row = locked.get(op.document_id)
if op.kind == "replace":
if row is not None and row.deleted_at is None:
row.deleted_at = now
continue
# reinforce
if op.document_id in stale_at_lock:
if op.fallback_document is not None:
fallbacks.append(op.fallback_document)
continue
if row is None or row.deleted_at is not None:
# An earlier op in this batch replaced this row.
continue
row.times_derived = max(row.times_derived + 1, op.incoming_times_derived)
await db.flush()
return fallbacks
class SemanticRejectionResult(Enum):
NOT_DUPLICATE = 0
REPLACED_EXISTING = 1
REJECTED = 2
async def _semantic_dup_decision(
db: AsyncSession,
doc: schemas.DocumentCreate,
workspace_name: str,
*,
observer: str,
observed: str,
candidate_document_ids: list[str] | None = None,
) -> tuple[SemanticRejectionResult, models.Document | None]:
"""Classify a semantic duplicate without writing."""
filters = _semantic_dup_filters(doc)
if filters is None:
return SemanticRejectionResult.NOT_DUPLICATE, None
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,
)
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,
query=doc.content,
observer=observer,
observed=observed,
filters=filters,
max_distance=_SEMANTIC_DUP_MAX_DISTANCE,
top_k=_SEMANTIC_DUP_TOP_K,
embedding=doc.embedding,
)
else:
return SemanticRejectionResult.NOT_DUPLICATE, None
if not similar_docs:
return SemanticRejectionResult.NOT_DUPLICATE, None
existing_doc = similar_docs[0]
tokens_new = set(embedding_client.encoding.encode(doc.content))
tokens_existing = set(embedding_client.encoding.encode(existing_doc.content))
unique_new = len(tokens_new - tokens_existing)
unique_existing = len(tokens_existing - tokens_new)
score_new = len(tokens_new) + (unique_new * 10)
score_existing = len(tokens_existing) + (unique_existing * 10)
if score_new >= score_existing:
return SemanticRejectionResult.REPLACED_EXISTING, existing_doc
return SemanticRejectionResult.REJECTED, existing_doc
async def is_rejected_duplicate(
db: AsyncSession,
doc: schemas.DocumentCreate,
@ -1165,90 +1408,29 @@ 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.
Uses: 1) Cosine similarity (>=0.95), 2) Token diff for retention.
Returns True if both:
- the document is deemed a duplicate of an existing document
- the existing document is deemed a superior duplicate
If the document is not a duplicate, returns False.
If the document is a duplicate AND the new document is superior,
deletes the existing document and returns False. In this case
``doc.times_derived`` is updated in place to carry the replaced
document's reinforcement count forward.
If the document is a duplicate AND the existing document is superior,
increments the existing document's ``times_derived`` to record the
reinforcement, then returns True.
Merges are scoped so they never cross document levels, and never cross
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).
"""
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
# Step 1: Find potential duplicates using cosine similarity
similar_docs = await query_documents(
db=db,
workspace_name=workspace_name,
query=doc.content,
"""Classify a semantic duplicate and apply the corresponding row write."""
result, existing_doc = await _semantic_dup_decision(
db,
doc,
workspace_name,
observer=observer,
observed=observed,
filters=filters,
max_distance=0.05,
top_k=1,
embedding=doc.embedding,
candidate_document_ids=candidate_document_ids,
)
if not similar_docs:
return SemanticRejectionResult.NOT_DUPLICATE
existing_doc = similar_docs[0]
# Step 2: Determine which has more information using token set difference
tokens_new = set(embedding_client.encoding.encode(doc.content))
tokens_existing = set(embedding_client.encoding.encode(existing_doc.content))
unique_new = len(tokens_new - tokens_existing)
unique_existing = len(tokens_existing - tokens_new)
score_new = len(tokens_new) + (unique_new * 10)
score_existing = len(tokens_existing) + (unique_existing * 10)
# If new document has more or equal information, keep it and delete existing
if score_new >= score_existing:
if existing_doc is None:
return result
if result is SemanticRejectionResult.REPLACED_EXISTING:
logger.debug(
"[DUPLICATE DETECTION] Deleting existing in favor of new. new=%r, existing=%r.",
doc.content,
existing_doc.content,
)
# Carry the reinforcement count forward so replacing a duplicate counts as
# 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
existing_doc.deleted_at = datetime.datetime.now(datetime.timezone.utc)
existing_doc.deleted_at = datetime.datetime.now(datetime.UTC)
await db.flush()
return (
SemanticRejectionResult.REPLACED_EXISTING
) # Don't reject the new document
# Existing document has more information, reject the new one but record the
# reinforcement: a semantic duplicate was derived again. greatest(...) keeps
# 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).
return result
existing_doc.times_derived = func.greatest(
models.Document.times_derived + 1,
doc.times_derived,
@ -1259,7 +1441,7 @@ async def is_rejected_duplicate(
doc.content,
existing_doc.content,
)
return SemanticRejectionResult.REJECTED
return result
async def cleanup_soft_deleted_documents(
@ -1284,7 +1466,7 @@ async def cleanup_soft_deleted_documents(
Returns:
Count of documents cleaned up (only those where vector deletion succeeded).
"""
cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(
cutoff = datetime.datetime.now(datetime.UTC) - datetime.timedelta(
minutes=older_than_minutes
)

7
src/deriver/consumer.py
View File

@ -27,6 +27,7 @@ from src.telemetry.events import (
from src.telemetry.logging import log_performance_metrics
from src.utils import summarizer
from src.utils.queue_payload import (
RETRY_ATTEMPTS_PAYLOAD_KEY,
DeletionPayload,
DreamPayload,
ReconcilerPayload,
@ -44,7 +45,11 @@ logging.getLogger("sqlalchemy.engine.Engine").disabled = True
async def process_item(queue_item: models.QueueItem) -> None:
"""Process a single item from the queue."""
task_type = queue_item.task_type
queue_payload = queue_item.payload
# Drop the work-unit retry counter before payload validation: every payload
# model sets extra="forbid", so leaving it in burns the item as
# extra_forbidden on the reclaim that was supposed to retry it.
queue_payload = dict(queue_item.payload or {})
queue_payload.pop(RETRY_ATTEMPTS_PAYLOAD_KEY, None)
workspace_name = queue_item.workspace_name
# Handle reconciler first - it's the only task type that doesn't require workspace_name

7
src/deriver/deriver.py
View File

@ -25,6 +25,7 @@ from src.telemetry.sentry import with_sentry_transaction
from src.utils.config_helpers import get_configuration
from src.utils.formatting import format_new_turn_with_timestamp
from src.utils.representation import PromptRepresentation, Representation
from src.utils.retryable_errors import is_retryable_error
from src.utils.tokens import track_deriver_input_tokens
from .prompts import estimate_deriver_prompt_tokens, minimal_deriver_prompt
@ -344,6 +345,12 @@ async def process_representation_tasks_batch(
)
)
retryable = next(
(exc for _, exc in save_errors if is_retryable_error(exc)),
None,
)
if retryable is not None:
raise retryable
if save_errors and successful_observer_count == 0:
details = "; ".join(
f"{observer}: {exc.__class__.__name__}: {exc}"

160
src/deriver/queue_manager.py
View File

@ -6,7 +6,7 @@ import time
from asyncio import Task
from collections.abc import Iterable, Sequence
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from datetime import UTC, datetime, timedelta
from logging import getLogger
from typing import Any, NamedTuple, cast
@ -15,7 +15,7 @@ from dotenv import load_dotenv
from nanoid import generate as generate_nanoid
from sentry_sdk.integrations.asyncio import AsyncioIntegration
from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration
from sqlalchemy import and_, delete, or_, select, update
from sqlalchemy import Text, and_, delete, literal, or_, select, update
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.engine import CursorResult
from sqlalchemy.ext.asyncio import AsyncSession
@ -43,6 +43,8 @@ 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.queue_payload import RETRY_ATTEMPTS_PAYLOAD_KEY
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 +55,12 @@ logger = getLogger(__name__)
load_dotenv(override=True)
# Total processing attempts per work unit for transient errors. Count is
# stored on the oldest unprocessed queue item so every deriver instance
# shares one budget.
MAX_RETRYABLE_ATTEMPTS = 3
RETRY_BACKOFF_SECONDS = 1.0
class WorkerOwnership(NamedTuple):
"""Represents the instance of a work unit that a worker is processing."""
@ -301,7 +309,7 @@ class QueueManager:
async def cleanup_stale_work_units(self) -> None:
"""Clean up stale work units"""
async with tracked_db("cleanup_stale_work_units") as db:
cutoff = datetime.now(timezone.utc) - timedelta(
cutoff = datetime.now(UTC) - timedelta(
minutes=settings.DERIVER.STALE_SESSION_TIMEOUT_MINUTES
)
@ -591,11 +599,24 @@ 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.
The attempt count lives on the oldest unprocessed queue item so a
different deriver instance continues the same budget after reclaim.
Reprocessing is at-least-once, not idempotent: the batch is re-derived
by a fresh LLM call, so identical text collapses via exact dedup and
near-identical text via semantic dedup. Retries can therefore inflate
times_derived and double-count LLM telemetry -- acceptable because the
alternative is dropping the batch.
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
@ -603,12 +624,37 @@ 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):
try:
attempts = await self._get_work_unit_retry_attempts(work_unit_key) + 1
if attempts < MAX_RETRYABLE_ATTEMPTS:
await self._set_work_unit_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
except Exception: # noqa: BLE001
logger.exception(
"Retry-counter I/O failed for work unit %s; releasing %s without recording an attempt",
work_unit_key,
context,
)
return True
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}",
@ -621,6 +667,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."""
@ -686,12 +733,18 @@ class QueueManager:
)
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(
@ -710,12 +763,16 @@ class QueueManager:
)
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(
@ -1068,6 +1125,87 @@ class QueueManager:
batch_max_tokens=batch_max_tokens,
)
async def _oldest_unprocessed_item(
self,
db: AsyncSession,
work_unit_key: str,
*,
for_update: bool = False,
) -> models.QueueItem | None:
stmt = (
select(models.QueueItem)
.where(
models.QueueItem.work_unit_key == work_unit_key,
models.QueueItem.processed.is_(False),
)
.order_by(models.QueueItem.id)
.limit(1)
)
if for_update:
stmt = stmt.with_for_update()
result = await db.execute(stmt)
return result.scalar_one_or_none()
async def _get_work_unit_retry_attempts(self, work_unit_key: str) -> int:
"""Read the shared transient-failure attempt count for a work unit."""
async with tracked_db("get_work_unit_retry_attempts") as db:
item = await self._oldest_unprocessed_item(db, work_unit_key)
if item is None:
return 0
raw = (item.payload or {}).get(RETRY_ATTEMPTS_PAYLOAD_KEY, 0)
try:
return max(0, int(raw))
except (TypeError, ValueError):
return 0
async def _set_work_unit_retry_attempts(
self, work_unit_key: str, attempts: int
) -> None:
"""Persist the shared attempt count on the oldest unprocessed item."""
async with tracked_db("set_work_unit_retry_attempts") as db:
item = await self._oldest_unprocessed_item(
db, work_unit_key, for_update=True
)
if item is None:
await db.commit()
return
new_payload = dict(item.payload or {})
new_payload[RETRY_ATTEMPTS_PAYLOAD_KEY] = attempts
await db.execute(
update(models.QueueItem)
.where(models.QueueItem.id == item.id)
.values(payload=new_payload)
)
await db.commit()
async def _clear_work_unit_retry_attempts(self, work_unit_key: str) -> None:
"""Drop the shared attempt count from remaining unprocessed items.
One statement on purpose: a multi-row ``SELECT ... FOR UPDATE`` here
would take locks on ``queue`` in scan order, which is a deadlock partner
for any other multi-row writer on the same table. The JSONB ``-``
operator does the strip server-side, so no rows are locked ahead of the
write and there is no lock order to get wrong.
"""
async with tracked_db("clear_work_unit_retry_attempts") as db:
await db.execute(
update(models.QueueItem)
.where(
models.QueueItem.work_unit_key == work_unit_key,
models.QueueItem.processed.is_(False),
models.QueueItem.payload.has_key(RETRY_ATTEMPTS_PAYLOAD_KEY),
)
.values(
# literal(..., Text) is required: an untyped bind leaves
# Postgres unable to pick between jsonb - text and its
# integer/array siblings.
payload=models.QueueItem.payload.op("-")(
literal(RETRY_ATTEMPTS_PAYLOAD_KEY, Text)
)
)
)
await db.commit()
async def mark_queue_items_as_processed(
self, items: list[QueueItem], work_unit_key: str
) -> None:

8
src/utils/queue_payload.py
View File

@ -5,6 +5,14 @@ from pydantic import BaseModel, ConfigDict
from src.schemas import DreamType, ReconcilerType, ResolvedConfiguration
# Queue mechanics, not task data: the deriver stores a per-work-unit transient
# failure count under this key so a retry budget survives work-unit reclaim.
# Every payload model below forbids extras, so anything that reads a raw
# QueueItem.payload must strip this key before validating. Lives here rather
# than in the deriver because both the writer (queue_manager) and the stripper
# (consumer) need it, and queue_manager imports consumer.
RETRY_ATTEMPTS_PAYLOAD_KEY = "_retry_attempts"
class BasePayload(BaseModel):
"""Base payload with common fields."""

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

642
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
@ -195,7 +198,7 @@ class TestDocumentCRUD:
deleted_doc = docs["User likes pizza"]
kept_doc = docs["User dislikes vegetables"]
deleted_doc.deleted_at = datetime.datetime.now(datetime.timezone.utc)
deleted_doc.deleted_at = datetime.datetime.now(datetime.UTC)
await db_session.commit()
results = await crud.query_documents(
@ -290,7 +293,7 @@ class TestDocumentCRUD:
db_session, test_workspace, test_peer
)
base = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc)
base = datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC)
# Three conclusions, all reinforced once -- the real-world steady state
# before the fix -- inserted oldest-first.
for i in range(3):
@ -1374,3 +1377,636 @@ class TestSessionPurityInvariant:
)
assert rejected is SemanticRejectionResult.NOT_DUPLICATE
mock_query.assert_not_awaited()
class TestCreateDocumentsConcurrency:
"""Concurrent same-collection reinforcements lock rows in id order."""
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],
):
"""Opposing-order batches on one collection must not deadlock."""
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; per-document failures skip one document."""
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_on_row_update_flush_aborts_batch(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
):
"""A DB error while applying row updates raises and commits 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())
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_db_error_in_loop_aborts_batch(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
):
"""A DB error during per-document classification raises and commits nothing."""
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
class FakePGError(Exception):
sqlstate: str = "40P01"
deadlock = OperationalError("SELECT documents", {}, FakePGError())
with (
patch(
"src.crud.document._semantic_dup_decision",
AsyncMock(side_effect=deadlock),
),
pytest.raises(OperationalError),
):
await crud.create_documents(
db_session,
[
self._doc("a brand new fact", session_name),
self._doc("another new fact", session_name),
],
workspace_name=workspace_name,
observer=observer,
observed=observed,
deduplicate=True,
)
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 docs == []
@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",
]
@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.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"
@pytest.mark.asyncio
async def test_same_batch_replace_then_reinforce_does_not_resurrect(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
):
"""A reinforce after a same-batch replace must not insert the inferior copy."""
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
await crud.create_documents(
db_session,
[self._doc("shared fact", session_name)],
workspace_name=workspace_name,
observer=observer,
observed=observed,
)
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()
fallback = self._doc("shared fact", session_name)
ops = [
document_module._DocumentRowOp("replace", existing.id), # pyright: ignore[reportPrivateUsage]
document_module._DocumentRowOp( # pyright: ignore[reportPrivateUsage]
"reinforce",
existing.id,
fallback_document=fallback,
),
]
fallbacks = await document_module._apply_document_row_updates( # pyright: ignore[reportPrivateUsage]
db_session,
ops,
workspace_name=workspace_name,
observer=observer,
observed=observed,
)
assert fallbacks == []
await db_session.commit()
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 live == []
class TestExternalCandidateHoist:
"""External-store dup candidates resolve before the first DB statement."""
async def _setup(
self,
db_session: AsyncSession,
test_workspace: models.Workspace,
test_peer: models.Peer,
) -> tuple[models.Peer, models.Session]:
observed_peer = 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([observed_peer, test_session])
await db_session.flush()
db_session.add(
models.Collection(
workspace_name=test_workspace.name,
observer=test_peer.name,
observed=observed_peer.name,
)
)
await db_session.commit()
return observed_peer, 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_external_candidates_resolved_before_db(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
monkeypatch: pytest.MonkeyPatch,
):
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:
events.append("execute")
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,
),
patch(
"src.crud.document.get_external_vector_store",
return_value=None,
),
):
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[:2] == ["resolve", "resolve"]
assert "execute" in events
@pytest.mark.asyncio
async def test_resolve_failure_skips_semantic_without_query_documents(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
monkeypatch: pytest.MonkeyPatch,
):
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)
with (
patch(
"src.crud.document.query_external_vector_document_ids",
side_effect=RuntimeError("store down"),
),
patch(
"src.crud.document.get_external_vector_store",
return_value=None,
),
patch(
"src.crud.document.query_documents",
new_callable=AsyncMock,
) as mock_query,
):
result = await crud.create_documents(
db_session,
[self._doc("fact one", test_session.name)],
workspace_name=test_workspace.name,
observer=test_peer.name,
observed=observed_peer.name,
deduplicate=True,
)
assert len(result.created_documents) == 1
mock_query.assert_not_awaited()

71
tests/deriver/test_deriver_processing.py
View File

@ -1,5 +1,5 @@
import signal
from datetime import datetime, timezone
from datetime import UTC, datetime
from typing import Any
from unittest.mock import AsyncMock, Mock, patch
@ -32,7 +32,7 @@ class TestDeriverProcessing:
peer_name="alice",
content="hello",
token_count=5,
created_at=datetime.now(timezone.utc),
created_at=datetime.now(UTC),
)
configuration = Mock()
configuration.reasoning.enabled = True
@ -82,7 +82,7 @@ class TestDeriverProcessing:
peer_name="alice",
content="hello",
token_count=5,
created_at=datetime.now(timezone.utc),
created_at=datetime.now(UTC),
)
configuration = Mock()
configuration.reasoning.enabled = True
@ -136,7 +136,7 @@ class TestDeriverProcessing:
peer_name="alice",
content="hello",
token_count=5,
created_at=datetime.now(timezone.utc),
created_at=datetime.now(UTC),
)
configuration = Mock()
configuration.reasoning.enabled = True
@ -182,6 +182,61 @@ class TestDeriverProcessing:
assert event.observer_count == 1
assert event.failed_observer_count == 1
async def test_retryable_observer_save_reraises_after_telemetry(self):
"""A deadlock on one observer must propagate so the queue can retry."""
from sqlalchemy.exc import OperationalError
class FakePGError(Exception):
sqlstate: str = "40P01"
deadlock = OperationalError("UPDATE documents", {}, FakePGError())
message = Mock(
id=1,
public_id="msg_1",
session_name="session-1",
workspace_name="workspace-1",
peer_name="alice",
content="hello",
token_count=5,
created_at=datetime.now(UTC),
)
configuration = Mock()
configuration.reasoning.enabled = True
mock_response = HonchoLLMCallResponse(
content=PromptRepresentation(
explicit=[
ExplicitObservationBase(content="The user has a dog named Rover")
]
),
input_tokens=10,
output_tokens=5,
finish_reasons=["STOP"],
)
partial_save = AsyncMock(side_effect=[crud.CreateDocumentsResult(), deadlock])
emitted: list[Any] = []
with (
patch(
"src.deriver.deriver.honcho_llm_call",
new_callable=AsyncMock,
return_value=mock_response,
),
patch.object(RepresentationManager, "save_representation", partial_save),
patch("src.deriver.deriver.emit", side_effect=emitted.append),
pytest.raises(OperationalError),
):
await process_representation_tasks_batch(
messages=[message],
message_level_configuration=configuration,
observers=["bob", "carol"],
observed="alice",
queue_item_message_ids=[1],
)
assert emitted, "expected telemetry to be emitted before the raised failure"
assert emitted[-1].observer_count == 1
assert emitted[-1].failed_observer_count == 1
async def test_process_representation_tasks_batch_passes_custom_instructions_into_prompt(
self,
) -> None:
@ -193,7 +248,7 @@ class TestDeriverProcessing:
peer_name="alice",
content="hello",
token_count=5,
created_at=datetime.now(timezone.utc),
created_at=datetime.now(UTC),
)
configuration = Mock()
configuration.reasoning.enabled = True
@ -343,7 +398,7 @@ class TestDeriverProcessing:
peer_name="alice",
content="hello",
token_count=100,
created_at=datetime.now(timezone.utc),
created_at=datetime.now(UTC),
)
configuration = Mock()
configuration.reasoning.enabled = True
@ -394,7 +449,7 @@ class TestDeriverProcessing:
peer_name="alice",
content="hello",
token_count=5,
created_at=datetime.now(timezone.utc),
created_at=datetime.now(UTC),
)
configuration = Mock()
configuration.reasoning.enabled = True
@ -443,7 +498,7 @@ class TestDeriverProcessing:
peer_name="alice",
content="hello",
token_count=5,
created_at=datetime.now(timezone.utc),
created_at=datetime.now(UTC),
)
configuration = Mock()
configuration.reasoning.enabled = True

367
tests/deriver/test_queue_processing.py
View File

@ -1,17 +1,21 @@
import asyncio
from collections.abc import Callable
from datetime import datetime, timedelta, timezone
from datetime import UTC, datetime, timedelta
from typing import Any
from unittest.mock import patch
from unittest.mock import AsyncMock, patch
import pytest
from nanoid import generate as generate_nanoid
from pydantic import ValidationError
from sqlalchemy import select
from sqlalchemy.exc import OperationalError
from sqlalchemy.ext.asyncio import AsyncSession
from src import models
from src.config import settings
from src.deriver.consumer import process_item
from src.deriver.queue_manager import QueueManager, WorkerOwnership
from src.utils.queue_payload import RETRY_ATTEMPTS_PAYLOAD_KEY, SummaryPayload
from src.utils.work_unit import construct_work_unit_key
@ -1519,7 +1523,7 @@ class TestQueueProcessing:
monkeypatch.setattr(
settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800
)
old_timestamp = datetime.now(timezone.utc) - timedelta(hours=2)
old_timestamp = datetime.now(UTC) - timedelta(hours=2)
work_unit_key, queue_items = await self._add_representation_work_unit(
db_session=db_session,
@ -1552,7 +1556,7 @@ class TestQueueProcessing:
) -> None:
monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", False)
monkeypatch.setattr(settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 0)
old_timestamp = datetime.now(timezone.utc) - timedelta(hours=2)
old_timestamp = datetime.now(UTC) - timedelta(hours=2)
work_unit_key, _queue_items = await self._add_representation_work_unit(
db_session=db_session,
@ -1602,7 +1606,7 @@ class TestQueueProcessing:
monkeypatch.setattr(
settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800
)
now = datetime.now(timezone.utc)
now = datetime.now(UTC)
work_unit_key, _queue_items = await self._add_representation_work_unit(
db_session=db_session,
@ -1628,7 +1632,7 @@ class TestQueueProcessing:
monkeypatch.setattr(
settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800
)
now = datetime.now(timezone.utc)
now = datetime.now(UTC)
work_unit_key, queue_items = await self._add_representation_work_unit(
db_session=db_session,
@ -1874,3 +1878,354 @@ 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 _retry_attempts_on_items(
self, db_session: AsyncSession, work_unit_key: str
) -> int | None:
items = await self._fetch_items(db_session, work_unit_key)
unprocessed = [item for item in items if not item.processed]
if not unprocessed:
return None
raw = (unprocessed[0].payload or {}).get("_retry_attempts")
return None if raw is None else int(raw)
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 await self._retry_attempts_on_items(db_session, work_unit_key) == 1
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
)
await qm._set_work_unit_retry_attempts( # pyright: ignore[reportPrivateUsage]
work_unit_key, MAX_RETRYABLE_ATTEMPTS - 1
)
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 await self._retry_attempts_on_items(db_session, work_unit_key) is None
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 await self._retry_attempts_on_items(db_session, work_unit_key) is None
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
)
await qm._set_work_unit_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)
# Counter lives on the oldest unprocessed item; once that item is
# processed the budget is gone even if the payload key remains.
assert await self._retry_attempts_on_items(db_session, work_unit_key) is None
async def test_retry_budget_survives_reclaim_by_another_manager(
self,
db_session: AsyncSession,
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
create_queue_payload: Callable[..., Any],
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A second QueueManager continues the durable attempt budget."""
from src.deriver.queue_manager import MAX_RETRYABLE_ATTEMPTS
monkeypatch.setattr("src.deriver.queue_manager.RETRY_BACKOFF_SECONDS", 0.0)
qm1, 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=self._retryable_error(),
):
await qm1.process_work_unit(work_unit_key, worker_id)
assert await self._retry_attempts_on_items(db_session, work_unit_key) == 1
assert await self._aqs_rows(db_session, work_unit_key) == 0
# Seed the remaining budget so the next reclaim is the terminal attempt.
qm2 = QueueManager()
await qm2._set_work_unit_retry_attempts( # pyright: ignore[reportPrivateUsage]
work_unit_key, MAX_RETRYABLE_ATTEMPTS - 1
)
claimed = await qm2.claim_work_units(db_session, [work_unit_key])
worker_id_2 = "test_worker_2"
qm2.worker_ownership[worker_id_2] = WorkerOwnership(
work_unit_key=work_unit_key, aqs_id=claimed[work_unit_key]
)
await db_session.commit()
with patch(
"src.deriver.queue_manager.process_representation_batch",
side_effect=self._retryable_error(),
):
await qm2.process_work_unit(work_unit_key, worker_id_2)
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
async def test_process_item_strips_retry_counter_before_validation(self) -> None:
"""A reclaimed non-representation item must survive its own retry counter.
The counter is written onto an *unprocessed* item so the budget outlives
a work-unit reclaim -- which means the next claim re-reads it. Every
payload model sets ``extra="forbid"``, so without the strip in
``process_item`` the reclaim raises extra_forbidden -> ValueError ->
not retryable -> the item is burned terminally on the very attempt that
was supposed to retry it. Representation tasks never hit this: their
batch path reads the payload with ``.get()`` instead of validating,
which is why the rest of this class cannot catch it.
"""
raw: dict[str, Any] = {
"task_type": "summary",
"session_name": "s",
"message_seq_in_session": 1,
"message_public_id": "msg-public-id",
"configuration": {
"reasoning": {"enabled": True},
"peer_card": {"use": True, "create": True},
"summary": {
"enabled": True,
"messages_per_short_summary": 20,
"messages_per_long_summary": 60,
},
"dream": {"enabled": True},
},
RETRY_ATTEMPTS_PAYLOAD_KEY: 1,
}
# Pin the premise: the payload model must keep rejecting the key, so
# this fails loudly if someone "fixes" the burn with extra="allow"
# instead of stripping.
with pytest.raises(ValidationError) as exc_info:
SummaryPayload.model_validate(raw)
assert any(err["type"] == "extra_forbidden" for err in exc_info.value.errors())
queue_item = models.QueueItem(
task_type="summary",
work_unit_key="summary:test-workspace:test-session",
payload=raw,
processed=False,
workspace_name="test-workspace",
message_id=1,
)
with patch(
"src.deriver.consumer.summarizer.summarize_if_needed",
new_callable=AsyncMock,
) as mock_summarize:
await process_item(queue_item)
mock_summarize.assert_awaited_once()
# The strip must happen on a copy: the counter has to stay on the row so
# the budget still advances if this attempt fails again.
assert raw[RETRY_ATTEMPTS_PAYLOAD_KEY] == 1

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)