fix(examples): bound durable workflow retries

Author: bokelley

docs/production-seller.md

@@ -62,10 +62,15 @@ The reference's `IdempotencyStore` wrapping is intentionally paired with its
 inline terminal responses. Do not put the same method-level wrapper around a
 method that returns a raw `TaskHandoff` or `WorkflowHandoff`: the wrapper runs
 before framework task issuance and therefore cannot cache the projected
-`{status: "submitted", task_id}` envelope. A durable workflow adapter must
-deduplicate task issuance in its queue/business transaction (and advertise
-idempotency only when it does so) until that projection boundary is supported
-directly by the SDK.
+`{status: "submitted", task_id}` envelope. A durable workflow adapter must not
+advertise method-level idempotency for that method unless an external durable
+request-to-task mapping can reuse the prior task id. The reference queue's
+uniqueness constraint is only on the framework-issued `task_id`, not the
+buyer's idempotency key. A web-process
+crash after enqueue commits but before the `submitted` response reaches the
+buyer can therefore cause a retried request to issue a second task and queue
+row. Fully closing that window requires SDK support for looking up or reusing
+a workflow task id by buyer idempotency key.
 
 For a mixed adapter, make the split explicit with
 `method_level_idempotency_methods`: include only methods that return terminal
@@ -98,7 +103,10 @@ artifact.
 [`workflow_queue.py`](../examples/v3_reference_seller/src/workflow_queue.py)
 is a PostgreSQL-backed adopter queue with expiring leases. The enqueue callback
 stores the framework task id before `WorkflowHandoff` returns `submitted`; a
-replacement worker reclaims an expired lease after a crash.
+replacement worker reclaims an expired lease after a crash. Handler failures
+retry with capped exponential backoff; after the configured attempt limit the
+registry task fails and the queue row moves to `dead_lettered`. Jobs without a
+matching account-scoped registry task dead-letter immediately.
 
 ```python
 queue = task_wiring.workflow_queue
@@ -198,8 +206,9 @@ missing field names. The retry horizon is projected into capabilities and must
 match the outbox value.
 
 `DurableTaskWiring.startup()` calls `create_schema()` for a convenient local
-bootstrap. Those idempotent `CREATE TABLE IF NOT EXISTS` calls are not a schema
-migration system: they do not detect or evolve a table with the wrong shape.
+bootstrap. The workflow example performs the one additive upgrade shown here,
+but these runtime DDL calls are not a general schema migration system and do
+not detect or safely evolve an arbitrarily mismatched table.
 For production, copy the SDK-owned SQL files (`decisioning_tasks.sql` and
 `task_webhook_outbox.sql`), the `PgBackend.create_schema()` DDL, and the
 reference workflow-queue DDL into reviewed, versioned migrations and apply

examples/v3_reference_seller/README.md

@@ -96,9 +96,16 @@ The bundled mock upstream resolves approvals quickly, so
 `create_media_buy` completes inline. When adapting the example to a human or
 long-running approval system, return `ctx.handoff_to_workflow(...)` and let a
 durable queue consumer call `registry.complete()` or `registry.fail()`; do not
-move a long polling loop into an in-process `TaskHandoff`. That queue must also
-deduplicate workflow issuance before enabling the idempotency capability; the
-reference's method wrapper is for its inline terminal-response path.
+move a long polling loop into an in-process `TaskHandoff`. Leave method-level
+idempotency unadvertised for that method unless an external durable
+request-to-task mapping can reuse the prior task id. The reference's method
+wrapper is for its inline terminal-response path.
+The reference queue deduplicates only the framework-issued `task_id`. It does
+not close the crash window between committing enqueue and returning
+`submitted`, because a buyer retry can receive a newly issued task id. That
+requires SDK support for reusing a workflow task id by buyer idempotency key.
+Worker failures use capped exponential backoff and a bounded attempt count;
+exhausted or account-mismatched jobs move to `dead_lettered` for operations.
 The queue and restart-recovery test are runnable reference infrastructure;
 adopters supply the workflow handler that talks to their approval system.
 

examples/v3_reference_seller/src/durable_tasks.py

@@ -7,6 +7,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import base64
 import binascii
 import logging
@@ -187,6 +188,23 @@ async def startup(self) -> None:
             await self.workflow_queue.create_schema()
             if self.idempotency_backend is not None:
                 await self.idempotency_backend.create_schema()
+        except asyncio.CancelledError:
+            cleanup = asyncio.create_task(self.shutdown())
+            # Once startup owns resources, repeated cancellation must not
+            # strand them. Shield the cleanup task and preserve cancellation
+            # for the caller after every close has had a chance to run.
+            while not cleanup.done():
+                try:
+                    await asyncio.shield(cleanup)
+                except asyncio.CancelledError:
+                    continue
+            try:
+                cleanup.result()
+            except asyncio.CancelledError:
+                logger.exception("Durable resource cleanup was cancelled")
+            except Exception:
+                logger.exception("Durable resource cleanup failed after startup cancellation")
+            raise
         except Exception:
             try:
                 await self.shutdown()
@@ -196,14 +214,18 @@ async def startup(self) -> None:
 
     async def shutdown(self) -> None:
         """Best-effort close every resource owned by this process."""
-        first_error: Exception | None = None
+        first_error: BaseException | None = None
         closers = [self.sender.aclose]
         if self.lock_pool is not None:
             closers.append(self.lock_pool.close)
         closers.append(self.pool.close)
         for close in closers:
             try:
                 await close()
+            except asyncio.CancelledError as exc:
+                logger.exception("Durable task resource close was cancelled")
+                if first_error is None:
+                    first_error = exc
             except Exception as exc:
                 logger.exception("Failed to close a durable task resource")
                 if first_error is None:

examples/v3_reference_seller/src/workflow_queue.py

@@ -26,6 +26,15 @@
 
 _SAFE_IDENTIFIER = re.compile(r"^[a-z_][a-z0-9_]{0,62}$")
 DEFAULT_WORKFLOW_TABLE = "adcp_reference_workflows"
+DEFAULT_MAX_ATTEMPTS = 8
+DEFAULT_MAX_RETRY_SECONDS = 300.0
+DEFAULT_RETRY_BASE_SECONDS = 1.0
+
+_RETRY_EXHAUSTED_ERROR: dict[str, str] = {
+    "code": "INTERNAL_ERROR",
+    "message": "Workflow processing exhausted its retry budget.",
+    "recovery": "terminal",
+}
 
 
 @dataclass(frozen=True)
@@ -61,15 +70,29 @@ def __init__(
         registry: TaskRegistry,
         table: str = DEFAULT_WORKFLOW_TABLE,
         lease_seconds: int = 60,
+        max_attempts: int = DEFAULT_MAX_ATTEMPTS,
+        retry_base_seconds: float = DEFAULT_RETRY_BASE_SECONDS,
+        max_retry_seconds: float = DEFAULT_MAX_RETRY_SECONDS,
     ) -> None:
         if not _SAFE_IDENTIFIER.fullmatch(table):
             raise ValueError("workflow table must be a safe lowercase PostgreSQL identifier")
         if lease_seconds < 2:
             raise ValueError("workflow lease_seconds must be at least 2")
+        if max_attempts < 1:
+            raise ValueError("workflow max_attempts must be at least 1")
+        if retry_base_seconds <= 0:
+            raise ValueError("workflow retry_base_seconds must be positive")
+        if max_retry_seconds < retry_base_seconds:
+            raise ValueError("workflow max_retry_seconds must be at least retry_base_seconds")
         self._pool = pool
         self._registry = registry
         self._table = table
+        suffix = "_state_check"
+        self._state_constraint = f"{table[: 63 - len(suffix)]}{suffix}"
         self._lease_seconds = lease_seconds
+        self._max_attempts = max_attempts
+        self._retry_base_seconds = retry_base_seconds
+        self._max_retry_seconds = max_retry_seconds
 
         self._sql_claim = (  # noqa: S608 - table is validated above
             f"WITH candidate AS ("
@@ -108,13 +131,20 @@ async def create_schema(self) -> None:
                 created_at       TIMESTAMPTZ NOT NULL DEFAULT now(),
                 updated_at       TIMESTAMPTZ NOT NULL DEFAULT now(),
                 completed_at     TIMESTAMPTZ,
-                CHECK (state IN ('pending', 'in_flight', 'completed')),
+                dead_lettered_at TIMESTAMPTZ,
                 CHECK (attempt_count >= 0),
                 CHECK (
                     (state = 'in_flight') =
                     (lease_token IS NOT NULL AND lease_expires_at IS NOT NULL)
                 )
             )""",
+            f"""ALTER TABLE {self._table}
+                ADD COLUMN IF NOT EXISTS dead_lettered_at TIMESTAMPTZ""",
+            f"""ALTER TABLE {self._table}
+                DROP CONSTRAINT IF EXISTS {self._state_constraint}""",
+            f"""ALTER TABLE {self._table}
+                ADD CONSTRAINT {self._state_constraint}
+                CHECK (state IN ('pending', 'in_flight', 'completed', 'dead_lettered'))""",
             f"""CREATE INDEX IF NOT EXISTS {self._table}_work_idx
                 ON {self._table} (available_at, created_at)
                 WHERE state = 'pending'""",
@@ -123,8 +153,15 @@ async def create_schema(self) -> None:
                 WHERE state = 'in_flight'""",
         ]
         async with self._pool.connection() as conn:
-            for statement in statements:
-                await conn.execute(statement)
+            async with conn.transaction():
+                # Web and worker processes may bootstrap concurrently in the
+                # reference deployment. Serialize this example-owned DDL.
+                await conn.execute(
+                    "SELECT pg_advisory_xact_lock(hashtext(%s))",
+                    (f"adcp.workflow_queue.schema:{self._table}",),
+                )
+                for statement in statements:
+                    await conn.execute(statement)
 
     async def enqueue(
         self,
@@ -207,18 +244,79 @@ async def claim(self) -> WorkflowJob | None:
             attempt_count=int(attempts),
         )
 
+    def _retry_delay(self, attempt_count: int) -> float:
+        exponent = min(max(attempt_count - 1, 0), 30)
+        return min(
+            self._max_retry_seconds,
+            self._retry_base_seconds * (2**exponent),
+        )
+
     async def _release(self, job: WorkflowJob, exc: Exception) -> None:
         # Persist the class for operations without writing arbitrary exception
         # text (which may contain request data or credentials) into PostgreSQL.
         message = type(exc).__name__
+        retry_delay = self._retry_delay(job.attempt_count)
         async with self._pool.connection() as conn:
             await conn.execute(
                 f"UPDATE {self._table} SET state = 'pending', "  # noqa: S608
-                "available_at = now() + interval '1 second', lease_token = NULL, "
+                "available_at = now() + (%s * interval '1 second'), "
+                "lease_token = NULL, "
                 "lease_expires_at = NULL, last_error = %s, updated_at = now() "
                 "WHERE task_id = %s AND state = 'in_flight' AND lease_token = %s",
-                (message, job.task_id, job.lease_token),
+                (retry_delay, message, job.task_id, job.lease_token),
+            )
+
+    async def _dead_letter(self, job: WorkflowJob, reason: str) -> None:
+        async with self._pool.connection() as conn:
+            cursor = await conn.execute(
+                f"UPDATE {self._table} SET state = 'dead_lettered', "  # noqa: S608
+                "lease_token = NULL, lease_expires_at = NULL, last_error = %s, "
+                "dead_lettered_at = now(), updated_at = now() "
+                "WHERE task_id = %s AND state = 'in_flight' AND lease_token = %s",
+                (reason, job.task_id, job.lease_token),
+            )
+        if cursor.rowcount != 1:
+            raise RuntimeError("workflow lease was lost before dead-lettering")
+
+    async def _handle_failure(self, job: WorkflowJob, exc: Exception) -> None:
+        if job.attempt_count < self._max_attempts:
+            await self._release(job, exc)
+            logger.warning(
+                "Workflow job %s failed on attempt %d with %s; released for retry",
+                job.task_id,
+                job.attempt_count,
+                type(exc).__name__,
+            )
+            return
+
+        await self._registry.fail(job.task_id, dict(_RETRY_EXHAUSTED_ERROR))
+        await self._dead_letter(job, type(exc).__name__)
+        logger.error(
+            "Workflow job %s exhausted %d attempts with %s and was dead-lettered",
+            job.task_id,
+            job.attempt_count,
+            type(exc).__name__,
+        )
+
+    async def _handle_unverified_failure(self, job: WorkflowJob, exc: Exception) -> None:
+        """Bound failures before account ownership has been established."""
+        if job.attempt_count < self._max_attempts:
+            await self._release(job, exc)
+            logger.warning(
+                "Workflow task lookup failed on attempt %d with %s; released for retry",
+                job.attempt_count,
+                type(exc).__name__,
             )
+            return
+
+        # Never call the task-id-only registry mutation methods until the
+        # account-scoped get above has positively established ownership.
+        await self._dead_letter(job, type(exc).__name__)
+        logger.error(
+            "Workflow task lookup exhausted %d attempts with %s; dead-lettered",
+            job.attempt_count,
+            type(exc).__name__,
+        )
 
     async def _acknowledge(self, job: WorkflowJob) -> None:
         async with self._pool.connection() as conn:
@@ -242,22 +340,54 @@ async def process_one(self, handler: WorkflowHandler) -> bool:
                 job.task_id,
                 expected_account_id=job.account_id,
             )
-            if task is None:
-                raise ValueError("workflow job does not match an account-scoped task")
-            if task.get("state") in {"completed", "failed"}:
-                # A prior worker may have committed terminal task state and
-                # died before acknowledging this queue lease. Do not rerun the
-                # business effect; reconcile the queue row to terminal state.
-                await self._acknowledge(job)
-                return True
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            await self._handle_unverified_failure(job, exc)
+            return True
+
+        if task is None:
+            # Keep this transition outside the generic failure path. If the
+            # queue write fails, a later lease retries only dead-lettering and
+            # can never mutate a task belonging to another account.
+            await self._dead_letter(job, "task_missing_or_account_mismatch")
+            logger.error(
+                "Workflow job %s has no matching account-scoped task; dead-lettered",
+                job.task_id,
+            )
+            return True
+        if task.get("state") == "failed" and task.get("error") == _RETRY_EXHAUSTED_ERROR:
+            # Registry failure and queue dead-lettering are separate commits.
+            # Finish the second transition after a crash between them without
+            # running the handler again.
+            await self._dead_letter(job, "retry_budget_exhausted")
+            return True
+        if task.get("state") in {"completed", "failed"}:
+            # A prior worker may have committed terminal task state and died
+            # before acknowledging this queue lease. Do not rerun the business
+            # effect; reconcile the queue row to terminal state.
+            await self._acknowledge(job)
+            return True
+        if job.attempt_count > self._max_attempts:
+            # A prior attempt exhausted the handler budget but could not
+            # persist terminal registry state. Retry only finalization; never
+            # execute the business handler beyond max_attempts.
+            await self._registry.fail(job.task_id, dict(_RETRY_EXHAUSTED_ERROR))
+            await self._dead_letter(job, "retry_budget_exhausted")
+            return True
+
+        try:
             result = await handler(job)
             await self._registry.complete(job.task_id, result)
-            await self._acknowledge(job)
         except asyncio.CancelledError:
             raise
         except Exception as exc:
-            await self._release(job, exc)
-            logger.exception("Workflow job %s failed; released for retry", job.task_id)
+            await self._handle_failure(job, exc)
+            return True
+
+        # A failed acknowledgement must not enter the handler retry path: the
+        # registry is already terminal, and lease recovery reconciles the row.
+        await self._acknowledge(job)
         return True
 
     async def run_worker(
@@ -302,6 +432,9 @@ async def get(self, task_id: str) -> dict[str, Any] | None:
 
 
 __all__ = [
+    "DEFAULT_MAX_ATTEMPTS",
+    "DEFAULT_MAX_RETRY_SECONDS",
+    "DEFAULT_RETRY_BASE_SECONDS",
     "DEFAULT_WORKFLOW_TABLE",
     "PgWorkflowQueue",
     "WorkflowHandler",