From 99bd196129849dcf6d2801c8f290392a4ba1507c Mon Sep 17 00:00:00 2001 From: Aakash Kattelu Date: Tue, 1 Sep 2026 12:05:58 -0400 Subject: [PATCH] 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. --- src/crud/document.py | 19 ++++-- src/deriver/consumer.py | 4 +- src/deriver/deriver.py | 7 +++ src/deriver/queue_manager.py | 79 +++++++++--------------- tests/crud/test_document.py | 76 ++++++++++++++++++++++- tests/deriver/test_deriver_processing.py | 71 ++++++++++++++++++--- tests/deriver/test_queue_processing.py | 41 ++++++++++-- 7 files changed, 226 insertions(+), 71 deletions(-) diff --git a/src/crud/document.py b/src/crud/document.py index 4e09e6c6..7a3fc63a 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -1296,6 +1296,7 @@ async def _apply_document_row_updates( """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) @@ -1307,11 +1308,18 @@ async def _apply_document_row_updates( ) .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.timezone.utc) + 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": @@ -1319,10 +1327,13 @@ async def _apply_document_row_updates( row.deleted_at = now continue # reinforce - if row is None or row.deleted_at is not None: + 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 @@ -1417,7 +1428,7 @@ async def is_rejected_duplicate( existing_doc.content, ) doc.times_derived = max(doc.times_derived, existing_doc.times_derived + 1) - existing_doc.deleted_at = datetime.datetime.now(datetime.timezone.utc) + existing_doc.deleted_at = datetime.datetime.now(datetime.UTC) await db.flush() return result existing_doc.times_derived = func.greatest( @@ -1455,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..8335798b 100644 --- a/src/deriver/consumer.py +++ b/src/deriver/consumer.py @@ -44,7 +44,9 @@ 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. + queue_payload = dict(queue_item.payload or {}) + queue_payload.pop("_retry_attempts", 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 463e9579..58532472 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 @@ -309,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 ) @@ -608,8 +608,11 @@ class QueueManager: 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 safe because a retried batch re-derives the same - observations and exact dedup collapses them into reinforcement. + 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 @@ -622,16 +625,24 @@ class QueueManager: context: Context string describing what was being processed (e.g., "processing representation batch") """ if is_retryable_error(error): - 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, + 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, - attempts, - MAX_RETRYABLE_ATTEMPTS, - exc_info=error, + context, ) return True @@ -720,9 +731,6 @@ class QueueManager: await self.mark_queue_items_as_processed( items_to_process, work_unit_key ) - await self._clear_work_unit_retry_attempts( - work_unit_key - ) queue_item_count += len(items_to_process) except Exception as e: if await self._handle_processing_error( @@ -753,9 +761,6 @@ class QueueManager: await self.mark_queue_items_as_processed( [queue_item], work_unit_key ) - await self._clear_work_unit_retry_attempts( - work_unit_key - ) queue_item_count += 1 except Exception as e: if await self._handle_processing_error( @@ -1211,23 +1216,12 @@ class QueueManager: async with tracked_db("process_queue_item_batch") as db: work_unit = parse_work_unit_key(work_unit_key) item_ids = [item.id for item in items] - result = await db.execute( - select(models.QueueItem) + await db.execute( + update(models.QueueItem) .where(models.QueueItem.id.in_(item_ids)) .where(models.QueueItem.work_unit_key == work_unit_key) - .with_for_update() + .values(processed=True) ) - for queue_item in result.scalars(): - await db.execute( - update(models.QueueItem) - .where(models.QueueItem.id == queue_item.id) - .values( - processed=True, - payload=self._payload_without_retry_attempts( - queue_item.payload - ), - ) - ) await db.execute( update(models.ActiveQueueSession) .where(models.ActiveQueueSession.work_unit_key == work_unit_key) @@ -1253,24 +1247,11 @@ class QueueManager: if not item: return async with tracked_db("mark_queue_item_as_errored") as db: - result = await db.execute( - select(models.QueueItem) - .where(models.QueueItem.id == item.id) - .where(models.QueueItem.work_unit_key == work_unit_key) - .with_for_update() - ) - queue_item = result.scalar_one_or_none() - if queue_item is None: - await db.commit() - return await db.execute( update(models.QueueItem) - .where(models.QueueItem.id == queue_item.id) - .values( - processed=True, - error=error[:65535], # Truncate to TEXT limit - payload=self._payload_without_retry_attempts(queue_item.payload), - ) + .where(models.QueueItem.id == item.id) + .where(models.QueueItem.work_unit_key == work_unit_key) + .values(processed=True, error=error[:65535]) # Truncate to TEXT limit ) await db.execute( update(models.ActiveQueueSession) diff --git a/tests/crud/test_document.py b/tests/crud/test_document.py index 99618756..593ee52b 100644 --- a/tests/crud/test_document.py +++ b/tests/crud/test_document.py @@ -198,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( @@ -293,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): @@ -1770,7 +1770,7 @@ class TestCreateDocumentsErrorHandling: real_apply = document_module._apply_document_row_updates # pyright: ignore[reportPrivateUsage] async def delete_then_apply(*args: Any, **kwargs: Any) -> Any: - existing.deleted_at = datetime.datetime.now(datetime.timezone.utc) + existing.deleted_at = datetime.datetime.now(datetime.UTC) await db_session.flush() return await real_apply(*args, **kwargs) @@ -1806,6 +1806,76 @@ class TestCreateDocumentsErrorHandling: 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.""" 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 76e3e925..99f44931 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -1,6 +1,6 @@ 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 @@ -1520,7 +1520,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, @@ -1553,7 +1553,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, @@ -1603,7 +1603,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, @@ -1629,7 +1629,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, @@ -2118,7 +2118,8 @@ class TestQueueRetry: items = await self._fetch_items(db_session, work_unit_key) assert all(item.processed for item in items) assert all(item.error is None for item in items) - assert all("_retry_attempts" not in (item.payload or {}) for item in items) + # 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( @@ -2168,3 +2169,31 @@ class TestQueueRetry: assert items[0].processed assert items[0].error is not None assert "OperationalError" in items[0].error + + async def test_summary_payload_forbids_retry_attempts_key(self) -> None: + from pydantic import ValidationError + + from src.deriver.queue_manager import QueueManager + from src.utils.queue_payload import SummaryPayload + + raw = { + "task_type": "summary", + "session_name": "s", + "message_seq_in_session": 1, + "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": 1, + } + with pytest.raises(ValidationError) as ei: + SummaryPayload.model_validate(raw) + assert any(err["type"] == "extra_forbidden" for err in ei.value.errors()) + cleaned = QueueManager._payload_without_retry_attempts(raw) # pyright: ignore[reportPrivateUsage] + SummaryPayload.model_validate(cleaned)