diff --git a/CLAUDE.md b/CLAUDE.md index 6a83066d..ab842f99 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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 `.** 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 ` 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. diff --git a/src/crud/document.py b/src/crud/document.py index 37eb94b4..7a3fc63a 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -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 ) diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py index 118135ca..6251dc73 100644 --- a/src/deriver/consumer.py +++ b/src/deriver/consumer.py @@ -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 diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py index f76c4d52..4c1e4dd2 100644 --- a/src/deriver/deriver.py +++ b/src/deriver/deriver.py @@ -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}" diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 1c19c131..b98c0ef6 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -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: diff --git a/src/utils/queue_payload.py b/src/utils/queue_payload.py index 59815ca5..6f3f1105 100644 --- a/src/utils/queue_payload.py +++ b/src/utils/queue_payload.py @@ -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.""" diff --git a/src/utils/retryable_errors.py b/src/utils/retryable_errors.py new file mode 100644 index 00000000..357f594d --- /dev/null +++ b/src/utils/retryable_errors.py @@ -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) + ) diff --git a/tests/crud/test_document.py b/tests/crud/test_document.py index 6686e688..593ee52b 100644 --- a/tests/crud/test_document.py +++ b/tests/crud/test_document.py @@ -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() diff --git a/tests/deriver/test_deriver_processing.py b/tests/deriver/test_deriver_processing.py index 29a983b5..86ab9d33 100644 --- a/tests/deriver/test_deriver_processing.py +++ b/tests/deriver/test_deriver_processing.py @@ -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 diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index 81dd0632..ce8e3e13 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -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 diff --git a/tests/utils/test_retryable_errors.py b/tests/utils/test_retryable_errors.py new file mode 100644 index 00000000..bead2d2e --- /dev/null +++ b/tests/utils/test_retryable_errors.py @@ -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)