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.
This commit is contained in:
Aakash Kattelu 2026-09-01 12:05:58 -04:00
parent bef6274920
commit 99bd196129
7 changed files with 226 additions and 71 deletions

View File

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

View File

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

View File

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

View File

@ -6,7 +6,7 @@ import time
from asyncio import Task
from collections.abc import Iterable, Sequence
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from datetime import UTC, datetime, timedelta
from logging import getLogger
from typing import Any, NamedTuple, cast
@ -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)

View File

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

View File

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

View File

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