From 02125d58237865fcfc49e0117d80cb62a78b3afa Mon Sep 17 00:00:00 2001 From: Brian O'Kelley Date: Fri, 2 Oct 2026 17:52:35 +0000 Subject: [PATCH 1/9] fix(decisioning): serve account-scoped task polling and listing --- docs/task-registry.md | 43 +++ src/adcp/decisioning/__init__.py | 2 + src/adcp/decisioning/handler.py | 64 +++- src/adcp/decisioning/pg/task_registry.py | 71 +++- src/adcp/decisioning/task_queries.py | 195 +++++++++++ src/adcp/decisioning/task_registry.py | 40 +++ src/adcp/server/base.py | 6 +- tests/test_decisioning_task_polling.py | 357 +++++++++++++++++++++ tests/type_checks/task_registry_listing.py | 21 ++ 9 files changed, 781 insertions(+), 18 deletions(-) create mode 100644 docs/task-registry.md create mode 100644 src/adcp/decisioning/task_queries.py create mode 100644 tests/test_decisioning_task_polling.py create mode 100644 tests/type_checks/task_registry_listing.py diff --git a/docs/task-registry.md b/docs/task-registry.md new file mode 100644 index 000000000..e19e0c47b --- /dev/null +++ b/docs/task-registry.md @@ -0,0 +1,43 @@ +# Decisioning task polling + +`PlatformHandler` serves `get_task_status` for every `TaskRegistry`. It resolves +and authorizes the account before reading, passes `expected_account_id`, and +returns `REFERENCE_NOT_FOUND` with the same message for unknown, discarded, +expired, and cross-account references. Custom registries must enforce that scope +and return `None` for expired records; an optional epoch `expires_at` is also +recognized. SDK registries retain tasks until explicit discard; no TTL is +installed by this change. + +Polls use `status`, ISO UTC timestamps and protocol metadata. Progress, failure +summary and the original task's context are preserved. Set `include_result=true` +to retrieve a terminal artifact. A failed task exposes the stored fatal error +both in its convenience `error` summary and in `result.adcp_error`. + +`list_tasks` is advertised only for the separate optional `ListableTaskRegistry` +protocol. Existing third-party `TaskRegistry` implementations need no changes. +The SDK's memory and PostgreSQL implementations support canonical filters, +sort and pagination, with signed cursors bound to the resolved account and +query. Set `ADCP_PAGINATION_SECRET` consistently across workers if cursors must +survive restarts. Ordering uses task ID as a deterministic tie-breaker. Pages +reflect current state, rather than a frozen snapshot across requests. + +```python +from adcp.decisioning import ListableTaskRegistry, TaskRegistry + +async def reconcile(registry: TaskRegistry, account_id: str) -> None: + if isinstance(registry, ListableTaskRegistry): + page = await registry.list( + account_id=account_id, + filters={"statuses": ["submitted", "working"]}, + pagination={"max_results": 25}, + ) + print(page["tasks"]) +``` + +The pinned legacy list response admits only `media-buy`, `signals`, and `creative` +in its `domain` field. SDK listing includes those domains; other tasks remain +pollable. SDK registries do not store conversation history, so `include_history` +does not fabricate history. Unknown filter extensions are rejected. The current +PostgreSQL implementation filters and pages an account-scoped snapshot in Python; +large task histories may need a custom SQL-native implementation of the optional +protocol. No schema migration is required. diff --git a/src/adcp/decisioning/__init__.py b/src/adcp/decisioning/__init__.py index 2ca3b9c2c..f44143852 100644 --- a/src/adcp/decisioning/__init__.py +++ b/src/adcp/decisioning/__init__.py @@ -223,6 +223,7 @@ def create_media_buy( ) from adcp.decisioning.task_registry import ( InMemoryTaskRegistry, + ListableTaskRegistry, TaskHandoffContext, TaskRegistry, TaskState, @@ -404,6 +405,7 @@ def __init__(self, *args: object, **kwargs: object) -> None: "IncrementalGetProducts", "InMemoryMockAdServer", "InMemoryTaskRegistry", + "ListableTaskRegistry", "MEDIA_BUY_TRANSITIONS", "MaybeAsync", "MediaBuyNotFoundError", diff --git a/src/adcp/decisioning/handler.py b/src/adcp/decisioning/handler.py index 301f8bcb1..b70361d7b 100644 --- a/src/adcp/decisioning/handler.py +++ b/src/adcp/decisioning/handler.py @@ -90,6 +90,8 @@ has_refine_support, project_refine_response, ) +from adcp.decisioning.task_queries import project_task, task_is_expired +from adcp.decisioning.task_registry import ListableTaskRegistry from adcp.decisioning.time_budget import ( SyncExecutorAdmission, project_incomplete_response, @@ -184,6 +186,7 @@ class _PublicAccount(_DecisioningAccount[Any]): GetRightsSuccessResponse, GetSignalsRequest, GetSignalsResponse, + GetTaskStatusRequest, ListAccountsRequest, ListAccountsResponse, ListCollectionListsRequest, @@ -196,6 +199,7 @@ class _PublicAccount(_DecisioningAccount[Any]): ListProductsResponse, ListPropertyListsRequest, ListPropertyListsResponse, + ListTasksRequest, ProvidePerformanceFeedbackRequest, ProvidePerformanceFeedbackResponse, RefineProposalsRequest, @@ -1275,7 +1279,13 @@ class PlatformHandler(ADCPHandler[ToolContext]): | set(_CONTENT_STANDARDS_ADVERTISED_TOOLS) | set(_PROPERTY_LISTS_ADVERTISED_TOOLS) | set(_COLLECTION_LISTS_ADVERTISED_TOOLS) - | {"get_reporting_status", "sync_reporting_status", "sync_reporting_receipts"} + | { + "get_reporting_status", + "sync_reporting_status", + "sync_reporting_receipts", + "get_task_status", + "list_tasks", + } ) _agent_type = "decisioning platform" @@ -1357,6 +1367,13 @@ def advertised_tools_for_instance(self) -> frozenset[str]: reporting = getattr(self._platform, "_reliable_reporting_service", None) if reporting is not None: serving.update(reporting.reporting_tools) + if not serving: + # Preserve the existing class-universe fallback for novel specialisms. + serving.update(self.advertised_tools) + serving.discard("list_tasks") + serving.add("get_task_status") + if isinstance(self._registry, ListableTaskRegistry): + serving.add("list_tasks") return frozenset(serving) def _log_account_tool_dropped(self, tool_name: str, method_name: str) -> None: @@ -2331,6 +2348,51 @@ async def control_media_buy( # type: ignore[override] await self._invoke_compact_lifecycle("control_media_buy", params, context), ) + async def get_task_status( # type: ignore[override] + self, + params: GetTaskStatusRequest, + context: ToolContext | None = None, + ) -> dict[str, Any]: + """Read an authenticated account's application task; hide missing references.""" + account = await self._resolve_account( + params.account, context or ToolContext(), tool_name="get_task_status" + ) + record = await self._registry.get(params.task_id, expected_account_id=account.id) + if record is None or task_is_expired(record): + raise AdcpError( + "REFERENCE_NOT_FOUND", + message="Task reference not found", + recovery="terminal", + field="task_id", + ) + return project_task(record, include_result=bool(params.include_result)) + + async def list_tasks( # type: ignore[override] + self, + params: ListTasksRequest, + context: ToolContext | None = None, + ) -> dict[str, Any] | NotImplementedResponse: + """Reconcile tasks when the registry supports the optional list protocol.""" + if not isinstance(self._registry, ListableTaskRegistry): + return self._not_supported("list_tasks") + account = await self._resolve_account( + params.account, context or ToolContext(), tool_name="list_tasks" + ) + return await self._registry.list( + account_id=account.id, + filters=( + params.filters.model_dump(mode="json", exclude_none=True) + if params.filters + else None + ), + sort=params.sort.model_dump(mode="json", exclude_none=True) if params.sort else None, + pagination=( + params.pagination.model_dump(mode="json", exclude_none=True) + if params.pagination + else None + ), + ) + async def get_products( # type: ignore[override] self, params: GetProductsRequest, diff --git a/src/adcp/decisioning/pg/task_registry.py b/src/adcp/decisioning/pg/task_registry.py index 0e1129cdd..4f0ab50a0 100644 --- a/src/adcp/decisioning/pg/task_registry.py +++ b/src/adcp/decisioning/pg/task_registry.py @@ -165,33 +165,33 @@ def __init__( # Pre-format queries at construction so the hot path avoids f-strings per call. # _table is whitelisted by _SAFE_IDENTIFIER_RE above. self._sql_insert = ( # noqa: S608 — table name is whitelisted - f"INSERT INTO {self._table}" + f"INSERT INTO {self._table}" # nosec B608 — table identifier is constructor-validated f" (task_id, account_id, state, task_type, request_context," f" webhook_registration, webhook_registration_nonce, created_at, updated_at)" f" VALUES (%s, %s, 'submitted', %s, %s::jsonb, %s, %s, %s, %s)" ) self._sql_update_progress = ( # noqa: S608 - f"UPDATE {self._table}" + f"UPDATE {self._table}" # nosec B608 — table identifier is constructor-validated f" SET state = CASE state WHEN 'submitted' THEN 'working' ELSE state END," f" progress = %s::jsonb, updated_at = %s" f" WHERE task_id = %s AND state NOT IN ('completed', 'failed')" ) self._sql_complete = ( # noqa: S608 - f"UPDATE {self._table}" + f"UPDATE {self._table}" # nosec B608 — table identifier is constructor-validated f" SET state = 'completed', result = %s::jsonb, updated_at = %s" f" WHERE task_id = %s AND state NOT IN ('completed', 'failed')" f" RETURNING task_id, account_id, task_type, webhook_registration," f" webhook_registration_nonce" ) self._sql_fail = ( # noqa: S608 - f"UPDATE {self._table}" + f"UPDATE {self._table}" # nosec B608 — table identifier is constructor-validated f" SET state = 'failed', error = %s::jsonb, updated_at = %s" f" WHERE task_id = %s AND state NOT IN ('completed', 'failed')" f" RETURNING task_id, account_id, task_type, webhook_registration," f" webhook_registration_nonce" ) self._sql_clear_webhook_registration = ( # noqa: S608 - f"UPDATE {self._table} SET webhook_registration = NULL," + f"UPDATE {self._table} SET webhook_registration = NULL," # nosec B608 — table identifier is constructor-validated f" webhook_registration_nonce = NULL WHERE task_id = %s" ) # Explicit ``::text`` cast on the optional account-filter @@ -201,21 +201,18 @@ def __init__( # ``%s IS NULL`` predicate gives psycopg no type context # for the parameter and the query fails at prepare time. self._sql_get = ( # noqa: S608 - f"SELECT task_id, account_id, state, task_type," + f"SELECT task_id, account_id, state, task_type," # nosec B608 — table identifier is constructor-validated f" progress, result, error, request_context, created_at, updated_at" f" FROM {self._table}" f" WHERE task_id = %s AND (%s::text IS NULL OR account_id = %s)" ) - self._sql_get_state_result = ( # noqa: S608 - f"SELECT state, result, task_type FROM {self._table} WHERE task_id = %s" + self._sql_get_state_result = ( + f"SELECT state, result, task_type FROM {self._table}" # nosec B608 — validated identifier + " WHERE task_id = %s" ) - self._sql_get_task_type = ( # noqa: S608 - f"SELECT task_type FROM {self._table} WHERE task_id = %s" - ) - self._sql_get_state_error = ( # noqa: S608 - f"SELECT state, error FROM {self._table} WHERE task_id = %s" - ) - self._sql_discard = f"DELETE FROM {self._table} WHERE task_id = %s" # noqa: S608 + self._sql_get_task_type = f"SELECT task_type FROM {self._table} WHERE task_id = %s" # noqa: S608 # nosec B608 — table identifier is constructor-validated + self._sql_get_state_error = f"SELECT state, error FROM {self._table} WHERE task_id = %s" # noqa: S608 # nosec B608 — table identifier is constructor-validated + self._sql_discard = f"DELETE FROM {self._table} WHERE task_id = %s" # noqa: S608 # nosec B608 — table identifier is constructor-validated self._sql_ddl = ( # noqa: S608 f"CREATE TABLE IF NOT EXISTS {self._table} (" f' task_id TEXT COLLATE "C" NOT NULL PRIMARY KEY,' @@ -524,6 +521,50 @@ async def get( **({"context": row[7]} if row[7] is not None else {}), } + async def list( + self, + *, + account_id: str, + filters: dict[str, Any] | None = None, + sort: dict[str, Any] | None = None, + pagination: dict[str, Any] | None = None, + ) -> dict[str, Any]: + from adcp.decisioning.task_queries import list_task_records + + # Account isolation is a SQL predicate, before rows are materialized. + async with self._pool.connection() as conn: + cur = await conn.execute( + f"SELECT task_id, account_id, state, task_type, progress, result, error," # nosec B608 — table identifier is constructor-validated + f" request_context, created_at, updated_at," + f" (webhook_registration IS NOT NULL) FROM {self._table} WHERE account_id = %s", + (account_id,), + ) + rows = await cur.fetchall() + records = [ + dict( + zip( + ( + "task_id", + "account_id", + "state", + "task_type", + "progress", + "result", + "error", + "context", + "created_at", + "updated_at", + "has_webhook", + ), + row, + ) + ) + for row in rows + ] + return list_task_records( + records, account_id=account_id, filters=filters, sort=sort, pagination=pagination + ) + async def _enqueue_terminal_if_registered( self, conn: Any, diff --git a/src/adcp/decisioning/task_queries.py b/src/adcp/decisioning/task_queries.py new file mode 100644 index 000000000..438fa85d6 --- /dev/null +++ b/src/adcp/decisioning/task_queries.py @@ -0,0 +1,195 @@ +"""Shared application-task wire projection and reconciliation queries.""" + +from __future__ import annotations + +import json +import time +from collections import Counter +from datetime import datetime, timezone +from functools import lru_cache +from typing import Any + +from adcp.decisioning.pagination import _decode_cursor, _encode_cursor, _query_hash +from adcp.decisioning.types import AdcpError +from adcp.validation.schema_loader import get_named_schema_document + + +@lru_cache(maxsize=1) +def _task_protocols() -> dict[str, str]: + manifest = get_named_schema_document("manifest.json") or {} + # Manifest namespaces include administrative surfaces; task categorization + # uses the coarser adcp-protocol enum. + aliases = { + "account": "media-buy", + "content-standards": "governance", + "property": "governance", + "collection": "governance", + "compliance": "governance", + } + return { + name: aliases.get(tool["protocol"], tool["protocol"]) + for name, tool in manifest.get("tools", {}).items() + } + + +def task_protocol(record: dict[str, Any]) -> str: + protocol = record.get("protocol") or _task_protocols().get(record["task_type"]) + if not isinstance(protocol, str): + raise AdcpError("INTERNAL_ERROR", message="Task has no protocol metadata") + return protocol + + +def task_is_expired(record: dict[str, Any]) -> bool: + """Custom registries may expose an epoch expiry; SDK stores retain until discard.""" + expires = record.get("expires_at") + return expires is not None and float(expires) <= time.time() + + +def project_task(record: dict[str, Any], *, include_result: bool = False) -> dict[str, Any]: + response: dict[str, Any] = { + "task_id": record["task_id"], + "task_type": record["task_type"], + "protocol": task_protocol(record), + "status": record["state"], + "created_at": datetime.fromtimestamp(record["created_at"], timezone.utc).isoformat(), + "updated_at": datetime.fromtimestamp(record["updated_at"], timezone.utc).isoformat(), + } + for key in ("progress", "error", "context", "has_webhook"): + if record.get(key) is not None: + response[key] = record[key] + if record["state"] in {"completed", "failed", "canceled", "rejected"}: + response["completed_at"] = response["updated_at"] + if include_result: + if record.get("result") is not None: + response["result"] = record["result"] + elif record["state"] == "failed" and record.get("error") is not None: + response["result"] = {"status": "failed", "adcp_error": record["error"]} + return response + + +def list_task_records( + records: list[dict[str, Any]], + *, + account_id: str, + filters: dict[str, Any] | None = None, + sort: dict[str, Any] | None = None, + pagination: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Query an account-scoped snapshot. Cursors bind account, filters and sort. + + The legacy list schema admits media-buy, creative and signals domains only. + Other domains remain pollable through get_task_status. SDK registries do not + persist conversation history. Unknown filter extensions are rejected rather + than silently returning a broader result set. + """ + filters, sort, pagination = filters or {}, sort or {}, pagination or {} + supported = { + "protocol", + "protocols", + "status", + "statuses", + "task_type", + "task_types", + "created_after", + "created_before", + "updated_after", + "updated_before", + "task_ids", + "context_contains", + "has_webhook", + } + if filters.keys() - supported: + raise AdcpError("INVALID_REQUEST", message="Unsupported task filter", field="filters") + field, direction = sort.get("field", "created_at"), sort.get("direction", "desc") + if field not in { + "created_at", + "updated_at", + "status", + "task_type", + "protocol", + } or direction not in {"asc", "desc"}: + raise AdcpError("INVALID_REQUEST", message="Invalid task sort", field="sort") + max_results = pagination.get("max_results", 50) + if ( + not isinstance(max_results, int) + or isinstance(max_results, bool) + or not 1 <= max_results <= 100 + ): + raise AdcpError( + "INVALID_REQUEST", message="max_results must be 1–100", field="pagination.max_results" + ) + query_hash = _query_hash( + { + "account_id": account_id, + "filters": filters, + "sort": {"field": field, "direction": direction}, + } + ) + cursor = pagination.get("cursor") + offset = _decode_cursor(cursor, query_hash) if cursor else 0 + if offset < 0: + raise AdcpError("INVALID_REQUEST", message="Invalid task cursor", field="pagination.cursor") + tasks = [] + for record in records: + if record["account_id"] != account_id or task_is_expired(record): + continue + task = project_task(record) + if task["protocol"] not in {"media-buy", "creative", "signals"}: + continue + matches = True + for singular, plural in ( + ("protocol", "protocols"), + ("status", "statuses"), + ("task_type", "task_types"), + ): + if singular in filters and task[singular] != filters[singular]: + matches = False + if plural in filters and task[singular] not in filters[plural]: + matches = False + for stamp in ("created", "updated"): + for bound, op in (("after", "after"), ("before", "before")): + value = filters.get(f"{stamp}_{bound}") + if value is not None: + try: + date = datetime.fromisoformat(str(value).replace("Z", "+00:00")) + if date.tzinfo is None: + raise ValueError("missing timezone") + epoch = date.timestamp() + except (ValueError, TypeError): + raise AdcpError( + "INVALID_REQUEST", message="Invalid task date filter", field="filters" + ) from None + if (op == "after" and record[f"{stamp}_at"] <= epoch) or ( + op == "before" and record[f"{stamp}_at"] >= epoch + ): + matches = False + if "task_ids" in filters and task["task_id"] not in filters["task_ids"]: + matches = False + if "context_contains" in filters and filters["context_contains"] not in json.dumps( + record.get("context") or {}, sort_keys=True + ): + matches = False + if ( + "has_webhook" in filters + and bool(record.get("has_webhook", False)) != filters["has_webhook"] + ): + matches = False + if matches: + tasks.append(task) + tasks.sort(key=lambda task: (task[field], task["task_id"]), reverse=direction == "desc") + total = len(tasks) + page = tasks[offset : offset + max_results] + metadata: dict[str, Any] = {"has_more": offset + len(page) < total, "total_count": total} + if metadata["has_more"]: + metadata["cursor"] = _encode_cursor(offset + len(page), query_hash) + summary = { + "total_matching": total, + "returned": len(page), + "domain_breakdown": dict(Counter(task["protocol"] for task in tasks)), + "status_breakdown": dict(Counter(task["status"] for task in tasks)), + "filters_applied": list(filters), + "sort_applied": {"field": field, "direction": direction}, + } + for task in page: + task["domain"] = task.pop("protocol") + return {"status": "completed", "query_summary": summary, "tasks": page, "pagination": metadata} diff --git a/src/adcp/decisioning/task_registry.py b/src/adcp/decisioning/task_registry.py index c4479b259..62e6c0540 100644 --- a/src/adcp/decisioning/task_registry.py +++ b/src/adcp/decisioning/task_registry.py @@ -382,6 +382,26 @@ async def discard(self, task_id: str) -> None: ... +@runtime_checkable +class ListableTaskRegistry(Protocol): + """Optional account-scoped reconciliation extension; TaskRegistry stays minimal. + + Return a wire-shaped list_tasks page with query_summary, tasks and pagination. + Filters/sort/pagination use the canonical request's JSON shape. Implementations + must scope reads before materializing records and hide expired references. + Conversation history is not required because SDK registries do not store it. + """ + + async def list( + self, + *, + account_id: str, + filters: dict[str, Any] | None = None, + sort: dict[str, Any] | None = None, + pagination: dict[str, Any] | None = None, + ) -> dict[str, Any]: ... + + # --------------------------------------------------------------------------- # In-memory reference implementation — v6.0 ships this; v6.1 lands a # durable Postgres-backed counterpart that implements the same Protocol. @@ -557,6 +577,26 @@ async def get( return None return record.to_dict() + async def list( + self, + *, + account_id: str, + filters: dict[str, Any] | None = None, + sort: dict[str, Any] | None = None, + pagination: dict[str, Any] | None = None, + ) -> dict[str, Any]: + from adcp.decisioning.task_queries import list_task_records + + async with self._lock: + records = [ + record.to_dict() + for record in self._records.values() + if record.account_id == account_id + ] + return list_task_records( + records, account_id=account_id, filters=filters, sort=sort, pagination=pagination + ) + async def discard(self, task_id: str) -> None: async with self._lock: # Idempotent: pop with default. The Protocol contract diff --git a/src/adcp/server/base.py b/src/adcp/server/base.py index c804c1e5f..94fad3a99 100644 --- a/src/adcp/server/base.py +++ b/src/adcp/server/base.py @@ -707,7 +707,8 @@ async def get_task_status( ) -> Any: """Get task status. - Override this to expose persisted async task status. + PlatformHandler implements this using its TaskRegistry. + Other handlers override it to expose persisted async task status. """ return self._not_supported("get_task_status") @@ -718,7 +719,8 @@ async def list_tasks( ) -> Any: """List tasks. - Override this to expose persisted async tasks. + PlatformHandler implements this for ListableTaskRegistry stores. + Other handlers override it to expose persisted async tasks. """ return self._not_supported("list_tasks") diff --git a/tests/test_decisioning_task_polling.py b/tests/test_decisioning_task_polling.py new file mode 100644 index 000000000..5a508ddc6 --- /dev/null +++ b/tests/test_decisioning_task_polling.py @@ -0,0 +1,357 @@ +"""Application task polling/listing through the actual MCP HTTP transport.""" + +from __future__ import annotations + +import asyncio +import json +import os +import secrets +from collections.abc import AsyncIterator +from typing import Any, ClassVar + +import pytest + +from adcp.decisioning import ( + Account, + DecisioningCapabilities, + DecisioningPlatform, + InMemoryTaskRegistry, + ListableTaskRegistry, + SingletonAccounts, + TaskRegistry, +) +from adcp.decisioning.serve import create_adcp_server_from_platform +from adcp.decisioning.types import AdcpError +from adcp.server.mcp_tools import create_tool_caller, get_tools_for_handler +from adcp.testing import build_test_client +from adcp.validation.schema_loader import get_validator + + +class Accounts(SingletonAccounts): + def resolve(self, ref: dict[str, Any] | None, *, auth_info: Any = None) -> Account[Any]: + return Account(id=(ref or {}).get("account_id", "acct_1")) + + +class Seller(DecisioningPlatform): + capabilities = DecisioningCapabilities( + specialisms=["sales-non-guaranteed"], channels=["display"], pricing_models=["cpm"] + ) + accounts = Accounts(account_id="acct_1") + + def __init__(self) -> None: + self.started = asyncio.Event() + self.finish = asyncio.Event() + self.failed = False + + def get_products(self, req: Any, ctx: Any) -> Any: + return ctx.handoff_to_task(self.work) + + async def work(self, ctx: Any) -> dict[str, Any]: + await ctx.update({"percentage": 25, "current_step": "Review"}) + self.started.set() + await self.finish.wait() + if self.failed: + raise AdcpError("INTERNAL_ERROR", message="Background failed", recovery="terminal") + return {"products": []} + + def create_media_buy(self, req: Any, ctx: Any) -> Any: + raise NotImplementedError + + def update_media_buy(self, media_buy_id: Any, patch: Any, ctx: Any) -> Any: + raise NotImplementedError + + def sync_creatives(self, req: Any, ctx: Any) -> Any: + raise NotImplementedError + + def get_media_buy_delivery(self, req: Any, ctx: Any) -> Any: + raise NotImplementedError + + def get_media_buys(self, req: Any, ctx: Any) -> Any: + raise NotImplementedError + + +@pytest.fixture(params=["memory", "postgres"]) +async def registry(request: pytest.FixtureRequest) -> AsyncIterator[Any]: + if request.param == "memory": + yield InMemoryTaskRegistry() + return + from psycopg_pool import AsyncConnectionPool + + from adcp.decisioning.pg import PgTaskRegistry + + url = os.environ.get("ADCP_PG_TEST_URL") + if not url: + pytest.skip("ADCP_PG_TEST_URL required") + table = f"test_polling_{secrets.token_hex(6)}" + async with AsyncConnectionPool(url, open=False) as pool: + reg = PgTaskRegistry(pool=pool, _table=table) + await reg.create_schema() + yield reg + async with pool.connection() as conn: + await conn.execute(f"DROP TABLE {table}") + + +def rpc_data(response: Any) -> dict[str, Any]: + body = ( + response.json() + if response.headers.get("content-type", "").startswith("application/json") + else json.loads( + next(line[6:] for line in response.text.splitlines() if line.startswith("data: ")) + ) + ) + assert "error" not in body, body + result = body["result"] + if "artifacts" in result: + return next( + part["data"] + for artifact in result["artifacts"] + for part in artifact["parts"] + if "data" in part + ) + if "structuredContent" in result: + return result["structuredContent"] + if "content" in result: + return json.loads(result["content"][0]["text"]) + return result + + +@pytest.mark.parametrize("failed", [False, True]) +@pytest.mark.parametrize("transport", ["mcp", "a2a"]) +async def test_handoff_poll_progress_terminal_over_http( + registry: Any, failed: bool, transport: str +) -> None: + seller = Seller() + seller.failed = failed + async with build_test_client( + seller, + registry=registry, + validate_at_init=False, + stateless_http=True, + validation=None, + transport=transport, + ) as client: + + async def rpc(method: str, params: dict[str, Any]) -> dict[str, Any]: + if transport == "a2a": + if method == "tools/list": + card = (await client.get("/.well-known/agent-card.json")).json() + return {"tools": [{"name": skill["id"]} for skill in card["skills"]]} + body = { + "jsonrpc": "2.0", + "id": str(secrets.token_hex(4)), + "method": "message/send", + "params": { + "message": { + "messageId": secrets.token_hex(8), + "role": "user", + "parts": [ + { + "kind": "data", + "data": { + "skill": params["name"], + "parameters": params["arguments"], + }, + } + ], + } + }, + } + response = await client.post("/", json=body) + else: + response = await client.post( + "/mcp/", + headers={"accept": "application/json, text/event-stream"}, + json={"jsonrpc": "2.0", "id": 1, "method": method, "params": params}, + ) + assert response.status_code == 200, response.text + return rpc_data(response) + + tools = await rpc("tools/list", {}) + assert {"get_task_status", "list_tasks"} <= {t["name"] for t in tools["tools"]} + issued = await rpc( + "tools/call", + { + "name": "get_products", + "arguments": { + "brief": "Test", + "buying_mode": "brief", + "account": {"account_id": "acct_1"}, + }, + }, + ) + task_id = issued["task_id"] + await asyncio.wait_for(seller.started.wait(), 5) + working = await rpc( + "tools/call", {"name": "get_task_status", "arguments": {"task_id": task_id}} + ) + assert working["status"] == "working" + assert working["progress"]["percentage"] == 25 + assert working["protocol"] == "media-buy" + assert isinstance(working["created_at"], str) + assert "result" not in working + seller.finish.set() + for _ in range(100): + record = await registry.get(task_id) + if record["state"] in {"completed", "failed"}: + break + await asyncio.sleep(0.01) + terminal = await rpc( + "tools/call", + {"name": "get_task_status", "arguments": {"task_id": task_id, "include_result": True}}, + ) + assert terminal["status"] == ("failed" if failed else "completed") + assert "completed_at" in terminal + if failed: + assert terminal["error"]["message"] == "Background failed" + assert terminal["result"]["adcp_error"] == terminal["error"] + else: + assert terminal["result"]["products"] == [] + validator = get_validator("get_task_status", "sync") + assert validator is not None + validator.validate(terminal) + listed = await rpc( + "tools/call", {"name": "list_tasks", "arguments": {"filters": {"task_ids": [task_id]}}} + ) + assert listed["tasks"][0]["task_id"] == task_id + assert listed["tasks"][0]["domain"] == "media-buy" + list_validator = get_validator("list_tasks", "sync") + assert list_validator is not None + list_validator.validate(listed) + + +async def test_scoped_listing_filters_cursor_and_not_found(registry: Any) -> None: + handler, executor, _ = create_adcp_server_from_platform( + Seller(), registry=registry, validate_at_init=False + ) + try: + caller = create_tool_caller(handler, "get_task_status") + task = await registry.issue( + account_id="acct_1", task_type="get_products", request_context={"campaign": "alpha"} + ) + foreign = await registry.issue(account_id="acct_2", task_type="get_products") + await registry.complete(task, {"products": []}) + other = await registry.issue(account_id="acct_1", task_type="get_signals") + + async def missing_error(task_id: str) -> dict[str, Any]: + with pytest.raises(AdcpError) as caught: + await caller({"task_id": task_id}) + return caught.value.to_wire() + + missing = await missing_error("absent") + cross = await missing_error(foreign) + await registry.discard(other) + discarded = await missing_error(other) + assert missing == cross == discarded + assert missing["code"] == "REFERENCE_NOT_FOUND" + first = await registry.list( + account_id="acct_1", + filters={ + "context_contains": "alpha", + "statuses": ["completed"], + "protocols": ["media-buy"], + }, + pagination={"max_results": 1}, + ) + assert [t["task_id"] for t in first["tasks"]] == [task] + for _ in range(2): + await registry.issue(account_id="acct_1", task_type="get_products") + page = await registry.list(account_id="acct_1", pagination={"max_results": 1}) + seen = {page["tasks"][0]["task_id"]} + while page["pagination"]["has_more"]: + page = await registry.list( + account_id="acct_1", + pagination={"max_results": 1, "cursor": page["pagination"]["cursor"]}, + ) + seen.add(page["tasks"][0]["task_id"]) + assert len(seen) == 3 and foreign not in seen + cursor = (await registry.list(account_id="acct_1", pagination={"max_results": 1}))[ + "pagination" + ]["cursor"] + with pytest.raises(AdcpError): + await registry.list(account_id="acct_2", pagination={"cursor": cursor}) + with pytest.raises(AdcpError): + await registry.list( + account_id="acct_1", filters={"status": "completed"}, pagination={"cursor": cursor} + ) + finally: + executor.shutdown() + + +class MinimalRegistry: + is_durable: ClassVar[bool] = True + + def __init__(self) -> None: + self.delegate = InMemoryTaskRegistry() + + async def issue(self, **kwargs: Any) -> str: + return await self.delegate.issue(**kwargs) + + async def get(self, task_id: str, *, expected_account_id: str | None = None) -> Any: + record = await self.delegate.get(task_id, expected_account_id=expected_account_id) + if record is not None: + record["expires_at"] = 0 + return record + + async def update_progress(self, task_id: str, progress: dict[str, Any]) -> None: + await self.delegate.update_progress(task_id, progress) + + async def complete(self, task_id: str, result: dict[str, Any]) -> None: + await self.delegate.complete(task_id, result) + + async def fail(self, task_id: str, error: dict[str, Any]) -> None: + await self.delegate.fail(task_id, error) + + async def discard(self, task_id: str) -> None: + await self.delegate.discard(task_id) + + +async def test_minimal_third_party_registry_boots_and_expiry_is_hidden() -> None: + registry = MinimalRegistry() + assert isinstance(registry, TaskRegistry) + assert not isinstance(registry, ListableTaskRegistry) + handler, executor, _ = create_adcp_server_from_platform( + Seller(), registry=registry, validate_at_init=False + ) + try: + names = {t["name"] for t in get_tools_for_handler(handler)} + assert "get_task_status" in names and "list_tasks" not in names + task = await registry.issue(account_id="acct_1", task_type="get_products") + caller = create_tool_caller(handler, "get_task_status") + errors = [] + for reference in (task, "absent"): + with pytest.raises(AdcpError) as caught: + await caller({"task_id": reference}) + errors.append(caught.value.to_wire()) + assert errors[0] == errors[1] + assert (await create_tool_caller(handler, "list_tasks")({}))["error"][ + "code" + ] == "NOT_SUPPORTED" + finally: + executor.shutdown() + + +async def test_account_resolution_precedes_registry_reads() -> None: + from unittest.mock import AsyncMock + + class DeniedAccounts(Accounts): + def resolve(self, ref: dict[str, Any] | None, *, auth_info: Any = None) -> Account[Any]: + raise AdcpError("PERMISSION_DENIED", message="Denied") + + seller = Seller() + seller.accounts = DeniedAccounts(account_id="acct_1") + reg = InMemoryTaskRegistry() + handler, executor, _ = create_adcp_server_from_platform( + seller, registry=reg, validate_at_init=False + ) + try: + read = AsyncMock(wraps=reg.get) + listing = AsyncMock(wraps=reg.list) + reg.get = read + reg.list = listing + for tool, params in (("get_task_status", {"task_id": "probe"}), ("list_tasks", {})): + with pytest.raises(AdcpError, match="PERMISSION_DENIED"): + await create_tool_caller(handler, tool)(params) + read.assert_not_awaited() + listing.assert_not_awaited() + finally: + executor.shutdown() diff --git a/tests/type_checks/task_registry_listing.py b/tests/type_checks/task_registry_listing.py new file mode 100644 index 000000000..a04ccbbf2 --- /dev/null +++ b/tests/type_checks/task_registry_listing.py @@ -0,0 +1,21 @@ +"""Adopter contract: optional reconciliation never enlarges TaskRegistry.""" + +from typing import Any + +from adcp.decisioning import InMemoryTaskRegistry, ListableTaskRegistry, TaskRegistry +from adcp.decisioning.pg import PgTaskRegistry + + +async def reconcile(registry: TaskRegistry, account_id: str) -> dict[str, Any] | None: + if isinstance(registry, ListableTaskRegistry): + return await registry.list( + account_id=account_id, filters={"status": "working"}, pagination={"max_results": 10} + ) + return None + + +memory: TaskRegistry = InMemoryTaskRegistry() + + +def postgres_registry(registry: PgTaskRegistry) -> TaskRegistry: + return registry From e38261515f0d59249bd95b68275add19ac9eb4fb Mon Sep 17 00:00:00 2001 From: Brian O'Kelley Date: Fri, 2 Oct 2026 17:56:31 +0000 Subject: [PATCH 2/9] feat(decisioning): observe committed task lifecycle transitions --- docs/task-lifecycle-observers.md | 37 ++++ src/adcp/decisioning/__init__.py | 4 + src/adcp/decisioning/pg/task_registry.py | 107 +++++++---- src/adcp/decisioning/task_registry.py | 90 +++++++++- tests/test_decisioning_task_lifecycle.py | 166 ++++++++++++++++++ tests/type_checks/task_lifecycle_observers.py | 39 ++++ 6 files changed, 404 insertions(+), 39 deletions(-) create mode 100644 docs/task-lifecycle-observers.md create mode 100644 tests/test_decisioning_task_lifecycle.py create mode 100644 tests/type_checks/task_lifecycle_observers.py diff --git a/docs/task-lifecycle-observers.md b/docs/task-lifecycle-observers.md new file mode 100644 index 000000000..66ed5b9c9 --- /dev/null +++ b/docs/task-lifecycle-observers.md @@ -0,0 +1,37 @@ +# Task lifecycle metrics + +`InMemoryTaskRegistry` and `PgTaskRegistry` offer additive lifecycle hooks. The +minimal `TaskRegistry` protocol and all existing write-method return values stay +unchanged. Register a synchronous callback with `add_lifecycle_observer` and +remove it with `remove_lifecycle_observer` (returns whether removal succeeded). +Registering the same callback twice is a no-op. + +```python +from adcp.decisioning import InMemoryTaskRegistry, TaskTransition + +registry = InMemoryTaskRegistry() + +def record_transition(event: TaskTransition, *, task_id: str, account_id: str, + task_type: str, created_at: float, updated_at: float) -> None: + print(event, task_type, updated_at - created_at) + +registry.add_lifecycle_observer(record_transition) +``` + +Events are `submitted`, the first `working` transition, `completed`, `failed`, +and `discarded`. Repeated progress writes, idempotent terminal retries, unknown +progress/discard requests and failed transactions do not emit events. Terminal +states cannot be changed into another terminal state; memory and PostgreSQL +registries enforce the same rule, including concurrent complete/fail calls. + +Callbacks run after the PostgreSQL transaction and connection contexts exit, or +after the memory registry lock is released. Each event contains a snapshot of +its transition metadata with epoch timestamps; discard uses deletion time for +`updated_at`. Registration is thread-safe, and delivery uses a snapshot of the +registered callbacks. Callback failures are logged and swallowed without +preventing remaining callbacks or changing the write result. + +These hooks are best-effort metrics, not durable events: process exit can lose a +notification, and concurrent writers may notify out of order after commit. Keep +callbacks short and synchronous. Use the durable task webhook outbox for +buyer-facing delivery. No database migration is needed. diff --git a/src/adcp/decisioning/__init__.py b/src/adcp/decisioning/__init__.py index 2ca3b9c2c..5bb339608 100644 --- a/src/adcp/decisioning/__init__.py +++ b/src/adcp/decisioning/__init__.py @@ -224,8 +224,10 @@ def create_media_buy( from adcp.decisioning.task_registry import ( InMemoryTaskRegistry, TaskHandoffContext, + TaskLifecycleObserver, TaskRegistry, TaskState, + TaskTransition, TaskWebhookAuthentication, ) from adcp.decisioning.tenant_store import create_tenant_store @@ -473,6 +475,8 @@ def __init__(self, *args: object, **kwargs: object) -> None: "TaskHandoff", "TaskHandoffContext", "TaskRegistry", + "TaskLifecycleObserver", + "TaskTransition", "TaskState", "TaskWebhookAuthentication", "WebhookSenderResolver", diff --git a/src/adcp/decisioning/pg/task_registry.py b/src/adcp/decisioning/pg/task_registry.py index 0e1129cdd..a9fe431a5 100644 --- a/src/adcp/decisioning/pg/task_registry.py +++ b/src/adcp/decisioning/pg/task_registry.py @@ -71,7 +71,7 @@ async def main(): from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias from adcp.decisioning.account_projection import strip_credentials_from_wire_result -from adcp.decisioning.task_registry import TaskWebhookAuthentication +from adcp.decisioning.task_registry import TaskWebhookAuthentication, _TaskLifecycleObservers if TYPE_CHECKING: from psycopg_pool import AsyncConnectionPool @@ -102,7 +102,7 @@ async def main(): WebhookSigningScopeResolver: TypeAlias = Callable[["RequestContext[Any]"], str | Awaitable[str]] -class PgTaskRegistry: +class PgTaskRegistry(_TaskLifecycleObservers): """PostgreSQL-backed :class:`~adcp.decisioning.TaskRegistry` — v6.1. Durable counterpart to :class:`~adcp.decisioning.InMemoryTaskRegistry`. @@ -156,6 +156,7 @@ def __init__( "webhook_signing_scope_resolver is required exactly when the " "PgTaskWebhookOutbox uses sender_resolver" ) + self._init_lifecycle_observers() self._pool = pool self._table = _table self.task_webhook_outbox = task_webhook_outbox @@ -165,33 +166,36 @@ def __init__( # Pre-format queries at construction so the hot path avoids f-strings per call. # _table is whitelisted by _SAFE_IDENTIFIER_RE above. self._sql_insert = ( # noqa: S608 — table name is whitelisted - f"INSERT INTO {self._table}" + f"INSERT INTO {self._table}" # nosec B608 — table identifier is constructor-validated f" (task_id, account_id, state, task_type, request_context," f" webhook_registration, webhook_registration_nonce, created_at, updated_at)" f" VALUES (%s, %s, 'submitted', %s, %s::jsonb, %s, %s, %s, %s)" ) self._sql_update_progress = ( # noqa: S608 - f"UPDATE {self._table}" - f" SET state = CASE state WHEN 'submitted' THEN 'working' ELSE state END," - f" progress = %s::jsonb, updated_at = %s" - f" WHERE task_id = %s AND state NOT IN ('completed', 'failed')" + f"WITH previous AS (SELECT task_id, state FROM {self._table}" # nosec B608 — table identifier is constructor-validated + f" WHERE task_id = %s AND state NOT IN ('completed', 'failed') FOR UPDATE)" + f" UPDATE {self._table} AS task" + f" SET state = 'working', progress = %s::jsonb, updated_at = %s" + f" FROM previous WHERE task.task_id = previous.task_id" + f" RETURNING previous.state, task.task_id, task.account_id, task.task_type," + f" task.created_at, task.updated_at" ) self._sql_complete = ( # noqa: S608 - f"UPDATE {self._table}" + f"UPDATE {self._table}" # nosec B608 — table identifier is constructor-validated f" SET state = 'completed', result = %s::jsonb, updated_at = %s" f" WHERE task_id = %s AND state NOT IN ('completed', 'failed')" f" RETURNING task_id, account_id, task_type, webhook_registration," - f" webhook_registration_nonce" + f" webhook_registration_nonce, created_at, updated_at" ) self._sql_fail = ( # noqa: S608 - f"UPDATE {self._table}" + f"UPDATE {self._table}" # nosec B608 — table identifier is constructor-validated f" SET state = 'failed', error = %s::jsonb, updated_at = %s" f" WHERE task_id = %s AND state NOT IN ('completed', 'failed')" f" RETURNING task_id, account_id, task_type, webhook_registration," - f" webhook_registration_nonce" + f" webhook_registration_nonce, created_at, updated_at" ) self._sql_clear_webhook_registration = ( # noqa: S608 - f"UPDATE {self._table} SET webhook_registration = NULL," + f"UPDATE {self._table} SET webhook_registration = NULL," # nosec B608 — table identifier is constructor-validated f" webhook_registration_nonce = NULL WHERE task_id = %s" ) # Explicit ``::text`` cast on the optional account-filter @@ -201,21 +205,21 @@ def __init__( # ``%s IS NULL`` predicate gives psycopg no type context # for the parameter and the query fails at prepare time. self._sql_get = ( # noqa: S608 - f"SELECT task_id, account_id, state, task_type," + f"SELECT task_id, account_id, state, task_type," # nosec B608 — table identifier is constructor-validated f" progress, result, error, request_context, created_at, updated_at" f" FROM {self._table}" f" WHERE task_id = %s AND (%s::text IS NULL OR account_id = %s)" ) self._sql_get_state_result = ( # noqa: S608 - f"SELECT state, result, task_type FROM {self._table} WHERE task_id = %s" - ) - self._sql_get_task_type = ( # noqa: S608 - f"SELECT task_type FROM {self._table} WHERE task_id = %s" + f"SELECT state, result, task_type FROM {self._table}" # nosec B608 — validated identifier + " WHERE task_id = %s" ) - self._sql_get_state_error = ( # noqa: S608 - f"SELECT state, error FROM {self._table} WHERE task_id = %s" + self._sql_get_task_type = f"SELECT task_type FROM {self._table} WHERE task_id = %s" # noqa: S608 # nosec B608 — table identifier is constructor-validated + self._sql_get_state_error = f"SELECT state, error FROM {self._table} WHERE task_id = %s" # noqa: S608 # nosec B608 — table identifier is constructor-validated + self._sql_discard = ( # noqa: S608 + f"DELETE FROM {self._table} WHERE task_id = %s" # nosec B608 — table identifier is constructor-validated + f" RETURNING task_id, account_id, task_type, created_at, updated_at" ) - self._sql_discard = f"DELETE FROM {self._table} WHERE task_id = %s" # noqa: S608 self._sql_ddl = ( # noqa: S608 f"CREATE TABLE IF NOT EXISTS {self._table} (" f' task_id TEXT COLLATE "C" NOT NULL PRIMARY KEY,' @@ -331,6 +335,16 @@ async def issue( now, ), ) + self._notify_lifecycle_observers( + "submitted", + { + "task_id": task_id, + "account_id": account_id, + "task_type": task_type, + "created_at": now, + "updated_at": now, + }, + ) return task_id async def resolve_webhook_signing_scope( @@ -390,21 +404,29 @@ async def update_progress( straggler progress event. """ async with self._pool.connection() as conn: - await conn.execute( + cur = await conn.execute( self._sql_update_progress, - (json.dumps(progress), time.time(), task_id), + (task_id, json.dumps(progress), time.time()), ) - # Zero rows updated means unknown task_id or terminal state — silent - # no-op per Protocol contract. The InMemoryTaskRegistry logs a - # WARNING on terminal-state drops; we omit the extra SELECT needed - # to distinguish the two cases since the dispatch wrapper swallows - # the result either way. + row = await cur.fetchone() + if row is not None and row[0] == "submitted": + self._notify_lifecycle_observers("working", self._lifecycle_record(row[1:])) - async def complete( + @staticmethod + def _lifecycle_record(row: tuple[Any, ...]) -> dict[str, Any]: + return dict(zip(("task_id", "account_id", "task_type", "created_at", "updated_at"), row)) + + async def complete(self, task_id: str, result: dict[str, Any]) -> None: + """Complete once; notify metrics only after the transaction commits.""" + row = await self._complete(task_id, result) + if row is not None: + self._notify_lifecycle_observers("completed", self._lifecycle_record(row[:3] + row[5:])) + + async def _complete( self, task_id: str, result: dict[str, Any], - ) -> None: + ) -> tuple[Any, ...] | None: """Mark the task ``completed`` with ``result`` as the terminal artifact. Idempotent on repeated calls with an equal ``result``; raises @@ -438,7 +460,7 @@ async def complete( status="completed", payload=safe_result, ) - return # updated successfully + return row # connection and transaction exit before the caller notifies # Zero rows in RETURNING — task is unknown or already terminal. cur2 = await conn.execute(self._sql_get_state_result, (task_id,)) @@ -449,15 +471,21 @@ async def complete( safe_result = strip_credentials_from_wire_result(task_type, result) if state == "completed": if existing_result == safe_result: - return # idempotent + return None # idempotent raise ValueError(f"Task {task_id!r} already completed with a different result") raise ValueError(f"Task {task_id!r} already in terminal state {state!r}") - async def fail( + async def fail(self, task_id: str, error: dict[str, Any]) -> None: + """Fail once; notify metrics only after the transaction commits.""" + row = await self._fail(task_id, error) + if row is not None: + self._notify_lifecycle_observers("failed", self._lifecycle_record(row[:3] + row[5:])) + + async def _fail( self, task_id: str, error: dict[str, Any], - ) -> None: + ) -> tuple[Any, ...] | None: """Mark the task ``failed`` with ``error`` as the terminal payload. Idempotent on repeated calls with an equal ``error``; raises @@ -477,7 +505,7 @@ async def fail( status="failed", payload=error, ) - return # updated successfully + return row # notify only after connection/transaction exit # Zero rows in RETURNING — task is unknown or already terminal. cur2 = await conn.execute(self._sql_get_state_error, (task_id,)) @@ -487,7 +515,7 @@ async def fail( state, existing_error = row if state == "failed": if existing_error == error: - return # idempotent + return None # idempotent raise ValueError(f"Task {task_id!r} already failed with a different error") raise ValueError(f"Task {task_id!r} already in terminal state {state!r}") @@ -533,7 +561,7 @@ async def _enqueue_terminal_if_registered( payload: dict[str, Any], ) -> None: """Enqueue on ``conn`` so task state and webhook commit atomically.""" - task_id, account_id, task_type, encrypted_registration, registration_nonce = row + task_id, account_id, task_type, encrypted_registration, registration_nonce = row[:5] if encrypted_registration is None: return if self.task_webhook_outbox is None: @@ -576,7 +604,12 @@ async def discard(self, task_id: str) -> None: matching the :class:`~adcp.decisioning.InMemoryTaskRegistry` contract. """ async with self._pool.connection() as conn: - await conn.execute(self._sql_discard, (task_id,)) + cur = await conn.execute(self._sql_discard, (task_id,)) + row = await cur.fetchone() + if row is not None: + event_record = self._lifecycle_record(row) + event_record["updated_at"] = time.time() + self._notify_lifecycle_observers("discarded", event_record) #: Backwards-compat alias. Renamed to :class:`PgTaskRegistry` in diff --git a/src/adcp/decisioning/task_registry.py b/src/adcp/decisioning/task_registry.py index c4479b259..6d819beeb 100644 --- a/src/adcp/decisioning/task_registry.py +++ b/src/adcp/decisioning/task_registry.py @@ -37,6 +37,7 @@ import asyncio import logging +import threading import time import uuid from collections.abc import Awaitable, Callable @@ -51,6 +52,71 @@ TaskState = Literal["submitted", "working", "completed", "failed"] +TaskTransition = Literal["submitted", "working", "completed", "failed", "discarded"] + + +class TaskLifecycleObserver(Protocol): + """Synchronous best-effort metrics hook; no durable delivery guarantee.""" + + def __call__( + self, + event: TaskTransition, + *, + task_id: str, + account_id: str, + task_type: str, + created_at: float, + updated_at: float, + ) -> None: ... + + +class _TaskLifecycleObservers: + """Thread-safe observer registration, shared by SDK registry implementations.""" + + def _init_lifecycle_observers(self) -> None: + self._lifecycle_observers: list[TaskLifecycleObserver] = [] + self._lifecycle_observers_lock = threading.Lock() + + def add_lifecycle_observer(self, observer: TaskLifecycleObserver) -> None: + """Register a synchronous hook fired after commit/outside registry locks. + + Re-registering the same callback is a no-op. Notifications use a snapshot; + registration changes during delivery apply to the next notification. + Observer exceptions are logged and swallowed. Hooks are best-effort metrics, + not a durable event stream, and concurrent transitions may notify out of order. + """ + with self._lifecycle_observers_lock: + if observer not in self._lifecycle_observers: + self._lifecycle_observers.append(observer) + + def remove_lifecycle_observer(self, observer: TaskLifecycleObserver) -> bool: + """Remove a hook, returning whether it was registered.""" + with self._lifecycle_observers_lock: + try: + self._lifecycle_observers.remove(observer) + except ValueError: + return False + return True + + def _notify_lifecycle_observers(self, event: TaskTransition, record: dict[str, Any]) -> None: + with self._lifecycle_observers_lock: + observers = tuple(self._lifecycle_observers) + for observer in observers: + try: + observer( + event, + task_id=record["task_id"], + account_id=record["account_id"], + task_type=record["task_type"], + created_at=record["created_at"], + updated_at=record["updated_at"], + ) + except Exception: + logger.exception( + "Task lifecycle observer failed for %s task %s", event, record["task_id"] + ) + + @dataclass(frozen=True) class TaskWebhookAuthentication: """Buyer-selected authentication for a durable task webhook. @@ -388,7 +454,7 @@ async def discard(self, task_id: str) -> None: # --------------------------------------------------------------------------- -class InMemoryTaskRegistry: +class InMemoryTaskRegistry(_TaskLifecycleObservers): """Process-local task registry — v6.0 reference implementation. Storage is a plain ``dict[str, TaskRecord]`` guarded by an @@ -411,6 +477,7 @@ class InMemoryTaskRegistry: is_durable: ClassVar[bool] = False def __init__(self) -> None: + self._init_lifecycle_observers() self._records: dict[str, TaskRecord] = {} self._lock = asyncio.Lock() @@ -454,6 +521,8 @@ async def issue( task_type=task_type, request_context=(dict(request_context) if request_context is not None else None), ) + event_record = self._records[task_id].to_dict() + self._notify_lifecycle_observers("submitted", event_record) return task_id async def update_progress( @@ -483,10 +552,14 @@ async def update_progress( record.state, ) return + transitioned = record.state == "submitted" record.progress = dict(progress) if record.state == "submitted": record.state = "working" record.updated_at = time.time() + event_record = record.to_dict() + if transitioned: + self._notify_lifecycle_observers("working", event_record) async def complete( self, @@ -515,9 +588,13 @@ async def complete( if record.result == stripped: return # idempotent raise ValueError(f"Task {task_id!r} already completed with a different result") + if record.state == "failed": + raise ValueError(f"Task {task_id!r} already in terminal state {record.state!r}") record.state = "completed" record.result = dict(stripped) if isinstance(stripped, dict) else stripped record.updated_at = time.time() + event_record = record.to_dict() + self._notify_lifecycle_observers("completed", event_record) async def fail( self, @@ -532,9 +609,13 @@ async def fail( if record.error == error: return # idempotent raise ValueError(f"Task {task_id!r} already failed with a different error") + if record.state == "completed": + raise ValueError(f"Task {task_id!r} already in terminal state {record.state!r}") record.state = "failed" record.error = dict(error) record.updated_at = time.time() + event_record = record.to_dict() + self._notify_lifecycle_observers("failed", event_record) async def get( self, @@ -562,7 +643,12 @@ async def discard(self, task_id: str) -> None: # Idempotent: pop with default. The Protocol contract # tolerates discarding an unknown id (no raise) so the # WorkflowHandoff projection's rollback can be unconditional. - self._records.pop(task_id, None) + record = self._records.pop(task_id, None) + if record is not None: + event_record = record.to_dict() + event_record["updated_at"] = time.time() + if record is not None: + self._notify_lifecycle_observers("discarded", event_record) # --------------------------------------------------------------------------- diff --git a/tests/test_decisioning_task_lifecycle.py b/tests/test_decisioning_task_lifecycle.py new file mode 100644 index 000000000..aa1b5190d --- /dev/null +++ b/tests/test_decisioning_task_lifecycle.py @@ -0,0 +1,166 @@ +"""Best-effort metrics hooks count committed transitions, including races.""" + +from __future__ import annotations + +import asyncio +import os +import secrets +from collections.abc import AsyncIterator +from typing import Any + +import pytest + +from adcp.decisioning import InMemoryTaskRegistry, TaskTransition + + +@pytest.fixture(params=["memory", "postgres", "postgres-autocommit"]) +async def registry(request: pytest.FixtureRequest) -> AsyncIterator[Any]: + if request.param == "memory": + yield InMemoryTaskRegistry() + return + from psycopg_pool import AsyncConnectionPool + + from adcp.decisioning.pg import PgTaskRegistry + + url = os.environ.get("ADCP_PG_TEST_URL") + if not url: + pytest.skip("ADCP_PG_TEST_URL required") + table = f"test_lifecycle_{secrets.token_hex(6)}" + async with AsyncConnectionPool( + url, + open=False, + min_size=1, + max_size=4, + kwargs={"autocommit": request.param == "postgres-autocommit"}, + ) as pool: + reg = PgTaskRegistry(pool=pool, _table=table) + await reg.create_schema() + yield reg + async with pool.connection() as conn: + await conn.execute(f"DROP TABLE {table}") + + +async def test_observe_committed_transitions_and_suppress_noops(registry: Any) -> None: + events: list[tuple[str, dict[str, Any]]] = [] + + def observe(event: TaskTransition, **metadata: Any) -> None: + # A synchronous callback can inspect committed state on another connection. + # It cannot observe an uncommitted row or run under the memory registry lock. + if isinstance(registry, InMemoryTaskRegistry): + assert not registry._lock.locked() + record = registry._records.get(metadata["task_id"]) + assert record is None if event == "discarded" else record.state == event + else: + import psycopg + + with psycopg.connect(os.environ["ADCP_PG_TEST_URL"]) as conn: + row = conn.execute( + f"SELECT state FROM {registry._table} WHERE task_id = %s", + (metadata["task_id"],), + ).fetchone() + assert row is None if event == "discarded" else row[0] == event + events.append((event, metadata)) + + registry.add_lifecycle_observer(observe) + registry.add_lifecycle_observer(observe) + task = await registry.issue(account_id="acct_1", task_type="get_products") + await asyncio.gather(*(registry.update_progress(task, {"percentage": n}) for n in range(10))) + await asyncio.gather(*(registry.complete(task, {"products": []}) for _ in range(10))) + await registry.update_progress(task, {"percentage": 100}) + await registry.update_progress("absent", {}) + await registry.discard("absent") + await registry.discard(task) + await registry.discard(task) + assert [event for event, _ in events] == ["submitted", "working", "completed", "discarded"] + for _, metadata in events: + assert metadata["task_id"] == task + assert metadata["account_id"] == "acct_1" + assert metadata["task_type"] == "get_products" + assert metadata["created_at"] <= metadata["updated_at"] + assert registry.remove_lifecycle_observer(observe) is True + assert registry.remove_lifecycle_observer(observe) is False + await registry.issue(account_id="acct_1", task_type="get_products") + assert len(events) == 4 + + +async def test_racing_terminal_transitions_emit_once(registry: Any) -> None: + events = [] + + def observe(event: TaskTransition, **metadata: Any) -> None: + events.append(event) + + registry.add_lifecycle_observer(observe) + task = await registry.issue(account_id="acct", task_type="get_products") + results = await asyncio.gather( + *(registry.complete(task, {"products": []}) for _ in range(10)), + *(registry.fail(task, {"code": "INTERNAL_ERROR", "message": "Failed"}) for _ in range(10)), + return_exceptions=True, + ) + assert events[0] == "submitted" + assert len(events) == 2 and events[1] in {"completed", "failed"} + assert any(isinstance(result, ValueError) for result in results) + record = await registry.get(task) + assert record["state"] == events[1] + if record["state"] == "completed": + with pytest.raises(ValueError): + await registry.fail(task, {"code": "CONFLICT", "message": "Conflict"}) + else: + with pytest.raises(ValueError): + await registry.complete(task, {"products": []}) + assert len(events) == 2 + + +async def test_observer_failure_and_snapshot_removal( + registry: Any, caplog: pytest.LogCaptureFixture +) -> None: + events = [] + + def broken(event: TaskTransition, **metadata: Any) -> None: + raise RuntimeError("Metrics unavailable") + + def removed(event: TaskTransition, **metadata: Any) -> None: + events.append(event) + + def remover(event: TaskTransition, **metadata: Any) -> None: + registry.remove_lifecycle_observer(removed) + + registry.add_lifecycle_observer(broken) + registry.add_lifecycle_observer(remover) + registry.add_lifecycle_observer(removed) + task = await registry.issue(account_id="acct", task_type="get_products") + await registry.fail(task, {"code": "INTERNAL_ERROR", "message": "Failed"}) + await registry.fail(task, {"code": "INTERNAL_ERROR", "message": "Failed"}) + assert events == ["submitted"] + assert "Task lifecycle observer failed" in caplog.text + assert (await registry.get(task))["state"] == "failed" + + +async def test_failed_database_transaction_never_notifies(registry: Any) -> None: + if isinstance(registry, InMemoryTaskRegistry): + with pytest.raises(ValueError): + await registry.complete("missing", {}) + return + import psycopg + + events = [] + registry.add_lifecycle_observer(lambda event, **metadata: events.append(event)) + task = await registry.issue(account_id="acct", task_type="get_products") + function = f"{registry._table}_reject" + async with registry._pool.connection() as conn: + await conn.execute( + f"CREATE FUNCTION {function}() RETURNS trigger LANGUAGE plpgsql AS $$ " + "BEGIN RAISE EXCEPTION 'fail commit'; END $$" + ) + await conn.execute( + f"CREATE CONSTRAINT TRIGGER reject_transition AFTER UPDATE ON {registry._table}" + f" DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION {function}()" + ) + try: + with pytest.raises(psycopg.errors.RaiseException): + await registry.complete(task, {"products": []}) + assert events == ["submitted"] + assert (await registry.get(task))["state"] == "submitted" + finally: + async with registry._pool.connection() as conn: + await conn.execute(f"DROP TRIGGER reject_transition ON {registry._table}") + await conn.execute(f"DROP FUNCTION {function}()") diff --git a/tests/type_checks/task_lifecycle_observers.py b/tests/type_checks/task_lifecycle_observers.py new file mode 100644 index 000000000..f628d38c9 --- /dev/null +++ b/tests/type_checks/task_lifecycle_observers.py @@ -0,0 +1,39 @@ +"""Public metrics hook works without changing the minimal registry protocol.""" + +from typing import Any + +from adcp.decisioning import ( + InMemoryTaskRegistry, + TaskLifecycleObserver, + TaskRegistry, + TaskTransition, +) +from adcp.decisioning.pg import PgTaskRegistry + + +class Metrics: + def __call__( + self, + event: TaskTransition, + *, + task_id: str, + account_id: str, + task_type: str, + created_at: float, + updated_at: float, + ) -> None: + print(event, task_id, account_id, task_type, updated_at - created_at) + + +observer: TaskLifecycleObserver = Metrics() + + +async def instrument(memory: InMemoryTaskRegistry, postgres: PgTaskRegistry) -> None: + for registry in (memory, postgres): + registry.add_lifecycle_observer(observer) + minimal: TaskRegistry = registry + task = await minimal.issue(account_id="account", task_type="get_products") + result: dict[str, Any] = {"products": []} + await minimal.complete(task, result) + removed: bool = registry.remove_lifecycle_observer(observer) + assert removed From d5319d7661a4d1044eda63bd3cf74c2176d9722f Mon Sep 17 00:00:00 2001 From: Brian O'Kelley Date: Fri, 2 Oct 2026 18:09:17 +0000 Subject: [PATCH 3/9] fix(decisioning): preserve optional gates and canonical task failures --- docs/task-registry.md | 4 +++- src/adcp/decisioning/handler.py | 8 ++++---- src/adcp/decisioning/task_queries.py | 6 +++++- tests/test_decisioning_advertised_per_specialism.py | 3 +++ tests/test_decisioning_task_polling.py | 5 ++++- 5 files changed, 19 insertions(+), 7 deletions(-) diff --git a/docs/task-registry.md b/docs/task-registry.md index e19e0c47b..3d3161ad8 100644 --- a/docs/task-registry.md +++ b/docs/task-registry.md @@ -11,7 +11,9 @@ installed by this change. Polls use `status`, ISO UTC timestamps and protocol metadata. Progress, failure summary and the original task's context are preserved. Set `include_result=true` to retrieve a terminal artifact. A failed task exposes the stored fatal error -both in its convenience `error` summary and in `result.adcp_error`. +in its convenience `error` summary, `result.adcp_error`, and the canonical +`result.errors` array. Successful results retain the adopter's original wire +response; the background callable must return a valid response for its task type. `list_tasks` is advertised only for the separate optional `ListableTaskRegistry` protocol. Existing third-party `TaskRegistry` implementations need no changes. diff --git a/src/adcp/decisioning/handler.py b/src/adcp/decisioning/handler.py index b70361d7b..cbbfb2715 100644 --- a/src/adcp/decisioning/handler.py +++ b/src/adcp/decisioning/handler.py @@ -1335,6 +1335,10 @@ def advertised_tools_for_instance(self) -> frozenset[str]: getattr(self._platform, tool_name, None) ): serving.add(tool_name) + if not serving: + # Preserve the existing class-universe fallback for novel specialisms. + serving.update(self.advertised_tools) + serving.discard("list_tasks") # Drop sync_accounts / list_accounts when the platform's # AccountStore doesn't expose the corresponding optional # Protocol method. ``sales-*`` claims union both tools in by @@ -1367,10 +1371,6 @@ def advertised_tools_for_instance(self) -> frozenset[str]: reporting = getattr(self._platform, "_reliable_reporting_service", None) if reporting is not None: serving.update(reporting.reporting_tools) - if not serving: - # Preserve the existing class-universe fallback for novel specialisms. - serving.update(self.advertised_tools) - serving.discard("list_tasks") serving.add("get_task_status") if isinstance(self._registry, ListableTaskRegistry): serving.add("list_tasks") diff --git a/src/adcp/decisioning/task_queries.py b/src/adcp/decisioning/task_queries.py index 438fa85d6..eb2c65fb7 100644 --- a/src/adcp/decisioning/task_queries.py +++ b/src/adcp/decisioning/task_queries.py @@ -63,7 +63,11 @@ def project_task(record: dict[str, Any], *, include_result: bool = False) -> dic if record.get("result") is not None: response["result"] = record["result"] elif record["state"] == "failed" and record.get("error") is not None: - response["result"] = {"status": "failed", "adcp_error": record["error"]} + response["result"] = { + "status": "failed", + "adcp_error": record["error"], + "errors": [record["error"]], + } return response diff --git a/tests/test_decisioning_advertised_per_specialism.py b/tests/test_decisioning_advertised_per_specialism.py index de76ba04b..6b0da7241 100644 --- a/tests/test_decisioning_advertised_per_specialism.py +++ b/tests/test_decisioning_advertised_per_specialism.py @@ -330,6 +330,9 @@ def get_products(self, req, ctx): # protocol families pre-filter — this is the documented # forward-compat fallback. assert "get_products" in tools + # The class-universe fallback must still pass through optional store gates. + assert "sync_accounts" not in tools + assert "list_accounts" not in tools def test_advertise_all_bypasses_per_specialism_filter(executor) -> None: diff --git a/tests/test_decisioning_task_polling.py b/tests/test_decisioning_task_polling.py index 5a508ddc6..2c7de51f4 100644 --- a/tests/test_decisioning_task_polling.py +++ b/tests/test_decisioning_task_polling.py @@ -52,7 +52,7 @@ async def work(self, ctx: Any) -> dict[str, Any]: await self.finish.wait() if self.failed: raise AdcpError("INTERNAL_ERROR", message="Background failed", recovery="terminal") - return {"products": []} + return {"products": [], "cache_scope": "account", "status": "completed"} def create_media_buy(self, req: Any, ctx: Any) -> Any: raise NotImplementedError @@ -209,6 +209,9 @@ async def rpc(method: str, params: dict[str, Any]) -> dict[str, Any]: validator = get_validator("get_task_status", "sync") assert validator is not None validator.validate(terminal) + result_validator = get_validator("get_products", "sync") + assert result_validator is not None + result_validator.validate(terminal["result"]) listed = await rpc( "tools/call", {"name": "list_tasks", "arguments": {"filters": {"task_ids": [task_id]}}} ) From b5c7d42eacecc20dc76368bc37b46b29e5945ae5 Mon Sep 17 00:00:00 2001 From: Brian O'Kelley Date: Fri, 2 Oct 2026 18:11:00 +0000 Subject: [PATCH 4/9] feat(pg): lazily resolve decisioning stores on the serving loop --- docs/lazy-pg-stores.md | 71 ++++ src/adcp/decisioning/pg/__init__.py | 14 + src/adcp/decisioning/pg/lazy.py | 470 +++++++++++++++++++++++ src/adcp/decisioning/webhook_emit.py | 5 + tests/test_decisioning_lazy_pg_stores.py | 449 ++++++++++++++++++++++ tests/test_decisioning_task_polling.py | 9 +- tests/type_checks/lazy_pg_stores.py | 66 ++++ 7 files changed, 1082 insertions(+), 2 deletions(-) create mode 100644 docs/lazy-pg-stores.md create mode 100644 src/adcp/decisioning/pg/lazy.py create mode 100644 tests/test_decisioning_lazy_pg_stores.py create mode 100644 tests/type_checks/lazy_pg_stores.py diff --git a/docs/lazy-pg-stores.md b/docs/lazy-pg-stores.md new file mode 100644 index 000000000..f464e6335 --- /dev/null +++ b/docs/lazy-pg-stores.md @@ -0,0 +1,71 @@ +# Lazy PostgreSQL decisioning stores + +`adcp.decisioning.pg` exports `LazyProposalStore`, `LazyTaskRegistry`, and +`LazyTaskWebhookOutbox`. Each accepts a synchronous or asynchronous zero-argument +factory returning its original concrete PostgreSQL store. Task factories may +also return an original `(PgTaskRegistry, PgTaskWebhookOutbox)` pair. + +Factories are called on the serving event loop, on the first async operation or +explicit `await wrapper.resolve()`. Concurrent calls share one successful +resolution. Failed or canceled initialization is not cached, so later calls +retry. Factories own pool opening and cleanup on failed initialization; wrappers +never open, close, replace or reopen their caller-owned pools. Schema setup stays +explicit: do it in the factory or call the wrapper's `create_schema()`. + +```python +from psycopg_pool import AsyncConnectionPool +from adcp.decisioning.pg import ( + LazyTaskRegistry, LazyTaskWebhookOutbox, PgTaskRegistry, PgTaskWebhookOutbox, +) + +pool = AsyncConnectionPool("postgresql://localhost/seller", open=False) + +async def task_stores() -> tuple[PgTaskRegistry, PgTaskWebhookOutbox]: + await pool.open() + outbox = PgTaskWebhookOutbox( + pool=pool, sender=signed_sender, encryption_key=encryption_key, + delivery_retry_horizon_seconds=86400, + ) + registry = PgTaskRegistry(pool=pool, task_webhook_outbox=outbox) + await registry.create_schema() + await outbox.create_schema() + return registry, outbox + +registry = LazyTaskRegistry(task_stores) +worker_outbox = LazyTaskWebhookOutbox.from_registry(registry) +# await worker_outbox.run_worker() and registry.issue(...) share the same pair. +# Cancel/join workers before closing pool during application shutdown. +``` + +Construct the original registry and outbox together, passing the same pool to +both. Their existing constructor checks (including signing-scope wiring) remain +authoritative, and a returned pair is checked again for identity. The worker +facade from `from_registry()` resolves the same original outbox, rather than +running an independent pool factory. Do not pass an unresolved worker facade to +the concrete `PgTaskRegistry` constructor. + +The wrappers declare the concrete stores' durability before resolution and +reject a factory returning a lossy store. `LazyTaskRegistry` retains the concrete +SDK registry's optional listing protocol; no arbitrary delegate attributes or +optional APIs are forwarded. Its additive observer methods can be registered or +removed before and after resolution. They receive the first submitted event and +retain the concrete registry's commit/no-op/failure semantics. + +`resolved` inspects the cached store without running the factory. A task registry +exposes its original `task_webhook_outbox` and atomic-outbox marker after +resolution. Servers advertising SDK task-webhook signing must await resolution +before synchronous server construction: the boot validator needs the actual +sender, retry horizon and shared-pool proof. Resolved exact SDK wrappers are +accepted by that proof; unaudited subclasses remain rejected. Polling-only +servers can resolve on first use. Outbox synchronous registration/crypto helpers +and retry-horizon inspection also require explicit resolution; they raise a +clear error when called too early. + +After success, a wrapper is bound to the event loop that resolved it. A different +loop raises an error, even if the previous loop has closed; create a fresh wrapper +and pool for the new loop. A failed initialization with no cached store may retry +on a new loop after the old loop closes. Closing the borrowed pool causes normal +pool errors on later operations, rather than rerunning the factory. + +No database migration or change to the existing `LazyBackend` is required. The +new wrappers reuse the concrete stores' current SQL and schema assets. diff --git a/src/adcp/decisioning/pg/__init__.py b/src/adcp/decisioning/pg/__init__.py index 14c586fde..cf4c7f254 100644 --- a/src/adcp/decisioning/pg/__init__.py +++ b/src/adcp/decisioning/pg/__init__.py @@ -35,6 +35,14 @@ PG_AVAILABLE, PgBuyerAgentRegistry, ) +from adcp.decisioning.pg.lazy import ( + LazyProposalStore, + LazyProposalStoreFactory, + LazyTaskRegistry, + LazyTaskRegistryFactory, + LazyTaskWebhookOutbox, + LazyTaskWebhookOutboxFactory, +) from adcp.decisioning.pg.proposal_store import PgProposalStore from adcp.decisioning.pg.task_registry import ( PgTaskRegistry, @@ -46,6 +54,12 @@ __all__ = [ "DEFAULT_TABLE_NAME", "PG_AVAILABLE", + "LazyProposalStore", + "LazyProposalStoreFactory", + "LazyTaskRegistry", + "LazyTaskRegistryFactory", + "LazyTaskWebhookOutbox", + "LazyTaskWebhookOutboxFactory", "PgBuyerAgentRegistry", "PgProposalStore", "PgTaskRegistry", diff --git a/src/adcp/decisioning/pg/lazy.py b/src/adcp/decisioning/pg/lazy.py new file mode 100644 index 000000000..542a3b1c1 --- /dev/null +++ b/src/adcp/decisioning/pg/lazy.py @@ -0,0 +1,470 @@ +"""Resolve-once wrappers for caller-owned PostgreSQL decisioning stores. + +Factories open infrastructure on the serving loop, then construct the original +concrete stores so constructor validation remains authoritative. A successful +resolution is cached; failed initialization can be retried. These wrappers never +open, close, replace or reopen pools themselves. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, Mapping +from datetime import datetime +from typing import Any, ClassVar, Generic, TypeAlias, TypeVar + +from adcp.decisioning.context import RequestContext +from adcp.decisioning.pg.proposal_store import DEFAULT_TABLE_NAME, PgProposalStore, _migration_sql +from adcp.decisioning.pg.task_registry import PgTaskRegistry +from adcp.decisioning.pg.task_webhook_outbox import PgTaskWebhookOutbox +from adcp.decisioning.proposal_store import ProposalRecord +from adcp.decisioning.recipe import Recipe +from adcp.decisioning.task_registry import ( + TaskTransition, + TaskWebhookAuthentication, + _TaskLifecycleObservers, +) + +_Store = TypeVar("_Store") +LazyProposalStoreFactory: TypeAlias = Callable[[], PgProposalStore | Awaitable[PgProposalStore]] +TaskRegistryStores: TypeAlias = PgTaskRegistry | tuple[PgTaskRegistry, PgTaskWebhookOutbox] +LazyTaskRegistryFactory: TypeAlias = Callable[ + [], TaskRegistryStores | Awaitable[TaskRegistryStores] +] +LazyTaskWebhookOutboxFactory: TypeAlias = Callable[ + [], PgTaskWebhookOutbox | Awaitable[PgTaskWebhookOutbox] +] + + +class _LazyStore(Generic[_Store]): + def __init__( + self, factory: Callable[[], _Store | Awaitable[_Store]], store_type: type[_Store] + ) -> None: + self._factory = factory + self._store_type = store_type + self._store: _Store | None = None + self._lock = asyncio.Lock() + self._loop: asyncio.AbstractEventLoop | None = None + + @property + def resolved(self) -> _Store | None: + """The cached concrete store, or None; inspection never invokes the factory.""" + return self._store + + def _bind_loop(self) -> None: + loop = asyncio.get_running_loop() + if self._loop is None: + self._loop = loop + elif self._loop is not loop: + # No successful resource was cached. A failed attempt on a now-dead + # loop may retry, with a fresh lock. The factory owns failed cleanup. + if self._store is None and self._loop.is_closed() and not self._lock.locked(): + self._loop = loop + self._lock = asyncio.Lock() + else: + raise RuntimeError( + "Lazy PostgreSQL stores cannot be reused on another event loop; " + "construct a new wrapper and pool for that loop" + ) + + async def resolve(self) -> _Store: + """Resolve once on the serving loop. Concurrent callers share success. + + Factory failures/cancellation are not cached. After successful resolution, + reusing this wrapper on another loop raises, even if the old loop has closed. + """ + self._bind_loop() + cached = self._store + if cached is not None: + return cached + async with self._lock: + cached = self._store + if cached is not None: + return cached + result = self._factory() + store = await result if isinstance(result, Awaitable) else result + if not isinstance(store, self._store_type): + raise TypeError(f"Factory must return {self._store_type.__name__}") + self._validate(store) + self._store = store + return store + + def _validate(self, store: _Store) -> None: + pass + + def _require_resolved(self) -> _Store: + store = self._store + if store is None: + raise RuntimeError("Await resolve() before using synchronous store helpers") + # Pure configuration/crypto helpers may run without a loop. When called + # on an active loop, reject accidental reuse just like async operations. + try: + asyncio.get_running_loop() + except RuntimeError: + return store + self._bind_loop() + return store + + +class LazyProposalStore(_LazyStore[PgProposalStore]): + """Deferred PgProposalStore; durable before and after resolution.""" + + is_durable: ClassVar[bool] = True + + def __init__(self, factory: LazyProposalStoreFactory) -> None: + super().__init__(factory, PgProposalStore) + + def _validate(self, store: PgProposalStore) -> None: + if store.is_durable is not True: + raise ValueError("LazyProposalStore requires a durable PostgreSQL store") + + @classmethod + def migration_sql(cls, table_name: str = DEFAULT_TABLE_NAME) -> dict[str, str]: + """Return the concrete store's migration SQL without opening infrastructure.""" + return _migration_sql(table_name) + + async def create_schema(self) -> None: + await (await self.resolve()).create_schema() + + async def put_draft( + self, + *, + proposal_id: str, + account_id: str, + recipes: Mapping[str, Recipe], + proposal_payload: Mapping[str, Any], + ) -> None: + await (await self.resolve()).put_draft( + proposal_id=proposal_id, + account_id=account_id, + recipes=recipes, + proposal_payload=proposal_payload, + ) + + async def get(self, proposal_id: str, *, expected_account_id: str) -> ProposalRecord | None: + return await (await self.resolve()).get( + proposal_id, expected_account_id=expected_account_id + ) + + async def commit( + self, + proposal_id: str, + *, + expires_at: datetime, + proposal_payload: Mapping[str, Any], + expected_account_id: str, + ) -> None: + await (await self.resolve()).commit( + proposal_id, + expires_at=expires_at, + proposal_payload=proposal_payload, + expected_account_id=expected_account_id, + ) + + async def try_reserve_consumption( + self, proposal_id: str, *, expected_account_id: str + ) -> ProposalRecord: + return await (await self.resolve()).try_reserve_consumption( + proposal_id, expected_account_id=expected_account_id + ) + + async def finalize_consumption( + self, proposal_id: str, *, media_buy_id: str, expected_account_id: str + ) -> None: + await (await self.resolve()).finalize_consumption( + proposal_id, media_buy_id=media_buy_id, expected_account_id=expected_account_id + ) + + async def release_consumption(self, proposal_id: str, *, expected_account_id: str) -> None: + await (await self.resolve()).release_consumption( + proposal_id, expected_account_id=expected_account_id + ) + + async def mark_consumed( + self, proposal_id: str, *, media_buy_id: str, expected_account_id: str + ) -> None: + await (await self.resolve()).mark_consumed( + proposal_id, media_buy_id=media_buy_id, expected_account_id=expected_account_id + ) + + async def discard(self, proposal_id: str, *, expected_account_id: str) -> None: + await (await self.resolve()).discard(proposal_id, expected_account_id=expected_account_id) + + async def get_by_media_buy_id( + self, media_buy_id: str, *, expected_account_id: str + ) -> ProposalRecord | None: + return await (await self.resolve()).get_by_media_buy_id( + media_buy_id, expected_account_id=expected_account_id + ) + + +class LazyTaskRegistry(_LazyStore[PgTaskRegistry], _TaskLifecycleObservers): + """Deferred PgTaskRegistry, including its listing and metrics contracts. + + A factory can return a registry, or an original (registry, outbox) pair. Build + them together using one caller-owned pool. Pair and signing-scope validation + are retained. Register observers before first use; they receive the first + submitted transition too. No arbitrary delegate attributes are forwarded. + + For a server advertising SDK task-webhook signing, await resolve() before + constructing the server: its synchronous boot validator needs the actual + outbox/sender/horizon configuration. Polling-only servers can resolve on the + first task operation. + """ + + is_durable: ClassVar[bool] = True + + def __init__(self, factory: LazyTaskRegistryFactory) -> None: + self._init_lifecycle_observers() + + async def resolve_registry() -> PgTaskRegistry: + result = factory() + stores = await result if isinstance(result, Awaitable) else result + if isinstance(stores, tuple): + registry, outbox = stores + if not isinstance(registry, PgTaskRegistry) or not isinstance( + outbox, PgTaskWebhookOutbox + ): + raise TypeError("Task factory must return a concrete registry/outbox pair") + if registry.task_webhook_outbox is not outbox or registry._pool is not outbox._pool: + raise ValueError( + "Task registry/outbox pair must share one pool and registration" + ) + return registry + return stores + + super().__init__(resolve_registry, PgTaskRegistry) + + def _validate(self, store: PgTaskRegistry) -> None: + if store.is_durable is not True: + raise ValueError("LazyTaskRegistry requires a durable PostgreSQL registry") + if not callable(store.list): + raise TypeError("LazyTaskRegistry requires PostgreSQL listing support") + outbox = store.task_webhook_outbox + if outbox is not None and outbox._pool is not store._pool: + raise ValueError("Task registry and outbox must share one pool") + store.add_lifecycle_observer(self._relay_transition) + + def _relay_transition( + self, + event: TaskTransition, + *, + task_id: str, + account_id: str, + task_type: str, + created_at: float, + updated_at: float, + ) -> None: + self._notify_lifecycle_observers( + event, + { + "task_id": task_id, + "account_id": account_id, + "task_type": task_type, + "created_at": created_at, + "updated_at": updated_at, + }, + ) + + @property + def task_webhook_outbox(self) -> PgTaskWebhookOutbox | None: + """Original coupled outbox, available after resolution; no eager opening.""" + store = self.resolved + return store.task_webhook_outbox if store is not None else None + + @property + def atomic_task_webhook_outbox(self) -> bool: + store = self.resolved + return store.atomic_task_webhook_outbox if store is not None else False + + async def create_schema(self) -> None: + await (await self.resolve()).create_schema() + + async def issue( + self, + *, + account_id: str, + task_type: str, + request_context: dict[str, Any] | None = None, + webhook_url: str | None = None, + webhook_operation_id: str | None = None, + webhook_token: str | None = None, + webhook_authentication: TaskWebhookAuthentication | None = None, + webhook_signing_scope_id: str | None = None, + **_extra: Any, + ) -> str: + return await (await self.resolve()).issue( + account_id=account_id, + task_type=task_type, + request_context=request_context, + webhook_url=webhook_url, + webhook_operation_id=webhook_operation_id, + webhook_token=webhook_token, + webhook_authentication=webhook_authentication, + webhook_signing_scope_id=webhook_signing_scope_id, + **_extra, + ) + + async def resolve_webhook_signing_scope(self, context: RequestContext[Any]) -> str | None: + return await (await self.resolve()).resolve_webhook_signing_scope(context) + + async def update_progress(self, task_id: str, progress: dict[str, Any]) -> None: + await (await self.resolve()).update_progress(task_id, progress) + + async def complete(self, task_id: str, result: dict[str, Any]) -> None: + await (await self.resolve()).complete(task_id, result) + + async def fail(self, task_id: str, error: dict[str, Any]) -> None: + await (await self.resolve()).fail(task_id, error) + + async def get( + self, task_id: str, *, expected_account_id: str | None = None + ) -> dict[str, Any] | None: + return await (await self.resolve()).get(task_id, expected_account_id=expected_account_id) + + async def list( + self, + *, + account_id: str, + filters: dict[str, Any] | None = None, + sort: dict[str, Any] | None = None, + pagination: dict[str, Any] | None = None, + ) -> dict[str, Any]: + return await (await self.resolve()).list( + account_id=account_id, filters=filters, sort=sort, pagination=pagination + ) + + async def discard(self, task_id: str) -> None: + await (await self.resolve()).discard(task_id) + + +class LazyTaskWebhookOutbox(_LazyStore[PgTaskWebhookOutbox]): + """Deferred concrete outbox. Synchronous crypto helpers require resolve(). + + Use from_registry() to share a registry factory's original outbox. Pass the + original objects, constructed inside that factory, to PgTaskRegistry; do not + substitute this worker facade into the concrete constructor's pool checks. + """ + + delivery_state_is_durable: ClassVar[bool] = True + supports_atomic_task_outbox: ClassVar[bool] = True + + def __init__(self, factory: LazyTaskWebhookOutboxFactory) -> None: + super().__init__(factory, PgTaskWebhookOutbox) + + def _validate(self, store: PgTaskWebhookOutbox) -> None: + if ( + store.delivery_state_is_durable is not True + or store.supports_atomic_task_outbox is not True + ): + raise ValueError("LazyTaskWebhookOutbox requires a durable atomic PostgreSQL outbox") + + @classmethod + def from_registry(cls, registry: LazyTaskRegistry) -> LazyTaskWebhookOutbox: + """Create a worker facade resolving the same original registry/outbox pair.""" + + async def resolve_outbox() -> PgTaskWebhookOutbox: + store = await registry.resolve() + if store.task_webhook_outbox is None: + raise ValueError("The registry factory did not configure a task webhook outbox") + return store.task_webhook_outbox + + return cls(resolve_outbox) + + @property + def delivery_retry_horizon_seconds(self) -> int: + return self._require_resolved().delivery_retry_horizon_seconds + + @property + def legacy_hmac_fallback(self) -> bool: + return self._require_resolved().legacy_hmac_fallback + + async def create_schema(self) -> None: + await (await self.resolve()).create_schema() + + async def enqueue_terminal( + self, + conn: Any, + *, + task_id: str, + account_id: str, + task_type: str, + status: str, + result: dict[str, Any], + url: str, + operation_id: str, + token: str | None, + authentication: TaskWebhookAuthentication | None = None, + signing_scope_id: str | None = None, + ) -> int: + return await (await self.resolve()).enqueue_terminal( + conn, + task_id=task_id, + account_id=account_id, + task_type=task_type, + status=status, + result=result, + url=url, + operation_id=operation_id, + token=token, + authentication=authentication, + signing_scope_id=signing_scope_id, + ) + + def validate_registration( + self, url: str, authentication: TaskWebhookAuthentication | None = None + ) -> None: + (self._require_resolved()).validate_registration(url, authentication) + + def protect_registration( + self, + *, + account_id: str, + task_id: str, + task_type: str, + url: str, + operation_id: str, + token: str | None, + authentication: TaskWebhookAuthentication | None = None, + signing_scope_id: str | None = None, + ) -> tuple[bytes, bytes]: + return (self._require_resolved()).protect_registration( + account_id=account_id, + task_id=task_id, + task_type=task_type, + url=url, + operation_id=operation_id, + token=token, + authentication=authentication, + signing_scope_id=signing_scope_id, + ) + + def open_registration( + self, + *, + account_id: str, + task_id: str, + task_type: str, + encrypted_registration: bytes, + nonce: bytes, + ) -> tuple[str, str, str | None]: + return (self._require_resolved()).open_registration( + account_id=account_id, + task_id=task_id, + task_type=task_type, + encrypted_registration=encrypted_registration, + nonce=nonce, + ) + + async def run_worker( + self, *, poll_interval: float = 1.0, purge_interval: float = 300.0 + ) -> None: + await (await self.resolve()).run_worker( + poll_interval=poll_interval, purge_interval=purge_interval + ) + + async def process_one(self) -> bool: + return await (await self.resolve()).process_one() + + async def purge_expired(self) -> None: + await (await self.resolve()).purge_expired() diff --git a/src/adcp/decisioning/webhook_emit.py b/src/adcp/decisioning/webhook_emit.py index b17159ad5..660a1bad9 100644 --- a/src/adcp/decisioning/webhook_emit.py +++ b/src/adcp/decisioning/webhook_emit.py @@ -44,10 +44,15 @@ def _sdk_task_outbox_pair_ready(registry: Any, task_outbox: Any) -> bool: if registry is None or task_outbox is None: return False try: + from adcp.decisioning.pg.lazy import LazyTaskRegistry from adcp.decisioning.pg.task_registry import PgTaskRegistry from adcp.decisioning.pg.task_webhook_outbox import PgTaskWebhookOutbox except ImportError: return False + if type(registry) is LazyTaskRegistry: + registry = registry.resolved + if registry is None: + return False return ( type(registry) is PgTaskRegistry and type(task_outbox) is PgTaskWebhookOutbox diff --git a/tests/test_decisioning_lazy_pg_stores.py b/tests/test_decisioning_lazy_pg_stores.py new file mode 100644 index 000000000..654c3b868 --- /dev/null +++ b/tests/test_decisioning_lazy_pg_stores.py @@ -0,0 +1,449 @@ +"""Lazy PostgreSQL infrastructure resolves once without weakening store contracts.""" + +from __future__ import annotations + +import asyncio +import os +import secrets +from collections.abc import AsyncIterator +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from adcp.decisioning import ListableTaskRegistry, TaskRegistry +from adcp.decisioning.pg import ( + LazyProposalStore, + LazyTaskRegistry, + LazyTaskWebhookOutbox, + PgProposalStore, + PgTaskRegistry, + PgTaskWebhookOutbox, +) +from adcp.decisioning.webhook_emit import _sdk_task_outbox_pair_ready +from adcp.webhook_sender import PreparedWebhook, WebhookDeliveryResult + + +def _sender() -> MagicMock: + import json + + sender = MagicMock() + sender._owns_client = True + sender._allow_private_destinations = False + sender._timeout = 10.0 + sender.signs_with_rfc9421 = True + sender._auth.alg = "ed25519" + + def prepare(**kwargs: Any) -> PreparedWebhook: + key = f"whk_{secrets.token_hex(16)}" + return PreparedWebhook( + url=kwargs["url"], + idempotency_key=key, + body=json.dumps({"idempotency_key": key, **kwargs}).encode(), + ) + + sender.prepare_mcp.side_effect = prepare + sender.send_prepared = AsyncMock( + side_effect=lambda prepared: WebhookDeliveryResult( + status_code=200, + idempotency_key=prepared.idempotency_key, + url=prepared.url, + response_headers={}, + response_body=b"{}", + sent_body=prepared.body, + ) + ) + return sender + + +def outbox(pool: Any) -> PgTaskWebhookOutbox: + return PgTaskWebhookOutbox( + pool=pool, sender=_sender(), encryption_key=b"e" * 32, delivery_retry_horizon_seconds=86400 + ) + + +@pytest.mark.parametrize("kind", ["proposal", "task", "outbox"]) +async def test_concurrent_resolution_and_retry(kind: str) -> None: + pool = MagicMock() + concrete: Any = { + "proposal": lambda: PgProposalStore(pool=pool), + "task": lambda: PgTaskRegistry(pool=pool), + "outbox": lambda: outbox(pool), + }[kind]() + classes = { + "proposal": LazyProposalStore, + "task": LazyTaskRegistry, + "outbox": LazyTaskWebhookOutbox, + } + calls = 0 + + async def factory() -> Any: + nonlocal calls + calls += 1 + await asyncio.sleep(0.01) + if calls == 1: + raise RuntimeError("Bootstrap unavailable") + return concrete + + lazy = classes[kind](factory) + assert lazy.resolved is None + assert not hasattr(lazy, "clear_all") + with pytest.raises(RuntimeError, match="Bootstrap"): + await lazy.resolve() + resolved = await asyncio.gather(*(lazy.resolve() for _ in range(20))) + assert calls == 2 + assert all(store is concrete for store in resolved) + assert concrete._pool is pool + + +@pytest.mark.parametrize("kind", ["proposal", "task", "outbox"]) +async def test_first_method_use_resolves_once_and_forwards_arguments(kind: str) -> None: + pool = MagicMock() + if kind == "proposal": + concrete = PgProposalStore(pool=pool) + method = AsyncMock(return_value=None) + concrete.get = method + lazy = LazyProposalStore(lambda: concrete) + assert lazy.is_durable is True + await lazy.get("proposal", expected_account_id="acct") + method.assert_awaited_once_with("proposal", expected_account_id="acct") + elif kind == "task": + concrete = PgTaskRegistry(pool=pool) + method = AsyncMock(return_value={"pagination": {"has_more": False}, "tasks": []}) + concrete.list = method + lazy = LazyTaskRegistry(lambda: concrete) + assert lazy.is_durable is True + assert isinstance(lazy, TaskRegistry) and isinstance(lazy, ListableTaskRegistry) + await lazy.list(account_id="acct", filters={"status": "working"}) + method.assert_awaited_once_with( + account_id="acct", filters={"status": "working"}, sort=None, pagination=None + ) + else: + concrete = outbox(pool) + method = AsyncMock(return_value=True) + concrete.process_one = method + lazy = LazyTaskWebhookOutbox(lambda: concrete) + assert lazy.delivery_state_is_durable and lazy.supports_atomic_task_outbox + with pytest.raises(RuntimeError, match="resolve"): + lazy.validate_registration("https://buyer.example/hooks") + assert await lazy.process_one() is True + lazy.validate_registration("https://buyer.example/hooks") + method.assert_awaited_once_with() + assert lazy.resolved is concrete + concrete.clear_all = AsyncMock() + assert not hasattr(lazy, "clear_all") + + +async def test_pair_shares_pool_resolution_and_observers() -> None: + pool = MagicMock() + calls = 0 + + async def factory() -> tuple[PgTaskRegistry, PgTaskWebhookOutbox]: + nonlocal calls + calls += 1 + await asyncio.sleep(0.01) + box = outbox(pool) + return PgTaskRegistry(pool=pool, task_webhook_outbox=box), box + + lazy = LazyTaskRegistry(factory) + facade = LazyTaskWebhookOutbox.from_registry(lazy) + events = [] + lazy.add_lifecycle_observer(lambda event, **metadata: events.append((event, metadata))) + assert lazy.task_webhook_outbox is None and lazy.atomic_task_webhook_outbox is False + registry, box = await asyncio.gather(lazy.resolve(), facade.resolve()) + assert calls == 1 + assert registry._pool is box._pool is pool + assert lazy.task_webhook_outbox is box + assert lazy.atomic_task_webhook_outbox is True + assert _sdk_task_outbox_pair_ready(lazy, box) + # The bridge is registered before any operation delegated by the facade. + registry._notify_lifecycle_observers( + "submitted", + { + "task_id": "task", + "account_id": "acct", + "task_type": "get_products", + "created_at": 1.0, + "updated_at": 1.0, + }, + ) + assert events[0][0] == "submitted" + + +async def test_original_pool_identity_and_signing_scope_checks_remain_authoritative() -> None: + pool = MagicMock() + other_pool = MagicMock() + attempts = 0 + + def factory() -> PgTaskRegistry: + nonlocal attempts + attempts += 1 + box = outbox(other_pool if attempts == 1 else pool) + return PgTaskRegistry(pool=pool, task_webhook_outbox=box) + + lazy = LazyTaskRegistry(factory) + with pytest.raises(ValueError, match="same connection pool"): + await lazy.resolve() + assert lazy.resolved is None + assert (await lazy.resolve()).task_webhook_outbox._pool is pool + assert attempts == 2 + + box = outbox(pool) + bad_pair = LazyTaskRegistry(lambda: (PgTaskRegistry(pool=pool), box)) + with pytest.raises(ValueError, match="pair must share"): + await bad_pair.resolve() + + def missing_resolver() -> PgTaskRegistry: + return PgTaskRegistry(pool=pool, webhook_signing_scope_resolver=lambda ctx: "tenant") + + with pytest.raises(ValueError, match="required exactly"): + await LazyTaskRegistry(missing_resolver).resolve() + + +async def test_cancelled_factory_retries_and_waiters_do_not_duplicate_success() -> None: + started, allow = asyncio.Event(), asyncio.Event() + pool = MagicMock() + attempts = 0 + + async def factory() -> PgTaskRegistry: + nonlocal attempts + attempts += 1 + started.set() + await allow.wait() + return PgTaskRegistry(pool=pool) + + lazy = LazyTaskRegistry(factory) + owner = asyncio.create_task(lazy.resolve()) + await started.wait() + waiter = asyncio.create_task(lazy.resolve()) + owner.cancel() + with pytest.raises(asyncio.CancelledError): + await owner + allow.set() + assert await waiter is await lazy.resolve() + assert attempts == 2 + + +def test_successful_store_refuses_cross_loop_reuse_without_reopening() -> None: + pool = MagicMock() + factory = MagicMock(return_value=PgTaskRegistry(pool=pool)) + lazy = LazyTaskRegistry(factory) + asyncio.run(lazy.resolve()) + with pytest.raises(RuntimeError, match="another event loop"): + asyncio.run(lazy.resolve()) + factory.assert_called_once_with() + + +def test_failed_initialization_can_retry_on_a_new_loop_after_old_loop_closes() -> None: + pool = MagicMock() + factory = MagicMock(side_effect=[RuntimeError("Bootstrap"), PgTaskRegistry(pool=pool)]) + lazy = LazyTaskRegistry(factory) + with pytest.raises(RuntimeError, match="Bootstrap"): + asyncio.run(lazy.resolve()) + assert asyncio.run(lazy.resolve())._pool is pool + assert factory.call_count == 2 + + +async def test_bad_factory_result_not_cached() -> None: + factory = MagicMock(side_effect=[object(), PgTaskRegistry(pool=MagicMock())]) + lazy = LazyTaskRegistry(factory) + with pytest.raises(TypeError, match="PgTaskRegistry"): + await lazy.resolve() + assert lazy.resolved is None + assert isinstance(await lazy.resolve(), PgTaskRegistry) + + +@pytest.fixture +async def real_stack() -> AsyncIterator[Any]: + from psycopg_pool import AsyncConnectionPool + + url = os.environ.get("ADCP_PG_TEST_URL") + if not url: + pytest.skip("ADCP_PG_TEST_URL required") + suffix = secrets.token_hex(6) + pool = AsyncConnectionPool(url, open=False, min_size=1, max_size=4) + sender = _sender() + calls = 0 + + async def factory() -> tuple[PgTaskRegistry, PgTaskWebhookOutbox]: + nonlocal calls + calls += 1 + await pool.open() + box = PgTaskWebhookOutbox( + pool=pool, + sender=sender, + encryption_key=b"e" * 32, + delivery_retry_horizon_seconds=86400, + table=f"test_lazy_box_{suffix}", + ) + registry = PgTaskRegistry( + pool=pool, task_webhook_outbox=box, _table=f"test_lazy_task_{suffix}" + ) + await registry.create_schema() + await box.create_schema() + return registry, box + + lazy = LazyTaskRegistry(factory) + facade = LazyTaskWebhookOutbox.from_registry(lazy) + try: + yield lazy, facade, pool, sender, lambda: calls + finally: + if not pool.closed: + async with pool.connection() as conn: + for table in ( + f"test_lazy_box_{suffix}", + f"test_lazy_task_{suffix}", + f"test_lazy_proposal_{suffix}", + ): + await conn.execute(f"DROP TABLE IF EXISTS {table}") + await pool.close() + + +async def test_real_lazy_pair_atomic_completion_listing_and_observer_commit( + real_stack: Any, +) -> None: + lazy, facade, pool, sender, calls = real_stack + events = [] + lazy.add_lifecycle_observer(lambda event, **metadata: events.append(event)) + task = await lazy.issue( + account_id="acct", + task_type="get_products", + webhook_url="https://buyer.example/hooks", + webhook_operation_id="buyer-operation", + ) + await lazy.update_progress(task, {"percentage": 25}) + await asyncio.gather(*(lazy.complete(task, {"products": []}) for _ in range(10))) + assert events == ["submitted", "working", "completed"] + page = await lazy.list(account_id="acct") + assert page["tasks"][0]["task_id"] == task + assert (await lazy.get(task, expected_account_id="acct"))["state"] == "completed" + assert await lazy.get(task, expected_account_id="other") is None + assert await facade.process_one() is True + assert sender.send_prepared.await_count == 1 + assert calls() == 1 + assert (await facade.resolve())._pool is (await lazy.resolve())._pool is pool + # A closed caller-owned pool is not implicitly reopened or replaced. + await pool.close() + from psycopg_pool import PoolClosed + + with pytest.raises(PoolClosed): + await lazy.get(task) + assert calls() == 1 + + +async def test_real_lazy_proposal_lifecycle(real_stack: Any) -> None: + from datetime import datetime, timedelta, timezone + + from adcp.decisioning.proposal_store import ProposalState + + lazy, facade, pool, sender, calls = real_stack + registry = await lazy.resolve() + table = registry._table.replace("test_lazy_task_", "test_lazy_proposal_") + store = LazyProposalStore(lambda: PgProposalStore(pool=pool, table_name=table)) + await store.create_schema() + await store.put_draft( + proposal_id="proposal", account_id="acct", recipes={}, proposal_payload={} + ) + assert await store.get("proposal", expected_account_id="other") is None + await store.commit( + "proposal", + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + proposal_payload={}, + expected_account_id="acct", + ) + reserved = await store.try_reserve_consumption("proposal", expected_account_id="acct") + assert reserved.state == ProposalState.CONSUMING + await store.release_consumption("proposal", expected_account_id="acct") + await store.try_reserve_consumption("proposal", expected_account_id="acct") + await store.finalize_consumption("proposal", media_buy_id="buy", expected_account_id="acct") + assert ( + await store.get_by_media_buy_id("buy", expected_account_id="acct") + ).state == ProposalState.CONSUMED + + +async def test_observer_registration_and_removal_survive_resolution(real_stack: Any) -> None: + lazy, facade, pool, sender, calls = real_stack + before, after = [], [] + + def observe_before(event: Any, **metadata: Any) -> None: + before.append(event) + + def observe_after(event: Any, **metadata: Any) -> None: + after.append(event) + + lazy.add_lifecycle_observer(observe_before) + task = await lazy.issue(account_id="acct", task_type="get_products") + assert lazy.remove_lifecycle_observer(observe_before) is True + lazy.add_lifecycle_observer(observe_after) + await lazy.update_progress(task, {"percentage": 50}) + await lazy.complete(task, {"products": []}) + assert before == ["submitted"] + assert after == ["working", "completed"] + lazy.remove_lifecycle_observer(observe_after) + await lazy.discard(task) + assert after == ["working", "completed"] + + +async def test_sync_outbox_crypto_and_worker_methods_delegate() -> None: + concrete = outbox(MagicMock()) + lazy = LazyTaskWebhookOutbox(lambda: concrete) + await lazy.resolve() + protected, nonce = lazy.protect_registration( + account_id="acct", + task_id="task", + task_type="get_products", + url="https://buyer.example/hooks", + operation_id="buyer", + token="token", + ) + assert lazy.open_registration( + account_id="acct", + task_id="task", + task_type="get_products", + encrypted_registration=protected, + nonce=nonce, + ) == ("https://buyer.example/hooks", "buyer", "token") + concrete.run_worker = AsyncMock() + concrete.purge_expired = AsyncMock() + await lazy.run_worker(poll_interval=0.25, purge_interval=10) + await lazy.purge_expired() + concrete.run_worker.assert_awaited_once_with(poll_interval=0.25, purge_interval=10) + concrete.purge_expired.assert_awaited_once_with() + with pytest.raises(ValueError): + lazy.validate_registration("http://buyer.example/hooks") + + +async def test_durability_is_not_silently_changed_by_factory() -> None: + class LossyRegistry(PgTaskRegistry): + is_durable = False + + lazy = LazyTaskRegistry(lambda: LossyRegistry(pool=MagicMock())) + assert lazy.is_durable is True + with pytest.raises(ValueError, match="durable"): + await lazy.resolve() + assert lazy.resolved is None + + +@pytest.mark.parametrize("flag", ["delivery_state_is_durable", "supports_atomic_task_outbox"]) +async def test_outbox_factory_cannot_weaken_declared_guarantees(flag: str) -> None: + box = outbox(MagicMock()) + setattr(box, flag, False) + lazy = LazyTaskWebhookOutbox(lambda: box) + with pytest.raises(ValueError, match="durable atomic"): + await lazy.resolve() + assert lazy.resolved is None + setattr(box, flag, True) + assert await lazy.resolve() is box + + +async def test_registry_factory_cannot_remove_declared_listing_support() -> None: + concrete = PgTaskRegistry(pool=MagicMock()) + original_list = concrete.list + concrete.list = None + lazy = LazyTaskRegistry(lambda: concrete) + assert isinstance(lazy, ListableTaskRegistry) + with pytest.raises(TypeError, match="listing support"): + await lazy.resolve() + assert lazy.resolved is None + concrete.list = original_list + assert await lazy.resolve() is concrete diff --git a/tests/test_decisioning_task_polling.py b/tests/test_decisioning_task_polling.py index 2c7de51f4..ba5754c94 100644 --- a/tests/test_decisioning_task_polling.py +++ b/tests/test_decisioning_task_polling.py @@ -70,7 +70,7 @@ def get_media_buys(self, req: Any, ctx: Any) -> Any: raise NotImplementedError -@pytest.fixture(params=["memory", "postgres"]) +@pytest.fixture(params=["memory", "postgres", "lazy-postgres"]) async def registry(request: pytest.FixtureRequest) -> AsyncIterator[Any]: if request.param == "memory": yield InMemoryTaskRegistry() @@ -86,7 +86,12 @@ async def registry(request: pytest.FixtureRequest) -> AsyncIterator[Any]: async with AsyncConnectionPool(url, open=False) as pool: reg = PgTaskRegistry(pool=pool, _table=table) await reg.create_schema() - yield reg + if request.param == "lazy-postgres": + from adcp.decisioning.pg import LazyTaskRegistry + + yield LazyTaskRegistry(lambda: reg) + else: + yield reg async with pool.connection() as conn: await conn.execute(f"DROP TABLE {table}") diff --git a/tests/type_checks/lazy_pg_stores.py b/tests/type_checks/lazy_pg_stores.py new file mode 100644 index 000000000..1c31bcee2 --- /dev/null +++ b/tests/type_checks/lazy_pg_stores.py @@ -0,0 +1,66 @@ +"""Typed adopter factories preserve concrete store, listing and observer APIs.""" + +from typing import Any + +from psycopg_pool import AsyncConnectionPool + +from adcp.decisioning import ListableTaskRegistry, TaskRegistry, TaskTransition +from adcp.decisioning.pg import ( + LazyProposalStore, + LazyProposalStoreFactory, + LazyTaskRegistry, + LazyTaskRegistryFactory, + LazyTaskWebhookOutbox, + LazyTaskWebhookOutboxFactory, + PgProposalStore, + PgTaskRegistry, + PgTaskWebhookOutbox, +) +from adcp.decisioning.proposal_store import ProposalStore + + +def example(pool: AsyncConnectionPool[Any], outbox: PgTaskWebhookOutbox) -> None: + async def task_factory() -> tuple[PgTaskRegistry, PgTaskWebhookOutbox]: + return PgTaskRegistry(pool=pool, task_webhook_outbox=outbox), outbox + + factory: LazyTaskRegistryFactory = task_factory + registry = LazyTaskRegistry(factory) + minimal: TaskRegistry = registry + listable: ListableTaskRegistry = registry + + def metrics( + event: TaskTransition, + *, + task_id: str, + account_id: str, + task_type: str, + created_at: float, + updated_at: float, + ) -> None: + print(event, task_id, account_id, task_type, updated_at - created_at) + + registry.add_lifecycle_observer(metrics) + registry.remove_lifecycle_observer(metrics) + facade: LazyTaskWebhookOutbox = LazyTaskWebhookOutbox.from_registry(registry) + + def proposal_store() -> PgProposalStore: + return PgProposalStore(pool=pool) + + proposal_factory: LazyProposalStoreFactory = proposal_store + proposals: ProposalStore = LazyProposalStore(proposal_factory) + + def task_outbox() -> PgTaskWebhookOutbox: + return outbox + + outbox_factory: LazyTaskWebhookOutboxFactory = task_outbox + other_facade = LazyTaskWebhookOutbox(outbox_factory) + print(minimal, listable, facade, proposals, other_facade) + + +async def consume(registry: LazyTaskRegistry, facade: LazyTaskWebhookOutbox) -> None: + concrete: PgTaskRegistry = await registry.resolve() + box: PgTaskWebhookOutbox = await facade.resolve() + original: PgTaskRegistry | None = registry.resolved + tasks: dict[str, Any] = await registry.list(account_id="acct") + assert concrete.task_webhook_outbox is box + print(original, tasks) From 37a69cd4b0dc9ce12a5f3bfd9d93cb91ad2eee15 Mon Sep 17 00:00:00 2001 From: Brian O'Kelley Date: Fri, 2 Oct 2026 21:30:04 +0000 Subject: [PATCH 5/9] fix(pg): type terminal task rows without optional driver imports --- src/adcp/decisioning/pg/task_registry.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/adcp/decisioning/pg/task_registry.py b/src/adcp/decisioning/pg/task_registry.py index a9fe431a5..497b97655 100644 --- a/src/adcp/decisioning/pg/task_registry.py +++ b/src/adcp/decisioning/pg/task_registry.py @@ -452,7 +452,7 @@ async def _complete( self._sql_complete, (json.dumps(safe_result), time.time(), task_id), ) - row = await cur.fetchone() + row: tuple[Any, ...] | None = await cur.fetchone() if row is not None: await self._enqueue_terminal_if_registered( conn, @@ -497,7 +497,7 @@ async def _fail( self._sql_fail, (json.dumps(error), time.time(), task_id), ) - row = await cur.fetchone() + row: tuple[Any, ...] | None = await cur.fetchone() if row is not None: await self._enqueue_terminal_if_registered( conn, From a022fb2a250e3b6406a2840aa1f6208fa174908e Mon Sep 17 00:00:00 2001 From: Brian O'Kelley Date: Sat, 3 Oct 2026 02:37:19 +0000 Subject: [PATCH 6/9] test(tasks): gate polling fixtures on configured PostgreSQL --- .github/workflows/ci.yml | 1 + tests/test_decisioning_handler_shims.py | 3 +++ tests/test_decisioning_task_polling.py | 6 +++--- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a72211d14..237830249 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -195,6 +195,7 @@ jobs: tests/conformance/signing/test_pg_replay_store_e2e.py \ tests/conformance/decisioning/test_pg_buyer_agent_registry.py \ tests/conformance/decisioning/test_pg_idempotency_backend.py \ + tests/test_decisioning_task_polling.py \ tests/conformance/decisioning/test_pg_task_webhook_outbox.py \ tests/conformance/decisioning/test_pg_reference_workflow_queue.py \ tests/conformance/reporting/ \ diff --git a/tests/test_decisioning_handler_shims.py b/tests/test_decisioning_handler_shims.py index 5d724845d..94d97f60e 100644 --- a/tests/test_decisioning_handler_shims.py +++ b/tests/test_decisioning_handler_shims.py @@ -63,6 +63,9 @@ def test_advertised_tools_covers_every_specialism_wire_tool() -> None: Protocols were dead code at runtime — buyers would 404 on ``build_creative``, ``get_signals``, etc.""" expected = { + # SDK task tools (listing is advertised only for listable registries) + "get_task_status", + "list_tasks", # Sales "get_products", "create_media_buy", diff --git a/tests/test_decisioning_task_polling.py b/tests/test_decisioning_task_polling.py index 2c7de51f4..c0f30a28c 100644 --- a/tests/test_decisioning_task_polling.py +++ b/tests/test_decisioning_task_polling.py @@ -75,13 +75,13 @@ async def registry(request: pytest.FixtureRequest) -> AsyncIterator[Any]: if request.param == "memory": yield InMemoryTaskRegistry() return + url = os.environ.get("ADCP_PG_TEST_URL") + if not url: + pytest.skip("ADCP_PG_TEST_URL required") from psycopg_pool import AsyncConnectionPool from adcp.decisioning.pg import PgTaskRegistry - url = os.environ.get("ADCP_PG_TEST_URL") - if not url: - pytest.skip("ADCP_PG_TEST_URL required") table = f"test_polling_{secrets.token_hex(6)}" async with AsyncConnectionPool(url, open=False) as pool: reg = PgTaskRegistry(pool=pool, _table=table) From 4b66cd58f564d83e5855a92f91c4226860e0fa08 Mon Sep 17 00:00:00 2001 From: Brian O'Kelley Date: Sat, 3 Oct 2026 02:39:58 +0000 Subject: [PATCH 7/9] test(tasks): configure PostgreSQL lifecycle coverage explicitly --- .github/workflows/ci.yml | 1 + tests/test_decisioning_task_lifecycle.py | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a72211d14..c6e3254a4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -195,6 +195,7 @@ jobs: tests/conformance/signing/test_pg_replay_store_e2e.py \ tests/conformance/decisioning/test_pg_buyer_agent_registry.py \ tests/conformance/decisioning/test_pg_idempotency_backend.py \ + tests/test_decisioning_task_lifecycle.py \ tests/conformance/decisioning/test_pg_task_webhook_outbox.py \ tests/conformance/decisioning/test_pg_reference_workflow_queue.py \ tests/conformance/reporting/ \ diff --git a/tests/test_decisioning_task_lifecycle.py b/tests/test_decisioning_task_lifecycle.py index aa1b5190d..26944ffdd 100644 --- a/tests/test_decisioning_task_lifecycle.py +++ b/tests/test_decisioning_task_lifecycle.py @@ -18,13 +18,13 @@ async def registry(request: pytest.FixtureRequest) -> AsyncIterator[Any]: if request.param == "memory": yield InMemoryTaskRegistry() return + url = os.environ.get("ADCP_PG_TEST_URL") + if not url: + pytest.skip("ADCP_PG_TEST_URL required") from psycopg_pool import AsyncConnectionPool from adcp.decisioning.pg import PgTaskRegistry - url = os.environ.get("ADCP_PG_TEST_URL") - if not url: - pytest.skip("ADCP_PG_TEST_URL required") table = f"test_lifecycle_{secrets.token_hex(6)}" async with AsyncConnectionPool( url, From 8247a3acb647e044b1fe05ec9b1c992a74d8bcbb Mon Sep 17 00:00:00 2001 From: Brian O'Kelley Date: Sat, 3 Oct 2026 02:42:42 +0000 Subject: [PATCH 8/9] test(pg): separate fake lazy delegates from real database fixtures --- .github/workflows/ci.yml | 1 + tests/test_decisioning_lazy_pg_stores.py | 64 +++++++++++++++++++++++- 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6a7714521..9201980e3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -197,6 +197,7 @@ jobs: tests/conformance/decisioning/test_pg_idempotency_backend.py \ tests/test_decisioning_task_polling.py \ tests/test_decisioning_task_lifecycle.py \ + tests/test_decisioning_lazy_pg_stores.py \ tests/conformance/decisioning/test_pg_task_webhook_outbox.py \ tests/conformance/decisioning/test_pg_reference_workflow_queue.py \ tests/conformance/reporting/ \ diff --git a/tests/test_decisioning_lazy_pg_stores.py b/tests/test_decisioning_lazy_pg_stores.py index 654c3b868..6c618119d 100644 --- a/tests/test_decisioning_lazy_pg_stores.py +++ b/tests/test_decisioning_lazy_pg_stores.py @@ -24,6 +24,18 @@ from adcp.webhook_sender import PreparedWebhook, WebhookDeliveryResult +@pytest.fixture +def fake_pg_driver(monkeypatch: pytest.MonkeyPatch) -> None: + """Fake pools never call psycopg; retain the real constructor validation. + + This explicit fixture enables only the optional-driver availability guards + for unit cases, following the existing task-outbox unit-test pattern. Real + database cases and the missing-extra regression do not use it. + """ + for module in ("proposal_store", "task_registry", "task_webhook_outbox"): + monkeypatch.setattr(f"adcp.decisioning.pg.{module}.PG_AVAILABLE", True) + + def _sender() -> MagicMock: import json @@ -63,6 +75,7 @@ def outbox(pool: Any) -> PgTaskWebhookOutbox: @pytest.mark.parametrize("kind", ["proposal", "task", "outbox"]) +@pytest.mark.usefixtures("fake_pg_driver") async def test_concurrent_resolution_and_retry(kind: str) -> None: pool = MagicMock() concrete: Any = { @@ -97,6 +110,7 @@ async def factory() -> Any: @pytest.mark.parametrize("kind", ["proposal", "task", "outbox"]) +@pytest.mark.usefixtures("fake_pg_driver") async def test_first_method_use_resolves_once_and_forwards_arguments(kind: str) -> None: pool = MagicMock() if kind == "proposal": @@ -134,6 +148,7 @@ async def test_first_method_use_resolves_once_and_forwards_arguments(kind: str) assert not hasattr(lazy, "clear_all") +@pytest.mark.usefixtures("fake_pg_driver") async def test_pair_shares_pool_resolution_and_observers() -> None: pool = MagicMock() calls = 0 @@ -170,6 +185,7 @@ async def factory() -> tuple[PgTaskRegistry, PgTaskWebhookOutbox]: assert events[0][0] == "submitted" +@pytest.mark.usefixtures("fake_pg_driver") async def test_original_pool_identity_and_signing_scope_checks_remain_authoritative() -> None: pool = MagicMock() other_pool = MagicMock() @@ -200,6 +216,7 @@ def missing_resolver() -> PgTaskRegistry: await LazyTaskRegistry(missing_resolver).resolve() +@pytest.mark.usefixtures("fake_pg_driver") async def test_cancelled_factory_retries_and_waiters_do_not_duplicate_success() -> None: started, allow = asyncio.Event(), asyncio.Event() pool = MagicMock() @@ -224,6 +241,7 @@ async def factory() -> PgTaskRegistry: assert attempts == 2 +@pytest.mark.usefixtures("fake_pg_driver") def test_successful_store_refuses_cross_loop_reuse_without_reopening() -> None: pool = MagicMock() factory = MagicMock(return_value=PgTaskRegistry(pool=pool)) @@ -234,6 +252,7 @@ def test_successful_store_refuses_cross_loop_reuse_without_reopening() -> None: factory.assert_called_once_with() +@pytest.mark.usefixtures("fake_pg_driver") def test_failed_initialization_can_retry_on_a_new_loop_after_old_loop_closes() -> None: pool = MagicMock() factory = MagicMock(side_effect=[RuntimeError("Bootstrap"), PgTaskRegistry(pool=pool)]) @@ -244,6 +263,7 @@ def test_failed_initialization_can_retry_on_a_new_loop_after_old_loop_closes() - assert factory.call_count == 2 +@pytest.mark.usefixtures("fake_pg_driver") async def test_bad_factory_result_not_cached() -> None: factory = MagicMock(side_effect=[object(), PgTaskRegistry(pool=MagicMock())]) lazy = LazyTaskRegistry(factory) @@ -255,11 +275,11 @@ async def test_bad_factory_result_not_cached() -> None: @pytest.fixture async def real_stack() -> AsyncIterator[Any]: - from psycopg_pool import AsyncConnectionPool - url = os.environ.get("ADCP_PG_TEST_URL") if not url: pytest.skip("ADCP_PG_TEST_URL required") + from psycopg_pool import AsyncConnectionPool + suffix = secrets.token_hex(6) pool = AsyncConnectionPool(url, open=False, min_size=1, max_size=4) sender = _sender() @@ -384,6 +404,7 @@ def observe_after(event: Any, **metadata: Any) -> None: assert after == ["working", "completed"] +@pytest.mark.usefixtures("fake_pg_driver") async def test_sync_outbox_crypto_and_worker_methods_delegate() -> None: concrete = outbox(MagicMock()) lazy = LazyTaskWebhookOutbox(lambda: concrete) @@ -413,6 +434,7 @@ async def test_sync_outbox_crypto_and_worker_methods_delegate() -> None: lazy.validate_registration("http://buyer.example/hooks") +@pytest.mark.usefixtures("fake_pg_driver") async def test_durability_is_not_silently_changed_by_factory() -> None: class LossyRegistry(PgTaskRegistry): is_durable = False @@ -425,6 +447,7 @@ class LossyRegistry(PgTaskRegistry): @pytest.mark.parametrize("flag", ["delivery_state_is_durable", "supports_atomic_task_outbox"]) +@pytest.mark.usefixtures("fake_pg_driver") async def test_outbox_factory_cannot_weaken_declared_guarantees(flag: str) -> None: box = outbox(MagicMock()) setattr(box, flag, False) @@ -436,6 +459,7 @@ async def test_outbox_factory_cannot_weaken_declared_guarantees(flag: str) -> No assert await lazy.resolve() is box +@pytest.mark.usefixtures("fake_pg_driver") async def test_registry_factory_cannot_remove_declared_listing_support() -> None: concrete = PgTaskRegistry(pool=MagicMock()) original_list = concrete.list @@ -447,3 +471,39 @@ async def test_registry_factory_cannot_remove_declared_listing_support() -> None assert lazy.resolved is None concrete.list = original_list assert await lazy.resolve() is concrete + + +@pytest.mark.parametrize("kind", ["proposal", "task", "outbox"]) +async def test_missing_pg_extra_remains_an_import_error_and_is_not_cached( + kind: str, monkeypatch: pytest.MonkeyPatch +) -> None: + modules = { + "proposal": "proposal_store", + "task": "task_registry", + "outbox": "task_webhook_outbox", + } + monkeypatch.setattr(f"adcp.decisioning.pg.{modules[kind]}.PG_AVAILABLE", False) + pool = MagicMock() + constructors = { + "proposal": lambda: PgProposalStore(pool=pool), + "task": lambda: PgTaskRegistry(pool=pool), + "outbox": lambda: outbox(pool), + } + wrappers = { + "proposal": LazyProposalStore, + "task": LazyTaskRegistry, + "outbox": LazyTaskWebhookOutbox, + } + calls = 0 + + def factory() -> Any: + nonlocal calls + calls += 1 + return constructors[kind]() + + lazy = wrappers[kind](factory) + for _ in range(2): + with pytest.raises(ImportError, match=r"adcp\[pg\]"): + await lazy.resolve() + assert lazy.resolved is None + assert calls == 2 From ead8d38fec3d6f3d6fea350d4618943d8e04cf17 Mon Sep 17 00:00:00 2001 From: Brian O'Kelley Date: Sat, 3 Oct 2026 22:46:13 +0000 Subject: [PATCH 9/9] perf(reporting): accelerate canonical surrogate detection --- src/adcp/reporting/canonical_json.py | 10 +++- tests/test_reporting_canonical_json.py | 70 ++++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 1 deletion(-) diff --git a/src/adcp/reporting/canonical_json.py b/src/adcp/reporting/canonical_json.py index d8bb4d26b..b71af1fe0 100644 --- a/src/adcp/reporting/canonical_json.py +++ b/src/adcp/reporting/canonical_json.py @@ -31,6 +31,7 @@ import hashlib import json +import re from typing import Any, Union __all__ = [ @@ -47,6 +48,8 @@ #: that survive a round trip through a JavaScript producer unchanged. MAX_SAFE_INTEGER = 9007199254740991 +_SURROGATE_CODE_POINT = re.compile(r"[\ud800-\udfff]") + JsonValue = Union[ None, bool, @@ -77,7 +80,12 @@ def _encode_string(value: str) -> str: then fail to UTF-8 encode, so they are escaped here the way well-formed ``JSON.stringify`` does. """ - if any(0xD800 <= ord(character) <= 0xDFFF for character in value): + has_surrogates = ( + _SURROGATE_CODE_POINT.search(value) is not None + if type(value) is str + else any(0xD800 <= ord(character) <= 0xDFFF for character in value) + ) + if has_surrogates: return _encode_string_with_lone_surrogates(value) return json.dumps(value, ensure_ascii=False, separators=(",", ":")) diff --git a/tests/test_reporting_canonical_json.py b/tests/test_reporting_canonical_json.py index 6ed0b9eda..d5f39d74a 100644 --- a/tests/test_reporting_canonical_json.py +++ b/tests/test_reporting_canonical_json.py @@ -13,6 +13,8 @@ from __future__ import annotations +from collections.abc import Iterator + import pytest from adcp.reporting.canonical_json import ( @@ -151,3 +153,71 @@ def test_lone_surrogates_are_escaped_not_emitted() -> None: # A raw lone surrogate cannot be UTF-8 encoded at all, so the only # byte-stable answer is the ES2019 well-formed ``JSON.stringify`` escape. assert canonical_json_utf8_v1({"k": "\ud800"}) == b'{"k":"\\ud800"}' + + +@pytest.mark.parametrize( + ("text", "encoded"), + [ + ("\ud7ff", b'"\xed\x9f\xbf"'), + ("\ud800", b'"\\ud800"'), + ("\udbff", b'"\\udbff"'), + ("\udc00", b'"\\udc00"'), + ("\udfff", b'"\\udfff"'), + ("\ue000", b'"\xee\x80\x80"'), + ("\ud800\udc00", b'"\\ud800\\udc00"'), + ("\udbff\udfff", b'"\\udbff\\udfff"'), + ("\udc00\ud800", b'"\\udc00\\ud800"'), + ("\U00010000\U0010ffff", b'"\xf0\x90\x80\x80\xf4\x8f\xbf\xbf"'), + ("\ud800\n\u00e9\udfff", b'"\\ud800\\n\xc3\xa9\\udfff"'), + ], +) +def test_surrogate_boundary_bytes_in_values_and_keys(text: str, encoded: bytes) -> None: + assert canonical_json_utf8_v1(text) == encoded + assert canonical_json_utf8_v1({text: text}) == b"{" + encoded + b":" + encoded + b"}" + + +def test_surrogate_keys_retain_utf16_order_and_explicit_pair_bytes() -> None: + entries = [("\ue000", 5), ("\udfff", 4), ("\ud800\udc00", 2), ("\ud800", 1), ("\udc00", 3)] + assert canonical_json_utf8_v1(dict(entries)) == ( + b'{"\\ud800":1,"\\ud800\\udc00":2,"\\udc00":3,"\\udfff":4,"\xee\x80\x80":5}' + ) + + +@pytest.mark.parametrize("as_key", [False, True]) +def test_string_subclass_uses_its_iterator_for_surrogate_detection(as_key: bool) -> None: + iterations: list[str] = [] + + class IteratedString(str): + def __iter__(self) -> Iterator[str]: + iterations.append("iterated") + return iter("\ud800") + + text = IteratedString("ordinary") + value = {text: 1} if as_key else text + assert canonical_json_utf8_v1(value) == (b'{"ordinary":1}' if as_key else b'"ordinary"') + assert iterations == ["iterated"] + + +@pytest.mark.parametrize("as_key", [False, True]) +def test_string_subclass_iterator_error_is_preserved(as_key: bool) -> None: + error = RuntimeError("iterator failed") + + class FailingString(str): + def __iter__(self) -> Iterator[str]: + raise error + + text = FailingString("ordinary") + with pytest.raises(RuntimeError) as caught: + canonical_json_utf8_v1({text: 1} if as_key else text) + assert caught.value is error + + +@pytest.mark.parametrize("as_key", [False, True]) +def test_string_subclass_empty_iterator_retains_utf8_error(as_key: bool) -> None: + class EmptyIteratorString(str): + def __iter__(self) -> Iterator[str]: + return iter(()) + + text = EmptyIteratorString("\ud800") + with pytest.raises(UnicodeEncodeError, match="surrogates not allowed"): + canonical_json_utf8_v1({text: 1} if as_key else text)