durable-workflow 2.0.2__tar.gz → 2.0.3__tar.gz

{durable_workflow-2.0.2/src/durable_workflow.egg-info → durable_workflow-2.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durable-workflow
-Version: 2.0.2
+Version: 2.0.3
 Summary: Python client and worker SDK for Durable Workflow Cloud and self-hosted Server
 Author: Durable Workflow Contributors
 License-Expression: MIT

{durable_workflow-2.0.2 → durable_workflow-2.0.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "durable-workflow"
-version = "2.0.2"
+version = "2.0.3"
 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.0.2"
-registry-version = "2.0.2"
+product-train = "2.0.3"
+registry-version = "2.0.3"
 supported-server-versions = "2.0.0"
 worker-protocol-version = "1.19"
 control-plane-version = "2"

{durable_workflow-2.0.2 → durable_workflow-2.0.3}/src/durable_workflow/errors.py

@@ -171,6 +171,32 @@ class ServerError(DurableWorkflowError):
             return self.body.get("reason")
         return None
 
+    def is_storage_admission_failure(self, poll_request_id: str | None = None) -> bool:
+        """Whether the runtime explicitly refused admission and requested an identity-preserving retry."""
+        body = self.body
+        if (
+            self.status != 503
+            or not isinstance(body, dict)
+            or self.reason() not in ("storage_pressure", "storage_admission_unavailable")
+            or body.get("retryable") is not True
+            or type(body.get("retry_after_seconds")) is not int
+            or body["retry_after_seconds"] <= 0
+            or body.get("storage_state") not in ("draining", "fenced")
+            or (self.reason() == "storage_admission_unavailable" and body["storage_state"] != "fenced")
+            or ("request_admitted" in body and body["request_admitted"] is not False)
+        ):
+            return False
+        if poll_request_id is None:
+            return body.get("request_admitted") is False
+        return (
+            bool(poll_request_id)
+            and "task" in body and body["task"] is None
+            and body.get("poll_status") == self.reason()
+            and body.get("poll_request_id") == poll_request_id
+            and body.get("retry_same_poll_request_id") is True
+            and body.get("claim_admitted") is False
+        )
+
 
 class NexusOperationFailed(DurableWorkflowError):
     """A Nexus service operation completed with a typed service failure."""

{durable_workflow-2.0.2 → durable_workflow-2.0.3}/src/durable_workflow/retry_policy.py

@@ -13,6 +13,9 @@
 from __future__ import annotations
 
 import asyncio
+import contextvars
+import json
+import logging
 import random
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
@@ -20,7 +23,40 @@ from typing import TypeVar
 
 import httpx
 
+from .errors import ServerError
+
 T = TypeVar("T")
+log = logging.getLogger("durable_workflow.worker")
+
+# Task-local so sharing a Client never changes unrelated client/control requests.
+_worker_storage_admission_stop: contextvars.ContextVar[Callable[[], bool] | None] = contextvars.ContextVar(
+    "worker_storage_admission_stop", default=None,
+)
+
+
+def _storage_refusal(exc: Exception) -> tuple[ServerError, str | None] | None:
+    if not isinstance(exc, httpx.HTTPStatusError):
+        return None
+    if "X-Durable-Workflow-Protocol-Version" not in exc.request.headers:
+        return None
+    try:
+        body = exc.response.json()
+    except ValueError:
+        return None
+    error = ServerError(exc.response.status_code, body)
+    if error.reason() not in ("storage_pressure", "storage_admission_unavailable"):
+        return None
+    poll_id = None
+    if exc.request.url.path.endswith("/poll"):
+        try:
+            request = json.loads(exc.request.content)
+            poll_id = request.get("poll_request_id") if isinstance(request, dict) else None
+        except ValueError:
+            pass
+        # An invalid submitted ID must not fall through to the non-poll contract.
+        if not isinstance(poll_id, str) or not poll_id:
+            poll_id = ""
+    return error, poll_id
 
 
 @dataclass
@@ -81,6 +117,7 @@ class TransportRetryPolicy:
         Raises the last exception if all retries are exhausted.
         """
         attempt = 0
+        storage_attempt = 0
         last_exc: Exception | None = None
 
         while attempt < self.max_attempts:
@@ -89,6 +126,32 @@ class TransportRetryPolicy:
                 return result
             except Exception as exc:
                 last_exc = exc
+                stop = _worker_storage_admission_stop.get()
+                refusal = _storage_refusal(exc) if stop is not None else None
+                if refusal is not None and stop is not None:
+                    error, poll_id = refusal
+                    if not error.is_storage_admission_failure(poll_id) or stop():
+                        raise
+                    storage_attempt += 1
+                    assert isinstance(error.body, dict)
+                    delay = min(
+                        5.0,
+                        max(
+                            self.backoff_seconds(min(storage_attempt - 1, 6)),
+                            error.body["retry_after_seconds"],
+                        ),
+                    )
+                    log.warning("storage admission paused; retrying the same worker request in %.2fs", delay)
+                    # Do not consume the finite transport budget or repeat serialization/uploads.
+                    while delay > 0:
+                        if stop():
+                            raise
+                        interval = min(0.1, delay)
+                        await asyncio.sleep(interval)
+                        delay -= interval
+                    if stop():
+                        raise
+                    continue
                 if not self.should_retry(exc, attempt):
                     raise
 

{durable_workflow-2.0.2 → durable_workflow-2.0.3}/src/durable_workflow/worker.py

@@ -27,10 +27,11 @@ import time
 import traceback
 import types
 import uuid
-from collections.abc import Awaitable, Callable, Iterable, Mapping
+from collections.abc import Awaitable, Callable, Coroutine, Iterable, Mapping
 from datetime import datetime, timezone
+from functools import wraps
 from types import FunctionType
-from typing import Annotated, Any, Literal, Union, get_args, get_origin, get_type_hints
+from typing import Annotated, Any, Concatenate, Literal, ParamSpec, TypeVar, Union, get_args, get_origin, get_type_hints
 
 from . import serializer
 from .activity import ActivityContext, ActivityInfo, _set_context
@@ -75,6 +76,7 @@ from .metrics import (
     WORKER_TASKS,
     MetricsRecorder,
 )
+from .retry_policy import _worker_storage_admission_stop
 from .workflow import (
     Command,
     NexusServiceCall,
@@ -137,6 +139,27 @@ _WORKFLOW_TASK_COMPLETION_MAX_ATTEMPTS = 3
 _WORKFLOW_TASK_COMPLETION_RETRY_DELAYS = (0.05, 0.2)
 _WORKFLOW_TASK_NEXUS_RESOLUTION_LIMIT = 100
 _WORKER_WORKFLOW_FINGERPRINTS: dict[tuple[str, str], str] = {}
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
+def _with_storage_admission_retries(
+    fn: Callable[Concatenate[Worker, _P], Coroutine[Any, Any, _R]],
+) -> Callable[Concatenate[Worker, _P], Coroutine[Any, Any, _R]]:
+    @wraps(fn)
+    async def run(self: Worker, /, *args: _P.args, **kwargs: _P.kwargs) -> _R:
+        token = _worker_storage_admission_stop.set(self._stop.is_set)
+        try:
+            return await fn(self, *args, **kwargs)
+        finally:
+            _worker_storage_admission_stop.reset(token)
+    return run
+
+
+def _is_storage_admission_error(error: BaseException) -> bool:
+    return isinstance(error, ServerError) and error.reason() in (
+        "storage_pressure", "storage_admission_unavailable",
+    )
 
 
 def _command_payload_codec(codec: object) -> str:
@@ -208,6 +231,8 @@ def _should_fail_workflow_task_after_completion_error(error: BaseException) -> b
 
 
 def _should_retry_workflow_task_completion_error(error: BaseException) -> bool:
+    if _is_storage_admission_error(error):
+        return False
     if isinstance(error, ServerError):
         return error.status >= 500 or error.status == 429
 
@@ -1424,6 +1449,8 @@ class Worker:
                 )
             except Exception as e:
                 log.warning("failed to complete workflow update task %s: %s", task_id, e)
+                if _is_storage_admission_error(e):
+                    return None
                 if _should_fail_workflow_task_after_completion_error(e):
                     await self._report_workflow_task_after_completion_error(task_id, attempt, e)
                     return None
@@ -1550,6 +1577,8 @@ class Worker:
             )
         except Exception as e:
             log.warning("failed to complete workflow task %s: %s", task_id, e)
+            if _is_storage_admission_error(e):
+                return None
             if _should_fail_workflow_task_after_completion_error(e):
                 await self._report_workflow_task_after_completion_error(task_id, attempt, e)
                 return None
@@ -1736,6 +1765,8 @@ class Worker:
                 log.warning("failed to report activity failure: %s", fe)
             return "failed_non_retryable"
         except Exception as e:
+            if _is_storage_admission_error(e):
+                raise
             log.exception("activity failed")
             try:
                 await self.client.fail_activity_task(
@@ -1971,6 +2002,9 @@ class Worker:
                 **self._external_storage_completion_kwargs(),
             )
         except ServerError as e:
+            if _is_storage_admission_error(e):
+                log.warning("query task %s acknowledgement paused: %s", query_task_id, e)
+                return "complete_error"
             if _is_final_query_task_rejection(e):
                 log.info(
                     "query task %s completion was rejected after the task ended server-side: %s",
@@ -2147,6 +2181,8 @@ class Worker:
                 self._release_workflow_capacity()
                 if self._stop.is_set():
                     return
+                if _is_storage_admission_error(e):
+                    raise
                 self._record_poll_metrics("workflow", "error", time.perf_counter() - poll_start)
                 log.warning("workflow poll error: %s", e)
                 await asyncio.sleep(1.0)
@@ -2236,10 +2272,15 @@ class Worker:
                     timeout=self._poll_http_timeout,
                     build_id=self.build_id,
                 )
+            except asyncio.CancelledError:
+                self._act_semaphore.release()
+                raise
             except Exception as e:
                 self._act_semaphore.release()
                 if self._stop.is_set():
                     return
+                if _is_storage_admission_error(e):
+                    raise
                 self._record_poll_metrics("activity", "error", time.perf_counter() - poll_start)
                 log.warning("activity poll error: %s", e)
                 await asyncio.sleep(1.0)
@@ -2289,6 +2330,8 @@ class Worker:
                     query_thread_stop is not None and query_thread_stop.is_set()
                 ):
                     return
+                if _is_storage_admission_error(e):
+                    raise
                 self._record_poll_metrics("query", "error", time.perf_counter() - poll_start)
                 log.warning("query poll error: %s", e)
                 await asyncio.sleep(1.0)
@@ -2525,6 +2568,7 @@ class Worker:
         except Exception:
             log.exception("query task poller thread stopped unexpectedly")
 
+    @_with_storage_admission_retries
     async def _query_task_thread_main(self) -> None:
         loop = asyncio.get_running_loop()
         task = asyncio.current_task()
@@ -2562,6 +2606,7 @@ class Worker:
                 "the worker registration remains active"
             )
 
+    @_with_storage_admission_retries
     async def run(self) -> None:
         """Register the worker and poll until `stop()` is called or the task is cancelled."""
         self._begin_run()
@@ -2622,6 +2667,8 @@ class Worker:
                     process_metrics=self._current_process_metrics(),
                 )
             except Exception as e:
+                if _is_storage_admission_error(e) and not self._stop.is_set():
+                    raise
                 log.warning("worker heartbeat failed: %s", e)
                 continue
             if isinstance(ack, dict):
@@ -2703,6 +2750,7 @@ class Worker:
 
         return metrics
 
+    @_with_storage_admission_retries
     async def run_until(
         self,
         *,
@@ -2714,17 +2762,19 @@ class Worker:
 
         This is intended for examples, smoke tests, and single-workflow scripts.
         Long-running workers should call :meth:`run` and coordinate shutdown from
-        their process supervisor.
+        their process supervisor. ``timeout`` includes registration and runtime
+        admission pauses; shutdown retains its separate drain timeout.
         """
         background_tasks: list[asyncio.Task[Any]] = []
+        deadline = asyncio.get_running_loop().time() + timeout
 
         self._begin_run()
         try:
-            await self._register()
-        finally:
-            self._registration_done.set()
+            try:
+                await asyncio.wait_for(self._register(), timeout=timeout)
+            finally:
+                self._registration_done.set()
 
-        try:
             if self._stop.is_set():
                 raise asyncio.CancelledError
             background_tasks.append(asyncio.create_task(self._heartbeat_loop()))
@@ -2742,7 +2792,11 @@ class Worker:
                 )
             )
             self._poller_tasks.add(run_until_loop)
-            return await run_until_loop
+            return await asyncio.wait_for(
+                run_until_loop, timeout=max(0.0, deadline - asyncio.get_running_loop().time()),
+            )
+        except asyncio.TimeoutError as error:
+            raise TimeoutError(f"workflow {workflow_id} not terminal after {timeout}s") from error
         finally:
             primary_error = sys.exc_info()[1]
             primary_traceback = primary_error.__traceback__ if primary_error is not None else None

{durable_workflow-2.0.2 → durable_workflow-2.0.3/src/durable_workflow.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durable-workflow
-Version: 2.0.2
+Version: 2.0.3
 Summary: Python client and worker SDK for Durable Workflow Cloud and self-hosted Server
 Author: Durable Workflow Contributors
 License-Expression: MIT

{durable_workflow-2.0.2 → durable_workflow-2.0.3}/src/durable_workflow.egg-info/SOURCES.txt

@@ -73,6 +73,7 @@ tests/test_serializer.py
 tests/test_signals.py
 tests/test_sleep.py
 tests/test_standalone_activity_client.py
+tests/test_storage_admission.py
 tests/test_sync.py
 tests/test_testing_harness.py
 tests/test_update_signal_condition_replay.py

{durable_workflow-2.0.2 → durable_workflow-2.0.3}/tests/test_release_metadata.py

@@ -41,9 +41,9 @@ def test_worker_release_identity_matches_supported_server_and_protocol() -> None
     project = manifest["project"]
     release = manifest["tool"]["durable-workflow"]
 
-    assert project["version"] == "2.0.2"
+    assert project["version"] == "2.0.3"
     assert release["product-train"] == project["version"]
-    assert release["registry-version"] == "2.0.2"
+    assert release["registry-version"] == "2.0.3"
     assert release["supported-server-versions"] == "2.0.0"
     assert release["worker-protocol-version"] == PROTOCOL_VERSION == "1.19"
     assert release["durable-selection"] is True

durable_workflow-2.0.3/tests/test_storage_admission.py

@@ -0,0 +1,590 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import threading
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import httpx
+import pytest
+
+import durable_workflow.retry_policy as retry_module
+from durable_workflow import activity, serializer, workflow
+from durable_workflow.client import Client
+from durable_workflow.errors import ServerError, Unauthorized
+from durable_workflow.retry_policy import TransportRetryPolicy, _worker_storage_admission_stop
+from durable_workflow.worker import PORTABLE_WORKER_AFFINITY_CAPABILITY_MANIFEST, Worker
+from tests.test_runtime_external_payload_transport import FakeRuntimePayloadServer, runtime_client
+from tests.test_worker import ValidatedUpdateWorkflow, compatible_cluster_info
+
+
+def pressure(poll_id: str | None = None, *, reason: str = "storage_pressure", entry: bool = True) -> dict[str, Any]:
+    body: dict[str, Any] = {
+        "reason": reason,
+        "storage_state": "draining" if reason == "storage_pressure" else "fenced",
+        "retryable": True,
+        "retry_after_seconds": 1,
+    }
+    if entry:
+        body["request_admitted"] = False
+    if poll_id is not None:
+        body.update({
+            "task": None, "poll_status": reason, "poll_request_id": poll_id,
+            "retry_same_poll_request_id": True, "claim_admitted": False,
+        })
+    return body
+
+
+@contextmanager
+def worker_scope(stop: Callable[[], bool] = lambda: False) -> Iterator[None]:
+    token = _worker_storage_admission_stop.set(stop)
+    try:
+        yield
+    finally:
+        _worker_storage_admission_stop.reset(token)
+
+
+@pytest.fixture
+def retry_sleeps(monkeypatch: pytest.MonkeyPatch) -> list[float]:
+    sleeps: list[float] = []
+
+    async def sleep(delay: float) -> None:
+        sleeps.append(delay)
+        assert len(sleeps) < 1000, "The simulated storage outage must recover or be interrupted."
+        await asyncio.sleep(0)
+
+    monkeypatch.setattr(retry_module, "asyncio", SimpleNamespace(sleep=sleep))
+    return sleeps
+
+
+def client_for(handler: Callable[..., Any]) -> Client:
+    client = Client(
+        "https://runtime.example", token="test-runtime-token",
+        retry_policy=TransportRetryPolicy(max_attempts=2, initial_backoff_seconds=0, jitter=False),
+    )
+    client._http = httpx.AsyncClient(base_url=client.base_url, transport=httpx.MockTransport(handler))
+    client._runtime_external_payload_transport_resolved = True
+    return client
+
+
+@pytest.mark.parametrize("kind", ["workflow", "activity", "query", "multiplexed"])
+@pytest.mark.parametrize("reason", ["storage_pressure", "storage_admission_unavailable"])
+@pytest.mark.parametrize("entry", [True, False])
+async def test_poll_keeps_identity_after_ambiguous_response_and_pressure(
+    kind: str, reason: str, entry: bool, retry_sleeps: list[float],
+) -> None:
+    requests: list[bytes] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        requests.append(request.content)
+        body = json.loads(request.content)
+        if len(requests) == 1:
+            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"}})
+
+    async with client_for(handler) as client:
+        with worker_scope():
+            if kind == "multiplexed":
+                result = await client.poll_workflow_task(
+                    worker_id="storage-worker", task_queue="orders",
+                    task_kinds=("workflow", "update_validation"),
+                )
+            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 len(requests) == 7
+    assert len(set(requests)) == 1
+    assert sum(retry_sleeps) == pytest.approx(5)
+
+
+@pytest.mark.parametrize("method,kwargs", [
+    ("register_worker", {
+        "worker_id": "storage-worker", "task_queue": "orders",
+        "capability_manifest": PORTABLE_WORKER_AFFINITY_CAPABILITY_MANIFEST,
+    }),
+    ("heartbeat_worker", {"worker_id": "storage-worker"}),
+    ("complete_workflow_task", {"task_id": "task", "workflow_task_attempt": 7, "commands": []}),
+    ("fail_workflow_task", {"task_id": "task", "workflow_task_attempt": 7, "message": "failed"}),
+    ("complete_activity_task", {"task_id": "task", "activity_attempt_id": "attempt-7", "result": "done"}),
+    ("fail_activity_task", {"task_id": "task", "activity_attempt_id": "attempt-7", "message": "failed"}),
+    ("heartbeat_activity_task", {"task_id": "task", "activity_attempt_id": "attempt-7", "details": {"step": 2}}),
+    ("complete_query_task", {"query_task_id": "task", "query_task_attempt": 7, "result": "done"}),
+    ("fail_query_task", {"query_task_id": "task", "query_task_attempt": 7, "message": "failed"}),
+    ("approve_update_validation_task", {"update_validation_task_id": "task", "update_validation_attempt": 7}),
+    ("reject_update_validation_task", {
+        "update_validation_task_id": "task", "update_validation_attempt": 7, "message": "rejected",
+        "reason": "update_rejected",
+    }),
+])
+async def test_worker_mutations_retry_the_prepared_request(
+    method: str, kwargs: dict[str, Any], retry_sleeps: list[float],
+) -> None:
+    requests: list[bytes] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        requests.append(request.content)
+        if len(requests) <= 4:
+            return httpx.Response(503, json=pressure(reason="storage_admission_unavailable"))
+        return httpx.Response(200, json={"recorded": True})
+
+    if method not in ("register_worker", "heartbeat_worker"):
+        kwargs = {**kwargs, "lease_owner": "storage-worker"}
+    async with client_for(handler) as client:
+        with worker_scope():
+            await getattr(client, method)(**kwargs)
+    assert len(requests) == 5
+    assert len(set(requests)) == 1
+    assert sum(retry_sleeps) == pytest.approx(4)
+
+
+@pytest.mark.parametrize("override", [
+    {"task": {"task_id": "already-claimed"}}, {"poll_request_id": "another-poll"},
+    {"poll_status": "empty"}, {"claim_admitted": True}, {"retry_same_poll_request_id": False},
+    {"request_admitted": True}, {"retryable": False}, {"storage_state": "normal"},
+    {"retry_after_seconds": 0}, {"retry_after_seconds": True}, {"retry_after_seconds": "1"},
+    {"reason": "storage_admission_unavailable", "poll_status": "storage_admission_unavailable"},
+])
+async def test_invalid_storage_poll_contract_is_not_retried(
+    override: dict[str, Any], retry_sleeps: list[float],
+) -> None:
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(503, json={**pressure(json.loads(request.content)["poll_request_id"]), **override})
+
+    async with client_for(handler) as client:
+        with worker_scope(), pytest.raises(ServerError):
+            await client.poll_activity_task(worker_id="storage-worker", task_queue="orders")
+    assert calls == 1
+    assert not retry_sleeps
+
+
+@pytest.mark.parametrize("scoped,worker_plane", [(False, True), (True, False)])
+async def test_direct_clients_and_control_plane_keep_bounded_retries(
+    scoped: bool, worker_plane: bool, retry_sleeps: list[float],
+) -> None:
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(503, json=pressure())
+
+    async with client_for(handler) as client:
+        async def call() -> None:
+            with pytest.raises(ServerError):
+                await client._request("POST", "/test", worker=worker_plane, json={})
+        if scoped:
+            with worker_scope():
+                await call()
+        else:
+            await call()
+    assert calls == 2
+
+
+async def test_retry_context_is_task_local_and_restored(retry_sleeps: list[float]) -> None:
+    async def check() -> None:
+        assert _worker_storage_admission_stop.get() is None
+
+    outside = asyncio.create_task(check())
+    with worker_scope():
+        await outside
+        assert _worker_storage_admission_stop.get() is not None
+    assert _worker_storage_admission_stop.get() is None
+
+
+async def test_shutdown_interrupts_pressure_without_another_mutation(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls = 0
+    stopped = False
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(503, json=pressure())
+
+    async def stop_on_sleep(delay: float) -> None:
+        nonlocal stopped
+        stopped = True
+
+    monkeypatch.setattr(retry_module, "asyncio", SimpleNamespace(sleep=stop_on_sleep))
+    async with client_for(handler) as client:
+        with worker_scope(lambda: stopped), pytest.raises(ServerError):
+            await client.complete_activity_task(
+                task_id="task", activity_attempt_id="attempt-7", lease_owner="worker", result="done",
+            )
+    assert calls == 1
+
+
+@pytest.mark.parametrize("kind", ["activity", "query", "workflow"])
+async def test_runtime_payload_upload_and_acknowledgement_are_not_repeated(
+    kind: str, retry_sleeps: list[float],
+) -> None:
+    server = FakeRuntimePayloadServer()
+    uploads: list[bytes] = []
+    acknowledgements: list[bytes] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        if request.method == "POST" and request.url.path == "/api/external-payloads/v1":
+            uploads.append(request.content)
+            if len(uploads) <= 3:
+                return httpx.Response(503, json=pressure())
+        if request.url.path.endswith("/complete"):
+            acknowledgements.append(request.content)
+            if len(acknowledgements) <= 4:
+                return httpx.Response(503, json=pressure())
+            return httpx.Response(200, json={"completed": True})
+        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(handler))
+        with worker_scope():
+            if kind == "activity":
+                await client.complete_activity_task(
+                    task_id="task", activity_attempt_id="attempt-7", lease_owner="worker", result="payload" * 30,
+                )
+            elif kind == "query":
+                await client.complete_query_task(
+                    query_task_id="task", query_task_attempt=7, lease_owner="worker", result="payload" * 30,
+                )
+            else:
+                await client.complete_workflow_task(
+                    task_id="task", workflow_task_attempt=7, lease_owner="worker",
+                    commands=[{"type": "complete_workflow", "result": serializer.envelope("payload" * 30)}],
+                )
+    assert len(uploads) == 4
+    assert len(set(uploads)) == 1
+    assert server.upload_count == 1
+    assert len(acknowledgements) == 5
+    assert len(set(acknowledgements)) == 1
+
+
+@pytest.mark.parametrize("fails", [False, True])
+async def test_running_worker_preserves_activity_outcome_without_reexecution(
+    fails: bool, retry_sleeps: list[float],
+) -> None:
+    calls = 0
+    claims = 0
+    heartbeats = 0
+    acks: list[bytes] = []
+    completed = asyncio.Event()
+
+    @activity.defn(name="storage.activity")
+    async def receipt() -> dict[str, bool]:
+        nonlocal calls
+        calls += 1
+        await activity.context().heartbeat({"step": "receipt"})
+        if fails:
+            raise ValueError("intentional activity failure")
+        return {"receipt": True}
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal claims, heartbeats
+        path = request.url.path
+        if path.endswith("/cluster/info"):
+            info = compatible_cluster_info()
+            info["worker_protocol"]["server_capabilities"]["query_tasks"] = False
+            return httpx.Response(200, json=info)
+        if path.endswith("/register") or request.method == "DELETE":
+            return httpx.Response(200, json={"registered": True})
+        if path.endswith("/activity-tasks/poll"):
+            claims += 1
+            if claims == 1:
+                return httpx.Response(200, json={"task": {
+                    "task_id": "storage-task", "activity_type": "storage.activity",
+                    "activity_attempt_id": "attempt-7", "attempt_number": 7,
+                    "payload_codec": "avro", "arguments": serializer.envelope([]),
+                }})
+        if path.endswith("/storage-task/heartbeat"):
+            heartbeats += 1
+            if heartbeats <= 3:
+                return httpx.Response(503, json=pressure())
+            return httpx.Response(200, json={"cancel_requested": False})
+        if path.endswith("/storage-task/fail" if fails else "/storage-task/complete"):
+            acks.append(request.content)
+            if len(acks) <= 4:
+                return httpx.Response(503, json=pressure())
+            completed.set()
+            return httpx.Response(200, json={"recorded": True})
+        assert path.endswith("/poll"), f"Unexpected request: {request.method} {path}"
+        return httpx.Response(200, json={"task": None})
+
+    async with client_for(handler) as client:
+        worker = Worker(client, task_queue="orders", activities=[receipt], max_concurrent_activity_tasks=1)
+        run = asyncio.create_task(worker.run())
+        try:
+            await asyncio.wait_for(completed.wait(), 2)
+        finally:
+            await worker.stop()
+            await run
+    assert calls == 1
+    assert heartbeats == 4
+    assert len(acks) == 5
+    assert len(set(acks)) == 1
+    assert worker._act_semaphore._value == 1
+
+
+async def test_query_thread_uses_pressure_retries(retry_sleeps: list[float]) -> None:
+    completed = threading.Event()
+    calls = 0
+    claims = 0
+    acks: list[bytes] = []
+
+    @workflow.defn(name="storage.query")
+    class QueryWorkflow:
+        def run(self, ctx):
+            yield ctx.sleep(30)
+
+        @workflow.query("status")
+        def status(self) -> str:
+            nonlocal calls
+            calls += 1
+            return "waiting"
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal claims
+        path = request.url.path
+        if path.endswith("/query-tasks/poll"):
+            claims += 1
+            if claims == 1:
+                return httpx.Response(200, json={"task": {
+                    "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)
+            if len(acks) <= 4:
+                return httpx.Response(503, json=pressure())
+            completed.set()
+            return httpx.Response(200, json={"completed": True})
+        raise AssertionError(f"Unexpected query request: {path}")
+
+    async with client_for(handler) as client:
+        worker = Worker(client, task_queue="orders", workflows=[QueryWorkflow])
+        worker._clone_client_for_query_tasks = lambda: client_for(handler)
+        worker._start_query_task_thread()
+        try:
+            assert await asyncio.to_thread(completed.wait, 2)
+        finally:
+            worker._request_query_task_thread_stop()
+            await worker._stop_query_task_thread(deadline=asyncio.get_running_loop().time() + 2)
+    assert calls == 1
+    assert len(acks) == 5
+    assert len(set(acks)) == 1
+
+
+@pytest.mark.parametrize("phase", ["registration", "poll", "completion"])
+async def test_run_until_timeout_interrupts_storage_pressure(phase: str) -> None:
+    refused = 0
+    deregistered = 0
+
+    @workflow.defn(name="storage.timeout")
+    class TimeoutWorkflow:
+        def run(self, ctx):
+            return "done"
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal refused, deregistered
+        path = request.url.path
+        if path.endswith("/cluster/info"):
+            info = compatible_cluster_info()
+            info["worker_protocol"]["server_capabilities"]["query_tasks"] = False
+            return httpx.Response(200, json=info)
+        if request.method == "DELETE":
+            deregistered += 1
+            return httpx.Response(200, json={})
+        if (
+            (phase == "registration" and path.endswith("/register"))
+            or (phase == "poll" and path.endswith("/poll"))
+            or (phase == "completion" and path.endswith("/complete"))
+        ):
+            refused += 1
+            poll_id = json.loads(request.content)["poll_request_id"] if phase == "poll" else None
+            return httpx.Response(503, json=pressure(poll_id))
+        if path.endswith("/register"):
+            return httpx.Response(200, json={})
+        if path.endswith("/workflows/storage-timeout"):
+            return httpx.Response(200, json={"status": "running"})
+        if path.endswith("/workflow-tasks/poll"):
+            return httpx.Response(200, json={"task": {
+                "task_id": "task", "workflow_task_attempt": 1, "workflow_type": "storage.timeout",
+                "payload_codec": "avro", "arguments": serializer.envelope([]), "history_events": [],
+            }})
+        raise AssertionError(f"Unexpected request: {path}")
+
+    async with client_for(handler) as client:
+        worker = Worker(client, task_queue="orders", workflows=[TimeoutWorkflow], shutdown_timeout=0.2)
+        with pytest.raises(TimeoutError, match="storage-timeout"):
+            await asyncio.wait_for(worker.run_until(workflow_id="storage-timeout", timeout=0.05), 1)
+    assert refused == 1
+    assert deregistered == (0 if phase == "registration" else 1)
+    assert worker._stop.is_set()
+    assert worker._registration_done.is_set()
+    assert worker._workflow_reserved == 0
+    assert not worker._in_flight
+    assert _worker_storage_admission_stop.get() is None
+
+
+async def test_stop_interrupts_registration_pressure() -> None:
+    refused = asyncio.Event()
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        if request.url.path.endswith("/cluster/info"):
+            return httpx.Response(200, json=compatible_cluster_info())
+        assert request.url.path.endswith("/register")
+        calls += 1
+        refused.set()
+        return httpx.Response(503, json=pressure())
+
+    async with client_for(handler) as client:
+        worker = Worker(client, task_queue="orders")
+        run = asyncio.create_task(worker.run())
+        try:
+            await asyncio.wait_for(refused.wait(), 1)
+        finally:
+            await asyncio.wait_for(worker.stop(), 1)
+        with pytest.raises(ServerError):
+            await run
+    assert calls == 1
+    assert worker._registration_done.is_set()
+
+
+@pytest.mark.parametrize("status,reason", [(401, "unauthenticated"), (409, "activity_attempt_expired")])
+async def test_recovery_does_not_retry_auth_or_lease_rejection(
+    status: int, reason: str, retry_sleeps: list[float],
+) -> None:
+    requests: list[bytes] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        requests.append(request.content)
+        if len(requests) <= 3:
+            return httpx.Response(503, json=pressure())
+        return httpx.Response(status, json={"reason": reason})
+
+    async with client_for(handler) as client:
+        with worker_scope(), pytest.raises(Unauthorized if status == 401 else ServerError) as error:
+            await client.complete_activity_task(
+                task_id="task", activity_attempt_id="attempt-7", lease_owner="worker", result="receipt",
+            )
+    if isinstance(error.value, ServerError):
+        assert error.value.status == status
+        assert error.value.reason() == reason
+    assert len(requests) == 4
+    assert len(set(requests)) == 1
+
+
+async def test_legacy_validation_poll_without_identity_fails_closed(retry_sleeps: list[float]) -> None:
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        body = pressure("")
+        body["poll_request_id"] = None
+        return httpx.Response(503, json=body)
+
+    async with client_for(handler) as client:
+        with worker_scope(), pytest.raises(ServerError):
+            await client.poll_update_validation_task(worker_id="worker", task_queue="orders")
+    assert calls == 1
+    assert not retry_sleeps
+
+
+@pytest.mark.parametrize("approved", [True, False])
+async def test_update_validator_is_not_reexecuted_for_acknowledgement(
+    approved: bool, retry_sleeps: list[float],
+) -> None:
+    ValidatedUpdateWorkflow.validator_calls = 0
+    ValidatedUpdateWorkflow.handler_calls = 0
+    requests: list[bytes] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        assert request.url.path.endswith("/approve" if approved else "/reject")
+        requests.append(request.content)
+        if len(requests) <= 4:
+            return httpx.Response(503, json=pressure())
+        return httpx.Response(200, json={})
+
+    async with client_for(handler) as client:
+        worker = Worker(client, task_queue="orders", workflows=[ValidatedUpdateWorkflow])
+        with worker_scope():
+            outcome = await worker._run_update_validation_task({
+                "update_validation_task_id": "validation-task", "update_validation_attempt": 7,
+                "workflow_type": "validated-update-wf", "update_name": "approve",
+                "history_events": [], "workflow_arguments": serializer.envelope([]),
+                "update_arguments": serializer.envelope([approved]), "payload_codec": "avro",
+            })
+    assert outcome == ("approved" if approved else "rejected")
+    assert ValidatedUpdateWorkflow.validator_calls == 1
+    assert ValidatedUpdateWorkflow.handler_calls == 0
+    assert len(requests) == 5
+    assert len(set(requests)) == 1
+
+
+async def test_cold_replay_completion_retries_only_its_prepared_acknowledgement(
+    retry_sleeps: list[float], monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from tests.test_golden_history_replay import GoldenSingleActivityWorkflow
+
+    path = Path(__file__).parent / "fixtures/replay_regressions/storage-paused-cold-completion.json"
+    fixture = json.loads(path.read_text())
+    replays = 0
+    original = GoldenSingleActivityWorkflow.run
+
+    def counted_run(self, *args):
+        nonlocal replays
+        replays += 1
+        return original(self, *args)
+
+    monkeypatch.setattr(GoldenSingleActivityWorkflow, "run", counted_run)
+    requests: list[bytes] = []
+    completed = False
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal completed
+        path = request.url.path
+        if path.endswith("/cluster/info"):
+            info = compatible_cluster_info()
+            info["worker_protocol"]["server_capabilities"]["query_tasks"] = False
+            return httpx.Response(200, json=info)
+        if path.endswith("/register") or request.method == "DELETE":
+            return httpx.Response(200, json={})
+        if path.endswith("/workflows/cold-replay"):
+            return httpx.Response(200, json={"status": "completed" if completed else "running"})
+        if path.endswith("/workflow-tasks/poll"):
+            return httpx.Response(200, json={"task": {
+                "task_id": "task", "workflow_task_attempt": 7, "workflow_type": fixture["workflow"]["type"],
+                "payload_codec": "avro", "arguments": serializer.envelope(fixture["workflow"]["input"]),
+                "history_events": fixture["history"],
+            }})
+        assert path.endswith("/workflow-tasks/task/complete"), f"Unexpected request: {path}"
+        requests.append(request.content)
+        if len(requests) <= 7:
+            return httpx.Response(503, json=pressure())
+        completed = True
+        return httpx.Response(200, json={})
+
+    async with client_for(handler) as client:
+        worker = Worker(client, task_queue="orders", workflows=[GoldenSingleActivityWorkflow])
+        result = await worker.run_until(workflow_id="cold-replay", timeout=2)
+    assert result.status == "completed"
+    assert replays == 1
+    assert len(requests) == 8
+    assert len(set(requests)) == 1
+    command, = json.loads(requests[0])["commands"]
+    assert command["type"] == "complete_workflow"
+    assert serializer.decode_envelope(command["result"]) == fixture["expected"]["result"]