From 7d033485265c626c852dc90985b88befe8d6adbc Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Wed, 14 Oct 2026 23:52:30 +0000 Subject: [PATCH 1/4] Fence query payload loading and handler dispatch to the original claim --- docs/query-task-authority.md | 37 +++ mkdocs.yml | 1 + src/durable_workflow/__init__.py | 2 + src/durable_workflow/_query_task_authority.py | 152 +++++++++++ src/durable_workflow/client.py | 79 +++++- src/durable_workflow/errors.py | 9 + src/durable_workflow/worker.py | 72 ++++- src/durable_workflow/workflow.py | 61 +++-- .../expired-query-handler-suppressed.json | 32 +++ tests/query_task_helpers.py | 15 + tests/test_client.py | 25 +- tests/test_condition_replay_drift.py | 6 +- tests/test_external_storage.py | 12 +- tests/test_query_task_authority.py | 257 ++++++++++++++++++ tests/test_replay_regression_corpus.py | 37 ++- tests/test_storage_admission.py | 22 +- tests/test_worker.py | 51 ++-- 17 files changed, 792 insertions(+), 78 deletions(-) create mode 100644 docs/query-task-authority.md create mode 100644 src/durable_workflow/_query_task_authority.py create mode 100644 tests/fixtures/replay_regressions/expired-query-handler-suppressed.json create mode 100644 tests/query_task_helpers.py create mode 100644 tests/test_query_task_authority.py diff --git a/docs/query-task-authority.md b/docs/query-task-authority.md new file mode 100644 index 0000000..3aec4b9 --- /dev/null +++ b/docs/query-task-authority.md @@ -0,0 +1,37 @@ +# Query loading and execution deadlines + +Query loading and replay consume the original query budget. They do not give a +worker extra time to invoke a handler. When the query lease expires, the worker +stops loading and suppresses the handler. Server still decides whether a result +can be accepted. + +Servers advertising `query_task_authority=true` expose a non-renewing query +status API for protocol 1.20. The SDK checks the original claim while loading +and again after replay, immediately before invoking the handler. Missing, +malformed, denied or unavailable observations suppress execution. + +To require that capability before registering a worker: + +```python +worker = Worker( + client, + task_queue="orders", + workflows=[OrderWorkflow], + require_query_task_authority=True, +) +``` + +The default retains compatibility with older Servers and explicit protocol +1.19. This path checks the supplied lease and any original caller deadline +against the local clock, then charges elapsed time monotonically. It cannot +observe an earlier caller timeout that Server did not expose. Use synchronized +clocks, or require the newer authority capability when that limitation is +unacceptable. A missing or malformed lease cannot authorize execution. + +An already-running synchronous Python handler cannot be safely interrupted. +Keep query handlers short and free of side effects. Expiry can cancel async +execution and payload I/O, but an observation never reserves future commit +authority. A result refused by Server does not restart the query or its budget. + +This change is undergoing source and published-runtime qualification under +[#143](https://github.com/durable-workflow/sdk-python/issues/143). diff --git a/mkdocs.yml b/mkdocs.yml index acf695a..7e890dc 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -77,6 +77,7 @@ nav: - Start here: index.md - SDK reference: sdk-reference.md - Sticky execution: guides/sticky-execution.md + - Query deadlines: query-task-authority.md - SDK guide: https://durable-workflow.com/docs/2.0/polyglot/python/ - PyPI: https://pypi.org/project/durable-workflow/ - API reference: diff --git a/src/durable_workflow/__init__.py b/src/durable_workflow/__init__.py index 7882bb8..7512447 100644 --- a/src/durable_workflow/__init__.py +++ b/src/durable_workflow/__init__.py @@ -88,6 +88,7 @@ NonDeterministicReplayError, NonRetryableError, QueryFailed, + QueryTaskAuthorityLost, RuntimeCapabilityUnsupported, RuntimeDiscoveryUnavailable, SagaCompensationFailed, @@ -375,6 +376,7 @@ "InvalidWorkflowDefinition", "NoopMetrics", "QueryFailed", + "QueryTaskAuthorityLost", "RuntimeCapabilityUnsupported", "RuntimeDiscoveryUnavailable", "RuntimeExternalPayloadReference", diff --git a/src/durable_workflow/_query_task_authority.py b/src/durable_workflow/_query_task_authority.py new file mode 100644 index 0000000..24ea377 --- /dev/null +++ b/src/durable_workflow/_query_task_authority.py @@ -0,0 +1,152 @@ +"""A query delivery's original, non-renewing execution budget.""" + +from __future__ import annotations + +import asyncio +import re +import time +from collections.abc import Awaitable, Callable, Mapping +from contextvars import ContextVar +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, TypeVar + +from .errors import QueryTaskAuthorityLost + +_T = TypeVar("_T") +_UTC6 = re.compile(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{6}Z\Z") +_OBSERVATION_INTERVAL_SECONDS = 5.0 + + +def _timestamp(value: Any, *, canonical: bool) -> datetime: + if not isinstance(value, str) or (canonical and _UTC6.fullmatch(value) is None): + raise ValueError("missing or malformed query timing") + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + if parsed.tzinfo is None: + raise ValueError("query timing must include a timezone") + return parsed + + +@dataclass +class QueryTaskAuthority: + client: object + namespace: str + task_id: str + attempt: int + owner: str + observed: bool + timing: dict[str, str] + until: float + identity: dict[str, Any] + observer: Callable[[], Awaitable[Mapping[str, Any]]] | None = field(default=None, repr=False) + + @classmethod + def delivery(cls, client: object, namespace: str, task: Mapping[str, Any], *, + owner: str, task_queue: str, observed: bool) -> QueryTaskAuthority: + task_id, attempt = task.get("query_task_id"), task.get("query_task_attempt") + if (not isinstance(owner, str) or not owner or not isinstance(task_queue, str) or not task_queue + or not isinstance(task_id, str) or not task_id or type(attempt) is not int or attempt < 1 + or task.get("lease_owner") != owner or task.get("task_queue") != task_queue + or task.get("namespace", namespace) != namespace): + raise QueryTaskAuthorityLost(str(task_id or ""), "invalid_query_claim") + try: + lease = _timestamp(task.get("lease_expires_at"), canonical=observed) + timing = {"lease_expires_at": task["lease_expires_at"]} + bounds = [lease] + if observed or "query_admitted_at" in task or "query_deadline_at" in task: + admitted = _timestamp(task.get("query_admitted_at"), canonical=True) + deadline = _timestamp(task.get("query_deadline_at"), canonical=True) + if deadline < admitted or (observed and lease > deadline): + raise ValueError("inconsistent original query timing") + timing.update({"query_admitted_at": task["query_admitted_at"], + "query_deadline_at": task["query_deadline_at"]}) + bounds.append(deadline) + # Legacy compatibility has no non-renewing observation. Its local + # clock check cannot establish an undisclosed caller deadline. + reference = (_timestamp(task.get("server_time"), canonical=True) if observed + else datetime.now(timezone.utc)) + until = time.monotonic() + (min(bounds) - reference).total_seconds() + except (ValueError, TypeError, OverflowError) as error: + raise QueryTaskAuthorityLost(task_id, "invalid_query_timing") from error + result = cls(client, namespace, task_id, attempt, owner, observed, timing, until, + {key: task.get(key) for key in ("workflow_id", "run_id", "task_queue")}) + result.remaining() + return result + + def matches(self, client: object, namespace: str, task: Mapping[str, Any], owner: str) -> bool: + return (self.client is client and self.namespace == namespace and self.owner == owner + and task.get("query_task_id") == self.task_id and task.get("query_task_attempt") == self.attempt + and task.get("lease_owner") == owner + and all(task.get(key) == value for key, value in self.identity.items()) + and all(task.get(key) == value for key, value in self.timing.items())) + + def remaining(self) -> float: + remaining = self.until - time.monotonic() + if remaining <= 0: + raise QueryTaskAuthorityLost(self.task_id, "query_execution_budget_expired") + return remaining + + async def confirm(self) -> None: + remaining = self.remaining() + if not self.observed: + return + started = time.monotonic() + try: + if self.observer is None: + raise ValueError("query authority observer is unavailable") + reply = await asyncio.wait_for(self.observer(), timeout=min(5.0, remaining)) + if (not isinstance(reply, Mapping) or reply.get("query_task_id") != self.task_id + or type(reply.get("query_task_attempt")) is not int + or reply.get("query_task_attempt") != self.attempt or reply.get("lease_owner") != self.owner + or any(reply.get(key) != value for key, value in self.timing.items())): + raise ValueError("query authority observation changed its original claim") + if (reply.get("can_continue") is not True or reply.get("query_task_status") != "leased" + or reply.get("reason") is not None): + raise QueryTaskAuthorityLost(self.task_id, str(reply.get("reason") or "query_claim_unavailable")) + server_time = _timestamp(reply.get("server_time"), canonical=True) + bound = min(_timestamp(value, canonical=True) for key, value in self.timing.items() + if key != "query_admitted_at") + # Charging the entire observation round trip is conservative. A + # later observation may shorten this bound but can never reset it. + self.until = min(self.until, started + (bound - server_time).total_seconds()) + except QueryTaskAuthorityLost: + raise + except Exception as error: + raise QueryTaskAuthorityLost(self.task_id, "query_authority_unavailable") from error + self.remaining() + + async def run(self, operation: Callable[[], Awaitable[_T]], *, observe: bool = True, + confirm_after: bool = True) -> _T: + await self.confirm() + execution = asyncio.ensure_future(operation()) + + async def watch() -> None: + while True: + await asyncio.sleep(min(_OBSERVATION_INTERVAL_SECONDS, self.remaining())) + await self.confirm() + + observation = asyncio.create_task(watch()) if observe and self.observed else None + pending: set[asyncio.Future[Any]] = {execution} + if observation is not None: + pending.add(observation) + try: + done, _ = await asyncio.wait(pending, timeout=self.remaining(), return_when=asyncio.FIRST_COMPLETED) + if not done: + raise QueryTaskAuthorityLost(self.task_id, "query_execution_budget_expired") + if observation is not None and observation in done: + await observation + raise QueryTaskAuthorityLost(self.task_id, "query_authority_unavailable") + result = await execution + if confirm_after: + await self.confirm() + return result + finally: + for future in pending: + future.cancel() + await asyncio.gather(*pending, return_exceptions=True) + + +# Async task creation copies this context. Each dispatched query keeps the +# loading budget without adding private fields to the public wire dictionary +# or retaining an unbounded task-ID registry on Client/Worker. +current_query_authority: ContextVar[QueryTaskAuthority | None] = ContextVar("query_task_authority", default=None) diff --git a/src/durable_workflow/client.py b/src/durable_workflow/client.py index 126cd3f..e3cc232 100644 --- a/src/durable_workflow/client.py +++ b/src/durable_workflow/client.py @@ -45,6 +45,7 @@ scope_timestamp, ) from ._cooperative_cancellation import CancellationDelivery, CancellationRequest +from ._query_task_authority import QueryTaskAuthority, current_query_authority from .errors import ( ExternalPayloadIntegrityMismatch, ExternalPayloadOversized, @@ -1809,6 +1810,9 @@ async def _request( if discovery and path != "/cluster/info": raise ValueError("cluster discovery credentials can only authorize /cluster/info") + query_poll = worker and path == "/worker/query-tasks/poll" + if query_poll: + current_query_authority.set(None) request_json = json if ( not discovery @@ -1872,6 +1876,10 @@ async def _do_request() -> httpx.Response: outcome = "ok" return None result = resp.json() + query_authority = None + if query_poll and isinstance(result, dict) and isinstance(result.get("task"), dict): + query_authority = await self._query_delivery_authority(result["task"], request_json) + current_query_authority.set(query_authority) if not discovery and _contains_runtime_payload_envelope(result): transport = await self._runtime_external_payload_transport() if transport is None: @@ -1879,14 +1887,17 @@ async def _do_request() -> httpx.Response: "runtime returned an external payload reference without advertising " "the authenticated namespace transport" ) - result = await self._resolve_workflow_claim_payloads( - result, - worker=worker, - transport=transport, - path=path, - request_json=request_json, - ) + if query_authority is not None: + result = await query_authority.run(lambda: self._resolve_runtime_payloads( + result, worker=worker, transport=transport, resolved={}, + )) + else: + result = await self._resolve_workflow_claim_payloads( + result, worker=worker, transport=transport, path=path, request_json=request_json, + ) result = self._hydrate_result_envelopes(result) + if query_authority is not None: + await query_authority.confirm() outcome = "ok" return result except Exception as exc: @@ -2332,6 +2343,7 @@ async def _resolve_runtime_payloads( *, worker: bool, transport: _RuntimeExternalPayloadTransport, + resolved: dict[tuple[str, str, int, str], str] | None = None, ) -> Any: if isinstance(value, dict): if "external_payload" in value: @@ -2347,6 +2359,9 @@ async def _resolve_runtime_payloads( "runtime external payload envelope codec does not match its reference", reference_id=reference.reference_id, ) + identity = (reference.reference_id, reference.codec, reference.size_bytes, reference.sha256) + if resolved is not None and identity in resolved: + return {"codec": reference.codec, "blob": resolved[identity]} data = await self._fetch_runtime_payload( reference, worker=worker, @@ -2359,12 +2374,15 @@ async def _resolve_runtime_payloads( "runtime external payload bytes are not a UTF-8 payload blob", reference_id=reference.reference_id, ) from exc + if resolved is not None: + resolved[identity] = blob return {"codec": reference.codec, "blob": blob} return { key: await self._resolve_runtime_payloads( item, worker=worker, transport=transport, + resolved=resolved, ) for key, item in value.items() } @@ -2374,6 +2392,7 @@ async def _resolve_runtime_payloads( item, worker=worker, transport=transport, + resolved=resolved, ) for item in value ] @@ -5625,6 +5644,48 @@ async def workflow_task_history( "POST", f"/worker/workflow-tasks/{task_id}/history", worker=True, json=body ) + async def _query_delivery_authority(self, task: dict[str, Any], request: Any) -> QueryTaskAuthority: + # Worker registration establishes discovery before it polls. Advanced + # callers without discovery retain only the exposed local lease checks. + info = self._cluster_info or {} + protocol = info.get("worker_protocol") + capabilities = protocol.get("server_capabilities") if isinstance(protocol, dict) else None + observed = (isinstance(capabilities, dict) and capabilities.get("query_task_authority") is True + and _supports_cooperative_cancellation_protocol(_protocol_version_from_env( + "DURABLE_WORKFLOW_WORKER_PROTOCOL_VERSION", PROTOCOL_VERSION))) + request = request if isinstance(request, dict) else {} + authority = QueryTaskAuthority.delivery(self, self.namespace, task, owner=request.get("worker_id", ""), + task_queue=request.get("task_queue", ""), observed=observed) + if observed: + authority.observer = lambda: self.query_task_status(query_task_id=authority.task_id, + lease_owner=authority.owner, query_task_attempt=authority.attempt) + await authority.confirm() + return authority + + async def query_task_status(self, *, query_task_id: str, lease_owner: str, + query_task_attempt: int) -> Any: + """Observe a protocol 1.20 query claim without renewing it. + + Requires ``query_task_authority=true`` discovery. A positive observation + does not reserve authority for later execution or completion. + """ + if not _supports_cooperative_cancellation_protocol(_protocol_version_from_env( + "DURABLE_WORKFLOW_WORKER_PROTOCOL_VERSION", PROTOCOL_VERSION, + )): + raise ValueError("query task observation requires worker protocol 1.20") + capability_path = "worker_protocol.server_capabilities.query_task_authority" + info = await self._runtime_discovery(operation="Client.query_task_status", required_path=capability_path) + protocol = info.get("worker_protocol") + capabilities = protocol.get("server_capabilities") if isinstance(protocol, dict) else None + supported = capabilities.get("query_task_authority") if isinstance(capabilities, dict) else None + if supported is not True: + error_type = RuntimeCapabilityUnsupported if supported is False else RuntimeDiscoveryUnavailable + raise error_type("Client.query_task_status", capability_path, + "Query task observation requires Server query_task_authority=true discovery.") + query = urlencode({"lease_owner": lease_owner, "query_task_attempt": query_task_attempt}) + return await self._request("GET", f"/worker/query-tasks/{quote(query_task_id, safe='')}/status?{query}", + worker=True, timeout=5.0) + async def poll_query_task( self, *, @@ -5640,6 +5701,10 @@ async def poll_query_task( must route a control-plane query to a non-PHP workflow runtime. ``timeout`` controls the server long-poll window; the HTTP request uses a small grace margin. + + Worker registration establishes capability discovery. Custom workers + should call ``get_cluster_info()`` first to enable advertised query + status checks; without discovery only supplied local lease bounds apply. """ body: dict[str, Any] = { "worker_id": worker_id, diff --git a/src/durable_workflow/errors.py b/src/durable_workflow/errors.py index 66a35f1..c6c3830 100644 --- a/src/durable_workflow/errors.py +++ b/src/durable_workflow/errors.py @@ -27,6 +27,15 @@ class DurableWorkflowError(Exception): """Base class for every exception raised by the SDK.""" +class QueryTaskAuthorityLost(DurableWorkflowError): + """A query delivery cannot safely load or invoke its handler.""" + + def __init__(self, query_task_id: str, reason: str) -> None: + super().__init__(f"query task {query_task_id!r} cannot continue: {reason}") + self.query_task_id = query_task_id + self.reason = reason + + class InvalidWorkflowDefinition(DurableWorkflowError, ValueError): """Workflow authoring uses an unsupported entry point or result.""" diff --git a/src/durable_workflow/worker.py b/src/durable_workflow/worker.py index db4b062..e94d3d6 100644 --- a/src/durable_workflow/worker.py +++ b/src/durable_workflow/worker.py @@ -41,6 +41,7 @@ from ._cancellation_scope_history import CancellationScopeBudget, CancellationScopeDeliveryReceipt from ._cooperative_cancellation import CancellationRequest, read_cancellation_history from ._prepared_local_activity import PreparedAttempt, PreparedCancellationObserved, PreparedLocalRunner +from ._query_task_authority import QueryTaskAuthority, current_query_authority from ._sticky_workflow_cache import CacheKey, StickyWorkflowCache, complete_history from .activity import ActivityContext, ActivityInfo, _set_context from .auth_composition import ( @@ -74,6 +75,7 @@ NonDeterministicReplayError, NonRetryableError, QueryFailed, + QueryTaskAuthorityLost, ServerError, UpdateRejected, UpdateValidationFailed, @@ -107,11 +109,11 @@ ScheduleActivity, StartChildWorkflow, UpsertMemo, + _prepare_query_state, _validate_workflow_entrypoint, _workflow_handler_registries, apply_update, commands_to_server_commands, - query_state, replay, validate_update, ) @@ -1035,8 +1037,11 @@ def __init__( external_storage: ExternalStorageDriver | None = None, external_storage_threshold_bytes: int | None = None, external_storage_cache: ExternalPayloadCache | None = None, + require_query_task_authority: bool = False, ) -> None: self.client = client + self.require_query_task_authority = require_query_task_authority + self._query_authority_supported = False self.task_queue = task_queue self.workflows: dict[str, type] = {} for definition in workflows: @@ -1300,6 +1305,14 @@ async def _register(self) -> None: ) self._validate_cooperative_activity_handlers() self._query_tasks_supported = _server_supports_query_tasks(info) + self._query_authority_supported = ( + isinstance(server_capabilities, Mapping) and server_capabilities.get("query_task_authority") is True + and _supports_cooperative_cancellation_protocol(_protocol_version_from_env( + "DURABLE_WORKFLOW_WORKER_PROTOCOL_VERSION", PROTOCOL_VERSION)) + ) + if self.require_query_task_authority and not self._query_authority_supported: + raise RuntimeError("query_task_authority_unavailable: require worker protocol 1.20 and " + "Server query_task_authority=true before registration") self._workflow_memo_updates_supported = _server_supports_workflow_memo_updates(info) has_update_validators = any( contract["update_validators"] for contract in self.workflow_command_contracts.values() @@ -3298,6 +3311,30 @@ async def call_interceptor( return await handler(context) async def _run_query_task(self, task: dict[str, Any], *, client: Client | None = None) -> str: + client = client or self.client + token = None + try: + namespace = getattr(client, "namespace", "default") + authority = current_query_authority.get() + if authority is None: + authority = QueryTaskAuthority.delivery(client, namespace, task, owner=self.worker_id, + task_queue=self.task_queue, observed=self._query_authority_supported) + if authority.observed: + authority.observer = lambda: client.query_task_status(query_task_id=authority.task_id, + lease_owner=authority.owner, query_task_attempt=authority.attempt) + token = current_query_authority.set(authority) + if not authority.matches(client, namespace, task, self.worker_id): + raise QueryTaskAuthorityLost(str(task.get("query_task_id", "")), "invalid_query_claim") + await authority.confirm() + return await self._run_authorized_query_task(task, client=client) + except QueryTaskAuthorityLost as error: + log.info("query delivery suppressed: %s", error) + return "expired" + finally: + if token is not None: + current_query_authority.reset(token) + + async def _run_authorized_query_task(self, task: dict[str, Any], *, client: Client) -> str: client = client or self.client codec = task.get("payload_codec") try: @@ -3341,6 +3378,11 @@ async def call_interceptor( async def _run_query_task_core(self, task: dict[str, Any], *, client: Client | None = None) -> str: client = client or self.client + authority = current_query_authority.get() + if authority is None or not authority.matches(client, getattr(client, "namespace", "default"), + task, self.worker_id): + raise QueryTaskAuthorityLost(str(task.get("query_task_id", "")), "invalid_query_claim") + await authority.confirm() query_task_id: str = task["query_task_id"] attempt: int = task.get("query_task_attempt", 1) wf_type: str = task.get("workflow_type", "") @@ -3390,7 +3432,7 @@ async def _run_query_task_core(self, task: dict[str, Any], *, client: Client | N codec=codec, context="workflow query arguments", ) - result = query_state( + invoke = _prepare_query_state( cls, history, start_input, @@ -3405,8 +3447,14 @@ async def _run_query_task_core(self, task: dict[str, Any], *, client: Client | N prepare_local_activity_groups=self._prepared_local_activity_groups_supported, local_activity_cancellation_policies=self._local_activity_cancellation_policies, ) - if inspect.isawaitable(result): - result = await result + await authority.confirm() + async def execute() -> Any: + value = invoke() + return await value if inspect.isawaitable(value) else value + + result = await authority.run(execute) + except QueryTaskAuthorityLost: + raise except AvroNotInstalledError as e: await self._fail_query_task( query_task_id, @@ -3452,7 +3500,7 @@ async def _run_query_task_core(self, task: dict[str, Any], *, client: Client | N return "failed" try: - await client.complete_query_task( + await authority.run(lambda: client.complete_query_task( query_task_id=query_task_id, lease_owner=self.worker_id, query_task_attempt=attempt, @@ -3462,7 +3510,9 @@ async def _run_query_task_core(self, task: dict[str, Any], *, client: Client | N run_id=task.get("run_id"), query_name=query_name, **self._external_storage_completion_kwargs(), - ) + ), observe=False, confirm_after=False) + except QueryTaskAuthorityLost: + raise except ServerError as e: if _is_storage_admission_error(e): log.warning("query task %s acknowledgement paused: %s", query_task_id, e) @@ -3552,6 +3602,12 @@ async def _fail_query_task( client: Client | None = None, ) -> None: client = client or self.client + authority = current_query_authority.get() + if authority is not None: + if (authority.client is not client or authority.task_id != query_task_id or authority.attempt != attempt + or authority.owner != self.worker_id): + raise QueryTaskAuthorityLost(query_task_id, "invalid_query_claim") + await authority.confirm() try: await client.fail_query_task( query_task_id=query_task_id, @@ -4013,7 +4069,7 @@ async def _dispatch_update_validation_task(self, task: dict[str, Any]) -> None: def _clone_client_for_query_tasks(self) -> Client: warning_config = self._payload_size_warning_config() - return Client( + client = Client( self.client.base_url, token=self.client.token, control_token=self.client.control_token, @@ -4035,6 +4091,8 @@ def _clone_client_for_query_tasks(self) -> Client: external_storage_threshold_bytes=self.external_storage_threshold_bytes, external_storage_cache=self.external_storage_cache, ) + client._cluster_info = self.client._cluster_info + return client def _can_clone_client_for_query_tasks(self) -> bool: return type(self.client) is Client diff --git a/src/durable_workflow/workflow.py b/src/durable_workflow/workflow.py index 9d62e30..c1f24a0 100644 --- a/src/durable_workflow/workflow.py +++ b/src/durable_workflow/workflow.py @@ -2991,7 +2991,7 @@ def replay( ).outcome -def query_state( +def _prepare_query_state( workflow_cls: type, history_events: Iterable[dict[str, Any]], start_input: list[Any], @@ -3006,17 +3006,8 @@ def query_state( prepare_local_activities: bool = False, prepare_local_activity_groups: bool = False, local_activity_cancellation_policies: tuple[str, ...] = (), -) -> Any: - """Replay a workflow to current state and invoke a registered query. - - This is the replay and invocation core used by the current server-routed - worker query path after it fetches durable history. That path is negotiated - through ``worker_protocol.server_capabilities.query_tasks``: workers - advertise ``query_tasks`` only after discovery reports it as available, - and clients fail with typed capability or discovery errors when it is not. - Unknown query names and handler exceptions are normalized to - :class:`~durable_workflow.errors.QueryFailed`. - """ +) -> Callable[[], Any]: + """Prepare the query so a worker can recheck authority after replay.""" try: state = _replay_state( workflow_cls, @@ -3047,12 +3038,46 @@ def query_state( if handler is None: raise QueryFailed(f"query handler {query_name!r} is not available") - try: - return handler(*(list(args) if args is not None else [])) - except QueryFailed: - raise - except Exception as exc: - raise QueryFailed(str(exc) or f"query {query_name!r} failed") from exc + def invoke() -> Any: + try: + return handler(*(list(args) if args is not None else [])) + except QueryFailed: + raise + except Exception as exc: + raise QueryFailed(str(exc) or f"query {query_name!r} failed") from exc + + return invoke + + +def query_state( + workflow_cls: type, + history_events: Iterable[dict[str, Any]], + start_input: list[Any], + query_name: str, + args: list[Any] | None = None, + *, + workflow_id: str | None = None, + run_id: str = "", + payload_codec: str | None = None, + external_storage: ExternalStorageDriver | None = None, + external_storage_cache: ExternalPayloadCache | None = None, + prepare_local_activities: bool = False, + prepare_local_activity_groups: bool = False, + local_activity_cancellation_policies: tuple[str, ...] = (), +) -> Any: + """Replay workflow state and invoke a query, normalizing handler failures. + + Server-routed Worker execution separately checks its original query claim + between replay and handler invocation. Direct local use has no Server lease. + """ + return _prepare_query_state( + workflow_cls, history_events, start_input, query_name, args, + workflow_id=workflow_id, run_id=run_id, payload_codec=payload_codec, + external_storage=external_storage, external_storage_cache=external_storage_cache, + prepare_local_activities=prepare_local_activities, + prepare_local_activity_groups=prepare_local_activity_groups, + local_activity_cancellation_policies=local_activity_cancellation_policies, + )() def apply_update( diff --git a/tests/fixtures/replay_regressions/expired-query-handler-suppressed.json b/tests/fixtures/replay_regressions/expired-query-handler-suppressed.json new file mode 100644 index 0000000..99c1d7a --- /dev/null +++ b/tests/fixtures/replay_regressions/expired-query-handler-suppressed.json @@ -0,0 +1,32 @@ +{ + "$schema": "https://durable-workflow.github.io/schemas/replay-regression-v1.json", + "fixture_schema": "durable-workflow.replay-regression/v1", + "id": "expired-query-handler-suppressed", + "protocol_version": "1.20", + "bindings": ["python"], + "workflow": { + "type": "sdk.query-task-dispatch", + "input": [{ + "query_task_id": "expired-query-original", + "query_task_attempt": 1, + "lease_owner": "query-worker", + "task_queue": "query-authority", + "lease_expires_at": "2000-01-01T00:00:00.000000Z", + "query_name": "current" + }] + }, + "command_sequence": [{ + "type": "query_task_dispatch", "outcome": "expired", + "handler_calls": 0, "completion_calls": 0, "failure_calls": 0 + }], + "expected": { + "command_sequence": [{ + "type": "query_task_dispatch", "outcome": "expired", + "handler_calls": 0, "completion_calls": 0, "failure_calls": 0 + }] + }, + "provenance": { + "issue": "https://github.com/durable-workflow/sdk-python/issues/143", + "scope": "minimal original lease-expiry worker-dispatch control, without durable-history qualification" + } +} diff --git a/tests/query_task_helpers.py b/tests/query_task_helpers.py new file mode 100644 index 0000000..9613886 --- /dev/null +++ b/tests/query_task_helpers.py @@ -0,0 +1,15 @@ +"""Ordinary ephemeral lease metadata for query transport and replay tests.""" + +from datetime import datetime, timedelta, timezone +from typing import Any + +from durable_workflow import Worker + + +def query_task_delivery(task: dict[str, Any], *, worker_id: str, task_queue: str) -> dict[str, Any]: + return {"query_task_attempt": 1, "lease_owner": worker_id, "task_queue": task_queue, + "lease_expires_at": (datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(), **task} + + +def leased_query_task(worker: Worker, task: dict[str, Any]) -> dict[str, Any]: + return query_task_delivery(task, worker_id=worker.worker_id, task_queue=worker.task_queue) diff --git a/tests/test_client.py b/tests/test_client.py index 0d2ba51..bceeb32 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -3761,16 +3761,20 @@ async def test_list_workers_rejects_non_object_response(self, client: Client) -> class TestQueryTasks: @pytest.mark.asyncio - async def test_poll_query_task_matches_polyglot_fixture(self, client: Client) -> None: + async def test_poll_keeps_fixture_request_but_refuses_missing_query_lease(self, client: Client) -> None: + from durable_workflow.errors import QueryTaskAuthorityLost + fixture_path = Path(__file__).parent / "fixtures" / "control-plane" / "query-task-poll-parity.json" fixture = json.loads(fixture_path.read_text()) resp = _mock_response(200, fixture["response_body"]) - with patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp) as mock: - task = await client.poll_query_task(**fixture["sdk_python"]["kwargs"]) + with ( + patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp) as mock, + pytest.raises(QueryTaskAuthorityLost) as refusal, + ): + await client.poll_query_task(**fixture["sdk_python"]["kwargs"]) - assert task == fixture["response_body"]["task"] - assert task["query_task_id"] == fixture["semantic_body"]["query_task_id"] + assert refusal.value.query_task_id == fixture["semantic_body"]["query_task_id"] assert mock.call_args.args[:2] == (fixture["request"]["method"], f"/api{fixture['request']['path']}") request_body = mock.call_args.kwargs["json"] assert request_body["worker_id"] == fixture["request"]["body"]["worker_id"] @@ -3808,11 +3812,14 @@ async def test_fail_query_task_matches_polyglot_fixture(self, client: Client) -> @pytest.mark.asyncio async def test_poll_query_task_returns_task_payload(self, client: Client) -> None: - resp = _mock_response(200, {"task": {"query_task_id": "qt1"}}) + from tests.query_task_helpers import query_task_delivery + + response_task = query_task_delivery({"query_task_id": "qt1"}, worker_id="w1", task_queue="q1") + resp = _mock_response(200, {"task": response_task}) with patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp) as mock: task = await client.poll_query_task(worker_id="w1", task_queue="q1", timeout=5.0) - assert task == {"query_task_id": "qt1"} + assert task == response_task assert mock.call_args.args[:2] == ("POST", "/api/worker/query-tasks/poll") body = mock.call_args.kwargs["json"] assert body["worker_id"] == "w1" @@ -3826,7 +3833,9 @@ async def test_poll_query_task_returns_task_payload(self, client: Client) -> Non async def test_poll_query_task_retries_once_with_same_poll_request_id_after_timeout( self, client: Client ) -> None: - response_task = {"task": {"query_task_id": "qt1"}} + from tests.query_task_helpers import query_task_delivery + + response_task = {"task": query_task_delivery({"query_task_id": "qt1"}, worker_id="w1", task_queue="q1")} with patch.object( client._http, diff --git a/tests/test_condition_replay_drift.py b/tests/test_condition_replay_drift.py index efd262f..39ea296 100644 --- a/tests/test_condition_replay_drift.py +++ b/tests/test_condition_replay_drift.py @@ -107,9 +107,11 @@ def predicate() -> bool: @pytest.mark.asyncio @pytest.mark.parametrize("field", ["condition_key", "condition_definition_fingerprint"]) async def test_query_worker_reports_unavailable_state_without_changing_the_run(field: str) -> None: + from tests.query_task_helpers import leased_query_task + client = AsyncMock(spec=Client) worker = Worker(client, task_queue="condition-drift", workflows=[ConditionIdentityWorkflow]) - outcome = await worker._run_query_task({ + outcome = await worker._run_query_task(leased_query_task(worker, { "query_task_id": "query-original", "query_task_attempt": 2, "workflow_id": "workflow-original", @@ -120,7 +122,7 @@ async def test_query_worker_reports_unavailable_state_without_changing_the_run(f "workflow_arguments": serializer.envelope([], codec="avro"), "query_arguments": serializer.envelope([], codec="avro"), "payload_codec": "avro", - }) + })) assert outcome == "failed" client.complete_query_task.assert_not_awaited() diff --git a/tests/test_external_storage.py b/tests/test_external_storage.py index c0305db..96806fb 100644 --- a/tests/test_external_storage.py +++ b/tests/test_external_storage.py @@ -262,6 +262,8 @@ def test_replayer_fetches_external_local_activity_result(tmp_path: Path) -> None @pytest.mark.asyncio async def test_worker_query_task_replay_fetches_external_history_payload(tmp_path: Path) -> None: + from tests.query_task_helpers import leased_query_task + storage = LocalFilesystemExternalStorage(tmp_path) client = _mock_worker_client(storage) worker = Worker(client, task_queue="q1", workflows=[ExternalStorageQueryWorkflow], activities=[]) @@ -279,7 +281,7 @@ async def test_worker_query_task_replay_fetches_external_history_payload(tmp_pat codec="avro", ) - outcome = await worker._run_query_task( + outcome = await worker._run_query_task(leased_query_task(worker, { "query_task_id": "qt-external-history", "query_task_attempt": 1, @@ -298,7 +300,7 @@ async def test_worker_query_task_replay_fetches_external_history_payload(tmp_pat "query_arguments": query_args, "payload_codec": "avro", } - ) + )) assert outcome == "completed" client.complete_query_task.assert_awaited_once() @@ -391,7 +393,9 @@ async def fake_request(*args: object, **kwargs: object) -> dict[str, str]: external_storage=storage, external_storage_threshold_bytes=1, ) - outcome = await worker._run_query_task( + from tests.query_task_helpers import leased_query_task + + outcome = await worker._run_query_task(leased_query_task(worker, { "query_task_id": "query-external-result", "query_task_attempt": 1, @@ -410,7 +414,7 @@ async def fake_request(*args: object, **kwargs: object) -> dict[str, str]: "query_arguments": serializer.envelope(["!"], codec="avro"), "payload_codec": "avro", } - ) + )) finally: await client.aclose() diff --git a/tests/test_query_task_authority.py b/tests/test_query_task_authority.py new file mode 100644 index 0000000..1df2727 --- /dev/null +++ b/tests/test_query_task_authority.py @@ -0,0 +1,257 @@ +from __future__ import annotations + +import asyncio +import time +from datetime import datetime, timedelta, timezone +from typing import Any +from unittest.mock import AsyncMock + +import httpx +import pytest + +from durable_workflow import Client, Worker, serializer, workflow +from durable_workflow._query_task_authority import QueryTaskAuthority, current_query_authority +from durable_workflow.errors import QueryTaskAuthorityLost + + +def timestamp(value: datetime) -> str: + return value.isoformat(timespec="microseconds").replace("+00:00", "Z") + + +def delivery(**changes: Any) -> dict[str, Any]: + now = datetime.now(timezone.utc) + return { + "query_task_id": "query-1", "query_task_attempt": 1, "lease_owner": "w1", "task_queue": "q1", + "query_admitted_at": timestamp(now - timedelta(seconds=2)), + "query_deadline_at": timestamp(now + timedelta(seconds=30)), + "lease_expires_at": timestamp(now + timedelta(seconds=20)), "server_time": timestamp(now), + "workflow_id": "workflow-1", "run_id": "run-1", "workflow_type": "authority-query", + "query_name": "current", "payload_codec": "avro", "history_events": [], + "workflow_arguments": serializer.envelope([], codec="avro"), + "query_arguments": serializer.envelope([], codec="avro"), **changes, + } + + +def authority(task: dict[str, Any], *, client: object | None = None, observed: bool = True) -> QueryTaskAuthority: + return QueryTaskAuthority.delivery(client or object(), "billing", task, owner="w1", task_queue="q1", + observed=observed) + + +def observation(task: dict[str, Any], **changes: Any) -> dict[str, Any]: + return {key: task[key] for key in ("query_task_id", "query_task_attempt", "lease_owner", "query_admitted_at", + "query_deadline_at", "lease_expires_at")} | { + "server_time": timestamp(datetime.now(timezone.utc)), "query_task_status": "leased", + "can_continue": True, "reason": None, **changes, + } + + +@pytest.mark.parametrize(("field", "value"), [ + ("query_task_attempt", True), ("lease_owner", "other"), ("task_queue", "other"), + ("lease_expires_at", None), ("lease_expires_at", "2026-10-14T00:00:00"), + ("query_admitted_at", None), ("query_deadline_at", "invalid"), ("server_time", "invalid"), +]) +def test_invalid_claim_cannot_grant_authority(field: str, value: Any) -> None: + with pytest.raises(QueryTaskAuthorityLost): + authority(delivery(**{field: value})) + + +@pytest.mark.parametrize("reason", ["query_task_timed_out", "lease_expired", "stale_worker_registration"]) +async def test_denied_observation_never_starts_operation(reason: str) -> None: + task = delivery() + budget = authority(task) + budget.observer = AsyncMock(return_value=observation(task, can_continue=False, reason=reason)) + operation = AsyncMock() + with pytest.raises(QueryTaskAuthorityLost, match=reason): + await budget.run(operation) + operation.assert_not_called() + + +async def test_observation_never_extends_original_budget_or_accepts_changed_deadline() -> None: + task = delivery() + budget = authority(task) + original_until = budget.until + budget.observer = AsyncMock(return_value=observation(task)) + await budget.confirm() + assert budget.until <= original_until + budget.observer.return_value = observation(task, query_deadline_at=timestamp( + datetime.now(timezone.utc) + timedelta(seconds=60))) + with pytest.raises(QueryTaskAuthorityLost): + await budget.confirm() + + +async def test_expiry_cancels_loading_without_starting_another_operation() -> None: + budget = authority(delivery(), observed=False) + budget.until = time.monotonic() + 0.03 + cancelled = asyncio.Event() + + async def load() -> None: + try: + await asyncio.sleep(60) + finally: + cancelled.set() + + with pytest.raises(QueryTaskAuthorityLost, match="budget_expired"): + await budget.run(load) + assert cancelled.is_set() + + +@workflow.defn(name="authority-query") +class QueryWorkflow: + calls = 0 + + def run(self, ctx: workflow.WorkflowContext) -> int: + return 7 + + @workflow.query("current") + def current(self) -> int: + type(self).calls += 1 + return 7 + + +@pytest.mark.parametrize("expire_after_replay", [False, True]) +async def test_worker_checks_after_replay_before_handler(monkeypatch: pytest.MonkeyPatch, + expire_after_replay: bool) -> None: + import durable_workflow.worker as worker_module + + client = AsyncMock(spec=Client) + client.namespace = "billing" + client.complete_query_task.return_value = {"outcome": "completed"} + worker = Worker(client, worker_id="w1", task_queue="q1", workflows=[QueryWorkflow]) + task = delivery() + budget = authority(task, client=client, observed=False) + original = worker_module._prepare_query_state + + def prepare(*args: Any, **kwargs: Any) -> Any: + invoke = original(*args, **kwargs) + if expire_after_replay: + budget.until = time.monotonic() - 1 + return invoke + + monkeypatch.setattr(worker_module, "_prepare_query_state", prepare) + QueryWorkflow.calls = 0 + token = current_query_authority.set(budget) + try: + outcome = await worker._run_query_task(task) + finally: + current_query_authority.reset(token) + assert outcome == ("expired" if expire_after_replay else "completed") + assert QueryWorkflow.calls == (0 if expire_after_replay else 1) + assert client.complete_query_task.await_count == (0 if expire_after_replay else 1) + client.fail_query_task.assert_not_called() + + +async def test_dispatch_context_keeps_its_own_original_budget() -> None: + first, second = authority(delivery()), authority(delivery(query_task_id="query-2")) + token = current_query_authority.set(first) + ready = asyncio.Event() + + async def dispatched() -> QueryTaskAuthority | None: + await ready.wait() + return current_query_authority.get() + + pending = asyncio.create_task(dispatched()) + current_query_authority.set(second) + ready.set() + try: + assert await pending is first + assert current_query_authority.get() is second + finally: + current_query_authority.reset(token) + + +async def test_query_loading_reuses_references_without_enlarging_persistent_cache() -> None: + from durable_workflow.external_storage import ExternalPayloadCache + from tests.test_runtime_external_payload_transport import FakeRuntimePayloadServer, runtime_client, runtime_envelope + + server = FakeRuntimePayloadServer() + first, second = runtime_envelope(server, ["first"]), runtime_envelope(server, ["second"]) + task = delivery(workflow_arguments=first, query_arguments=second, + history_events=[{"event_type": "WorkflowStarted", "payload": {"arguments": first}}]) + server.responses[("POST", "/api/worker/query-tasks/poll")] = {"task": task} + cache = ExternalPayloadCache(max_bytes=1) + async with runtime_client(server, cache=cache) as client: + token = current_query_authority.set(None) + try: + result = await client.poll_query_task(worker_id="w1", task_queue="q1") + finally: + current_query_authority.reset(token) + assert server.fetch_count == 2 + assert result["workflow_arguments"] == result["history_events"][0]["payload"]["arguments"] + assert result["workflow_arguments"] is not result["history_events"][0]["payload"]["arguments"] + assert cache.max_bytes == 1 + + +async def test_status_loss_cancels_active_query_fetch(monkeypatch: pytest.MonkeyPatch) -> None: + import durable_workflow._query_task_authority as authority_module + from tests.test_runtime_external_payload_transport import FakeRuntimePayloadServer, runtime_client, runtime_envelope + + monkeypatch.setattr(authority_module, "_OBSERVATION_INTERVAL_SECONDS", 0.01) + server = FakeRuntimePayloadServer() + task = delivery(workflow_arguments=runtime_envelope(server, ["large"])) + info = server.cluster_info() + info["worker_protocol"]["server_capabilities"]["query_task_authority"] = True + server.cluster_info = lambda: info # type: ignore[method-assign] + server.responses[("POST", "/api/worker/query-tasks/poll")] = {"task": task} + fetching, cancelled = asyncio.Event(), asyncio.Event() + + async def handle(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/status"): + assert request.method == "GET" + assert request.url.params["lease_owner"] == "w1" + assert request.url.params["query_task_attempt"] == "1" + return httpx.Response(200, json=observation(task, can_continue=not fetching.is_set(), + reason="stale_worker_registration" if fetching.is_set() else None)) + if request.method == "GET" and "/external-payloads/" in request.url.path: + fetching.set() + try: + await asyncio.sleep(60) + finally: + cancelled.set() + return server.handler(request) + + async with runtime_client(server) as client: + await client._http.aclose() + client._http = httpx.AsyncClient(base_url=client.base_url, transport=httpx.MockTransport(handle)) + await client.get_cluster_info() + token = current_query_authority.set(None) + try: + with pytest.raises(QueryTaskAuthorityLost, match="stale_worker_registration"): + await asyncio.wait_for(client.poll_query_task(worker_id="w1", task_queue="q1"), timeout=1) + finally: + current_query_authority.reset(token) + assert fetching.is_set() and cancelled.is_set() + + +@pytest.mark.parametrize(("version", "supported", "accepted"), [ + ("1.20", False, False), ("1.19", True, False), ("1.20", True, True), +]) +async def test_required_authority_is_checked_before_registration(monkeypatch: pytest.MonkeyPatch, + version: str, supported: bool, accepted: bool) -> None: + from tests.test_worker import compatible_cluster_info + + monkeypatch.setenv("DURABLE_WORKFLOW_WORKER_PROTOCOL_VERSION", version) + info = compatible_cluster_info() + protocol = info["worker_protocol"] + assert isinstance(protocol, dict) + protocol["server_capabilities"]["query_task_authority"] = supported + client = AsyncMock(spec=Client) + client.get_cluster_info.return_value = info + client.register_worker.return_value = {"worker_id": "w1", "registered": True} + worker = Worker(client, worker_id="w1", task_queue="q1", workflows=[QueryWorkflow], + require_query_task_authority=True) + if accepted: + await worker._register() + client.register_worker.assert_awaited_once() + else: + with pytest.raises(RuntimeError, match="query_task_authority_unavailable"): + await worker._register() + client.register_worker.assert_not_called() + + +async def test_explicit_legacy_status_does_not_send_a_request(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DURABLE_WORKFLOW_WORKER_PROTOCOL_VERSION", "1.19") + async with Client("https://runtime.example", namespace="billing") as client: + client._request = AsyncMock() # type: ignore[method-assign] + with pytest.raises(ValueError, match="protocol 1.20"): + await client.query_task_status(query_task_id="query-1", lease_owner="w1", query_task_attempt=1) + client._request.assert_not_called() diff --git a/tests/test_replay_regression_corpus.py b/tests/test_replay_regression_corpus.py index 144b6da..02217d7 100644 --- a/tests/test_replay_regression_corpus.py +++ b/tests/test_replay_regression_corpus.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import json from collections.abc import Mapping, Sequence from importlib import metadata @@ -45,6 +46,19 @@ FIXTURE_DIR = Path(__file__).parent / "fixtures" / "replay_regressions" +@workflow.defn(name="sdk.query-task-dispatch") +class QueryAuthorityWorkflow: + calls = 0 + + def run(self, ctx: WorkflowContext) -> int: + return 7 + + @workflow.query("current") + def current(self) -> int: + type(self).calls += 1 + return 7 + + @workflow.defn(name="tests.replay.parallel-metadata-producer") class ParallelMetadataProducerWorkflow: def run(self, ctx: WorkflowContext): # type: ignore[no-untyped-def] @@ -375,6 +389,27 @@ def _execute_fixture(fixture: dict[str, Any]) -> list[dict[str, Any]]: assert isinstance(workflow, dict) workflow_type = workflow.get("type") assert isinstance(workflow_type, str) and workflow_type + if workflow_type == "sdk.query-task-dispatch": + # The regression is a worker delivery, not an invented durable history. + assert "history" not in fixture + inputs = workflow.get("input") + assert isinstance(inputs, list) and len(inputs) == 1 and isinstance(inputs[0], dict) + client = AsyncMock(spec=Client) + client.namespace = "regression-corpus" + worker = Worker(client, task_queue="query-authority", worker_id="query-worker", + workflows=[QueryAuthorityWorkflow]) + task = {**inputs[0], "workflow_type": workflow_type, "payload_codec": "avro", + "workflow_arguments": serializer.envelope([], codec="avro"), + "query_arguments": serializer.envelope([], codec="avro"), "history_events": []} + QueryAuthorityWorkflow.calls = 0 + outcome = asyncio.run(worker._run_query_task(task)) + commands = [{"type": "query_task_dispatch", "outcome": outcome, + "handler_calls": QueryAuthorityWorkflow.calls, + "completion_calls": client.complete_query_task.await_count, + "failure_calls": client.fail_query_task.await_count}] + _assert_matches(fixture["command_sequence"], commands, f"{fixture['id']}.command_sequence") + _assert_matches(fixture["expected"], {"command_sequence": commands}, f"{fixture['id']}.expected") + return commands if workflow_type == "sdk.replay-conformance-report": # This public command executes the SDK's offline replay scenarios. Its # regression input is artifact declarations, not an invented history. @@ -487,7 +522,7 @@ def test_checked_in_replay_regression_corpus_uses_official_replayer( assert isinstance(workflow, dict) expected = fixture.get("expected") assert isinstance(expected, dict) and expected - if workflow.get("type") == "sdk.replay-conformance-report": + if workflow.get("type") in {"sdk.replay-conformance-report", "sdk.query-task-dispatch"}: _execute_fixture(fixture) return expected_error = expected.get("error") diff --git a/tests/test_storage_admission.py b/tests/test_storage_admission.py index 8542f04..63c4b82 100644 --- a/tests/test_storage_admission.py +++ b/tests/test_storage_admission.py @@ -109,7 +109,11 @@ def client_for(handler: Callable[..., Any], base_url: str = "https://runtime.exa async def test_poll_keeps_identity_after_ambiguous_response_and_pressure( kind: str, reason: str, entry: bool, retry_sleeps: list[float], ) -> None: + from tests.query_task_helpers import query_task_delivery + requests: list[bytes] = [] + response_task = (query_task_delivery({"query_task_id": "same-claim"}, worker_id="storage-worker", + task_queue="orders") if kind == "query" else {"task_id": "same-claim"}) def handler(request: httpx.Request) -> httpx.Response: requests.append(request.content) @@ -118,7 +122,7 @@ def handler(request: httpx.Request) -> httpx.Response: raise httpx.ReadTimeout("Response lost after a possible claim.", request=request) if len(requests) < 7: return httpx.Response(503, json=pressure(body["poll_request_id"], reason=reason, entry=entry)) - return httpx.Response(200, json={"task": {"task_id": "same-claim"}}) + return httpx.Response(200, json={"task": response_task}) async with client_for(handler) as client: with worker_scope(): @@ -130,7 +134,7 @@ def handler(request: httpx.Request) -> httpx.Response: else: poll = getattr(client, f"poll_{kind}_task") result = await poll(worker_id="storage-worker", task_queue="orders") - assert result == {"task_id": "same-claim"} + assert result == response_task assert len(requests) == 7 assert len(set(requests)) == 1 assert sum(retry_sleeps) == pytest.approx(5) @@ -140,7 +144,11 @@ def handler(request: httpx.Request) -> httpx.Response: async def test_backend_outage_preserves_logical_poll_after_transport_budget( kind: str, retry_sleeps: list[float], ) -> None: + from tests.query_task_helpers import query_task_delivery + requests: list[bytes] = [] + response_task = (query_task_delivery({"query_task_id": "reconciled-claim"}, worker_id="backend-worker", + task_queue="orders") if kind == "query" else {"task_id": "reconciled-claim"}) def handler(request: httpx.Request) -> httpx.Response: requests.append(request.content) @@ -148,7 +156,7 @@ def handler(request: httpx.Request) -> httpx.Response: raise httpx.ReadTimeout("Response lost after a possible claim.", request=request) if len(requests) <= 5: return httpx.Response(503, json=backend_unavailable(request)) - return httpx.Response(200, json={"task": {"task_id": "reconciled-claim"}}) + return httpx.Response(200, json={"task": response_task}) async with client_for(handler) as client: with worker_scope(): @@ -161,7 +169,7 @@ def handler(request: httpx.Request) -> httpx.Response: task = await getattr(client, f"poll_{kind}_task")( worker_id="backend-worker", task_queue="orders", ) - assert task == {"task_id": "reconciled-claim"} + assert task == response_task assert len(requests) == 6 assert len(set(requests)) == 1 assert sum(retry_sleeps) == pytest.approx(4) @@ -576,6 +584,8 @@ def handler(request: httpx.Request) -> httpx.Response: async def test_query_thread_uses_pressure_retries(retry_sleeps: list[float]) -> None: + from tests.query_task_helpers import leased_query_task + completed = threading.Event() calls = 0 claims = 0 @@ -598,12 +608,12 @@ def handler(request: httpx.Request) -> httpx.Response: if path.endswith("/query-tasks/poll"): claims += 1 if claims == 1: - return httpx.Response(200, json={"task": { + return httpx.Response(200, json={"task": leased_query_task(worker, { "query_task_id": "query-task", "query_task_attempt": 7, "workflow_type": "storage.query", "workflow_id": "workflow", "run_id": "run", "query_name": "status", "payload_codec": "avro", "history_events": [], "arguments": serializer.envelope([]), - }}) + })}) return httpx.Response(200, json={"task": None}) if path.endswith("/query-task/complete"): acks.append(request.content) diff --git a/tests/test_worker.py b/tests/test_worker.py index 871fbda..14cc2e9 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -46,6 +46,7 @@ _query_history_with_export_signal_arguments, _should_fail_workflow_task_after_completion_error, ) +from tests.query_task_helpers import leased_query_task class TypedCancelFlightError(Exception): @@ -2235,7 +2236,7 @@ async def test_query_task_executes_registered_query(self, mock_client: AsyncMock "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "completed" mock_client.complete_query_task.assert_awaited_once_with( @@ -2291,7 +2292,7 @@ async def test_query_task_replays_signal_arguments_from_history_export( "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "completed" mock_client.complete_query_task.assert_awaited_once_with( @@ -2339,7 +2340,7 @@ async def test_query_task_uses_export_history_when_inline_history_is_empty( "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "completed" mock_client.complete_query_task.assert_awaited_once_with( @@ -2426,7 +2427,7 @@ async def test_query_task_replays_repeated_condition_wait_signal_arguments( "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "completed" mock_client.complete_query_task.assert_awaited_once_with( @@ -2538,7 +2539,7 @@ async def test_query_task_replays_signal_woken_false_wait_reopens( "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "completed" mock_client.complete_query_task.assert_awaited_once_with( @@ -2611,7 +2612,7 @@ async def test_query_task_replays_history_from_export_after_worker_restart( "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "completed" mock_client.complete_query_task.assert_awaited_once_with( @@ -2711,7 +2712,7 @@ async def test_query_task_replaces_missing_history_payload_placeholders_from_exp "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "completed" mock_client.complete_query_task.assert_awaited_once_with( @@ -2767,7 +2768,7 @@ async def test_query_task_enriches_compact_activity_completion_from_export( "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "completed" mock_client.complete_query_task.assert_awaited_once_with( @@ -2800,7 +2801,7 @@ async def test_query_task_awaits_async_query_result(self, mock_client: AsyncMock "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "completed" mock_client.complete_query_task.assert_awaited_once_with( @@ -2829,7 +2830,7 @@ async def test_query_task_reports_unknown_query(self, mock_client: AsyncMock) -> "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "failed" mock_client.fail_query_task.assert_awaited_once() @@ -2852,7 +2853,7 @@ async def test_query_task_reports_state_unavailable_when_replay_fails(self, mock "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "failed" mock_client.fail_query_task.assert_awaited_once() @@ -2881,7 +2882,7 @@ async def test_query_task_completion_rejection_after_server_timeout_is_handled( "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "expired" mock_client.complete_query_task.assert_awaited_once() @@ -2904,7 +2905,7 @@ async def test_query_task_reports_query_result_completion_failure( "payload_codec": "avro", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "failed" mock_client.fail_query_task.assert_awaited_once() @@ -3585,7 +3586,7 @@ async def execute_query_task( interceptors=[Recorder()], ) - await worker._run_query_task( + await worker._run_query_task(leased_query_task(worker, { "query_task_id": "qt-intercept", "query_task_attempt": 1, @@ -3596,7 +3597,7 @@ async def execute_query_task( "query_arguments": serializer.envelope([], codec="avro"), "payload_codec": "avro", } - ) + )) assert events == ["before:qt-intercept", "after:completed"] @@ -3689,7 +3690,7 @@ async def execute_query_task( task["payload_codec"] = payload_codec if task_kind == "query": - assert await worker._run_query_task(task) == "failed" + assert await worker._run_query_task(leased_query_task(worker, task)) == "failed" failure = mock_client.fail_query_task else: assert await worker._run_workflow_task(task) is None @@ -3839,7 +3840,7 @@ async def test_task_root_codec_rejection_precedes_decode_and_user_work( activity_spy = AsyncMock(side_effect=AssertionError("activity handler must not run")) monkeypatch.setattr(serializer, "decode_envelope", decode_spy) monkeypatch.setattr(worker_module, "replay", replay_spy) - monkeypatch.setattr(worker_module, "query_state", query_spy) + monkeypatch.setattr(worker_module, "_prepare_query_state", query_spy) monkeypatch.setattr(worker_module, "apply_update", update_spy) monkeypatch.setattr(worker_module, "validate_update", validator_spy) monkeypatch.setattr(worker, "_execute_activity_callable", activity_spy) @@ -3849,7 +3850,7 @@ async def test_task_root_codec_rejection_precedes_decode_and_user_work( failure = mock_client.fail_activity_task assert outcome == "decode_error" elif task_kind == "query": - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) failure = mock_client.fail_query_task assert outcome == "failed" elif task_kind == "update_validation": @@ -3968,13 +3969,13 @@ async def test_exact_avro_task_completes_each_path_with_handler_work( decode_spy = Mock(wraps=serializer.decode_envelope) replay_spy = Mock(wraps=worker_module.replay) - query_spy = Mock(wraps=worker_module.query_state) + query_spy = Mock(wraps=worker_module._prepare_query_state) update_spy = Mock(wraps=worker_module.apply_update) validator_spy = Mock(wraps=worker_module.validate_update) activity_spy = AsyncMock(wraps=worker._execute_activity_callable) monkeypatch.setattr(serializer, "decode_envelope", decode_spy) monkeypatch.setattr(worker_module, "replay", replay_spy) - monkeypatch.setattr(worker_module, "query_state", query_spy) + monkeypatch.setattr(worker_module, "_prepare_query_state", query_spy) monkeypatch.setattr(worker_module, "apply_update", update_spy) monkeypatch.setattr(worker_module, "validate_update", validator_spy) monkeypatch.setattr(worker, "_execute_activity_callable", activity_spy) @@ -3985,7 +3986,7 @@ async def test_exact_avro_task_completes_each_path_with_handler_work( mock_client.fail_activity_task.assert_not_awaited() activity_spy.assert_awaited_once() elif task_kind == "query": - assert await worker._run_query_task(task) == "completed" + assert await worker._run_query_task(leased_query_task(worker, task)) == "completed" mock_client.complete_query_task.assert_awaited_once() mock_client.fail_query_task.assert_not_awaited() query_spy.assert_called_once() @@ -4156,7 +4157,7 @@ async def test_query_unsupported_payload_codec_fails_before_query_handler( "payload_codec": "zstd", } - outcome = await worker._run_query_task(task) + outcome = await worker._run_query_task(leased_query_task(worker, task)) assert outcome == "failed" mock_client.fail_query_task.assert_called_once() @@ -4677,7 +4678,7 @@ async def poll_query_task(self, **_: object) -> dict[str, object] | None: await asyncio.Event().wait() except asyncio.CancelledError: events.append("query_returned_during_shutdown") - return query_task + return leased_query_task(worker, query_task) async def complete_query_task(self, **_: object) -> dict[str, str]: events.append("query_dispatched") @@ -4864,7 +4865,7 @@ async def __aexit__(self, *_: object) -> None: async def poll_query_task(self, **_: object) -> dict[str, object] | None: if not self.polled: self.polled = True - return query_task + return leased_query_task(worker, query_task) await asyncio.sleep(0.01) return None @@ -5344,7 +5345,7 @@ async def poll_query_task(**_: object) -> dict[str, object] | None: nonlocal poll_count poll_count += 1 if poll_count == 1: - return query_task + return leased_query_task(worker, query_task) await asyncio.sleep(0) return None From 778f7c01fddb143be613af43dc62057e249ec2c5 Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Thu, 15 Oct 2026 00:08:14 +0000 Subject: [PATCH 2/4] Prepare compatible query authority patch 2.5.14 --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0158f0c..bf96aaf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "durable-workflow" -version = "2.5.13" +version = "2.5.14" description = "Python client and worker SDK for Durable Workflow Cloud and self-hosted Server" readme = "README.md" requires-python = ">=3.10" @@ -71,8 +71,8 @@ durable-workflow-replay-conformance = "durable_workflow.replay_conformance:main" durable-workflow-workflow-updates-conformance = "durable_workflow.workflow_updates_conformance:main" [tool.durable-workflow] -product-train = "2.5.13" -registry-version = "2.5.13" +product-train = "2.5.14" +registry-version = "2.5.14" supported-server-versions = "2.5.16" worker-protocol-version = "1.20" control-plane-version = "2" From 21e1f268ccc38d71551e12f9d6946be41604c10f Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Thu, 15 Oct 2026 00:11:08 +0000 Subject: [PATCH 3/4] Bound query transport discovery and preserve typed attempt identity --- src/durable_workflow/_query_task_authority.py | 3 +- src/durable_workflow/client.py | 3 +- tests/test_query_task_authority.py | 48 +++++++++++++++++++ 3 files changed, 52 insertions(+), 2 deletions(-) diff --git a/src/durable_workflow/_query_task_authority.py b/src/durable_workflow/_query_task_authority.py index 24ea377..7d04665 100644 --- a/src/durable_workflow/_query_task_authority.py +++ b/src/durable_workflow/_query_task_authority.py @@ -75,7 +75,8 @@ def delivery(cls, client: object, namespace: str, task: Mapping[str, Any], *, def matches(self, client: object, namespace: str, task: Mapping[str, Any], owner: str) -> bool: return (self.client is client and self.namespace == namespace and self.owner == owner - and task.get("query_task_id") == self.task_id and task.get("query_task_attempt") == self.attempt + and task.get("query_task_id") == self.task_id + and type(task.get("query_task_attempt")) is int and task.get("query_task_attempt") == self.attempt and task.get("lease_owner") == owner and all(task.get(key) == value for key, value in self.identity.items()) and all(task.get(key) == value for key, value in self.timing.items())) diff --git a/src/durable_workflow/client.py b/src/durable_workflow/client.py index e3cc232..e1273dd 100644 --- a/src/durable_workflow/client.py +++ b/src/durable_workflow/client.py @@ -1881,7 +1881,8 @@ async def _do_request() -> httpx.Response: query_authority = await self._query_delivery_authority(result["task"], request_json) current_query_authority.set(query_authority) if not discovery and _contains_runtime_payload_envelope(result): - transport = await self._runtime_external_payload_transport() + transport = (await query_authority.run(self._runtime_external_payload_transport) + if query_authority is not None else await self._runtime_external_payload_transport()) if transport is None: raise ExternalPayloadUnsupported( "runtime returned an external payload reference without advertising " diff --git a/tests/test_query_task_authority.py b/tests/test_query_task_authority.py index 1df2727..5b45cfc 100644 --- a/tests/test_query_task_authority.py +++ b/tests/test_query_task_authority.py @@ -159,6 +159,54 @@ async def dispatched() -> QueryTaskAuthority | None: current_query_authority.reset(token) +async def test_boolean_attempt_mutation_cannot_invoke_query_handler() -> None: + client = AsyncMock(spec=Client) + client.namespace = "billing" + task = delivery() + budget = authority(task, client=client, observed=False) + worker = Worker(client, worker_id="w1", task_queue="q1", workflows=[QueryWorkflow]) + task["query_task_attempt"] = True + QueryWorkflow.calls = 0 + token = current_query_authority.set(budget) + try: + assert await worker._run_query_task(task) == "expired" + finally: + current_query_authority.reset(token) + assert QueryWorkflow.calls == 0 + client.complete_query_task.assert_not_called() + client.fail_query_task.assert_not_called() + + +async def test_query_transport_discovery_uses_original_loading_budget() -> None: + from tests.test_runtime_external_payload_transport import FakeRuntimePayloadServer, runtime_client, runtime_envelope + + server = FakeRuntimePayloadServer() + task = delivery(workflow_arguments=runtime_envelope(server, ["large"])) + task["lease_expires_at"] = timestamp(datetime.now(timezone.utc) + timedelta(seconds=0.1)) + server.responses[("POST", "/api/worker/query-tasks/poll")] = {"task": task} + cancelled = asyncio.Event() + + async def handle(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/cluster/info"): + try: + await asyncio.sleep(60) + finally: + cancelled.set() + return server.handler(request) + + async with runtime_client(server) as client: + await client._http.aclose() + client._http = httpx.AsyncClient(base_url=client.base_url, transport=httpx.MockTransport(handle)) + token = current_query_authority.set(None) + try: + with pytest.raises(QueryTaskAuthorityLost, match="budget_expired"): + await asyncio.wait_for(client.poll_query_task(worker_id="w1", task_queue="q1"), timeout=1) + finally: + current_query_authority.reset(token) + assert cancelled.is_set() + assert server.fetch_count == 0 + + async def test_query_loading_reuses_references_without_enlarging_persistent_cache() -> None: from durable_workflow.external_storage import ExternalPayloadCache from tests.test_runtime_external_payload_transport import FakeRuntimePayloadServer, runtime_client, runtime_envelope From 1bcf50b48c17ee75e676abe7f6f86cf331425c96 Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Thu, 15 Oct 2026 00:54:15 +0000 Subject: [PATCH 4/4] Qualify query authority against published Server 2.8.0 --- docs/query-task-authority.md | 7 +- pyproject.toml | 2 +- .../integration/test_query_task_authority.py | 163 ++++++++++++++++++ 3 files changed, 169 insertions(+), 3 deletions(-) create mode 100644 tests/integration/test_query_task_authority.py diff --git a/docs/query-task-authority.md b/docs/query-task-authority.md index 3aec4b9..d068777 100644 --- a/docs/query-task-authority.md +++ b/docs/query-task-authority.md @@ -33,5 +33,8 @@ Keep query handlers short and free of side effects. Expiry can cancel async execution and payload I/O, but an observation never reserves future commit authority. A result refused by Server does not restart the query or its budget. -This change is undergoing source and published-runtime qualification under -[#143](https://github.com/durable-workflow/sdk-python/issues/143). +The non-renewing authority path is qualified against Server 2.8.0, including +expiry during external payload loading and async handler execution. Earlier +Server 2.6.3 remains compatible in default 1.20 and explicit 1.19 mode, with the +legacy limits described above. Qualification and release follow-through are +recorded in [#143](https://github.com/durable-workflow/sdk-python/issues/143). diff --git a/pyproject.toml b/pyproject.toml index bf96aaf..a0b2ef1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ durable-workflow-workflow-updates-conformance = "durable_workflow.workflow_updat [tool.durable-workflow] product-train = "2.5.14" registry-version = "2.5.14" -supported-server-versions = "2.5.16" +supported-server-versions = "2.8.0" worker-protocol-version = "1.20" control-plane-version = "2" durable-selection = true diff --git a/tests/integration/test_query_task_authority.py b/tests/integration/test_query_task_authority.py new file mode 100644 index 0000000..7594476 --- /dev/null +++ b/tests/integration/test_query_task_authority.py @@ -0,0 +1,163 @@ +"""Actual query authority with a locally configured, isolated Server. + +Run with DW_QUERY_TASK_TIMEOUT=4 on Server and +DURABLE_WORKFLOW_QUERY_AUTHORITY_QUALIFICATION=1 on pytest. The token must be +allowed to create disposable namespaces. Payload delay is injected after real +HTTP response headers; responses, claims and bytes are not rewritten. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import threading +import uuid +from datetime import datetime +from typing import Any + +import httpx +import pytest + +from durable_workflow import Client, Worker, workflow +from durable_workflow.errors import QueryFailed +from durable_workflow.external_storage import ExternalPayloadCache + +pytestmark = pytest.mark.skipif( + os.environ.get("DURABLE_WORKFLOW_QUERY_AUTHORITY_QUALIFICATION") != "1", + reason="requires an isolated query-authority Server with a four-second original query budget", +) + + +@workflow.defn(name="tests.query-authority.http") +class QueryWorkflow: + calls = 0 + cancelled = 0 + returned = 0 + + def run(self, ctx: Any, value: Any) -> Any: + self.value = value + yield ctx.sleep(60) + return 7 + + @workflow.query("value") + async def current(self, slow: bool = False) -> Any: + type(self).calls += 1 + assert self.value == {"body": "x" * 65536} + try: + if slow: + await asyncio.sleep(60) + type(self).returned += 1 + return {"long": 9223372036854775807, "nil": None, "text": "λ"} + except asyncio.CancelledError: + type(self).cancelled += 1 + raise + + +class ObservedWorker(Worker): + """Keep normal Worker.run/client cloning and observe its real HTTP path.""" + + def __init__(self, *args: Any, delay_payload: bool, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.delay_payload = delay_payload + self.claims: list[dict[str, Any]] = [] + self.status_reads = 0 + self.payload_reads = 0 + self.completions = 0 + self.payload_cancelled = threading.Event() + + def _clone_client_for_query_tasks(self) -> Client: + client = super()._clone_client_for_query_tasks() + + async def observe(response: httpx.Response) -> None: + path = response.request.url.path + if path.endswith("/query-tasks/poll") and response.status_code == 200: + await response.aread() + task = response.json().get("task") + if isinstance(task, dict): + admitted = datetime.fromisoformat(task["query_admitted_at"].replace("Z", "+00:00")) + deadline = datetime.fromisoformat(task["query_deadline_at"].replace("Z", "+00:00")) + assert (deadline - admitted).total_seconds() == 4, "configure the original Server budget" + self.claims.append(task) + elif "/query-tasks/" in path and path.endswith("/status"): + self.status_reads += 1 + elif "/query-tasks/" in path and path.endswith("/complete"): + self.completions += 1 + elif response.request.method == "GET" and "/external-payloads/v1/" in path: + self.payload_reads += 1 + if self.delay_payload: + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + self.payload_cancelled.set() + raise + + client._http.event_hooks["response"].append(observe) + return client + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["timely", "loading-expiry", "handler-expiry"]) +async def test_installed_query_authority_http(server_url: str, server_token: str, mode: str) -> None: + suffix = uuid.uuid4().hex[:12] + namespace = f"query-authority-{suffix}" + QueryWorkflow.calls = QueryWorkflow.cancelled = QueryWorkflow.returned = 0 + async with Client(server_url, token=server_token) as admin: + await admin.create_namespace(namespace) + try: + await admin.set_namespace_external_storage(namespace, driver="local", threshold_bytes=1024) + async with Client(server_url, token=server_token, namespace=namespace, + external_storage_cache=ExternalPayloadCache(max_bytes=1)) as client: + info = await client.get_cluster_info() + assert info["worker_protocol"]["server_capabilities"]["query_task_authority"] is True + assert info["namespace"]["external_payload_storage"]["status"] == "available" + worker = ObservedWorker(client, workflows=[QueryWorkflow], task_queue=namespace, + poll_timeout=0.5, shutdown_timeout=10, require_query_task_authority=True, + delay_payload=mode == "loading-expiry") + runner = asyncio.create_task(worker.run()) + handle = None + try: + await asyncio.wait_for(worker._registration_done.wait(), timeout=10) + assert not runner.done(), runner.exception() if runner.done() else "" + handle = await client.start_workflow(workflow_type="tests.query-authority.http", + task_queue=namespace, workflow_id=namespace, input=[{"body": "x" * 65536}]) + for _ in range(100): + if (await handle.describe()).status.lower() == "waiting": + break + await asyncio.sleep(0.05) + else: + raise AssertionError("original workflow did not reach waiting") + before = (await handle.get_history())["events"] + if mode == "timely": + reply = await client.query_workflow(handle.workflow_id, "value", args=[False]) + assert reply["result"] == {"long": 9223372036854775807, "nil": None, "text": "λ"} + assert type(reply["result"]["long"]) is int + else: + with pytest.raises(QueryFailed) as refused: + await client.query_workflow(handle.workflow_id, "value", args=[mode == "handler-expiry"]) + assert refused.value.status == 504 + assert refused.value.reason == "query_worker_execution_timeout" + for _ in range(100): + if worker.payload_cancelled.is_set() or QueryWorkflow.cancelled: + break + await asyncio.sleep(0.05) + assert (await handle.get_history())["events"] == before + assert len({task["query_task_id"] for task in worker.claims}) == 1 + assert worker.status_reads > 0 + assert worker.payload_reads == 1 + assert worker.completions == (1 if mode == "timely" else 0) + assert QueryWorkflow.calls == (0 if mode == "loading-expiry" else 1) + assert QueryWorkflow.returned == (1 if mode == "timely" else 0) + assert QueryWorkflow.cancelled == (1 if mode == "handler-expiry" else 0) + assert worker.payload_cancelled.is_set() is (mode == "loading-expiry") + print(json.dumps({"outcome": "pass", "mode": mode, "server": info.get("version"), + "workflow_id": handle.workflow_id, "run_id": handle.run_id, + "query_task_id": worker.claims[0]["query_task_id"], "status_reads": worker.status_reads, + "payload_reads": worker.payload_reads, "completions": worker.completions}), flush=True) + finally: + await worker.stop() + await asyncio.wait_for(runner, timeout=10) + if handle is not None: + await client.terminate_workflow(handle.workflow_id, reason="qualification complete") + finally: + await admin.delete_namespace(namespace)