diff --git a/README.md b/README.md index 649e0f03c..2d54cebe1 100644 --- a/README.md +++ b/README.md @@ -354,7 +354,7 @@ forward traffic degrades gracefully rather than failing. - **[API Reference](https://adcontextprotocol.github.io/adcp-client-python/)** - Complete API documentation with type signatures and examples - **[Protocol Spec](https://github.com/adcontextprotocol/adcp)** - Ad Context Protocol specification - **[Handler authoring](docs/handler-authoring.md)** - Building an AdCP-compliant agent on `adcp.server` -- **[Production seller path](docs/production-seller.md)** - Choose the server abstraction and wire durable multi-tenant tasks, idempotency, and webhook delivery +- **[Production seller path](docs/production-seller.md)** - Choose the server abstraction and wire durable multi-tenant tasks and webhook delivery - **[Migrating from SDK 6 to 7](https://github.com/adcontextprotocol/adcp-client-python/blob/main/MIGRATION_v6_to_v7.md)** - Breaking API, security, concurrency, and webhook changes - **[Migrating from SDK 7 to 8](https://github.com/adcontextprotocol/adcp-client-python/blob/main/MIGRATION_v7_to_v8.md)** - Secure webhook defaults and telemetry changes - **[Migrating from AdCP 3.1 to 3.2 beta](MIGRATION_ADCP_3.1_TO_3.2.md)** - Compact lifecycle adoption and old/new compatibility matrix @@ -1058,6 +1058,12 @@ class MySeller(ADCPHandler): serve(MySeller(), name="my-seller") ``` +The advertised capability applies globally, not just to the decorated example +method. Declare it only after every mutating operation the seller supports is +covered. Production integrations also need operation-aware handling for both +inline responses and durable task handoffs; a partial set of method wrappers +must leave `adcp.idempotency.supported` false. + **What the middleware does for you:** - Extracts `idempotency_key` from `params`, scopes lookups by `context.caller_identity` (per-principal — a security requirement from AdCP §2315) diff --git a/docs/production-seller.md b/docs/production-seller.md index d03aa31b8..e3d4936b9 100644 --- a/docs/production-seller.md +++ b/docs/production-seller.md @@ -25,12 +25,12 @@ and webhook delivery so adopters can replace one boundary at a time. buyer request │ ▼ -auth → tenant router → buyer registry → idempotency lock +auth → tenant router → buyer registry │ ▼ DecisioningPlatform method → upstream ad server │ - ├─ terminal result ───────────────► return inline and cache result + ├─ terminal result ───────────────► return inline │ └─ TaskHandoff / WorkflowHandoff ─► persist submitted task │ @@ -58,31 +58,29 @@ not make arbitrary in-process work durable. The reference includes a leased PostgreSQL queue and restart-recovery test; adopters still provide the business-specific approval handler. -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 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 -responses from the method-level cache, and implement the workflow method's -deduplication in the durable queue. The default reference includes -`create_media_buy` because its mock path is inline. +The reference intentionally advertises `adcp.idempotency.supported=false` and +does not install request-idempotency middleware. That capability is global: +wrapping only selected methods would overstate support while mutations such as +`sync_accounts` remain uncovered. A production implementation needs a future +operation-aware integration that covers every supported mutating operation, +including both inline results and task handoffs, before it can truthfully +advertise request idempotency. + +Workflow business effects still require independent idempotency. 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. Closing that window +requires a durable request-to-task mapping that can look up or reuse the prior +task id by buyer idempotency key. The workflow worker must also deduplicate +every external write using the stored buyer key because a crash can occur +after the write but before the handler outcome is staged. ## Task transitions | Handler outcome | Work owner | Persistence owner | What the buyer does next | |---|---|---|---| -| Return a result | Request handler | Idempotency backend caches the terminal response | Consume the inline result | +| Return a result | Request handler | The reference does not cache request results | Consume the inline result | | Raise `AdcpError` | Request handler | Framework projects the structured error; no task is created | Follow its recovery guidance | | `ctx.handoff_to_task(fn)` | SDK runs `fn` in the web process | `PgTaskRegistry` records submitted, progress, and terminal state | Poll `tasks/get` or await a webhook | | `ctx.handoff_to_workflow(enqueue)` | The adopter's queue/worker/HITL system | The enqueue callback stores the task id; that system later calls `registry.complete()` or `registry.fail()` | Poll `tasks/get` or await a webhook | @@ -137,13 +135,16 @@ await run_with_signals(workflow_handler=handle_approval) ``` The queue completes the original `PgTaskRegistry` record only after the -handler returns. A crash after an external side effect but before queue -acknowledgement causes deliberate re-execution after lease expiry, so the -business effect must be independently idempotent. The PostgreSQL conformance +handler returns. A crash after an external side effect but before the handler +outcome is staged causes deliberate re-execution after lease expiry, so the +business effect must be independently idempotent. Once staging commits, a +replacement retries only registry publication. The PostgreSQL conformance test kills the first logical worker after claim, creates fresh queue/registry objects, and verifies that the replacement completes the same task id. Queue payloads are ordinary JSONB: store only minimal continuation state and -never copy push-notification credentials or other secrets into them. +never copy push-notification credentials or other secrets into them. Scrubbed +terminal results are staged there until registry finalization succeeds, then +the queue clears that staged payload. ## Run the reference deployment @@ -210,10 +211,10 @@ 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 -them before either process starts. Runtime bootstrap can remain a safety net, -but migrations own schema evolution and rollback. +`task_webhook_outbox.sql`) and the reference workflow-queue DDL into reviewed, +versioned migrations and apply them before either process starts. Runtime +bootstrap can remain a safety net, but migrations own schema evolution and +rollback. The worker installs `SIGTERM` and `SIGINT` handlers, cancels its polling loops, awaits their cleanup, and then closes the sender and PostgreSQL pools. This is @@ -231,11 +232,10 @@ the shutdown path used by ordinary container and process supervisors. - Route human or long-running approval work through `WorkflowHandoff` and a durable queue; reserve `TaskHandoff` for bounded work that may safely fail on web-process restart. -- Put a uniqueness constraint on business effects keyed by the buyer's - idempotency key. The SDK cache cannot make a separate upstream transaction - atomic. -- Schedule `PgBackend.delete_expired()` (or equivalent SQL/pg_cron cleanup) so - expired idempotency rows do not accumulate. +- Make every external workflow effect independently idempotent using the + stored buyer key; queue lease recovery deliberately permits re-execution. +- Keep request idempotency unadvertised until an operation-aware integration + covers every supported mutation and durable task handoff. - Apply URL challenge and SSRF validation before accepting durable callback destinations. - Run the in-process tests and the media-buy seller storyboard before deploy. diff --git a/examples/v3_reference_seller/README.md b/examples/v3_reference_seller/README.md index ebdf84210..d8327e1f7 100644 --- a/examples/v3_reference_seller/README.md +++ b/examples/v3_reference_seller/README.md @@ -25,7 +25,6 @@ salesagent), see [MIGRATION.md](MIGRATION.md). | Durable task state | `src/durable_tasks.py` | `adcp.decisioning.PgTaskRegistry` | | Restart-safe external workflows | `src/workflow_queue.py` | `WorkflowHandoff` + leased PostgreSQL queue | | Atomic signed webhook delivery | `src/durable_tasks.py` + `src/worker.py` | `PgTaskWebhookOutbox` | -| Durable request idempotency | `src/durable_tasks.py` | `adcp.server.idempotency.PgBackend` | | Account v3 projection on read | `src/platform.py::list_accounts` | `adcp.types.project_account_for_response` | ## Architecture @@ -88,22 +87,28 @@ The seller binds `0.0.0.0:3001` and serves both transports. The commands above use the lightweight local path: async tasks are pollable but intentionally in-memory, and push-configured handoffs are rejected. For -the production-shaped PostgreSQL registry, idempotency cache, signed atomic -outbox, and separately supervised worker, follow the +the production-shaped PostgreSQL registry, signed atomic outbox, and +separately supervised worker, follow the [production seller path](../../docs/production-seller.md#run-the-reference-deployment). 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`. 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. +move a long polling loop into an in-process `TaskHandoff`. + +The reference advertises `adcp.idempotency.supported=false` and does not claim +durable request idempotency. The capability is global, so wrapping only a few +methods would be misleading while other mutations such as `sync_accounts` +remain uncovered. Truthful support requires a future operation-aware +integration across every supported mutation, including task handoffs. + 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. +requires a durable request-to-task mapping that reuses the prior task id by +buyer idempotency key. External workflow effects must independently +deduplicate on that stored key because lease recovery may re-execute a job. 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; diff --git a/examples/v3_reference_seller/src/app.py b/examples/v3_reference_seller/src/app.py index 85705276f..20adba236 100644 --- a/examples/v3_reference_seller/src/app.py +++ b/examples/v3_reference_seller/src/app.py @@ -309,7 +309,6 @@ def main() -> None: webhook_retry_horizon_seconds=( task_wiring.retry_horizon_seconds if task_wiring else 86_400 ), - idempotency=task_wiring.idempotency if task_wiring else None, ) logger.info( diff --git a/examples/v3_reference_seller/src/durable_tasks.py b/examples/v3_reference_seller/src/durable_tasks.py index fda404eb2..691fe2f33 100644 --- a/examples/v3_reference_seller/src/durable_tasks.py +++ b/examples/v3_reference_seller/src/durable_tasks.py @@ -17,7 +17,6 @@ from typing import TYPE_CHECKING from adcp.decisioning import PgTaskRegistry, PgTaskWebhookOutbox -from adcp.server.idempotency import IdempotencyStore, PgBackend from adcp.webhook_sender import WebhookSender from .workflow_queue import PgWorkflowQueue @@ -78,13 +77,10 @@ class DurableTaskWiring: """One process's validated durable task/webhook resources.""" pool: AsyncConnectionPool - lock_pool: AsyncConnectionPool | None sender: WebhookSender outbox: PgTaskWebhookOutbox registry: PgTaskRegistry workflow_queue: PgWorkflowQueue - idempotency_backend: PgBackend | None - idempotency: IdempotencyStore | None retry_horizon_seconds: int signing_algorithm: str @@ -94,7 +90,6 @@ def from_env( environ: Mapping[str, str] | None = None, *, required: bool | None = None, - include_idempotency: bool = True, ) -> DurableTaskWiring | None: """Validate configuration and build resources without opening sockets. @@ -132,18 +127,6 @@ def from_env( max_size=10, open=False, ) - # Advisory-lock operations reserve a connection for the whole handler - # call, so the SDK requires a distinct pool to prevent self-deadlock. - lock_pool = ( - AsyncConnectionPool( - values["ADCP_TASK_DATABASE_URL"], - min_size=1, - max_size=10, - open=False, - ) - if include_idempotency - else None - ) outbox = PgTaskWebhookOutbox( pool=pool, sender=sender, @@ -152,27 +135,12 @@ def from_env( ) registry = PgTaskRegistry(pool=pool, task_webhook_outbox=outbox) workflow_queue = PgWorkflowQueue(pool=pool, registry=registry) - idempotency_backend = ( - PgBackend(pool=pool, lock_pool=lock_pool) if lock_pool is not None else None - ) - idempotency = ( - IdempotencyStore( - backend=idempotency_backend, - ttl_seconds=retry_horizon, - raise_on_persist_error=True, - ) - if idempotency_backend is not None - else None - ) return cls( pool=pool, - lock_pool=lock_pool, sender=sender, outbox=outbox, registry=registry, workflow_queue=workflow_queue, - idempotency_backend=idempotency_backend, - idempotency=idempotency, retry_horizon_seconds=retry_horizon, signing_algorithm=signing_algorithm, ) @@ -181,13 +149,9 @@ async def startup(self) -> None: """Open the process-local pool and idempotently create SDK tables.""" try: await self.pool.open() - if self.lock_pool is not None: - await self.lock_pool.open() await self.registry.create_schema() await self.outbox.create_schema() 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 @@ -215,10 +179,7 @@ async def startup(self) -> None: async def shutdown(self) -> None: """Best-effort close every resource owned by this process.""" 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) + closers = [self.sender.aclose, self.pool.close] for close in closers: try: await close() diff --git a/examples/v3_reference_seller/src/platform.py b/examples/v3_reference_seller/src/platform.py index 96f68c1dc..0680e348a 100644 --- a/examples/v3_reference_seller/src/platform.py +++ b/examples/v3_reference_seller/src/platform.py @@ -66,7 +66,6 @@ import random from dataclasses import replace as _dc_replace from datetime import datetime, timezone -from types import MethodType from typing import TYPE_CHECKING, Any, ClassVar, cast from urllib.parse import urlsplit @@ -109,7 +108,6 @@ from adcp.decisioning.specialisms import SalesPlatform from adcp.server import current_tenant from adcp.server.helpers import valid_actions_for_status -from adcp.server.idempotency import IdempotencyStore from adcp.server.responses import list_creatives_response from adcp.types import ( BusinessEntity, @@ -133,7 +131,6 @@ UpdateMediaBuyRequest, UpdateMediaBuySuccessResponse, ) -from adcp.types.capabilities import IdempotencySupported from adcp.types.legacy import ( LegacyFormat, LegacyFormatId, @@ -159,13 +156,6 @@ "https://your-platform.example.com", } -_DEFAULT_IDEMPOTENCY_METHODS = ( - "create_media_buy", - "update_media_buy", - "sync_creatives", - "provide_performance_feedback", -) - def _legacy_format_converter(context: LegacyFormatConversionContext) -> dict[str, Any] | None: """Explicitly map the reference fixture catalogs into canonical kinds.""" @@ -827,8 +817,6 @@ def __init__( approval_poll_max_iterations: int = 60, webhook_signing_alg: str | None = None, webhook_retry_horizon_seconds: int = 86_400, - idempotency: IdempotencyStore | None = None, - method_level_idempotency_methods: tuple[str, ...] | None = None, ) -> None: """Construct the reference seller. @@ -864,12 +852,6 @@ def __init__( no signing advertised, sender wiring optional. :param webhook_retry_horizon_seconds: Durable retry horizon advertised to webhook receivers. Must match the configured task outbox. - :param idempotency: Optional durable idempotency store. When supplied, - mutating methods are wrapped and the replay window is advertised. - :param method_level_idempotency_methods: Methods whose inline terminal - responses are cached by ``idempotency``. Workflow adapters must - exclude methods returning handoff markers and deduplicate their - durable queue issuance separately. """ # Override the class-level capabilities iff signing is wired. # ``dataclasses.replace`` preserves every other field from the @@ -886,20 +868,6 @@ def __init__( delivery_retry_horizon_seconds=webhook_retry_horizon_seconds, ), ) - if idempotency is not None: - assert self.capabilities.adcp is not None - self.capabilities = _dc_replace( - self.capabilities, - adcp=self.capabilities.adcp.model_copy( - update={ - "idempotency": IdempotencySupported( - supported=True, - replay_ttl_seconds=idempotency.ttl_seconds, - ) - } - ), - ) - self._sessionmaker = sessionmaker # Single auth instance shared across every upstream_for() call. # The framework's client cache keys on (base_url, id(auth)), @@ -938,22 +906,6 @@ def __init__( sessionmaker, mock_upstream_url=mock_upstream_url, ) - if idempotency is not None: - selected_methods = ( - _DEFAULT_IDEMPOTENCY_METHODS - if method_level_idempotency_methods is None - else method_level_idempotency_methods - ) - for method_name in selected_methods: - method = getattr(type(self), method_name, None) - if not callable(method): - raise ValueError(f"Unknown method-level idempotency method: {method_name}") - # Bind the decorated unbound method back to this instance. - # Wrapping an already-bound method and assigning the plain - # closure would shift ``req`` into the decorator's ``self`` - # slot, causing idempotency-key extraction to silently miss. - wrapped = idempotency.wrap(method) - setattr(self, method_name, MethodType(wrapped, self)) def _client(self, ctx: RequestContext) -> UpstreamHttpClient: """Resolve the pooled :class:`UpstreamHttpClient` for this diff --git a/examples/v3_reference_seller/src/worker.py b/examples/v3_reference_seller/src/worker.py index 75e612d07..960da5c3e 100644 --- a/examples/v3_reference_seller/src/worker.py +++ b/examples/v3_reference_seller/src/worker.py @@ -27,7 +27,7 @@ async def run( workflow_handler: WorkflowHandler | None = None, ) -> None: """Run durable workers until a shutdown signal sets ``stop_event``.""" - wiring = DurableTaskWiring.from_env(required=True, include_idempotency=False) + wiring = DurableTaskWiring.from_env(required=True) assert wiring is not None # required=True makes None impossible await wiring.startup() stop = stop_event or asyncio.Event() @@ -53,9 +53,7 @@ async def run( ) completed_workers = done.intersection(workers) for task in completed_workers: - exc = task.exception() - if exc is not None: - raise exc + task.result() if completed_workers: raise RuntimeError("A durable worker exited unexpectedly") logger.info("Shutdown requested; stopping durable workers") diff --git a/examples/v3_reference_seller/src/workflow_queue.py b/examples/v3_reference_seller/src/workflow_queue.py index f019cd536..c8c596184 100644 --- a/examples/v3_reference_seller/src/workflow_queue.py +++ b/examples/v3_reference_seller/src/workflow_queue.py @@ -1,9 +1,9 @@ """PostgreSQL queue for adopter-owned ``WorkflowHandoff`` work. This is deliberately example code, not an SDK primitive. It demonstrates the -minimum production invariants: durable enqueue, single-worker leases, -restart recovery after lease expiry, account scoping, and completion through -the SDK's durable task registry. +production invariants needed around a durable task registry: leases with +heartbeats, bounded handler attempts, durable finalization, account scoping, +and restart-safe reconciliation. """ from __future__ import annotations @@ -14,8 +14,11 @@ import re import uuid from collections.abc import Awaitable, Callable -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any, Literal + +from adcp.decisioning import AdcpError +from adcp.decisioning.account_projection import strip_credentials_from_wire_result if TYPE_CHECKING: from psycopg_pool import AsyncConnectionPool @@ -36,6 +39,12 @@ "recovery": "terminal", } +FinalizationAction = Literal["complete", "fail"] + + +class WorkflowLeaseLostError(RuntimeError): + """The row is no longer owned by this worker's lease token.""" + @dataclass(frozen=True) class WorkflowJob: @@ -47,6 +56,8 @@ class WorkflowJob: payload: dict[str, Any] lease_token: str attempt_count: int + finalization_action: FinalizationAction | None = None + finalization_payload: dict[str, Any] | None = None WorkflowHandler = Callable[[WorkflowJob], Awaitable[dict[str, Any]]] @@ -55,9 +66,10 @@ class WorkflowJob: class PgWorkflowQueue: """A small durable queue that completes SDK-managed workflow tasks. - Handlers must make external side effects idempotent. A worker can finish - the side effect and die before acknowledging the queue row; after the - lease expires, a replacement worker intentionally runs the job again. + A handler's result (or terminal error) is written to the queue before the + task registry is updated. Registry outages therefore retry finalization, + never the business handler. Handlers must still make external side effects + idempotent because a process can die before that staging write commits. ``payload`` is ordinary JSONB, not an encrypted secret store. Enqueue only the minimum continuation data and never include callback credentials. @@ -104,44 +116,43 @@ def __init__( f"UPDATE {table} AS jobs SET" " state = 'in_flight', lease_token = %s," " lease_expires_at = now() + (%s * interval '1 second')," - " attempt_count = attempt_count + 1, updated_at = now()" + " updated_at = now()" " FROM candidate WHERE jobs.task_id = candidate.task_id" " RETURNING jobs.task_id, jobs.account_id, jobs.workflow_type," - " jobs.payload, jobs.lease_token, jobs.attempt_count" + " jobs.payload, jobs.lease_token, jobs.attempt_count," + " jobs.finalization_action, jobs.finalization_payload" ) async def create_schema(self) -> None: - """Bootstrap the example table for local development and tests. - - Production deployments should apply the equivalent DDL through their - migration system before starting web or worker processes. - """ + """Bootstrap or forward-upgrade the example table.""" statements = [ f"""CREATE TABLE IF NOT EXISTS {self._table} ( - task_id TEXT COLLATE "C" PRIMARY KEY, - account_id TEXT COLLATE "C" NOT NULL, - workflow_type TEXT NOT NULL, - payload JSONB NOT NULL, - state TEXT NOT NULL DEFAULT 'pending', - attempt_count INTEGER NOT NULL DEFAULT 0, - available_at TIMESTAMPTZ NOT NULL DEFAULT now(), - lease_token TEXT COLLATE "C", - lease_expires_at TIMESTAMPTZ, - last_error TEXT, - created_at TIMESTAMPTZ NOT NULL DEFAULT now(), - updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), - completed_at TIMESTAMPTZ, - dead_lettered_at TIMESTAMPTZ, + task_id TEXT COLLATE "C" PRIMARY KEY, + account_id TEXT COLLATE "C" NOT NULL, + workflow_type TEXT NOT NULL, + payload JSONB NOT NULL, + state TEXT NOT NULL DEFAULT 'pending', + attempt_count INTEGER NOT NULL DEFAULT 0, + available_at TIMESTAMPTZ NOT NULL DEFAULT now(), + lease_token TEXT COLLATE "C", + lease_expires_at TIMESTAMPTZ, + finalization_action TEXT, + finalization_payload JSONB, + last_error TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + completed_at TIMESTAMPTZ, + 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 COLUMN IF NOT EXISTS dead_lettered_at TIMESTAMPTZ", + f"ALTER TABLE {self._table} ADD COLUMN IF NOT EXISTS finalization_action TEXT", + f"ALTER TABLE {self._table} ADD COLUMN IF NOT EXISTS finalization_payload JSONB", + 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'))""", @@ -154,8 +165,6 @@ async def create_schema(self) -> None: ] async with self._pool.connection() as conn: 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}",), @@ -183,10 +192,7 @@ async def enqueue( ) async with self._pool.connection() as conn: row = await ( - await conn.execute( - sql, - (task_id, account_id, workflow_type, serialized), - ) + await conn.execute(sql, (task_id, account_id, workflow_type, serialized)) ).fetchone() if row is not None: return @@ -217,7 +223,11 @@ async def enqueue_from_handoff( ) async def claim(self) -> WorkflowJob | None: - """Recover expired leases and claim one eligible job.""" + """Recover expired leases and claim one eligible job. + + Claiming does not consume a handler attempt. The attempt counter is + incremented only after the account-scoped registry lookup succeeds. + """ lease_token = uuid.uuid4().hex async with self._pool.connection() as conn: await conn.execute( @@ -226,168 +236,326 @@ async def claim(self) -> WorkflowJob | None: "WHERE state = 'in_flight' AND lease_expires_at <= now()" ) row = await ( - await conn.execute( - self._sql_claim, - (lease_token, self._lease_seconds), - ) + await conn.execute(self._sql_claim, (lease_token, self._lease_seconds)) ).fetchone() if row is None: return None - task_id, account_id, workflow_type, payload, token, attempts = row - payload_dict = payload if isinstance(payload, dict) else json.loads(payload) + payload = row[3] if isinstance(row[3], dict) else json.loads(row[3]) + finalization_payload = row[7] + if finalization_payload is not None and not isinstance(finalization_payload, dict): + finalization_payload = json.loads(finalization_payload) return WorkflowJob( - task_id=str(task_id), - account_id=str(account_id), - workflow_type=str(workflow_type), - payload=payload_dict, - lease_token=str(token), - attempt_count=int(attempts), + task_id=str(row[0]), + account_id=str(row[1]), + workflow_type=str(row[2]), + payload=payload, + lease_token=str(row[4]), + attempt_count=int(row[5]), + finalization_action=row[6], + finalization_payload=finalization_payload, ) 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), - ) + return float(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 def _token_update( + self, + job: WorkflowJob, + sql: str, + params: tuple[Any, ...], + operation: str, + ) -> None: async with self._pool.connection() as conn: - await conn.execute( - f"UPDATE {self._table} SET state = 'pending', " # noqa: S608 - "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", - (retry_delay, message, job.task_id, job.lease_token), - ) + cursor = await conn.execute(sql, params) + if cursor.rowcount != 1: + raise WorkflowLeaseLostError(f"workflow lease was lost before {operation}") + + async def _release(self, job: WorkflowJob, reason: str) -> None: + retry_delay = self._retry_delay(job.attempt_count) + await self._token_update( + job, + f"UPDATE {self._table} SET state = 'pending', " # noqa: S608 + "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", + (retry_delay, reason, job.task_id, job.lease_token), + "release", + ) async def _dead_letter(self, job: WorkflowJob, reason: str) -> None: + await self._token_update( + job, + f"UPDATE {self._table} SET state = 'dead_lettered', " # noqa: S608 + "lease_token = NULL, lease_expires_at = NULL, " + "finalization_payload = NULL, " + "last_error = COALESCE(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), + "dead-lettering", + ) + + async def _acknowledge(self, job: WorkflowJob) -> None: + await self._token_update( + job, + f"UPDATE {self._table} SET state = 'completed', " # noqa: S608 + "lease_token = NULL, lease_expires_at = NULL, last_error = NULL, " + "finalization_payload = NULL, " + "completed_at = now(), updated_at = now() " + "WHERE task_id = %s AND state = 'in_flight' AND lease_token = %s", + (job.task_id, job.lease_token), + "acknowledgement", + ) + + async def _renew_lease(self, job: WorkflowJob) -> None: + await self._token_update( + job, + f"UPDATE {self._table} SET " # noqa: S608 + "lease_expires_at = now() + (%s * interval '1 second'), updated_at = now() " + "WHERE task_id = %s AND state = 'in_flight' AND lease_token = %s", + (self._lease_seconds, job.task_id, job.lease_token), + "lease renewal", + ) + + async def _heartbeat(self, job: WorkflowJob) -> None: + interval = self._lease_seconds / 3 + while True: + await asyncio.sleep(interval) + await self._renew_lease(job) + + async def _begin_handler_attempt(self, job: WorkflowJob) -> WorkflowJob | 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") + row = await ( + await conn.execute( + f"UPDATE {self._table} SET attempt_count = attempt_count + 1, " # noqa: S608 + "updated_at = now() WHERE task_id = %s AND state = 'in_flight' " + "AND lease_token = %s AND finalization_action IS NULL " + "AND attempt_count < %s RETURNING attempt_count", + (job.task_id, job.lease_token, self._max_attempts), + ) + ).fetchone() + if row is None: + return None + return replace(job, attempt_count=int(row[0])) - 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", + async def _stage_finalization( + self, + job: WorkflowJob, + action: FinalizationAction, + payload: dict[str, Any], + *, + last_error: str | None = None, + ) -> WorkflowJob: + serialized = json.dumps(payload, separators=(",", ":"), sort_keys=True) + await self._token_update( + job, + f"UPDATE {self._table} SET finalization_action = %s, " # noqa: S608 + "finalization_payload = %s::jsonb, last_error = %s, updated_at = now() " + "WHERE task_id = %s AND state = 'in_flight' AND lease_token = %s " + "AND finalization_action IS NULL", + (action, serialized, last_error, job.task_id, job.lease_token), + "staging finalization", + ) + return replace( + job, + finalization_action=action, + finalization_payload=dict(payload), + ) + + async def _stage_despite_cancellation( + self, + job: WorkflowJob, + action: FinalizationAction, + payload: dict[str, Any], + *, + last_error: str | None = None, + ) -> WorkflowJob: + """Finish the short DB write even if worker shutdown races it.""" + staging = asyncio.create_task( + self._stage_finalization(job, action, payload, last_error=last_error) + ) + try: + return await asyncio.shield(staging) + except asyncio.CancelledError: + staged_job = await staging + await self._release_despite_cancellation(staged_job) + raise + + async def _release_despite_cancellation(self, job: WorkflowJob) -> None: + """Best-effort lease release after cancellation has already begun.""" + try: + await asyncio.shield(self._release(job, "CancelledError")) + except WorkflowLeaseLostError: + logger.debug( + "Workflow task %s lease was already lost during cancellation cleanup", 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__, + async def _run_handler(self, job: WorkflowJob, handler: WorkflowHandler) -> dict[str, Any]: + handler_task: asyncio.Future[dict[str, Any]] = asyncio.ensure_future(handler(job)) + heartbeat_task = asyncio.create_task(self._heartbeat(job)) + waiters: set[asyncio.Future[Any]] = {handler_task, heartbeat_task} + try: + done, _ = await asyncio.wait( + waiters, + return_when=asyncio.FIRST_COMPLETED, ) - 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__, - ) + if heartbeat_task in done: + heartbeat_error = heartbeat_task.exception() + if heartbeat_error is not None: + handler_task.cancel() + await asyncio.gather(handler_task, return_exceptions=True) + raise heartbeat_error + return await handler_task + except asyncio.CancelledError: + handler_task.cancel() + heartbeat_task.cancel() + await asyncio.gather(handler_task, heartbeat_task, return_exceptions=True) + await self._release_despite_cancellation(job) + raise + finally: + heartbeat_task.cancel() + await asyncio.gather(heartbeat_task, return_exceptions=True) + + async def _apply_finalization(self, job: WorkflowJob) -> None: + action = job.finalization_action + payload = job.finalization_payload + if action is None or payload is None: + raise RuntimeError("workflow finalization is incomplete") + if action == "complete": + await self._registry.complete(job.task_id, payload) + await self._acknowledge(job) + else: + await self._registry.fail(job.task_id, payload) + await self._dead_letter(job, str(payload.get("code", "workflow_failed"))) - async def _acknowledge(self, job: WorkflowJob) -> None: - async with self._pool.connection() as conn: - cursor = await conn.execute( - f"UPDATE {self._table} SET state = 'completed', " # noqa: S608 - "lease_token = NULL, lease_expires_at = NULL, last_error = NULL, " - "completed_at = now(), updated_at = now() " - "WHERE task_id = %s AND state = 'in_flight' AND lease_token = %s", - (job.task_id, job.lease_token), + async def _retry_finalization(self, job: WorkflowJob) -> None: + try: + await self._apply_finalization(job) + except asyncio.CancelledError: + await self._release_despite_cancellation(job) + raise + except Exception as exc: + await self._release(job, type(exc).__name__) + logger.warning( + "Workflow task %s finalization failed with %s; released for retry", + job.task_id, + type(exc).__name__, ) - if cursor.rowcount != 1: - raise RuntimeError("workflow lease was lost before acknowledgement") async def process_one(self, handler: WorkflowHandler) -> bool: - """Run one leased job and complete the matching SDK task.""" + """Run one leased job and reconcile it with the matching SDK task.""" job = await self.claim() if job is None: return False + try: task = await self._registry.get( job.task_id, expected_account_id=job.account_id, ) except asyncio.CancelledError: + await self._release_despite_cancellation(job) raise except Exception as exc: - await self._handle_unverified_failure(job, exc) + # Registry availability is not a business-handler failure. Retry + # indefinitely without consuming the handler attempt budget. + await self._release(job, type(exc).__name__) + logger.warning( + "Workflow task lookup failed with %s; released for retry", + type(exc).__name__, + ) 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. + # Unknown and cross-account task ids are intentionally + # indistinguishable. Never mutate the registry by task id alone. 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) + if task.get("state") == "failed": + reason = ( + "retry_budget_exhausted" + if task.get("error") == _RETRY_EXHAUSTED_ERROR + else "registry_task_already_failed" + ) + await self._dead_letter(job, reason) + else: + await self._acknowledge(job) + return True + + if job.finalization_action is not None: + await self._retry_finalization(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") + + attempted_job = await self._begin_handler_attempt(job) + if attempted_job is None: + # A crash after the last handler attempt but before outcome staging + # must not permit an extra execution. + staged = await self._stage_despite_cancellation( + job, + "fail", + dict(_RETRY_EXHAUSTED_ERROR), + last_error="retry_budget_exhausted", + ) + await self._retry_finalization(staged) return True + job = attempted_job try: - result = await handler(job) - await self._registry.complete(job.task_id, result) + result = await self._run_handler(job, handler) except asyncio.CancelledError: raise + except AdcpError as exc: + # Correctable errors require a changed buyer request, so retrying + # this unchanged queued job cannot resolve them. Only transient + # failures consume the workflow retry budget. + if exc.recovery == "transient" and job.attempt_count < self._max_attempts: + await self._release(job, type(exc).__name__) + return True + staged = await self._stage_despite_cancellation( + job, + "fail", + exc.to_wire(), + last_error=exc.code, + ) + await self._retry_finalization(staged) + return True + except WorkflowLeaseLostError: + # The handler was cancelled by _run_handler before this escapes. + # A replacement owner is responsible for the row. + raise except Exception as exc: - await self._handle_failure(job, exc) + if job.attempt_count < self._max_attempts: + await self._release(job, type(exc).__name__) + logger.warning( + "Workflow job %s failed on attempt %d with %s; released for retry", + job.task_id, + job.attempt_count, + type(exc).__name__, + ) + return True + staged = await self._stage_despite_cancellation( + job, + "fail", + dict(_RETRY_EXHAUSTED_ERROR), + last_error=type(exc).__name__, + ) + await self._retry_finalization(staged) 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) + safe_result = strip_credentials_from_wire_result(str(task["task_type"]), result) + if not isinstance(safe_result, dict): + raise TypeError("workflow handler result must serialize to an object") + staged = await self._stage_despite_cancellation(job, "complete", safe_result) + await self._retry_finalization(staged) return True async def run_worker( @@ -415,8 +583,9 @@ async def get(self, task_id: str) -> dict[str, Any] | None: async with self._pool.connection() as conn: row = await ( await conn.execute( - f"SELECT account_id, workflow_type, state, attempt_count, last_error " # noqa: S608 - f"FROM {self._table} WHERE task_id = %s", + f"SELECT account_id, workflow_type, state, attempt_count, last_error, " # noqa: S608 + f"finalization_action, finalization_payload FROM {self._table} " + "WHERE task_id = %s", (task_id,), ) ).fetchone() @@ -428,6 +597,8 @@ async def get(self, task_id: str) -> dict[str, Any] | None: "state": str(row[2]), "attempt_count": int(row[3]), "last_error": row[4], + "finalization_action": row[5], + "finalization_payload": row[6], } @@ -439,4 +610,5 @@ async def get(self, task_id: str) -> dict[str, Any] | None: "PgWorkflowQueue", "WorkflowHandler", "WorkflowJob", + "WorkflowLeaseLostError", ] diff --git a/examples/v3_reference_seller/tests/test_durable_tasks.py b/examples/v3_reference_seller/tests/test_durable_tasks.py index 4ab5772e2..929c37386 100644 --- a/examples/v3_reference_seller/tests/test_durable_tasks.py +++ b/examples/v3_reference_seller/tests/test_durable_tasks.py @@ -129,34 +129,12 @@ def test_complete_bundle_builds_atomic_registry_outbox_pair(tmp_path: Path) -> N assert wiring.registry.task_webhook_outbox is wiring.outbox assert wiring.workflow_queue._registry is wiring.registry # noqa: SLF001 assert wiring.registry.atomic_task_webhook_outbox is True - assert wiring.idempotency_backend is not None - assert wiring.idempotency_backend._lock_pool is wiring.lock_pool # noqa: SLF001 - assert wiring.idempotency is not None - assert wiring.idempotency.raise_on_persist_error is True assert wiring.retry_horizon_seconds == 172800 assert wiring.signing_algorithm == "ed25519" finally: asyncio.run(wiring.shutdown()) -def test_worker_bundle_omits_unused_idempotency_pool(tmp_path: Path) -> None: - key_path = tmp_path / "webhook-signing.pem" - _write_signing_key(key_path) - - wiring = DurableTaskWiring.from_env( - _environment(key_path), - include_idempotency=False, - ) - - assert wiring is not None - try: - assert wiring.lock_pool is None - assert wiring.idempotency_backend is None - assert wiring.idempotency is None - finally: - asyncio.run(wiring.shutdown()) - - def test_complete_bundle_passes_capability_wiring_preflight(tmp_path: Path) -> None: from src.platform import V3ReferenceSeller @@ -166,7 +144,6 @@ def test_complete_bundle_passes_capability_wiring_preflight(tmp_path: Path) -> N _write_signing_key(key_path) wiring = DurableTaskWiring.from_env(_environment(key_path)) assert wiring is not None - assert wiring.idempotency is not None executor = None try: seller = V3ReferenceSeller( @@ -175,8 +152,9 @@ def test_complete_bundle_passes_capability_wiring_preflight(tmp_path: Path) -> N mock_upstream_url=None, webhook_signing_alg=wiring.signing_algorithm, webhook_retry_horizon_seconds=wiring.retry_horizon_seconds, - idempotency=wiring.idempotency, ) + assert seller.capabilities.adcp is not None + assert seller.capabilities.adcp.idempotency.supported is False handler, executor, registry = create_adcp_server_from_platform( seller, @@ -193,21 +171,16 @@ def test_complete_bundle_passes_capability_wiring_preflight(tmp_path: Path) -> N @pytest.mark.asyncio async def test_startup_failure_closes_every_resource() -> None: pool = MagicMock(open=AsyncMock(), close=AsyncMock()) - lock_pool = MagicMock(open=AsyncMock(), close=AsyncMock()) sender = MagicMock(aclose=AsyncMock()) registry = MagicMock(create_schema=AsyncMock(side_effect=RuntimeError("DDL failed"))) outbox = MagicMock(create_schema=AsyncMock()) workflow_queue = MagicMock(create_schema=AsyncMock()) - backend = MagicMock(create_schema=AsyncMock()) wiring = DurableTaskWiring( pool=pool, - lock_pool=lock_pool, sender=sender, outbox=outbox, registry=registry, workflow_queue=workflow_queue, - idempotency_backend=backend, - idempotency=None, retry_horizon_seconds=86400, signing_algorithm="ed25519", ) @@ -216,14 +189,12 @@ async def test_startup_failure_closes_every_resource() -> None: await wiring.startup() sender.aclose.assert_awaited_once() - lock_pool.close.assert_awaited_once() pool.close.assert_awaited_once() @pytest.mark.asyncio async def test_startup_cancellation_closes_every_resource_and_reraises() -> None: pool = MagicMock(open=AsyncMock(), close=AsyncMock()) - lock_pool = MagicMock(open=AsyncMock(), close=AsyncMock()) sender_close_started = asyncio.Event() release_sender_close = asyncio.Event() @@ -241,13 +212,10 @@ async def wait_forever() -> None: registry = MagicMock(create_schema=AsyncMock(side_effect=wait_forever)) wiring = DurableTaskWiring( pool=pool, - lock_pool=lock_pool, sender=sender, outbox=MagicMock(create_schema=AsyncMock()), registry=registry, workflow_queue=MagicMock(create_schema=AsyncMock()), - idempotency_backend=MagicMock(create_schema=AsyncMock()), - idempotency=None, retry_horizon_seconds=86400, signing_algorithm="ed25519", ) @@ -265,24 +233,19 @@ async def wait_forever() -> None: assert startup.cancelled() sender.aclose.assert_awaited_once() - lock_pool.close.assert_awaited_once() pool.close.assert_awaited_once() @pytest.mark.asyncio async def test_shutdown_continues_after_close_failure() -> None: pool = MagicMock(close=AsyncMock()) - lock_pool = MagicMock(close=AsyncMock()) sender = MagicMock(aclose=AsyncMock(side_effect=RuntimeError("sender close failed"))) wiring = DurableTaskWiring( pool=pool, - lock_pool=lock_pool, sender=sender, outbox=MagicMock(), registry=MagicMock(), workflow_queue=MagicMock(), - idempotency_backend=None, - idempotency=None, retry_horizon_seconds=86400, signing_algorithm="ed25519", ) @@ -290,5 +253,4 @@ async def test_shutdown_continues_after_close_failure() -> None: with pytest.raises(RuntimeError, match="sender close failed"): await wiring.shutdown() - lock_pool.close.assert_awaited_once() pool.close.assert_awaited_once() diff --git a/examples/v3_reference_seller/tests/test_smoke.py b/examples/v3_reference_seller/tests/test_smoke.py index d607f607c..a32af3e6b 100644 --- a/examples/v3_reference_seller/tests/test_smoke.py +++ b/examples/v3_reference_seller/tests/test_smoke.py @@ -232,60 +232,6 @@ def test_platform_advertises_webhook_signing_when_alg_passed() -> None: assert [a.value for a in ws.algorithms] == ["ed25519"] -@pytest.mark.asyncio -async def test_platform_runtime_idempotency_wrapper_replays() -> None: - from src.platform import V3ReferenceSeller - - from adcp.decisioning import RequestContext - from adcp.server.idempotency import IdempotencyStore, MemoryBackend, is_wrapped - - class _Seller(V3ReferenceSeller): - calls = 0 - - async def create_media_buy(self, req, ctx): - del req, ctx - self.calls += 1 - return {"status": "completed", "media_buy_id": "mb_1"} - - seller = _Seller( - sessionmaker=lambda: None, # type: ignore[arg-type] - upstream_api_key="test-key", - mock_upstream_url=None, - idempotency=IdempotencyStore(MemoryBackend(), ttl_seconds=86400), - ) - assert is_wrapped(seller.create_media_buy) - context = RequestContext(caller_identity="buyer", tenant_id="tenant") - request = {"idempotency_key": "same-key", "budget": 100} - - first = await seller.create_media_buy(request, context) - replay = await seller.create_media_buy(request, context) - - assert first == {"status": "completed", "media_buy_id": "mb_1"} - assert replay == { - "status": "completed", - "media_buy_id": "mb_1", - "replayed": True, - } - assert seller.calls == 1 - - -def test_platform_idempotency_method_selection_is_explicit() -> None: - from src.platform import V3ReferenceSeller - - from adcp.server.idempotency import IdempotencyStore, MemoryBackend, is_wrapped - - seller = V3ReferenceSeller( - sessionmaker=lambda: None, # type: ignore[arg-type] - upstream_api_key="test-key", - mock_upstream_url=None, - idempotency=IdempotencyStore(MemoryBackend(), ttl_seconds=86400), - method_level_idempotency_methods=("sync_creatives",), - ) - - assert not is_wrapped(seller.create_media_buy) - assert is_wrapped(seller.sync_creatives) - - @pytest.mark.asyncio @pytest.mark.parametrize("authenticated_tenant", ["tenant-a", None]) async def test_bearer_context_rejects_cross_tenant_rebinding( diff --git a/examples/v3_reference_seller/tests/test_worker.py b/examples/v3_reference_seller/tests/test_worker.py index d2469615a..0e098ac94 100644 --- a/examples/v3_reference_seller/tests/test_worker.py +++ b/examples/v3_reference_seller/tests/test_worker.py @@ -16,24 +16,37 @@ @pytest.mark.asyncio -async def test_stop_event_cancels_worker_before_resource_shutdown(monkeypatch) -> None: +async def test_stop_event_cancels_workers_before_resource_shutdown(monkeypatch) -> None: import src.worker as worker_module - started = asyncio.Event() - cancelled = asyncio.Event() + outbox_started = asyncio.Event() + workflow_started = asyncio.Event() + shutdown_order = [] async def run_outbox() -> None: - started.set() + outbox_started.set() try: await asyncio.Event().wait() finally: - cancelled.set() + await asyncio.sleep(0) + shutdown_order.append("outbox") + + async def run_workflow(_handler) -> None: + workflow_started.set() + try: + await asyncio.Event().wait() + finally: + await asyncio.sleep(0) + shutdown_order.append("workflow") + + async def shutdown() -> None: + shutdown_order.append("wiring") wiring = SimpleNamespace( startup=AsyncMock(), - shutdown=AsyncMock(), + shutdown=AsyncMock(side_effect=shutdown), outbox=SimpleNamespace(run_worker=run_outbox), - workflow_queue=SimpleNamespace(), + workflow_queue=SimpleNamespace(run_worker=run_workflow), ) monkeypatch.setattr( worker_module.DurableTaskWiring, @@ -43,14 +56,18 @@ async def run_outbox() -> None: stop_event = asyncio.Event() async def request_stop() -> None: - await started.wait() + await asyncio.gather(outbox_started.wait(), workflow_started.wait()) stop_event.set() stopper = asyncio.create_task(request_stop()) - await worker_module.run(stop_event=stop_event) + await worker_module.run( + stop_event=stop_event, + workflow_handler=AsyncMock(), + ) + assert await stopper is None - assert stopper.done() - assert cancelled.is_set() + assert set(shutdown_order[:2]) == {"outbox", "workflow"} + assert shutdown_order[2] == "wiring" wiring.startup.assert_awaited_once() wiring.shutdown.assert_awaited_once() @@ -60,10 +77,11 @@ async def test_worker_failure_wins_when_shutdown_is_also_ready(monkeypatch) -> N import src.worker as worker_module stop_event = asyncio.Event() + failure = RuntimeError("outbox failed") async def fail_outbox() -> None: stop_event.set() - raise RuntimeError("outbox failed") + raise failure wiring = SimpleNamespace( startup=AsyncMock(), @@ -84,9 +102,10 @@ async def wait_for_all(awaitables, *, return_when): monkeypatch.setattr(worker_module.asyncio, "wait", wait_for_all) - with pytest.raises(RuntimeError, match="outbox failed"): + with pytest.raises(RuntimeError, match="outbox failed") as exc_info: await worker_module.run(stop_event=stop_event) + assert exc_info.value is failure wiring.shutdown.assert_awaited_once() diff --git a/tests/conformance/decisioning/test_pg_reference_workflow_queue.py b/tests/conformance/decisioning/test_pg_reference_workflow_queue.py index 62da421d7..2604cb738 100644 --- a/tests/conformance/decisioning/test_pg_reference_workflow_queue.py +++ b/tests/conformance/decisioning/test_pg_reference_workflow_queue.py @@ -5,6 +5,7 @@ from __future__ import annotations +import asyncio import os import secrets import sys @@ -26,10 +27,11 @@ _EXAMPLE = Path(__file__).resolve().parents[3] / "examples" / "v3_reference_seller" sys.path.insert(0, str(_EXAMPLE)) -from adcp.decisioning import Account, PgTaskRegistry, RequestContext # noqa: E402 -from adcp.decisioning.dispatch import _project_workflow_handoff # noqa: E402 from src.workflow_queue import PgWorkflowQueue # noqa: E402 +from adcp.decisioning import Account, AdcpError, PgTaskRegistry, RequestContext # noqa: E402 +from adcp.decisioning.dispatch import _project_workflow_handoff # noqa: E402 + class _Request(BaseModel): context: dict[str, str] | None = None @@ -103,7 +105,7 @@ async def enqueue(task_ctx) -> None: first_claim = await first_worker_queue.claim() assert first_claim is not None assert first_claim.task_id == task_id - assert first_claim.attempt_count == 1 + assert first_claim.attempt_count == 0 # Advance only the database lease, avoiding a wall-clock sleep. async with web_pool.connection() as conn: @@ -135,7 +137,7 @@ async def enqueue(task_ctx) -> None: ) async def complete(job): - assert job.attempt_count == 2 + assert job.attempt_count == 1 return job.payload["result"] assert await replacement_queue.process_one(complete) is True @@ -147,7 +149,7 @@ async def complete(job): assert queue_record is not None assert queue_record["state"] == "completed" - assert queue_record["attempt_count"] == 2 + assert queue_record["attempt_count"] == 1 assert task_record is not None assert task_record["state"] == "completed" assert task_record["result"] == expected_result @@ -302,7 +304,7 @@ async def handler(_job): assert handler_called is False assert queue_record is not None assert queue_record["state"] == "dead_lettered" - assert queue_record["attempt_count"] == 1 + assert queue_record["attempt_count"] == 0 assert queue_record["last_error"] == "task_missing_or_account_mismatch" if issue_task: assert task_record is not None @@ -453,7 +455,7 @@ async def handler(_job): @pytest.mark.asyncio -async def test_registry_finalization_failure_never_reruns_exhausted_handler() -> None: +async def test_exhausted_transient_error_survives_finalization_retry() -> None: suffix = secrets.token_hex(6) task_table = f"test_dtasks_{suffix}" workflow_table = f"test_workflows_{suffix}" @@ -509,19 +511,28 @@ async def fail(self, task_id, error): payload={"upstream_order_id": "order-1"}, ) handler_calls = 0 + transient_error = AdcpError( + "SERVICE_UNAVAILABLE", + message="Approval service is unavailable", + recovery="transient", + retry_after=30, + ) async def fail_handler(_job): nonlocal handler_calls handler_calls += 1 - raise RuntimeError("upstream failed") + raise transient_error - with pytest.raises(RuntimeError, match="registry temporarily unavailable"): - await queue.process_one(fail_handler) + assert await queue.process_one(fail_handler) is True + staged_record = await queue.get(task_id) + assert staged_record is not None + assert staged_record["state"] == "pending" + assert staged_record["finalization_action"] == "fail" + assert staged_record["finalization_payload"] == transient_error.to_wire() async with pool.connection() as conn: await conn.execute( - f"UPDATE {workflow_table} " # noqa: S608 - "SET lease_expires_at = now() - interval '1 second' " + f"UPDATE {workflow_table} SET available_at = now() " # noqa: S608 "WHERE task_id = %s", (task_id,), ) @@ -539,6 +550,7 @@ async def fail_handler(_job): assert queue_record["state"] == "dead_lettered" assert task_record is not None assert task_record["state"] == "failed" + assert task_record["error"] == transient_error.to_wire() finally: async with pool.connection() as conn: await conn.execute(f"DROP TABLE IF EXISTS {workflow_table}") # noqa: S608 @@ -643,14 +655,19 @@ async def test_registry_lookup_outage_never_mutates_unverified_task() -> None: registry = PgTaskRegistry(pool=pool, _table=task_table) class LookupOutageRegistry: + lookup_calls = 0 + async def get(self, task_id, *, expected_account_id=None): - raise RuntimeError("registry lookup unavailable") + self.lookup_calls += 1 + if self.lookup_calls <= 2: + raise RuntimeError("registry lookup unavailable") + return await registry.get(task_id, expected_account_id=expected_account_id) async def complete(self, task_id, result): - raise AssertionError("unverified task must not be completed") + await registry.complete(task_id, result) async def fail(self, task_id, error): - raise AssertionError("unverified task must not be failed") + await registry.fail(task_id, error) queue = PgWorkflowQueue( pool=pool, @@ -695,11 +712,309 @@ async def handler(_job): ) assert handler_called is False assert queue_record is not None - assert queue_record["state"] == "dead_lettered" - assert queue_record["attempt_count"] == 2 + assert queue_record["state"] == "pending" + assert queue_record["attempt_count"] == 0 assert queue_record["last_error"] == "RuntimeError" assert task_record is not None assert task_record["state"] == "submitted" + + async with pool.connection() as conn: + await conn.execute( + f"UPDATE {workflow_table} SET available_at = now() " # noqa: S608 + "WHERE task_id = %s", + (task_id,), + ) + assert await queue.process_one(handler) is True + recovered_queue_record = await queue.get(task_id) + recovered_task_record = await registry.get( + task_id, + expected_account_id=account_id, + ) + assert handler_called is True + assert recovered_queue_record is not None + assert recovered_queue_record["state"] == "completed" + assert recovered_queue_record["attempt_count"] == 1 + assert recovered_task_record is not None + assert recovered_task_record["state"] == "completed" + finally: + async with pool.connection() as conn: + await conn.execute(f"DROP TABLE IF EXISTS {workflow_table}") # noqa: S608 + await conn.execute(f"DROP TABLE IF EXISTS {task_table}") # noqa: S608 + + +@pytest.mark.asyncio +async def test_completed_payload_is_staged_before_registry_retry() -> None: + suffix = secrets.token_hex(6) + task_table = f"test_dtasks_{suffix}" + workflow_table = f"test_workflows_{suffix}" + account_id = "tenant-a:account-a" + unsafe_result = { + "media_buy_id": "mb_staged", + "status": "active", + "account": { + "billing_entity": { + "legal_name": "Example Buyer", + "bank": {"account_number": "secret-bank-account"}, + }, + "governance_agents": [ + { + "agent_url": "https://governance.example", + "authentication": {"credentials": "secret-token"}, + } + ], + }, + } + safe_result = { + "media_buy_id": "mb_staged", + "status": "active", + "account": { + "billing_entity": {"legal_name": "Example Buyer"}, + "governance_agents": [{"agent_url": "https://governance.example"}], + }, + } + + async with psycopg_pool.AsyncConnectionPool( + TEST_URL, min_size=1, max_size=4, open=False + ) as pool: + await pool.open() + registry = PgTaskRegistry(pool=pool, _table=task_table) + + class CompleteOnceUnavailableRegistry: + complete_calls = 0 + + async def get(self, task_id, *, expected_account_id=None): + return await registry.get(task_id, expected_account_id=expected_account_id) + + async def complete(self, task_id, payload): + self.complete_calls += 1 + if self.complete_calls == 1: + raise RuntimeError("registry temporarily unavailable") + await registry.complete(task_id, payload) + + async def fail(self, task_id, error): + await registry.fail(task_id, error) + + flaky_registry = CompleteOnceUnavailableRegistry() + queue = PgWorkflowQueue( + pool=pool, + registry=flaky_registry, # type: ignore[arg-type] + table=workflow_table, + retry_base_seconds=1, + ) + await registry.create_schema() + await queue.create_schema() + try: + task_id = await registry.issue(account_id=account_id, task_type="create_media_buy") + await queue.enqueue( + task_id=task_id, + account_id=account_id, + workflow_type="manual_media_buy_approval", + payload={"upstream_order_id": "order-1"}, + ) + handler_calls = 0 + + async def handler(_job): + nonlocal handler_calls + handler_calls += 1 + return unsafe_result + + assert await queue.process_one(handler) is True + staged = await queue.get(task_id) + assert staged is not None + assert staged["state"] == "pending" + assert staged["finalization_action"] == "complete" + assert staged["finalization_payload"] == safe_result + assert unsafe_result["account"]["billing_entity"]["bank"] == { + "account_number": "secret-bank-account" + } + + async with pool.connection() as conn: + await conn.execute( + f"UPDATE {workflow_table} SET available_at = now() WHERE task_id = %s", # noqa: S608 + (task_id,), + ) + assert await queue.process_one(handler) is True + task = await registry.get(task_id, expected_account_id=account_id) + queue_record = await queue.get(task_id) + assert handler_calls == 1 + assert flaky_registry.complete_calls == 2 + assert task is not None and task["result"] == safe_result + assert queue_record is not None and queue_record["state"] == "completed" + assert queue_record["finalization_payload"] is None + finally: + async with pool.connection() as conn: + await conn.execute(f"DROP TABLE IF EXISTS {workflow_table}") # noqa: S608 + await conn.execute(f"DROP TABLE IF EXISTS {task_table}") # noqa: S608 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("recovery", ["terminal", "correctable"]) +async def test_non_transient_adcp_error_round_trips_exact_wire_payload(recovery) -> None: + suffix = secrets.token_hex(6) + task_table = f"test_dtasks_{suffix}" + workflow_table = f"test_workflows_{suffix}" + account_id = "tenant-a:account-a" + error = AdcpError( + "POLICY_VIOLATION", + message="Approval was denied", + recovery=recovery, + field="packages[0]", + suggestion="Choose a different package", + details={"policy_id": "policy-7"}, + ) + + async with psycopg_pool.AsyncConnectionPool( + TEST_URL, min_size=1, max_size=4, open=False + ) as pool: + await pool.open() + registry = PgTaskRegistry(pool=pool, _table=task_table) + queue = PgWorkflowQueue(pool=pool, registry=registry, table=workflow_table) + await registry.create_schema() + await queue.create_schema() + try: + task_id = await registry.issue(account_id=account_id, task_type="create_media_buy") + await queue.enqueue( + task_id=task_id, + account_id=account_id, + workflow_type="manual_media_buy_approval", + payload={"upstream_order_id": "order-1"}, + ) + + async def handler(_job): + raise error + + assert await queue.process_one(handler) is True + task = await registry.get(task_id, expected_account_id=account_id) + queue_record = await queue.get(task_id) + assert task is not None + assert task["state"] == "failed" + assert task["error"] == error.to_wire() + assert queue_record is not None + assert queue_record["state"] == "dead_lettered" + assert queue_record["attempt_count"] == 1 + assert queue_record["finalization_payload"] is None + finally: + async with pool.connection() as conn: + await conn.execute(f"DROP TABLE IF EXISTS {workflow_table}") # noqa: S608 + await conn.execute(f"DROP TABLE IF EXISTS {task_table}") # noqa: S608 + + +@pytest.mark.asyncio +async def test_heartbeat_prevents_reclaim_while_handler_is_live() -> None: + suffix = secrets.token_hex(6) + task_table = f"test_dtasks_{suffix}" + workflow_table = f"test_workflows_{suffix}" + account_id = "tenant-a:account-a" + + async with psycopg_pool.AsyncConnectionPool( + TEST_URL, min_size=1, max_size=6, open=False + ) as pool: + await pool.open() + registry = PgTaskRegistry(pool=pool, _table=task_table) + queue = PgWorkflowQueue(pool=pool, registry=registry, table=workflow_table, lease_seconds=2) + contender = PgWorkflowQueue( + pool=pool, registry=registry, table=workflow_table, lease_seconds=2 + ) + await registry.create_schema() + await queue.create_schema() + try: + task_id = await registry.issue(account_id=account_id, task_type="create_media_buy") + await queue.enqueue( + task_id=task_id, + account_id=account_id, + workflow_type="manual_media_buy_approval", + payload={"upstream_order_id": "order-1"}, + ) + started = asyncio.Event() + finish = asyncio.Event() + + async def handler(_job): + started.set() + await finish.wait() + return {"status": "active"} + + processing = asyncio.create_task(queue.process_one(handler)) + await asyncio.wait_for(started.wait(), timeout=2) + # Make the initial lease expire before the first heartbeat. The + # heartbeat must renew it before the contender's claim. + async with pool.connection() as conn: + await conn.execute( + f"UPDATE {workflow_table} SET lease_expires_at = " # noqa: S608 + "now() + interval '0.2 seconds' WHERE task_id = %s", + (task_id,), + ) + await asyncio.sleep(0.9) + assert await contender.claim() is None + finish.set() + assert await asyncio.wait_for(processing, timeout=2) is True + finally: + async with pool.connection() as conn: + await conn.execute(f"DROP TABLE IF EXISTS {workflow_table}") # noqa: S608 + await conn.execute(f"DROP TABLE IF EXISTS {task_table}") # noqa: S608 + + +@pytest.mark.asyncio +async def test_worker_cancellation_releases_lease_and_cancels_handler() -> None: + suffix = secrets.token_hex(6) + task_table = f"test_dtasks_{suffix}" + workflow_table = f"test_workflows_{suffix}" + account_id = "tenant-a:account-a" + + async with psycopg_pool.AsyncConnectionPool( + TEST_URL, min_size=1, max_size=4, open=False + ) as pool: + await pool.open() + registry = PgTaskRegistry(pool=pool, _table=task_table) + queue = PgWorkflowQueue( + pool=pool, + registry=registry, + table=workflow_table, + lease_seconds=2, + retry_base_seconds=1, + ) + await registry.create_schema() + await queue.create_schema() + try: + task_id = await registry.issue(account_id=account_id, task_type="create_media_buy") + await queue.enqueue( + task_id=task_id, + account_id=account_id, + workflow_type="manual_media_buy_approval", + payload={"upstream_order_id": "order-1"}, + ) + started = asyncio.Event() + handler_cancelled = asyncio.Event() + + async def blocked_handler(_job): + started.set() + try: + await asyncio.Event().wait() + finally: + handler_cancelled.set() + + processing = asyncio.create_task(queue.process_one(blocked_handler)) + await asyncio.wait_for(started.wait(), timeout=2) + processing.cancel() + with pytest.raises(asyncio.CancelledError): + _ = await processing + assert handler_cancelled.is_set() + released = await queue.get(task_id) + assert released is not None + assert released["state"] == "pending" + assert released["attempt_count"] == 1 + + async with pool.connection() as conn: + await conn.execute( + f"UPDATE {workflow_table} SET available_at = now() WHERE task_id = %s", # noqa: S608 + (task_id,), + ) + + async def replacement_handler(_job): + return {"status": "active"} + + assert await queue.process_one(replacement_handler) is True + task = await registry.get(task_id, expected_account_id=account_id) + assert task is not None and task["state"] == "completed" finally: async with pool.connection() as conn: await conn.execute(f"DROP TABLE IF EXISTS {workflow_table}") # noqa: S608