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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durable-workflow
-Version: 2.0.2
+Version: 2.0.4
 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.4}/pyproject.toml

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

{durable_workflow-2.0.2 → durable_workflow-2.0.4}/src/durable_workflow/client.py

@@ -20,8 +20,10 @@ from __future__ import annotations
 
 import asyncio
 import hashlib
+import json as json_module
 import math
 import os
+import re
 import time
 import uuid
 import warnings
@@ -30,7 +32,7 @@ from dataclasses import dataclass
 from importlib.metadata import PackageNotFoundError
 from importlib.metadata import version as _pkg_version
 from typing import Any
-from urllib.parse import quote, urlencode, urlsplit
+from urllib.parse import quote, unquote, urlencode, urlsplit
 
 import httpx
 
@@ -101,6 +103,26 @@ _RUNTIME_EXTERNAL_PAYLOAD_FETCH_PATH_TEMPLATE = (
     "/external-payloads/v1/{referenceId}"
 )
 _RUNTIME_EXTERNAL_PAYLOAD_ERROR_BODY_LIMIT = 64 * 1024
+_PAYLOAD_COMPLETION_SCHEMA = "durable-workflow.v2.payload-completion-context.v1"
+_PAYLOAD_COMPLETION_HEADER = "X-Durable-Workflow-Payload-Completion"
+
+
+def _payload_completion_context(path: str, body: Any) -> dict[str, Any] | None:
+    match = re.fullmatch(r"/worker/(activity|workflow|query)-tasks/([^/]+)/(complete|fail)", path.split("?")[0])
+    if match is None or not isinstance(body, dict):
+        return None
+    kind, task_id, operation = match.groups()
+    attempt = body.get("activity_attempt_id" if kind == "activity" else f"{kind}_task_attempt")
+    owner = body.get("lease_owner")
+    if not isinstance(owner, str) or not owner:
+        return None
+    if kind == "activity":
+        if not isinstance(attempt, str) or not attempt:
+            return None
+    elif type(attempt) is not int or attempt < 1:
+        return None
+    return {"schema": _PAYLOAD_COMPLETION_SCHEMA, "kind": kind, "task_id": unquote(task_id),
+            "attempt": attempt, "lease_owner": owner, "operation": operation}
 
 
 def _default_sdk_version() -> str:
@@ -277,6 +299,7 @@ class _RuntimeExternalPayloadTransport:
     max_payload_bytes: int
     request_timeout_seconds: float
     status: str
+    completion_context: bool = False
 
 
 @dataclass
@@ -1734,6 +1757,9 @@ class Client:
                     worker=worker,
                     transport=transport,
                     uploaded={},
+                    completion=(
+                        _payload_completion_context(path, json) if worker and transport.completion_context else None
+                    ),
                 )
 
         start = time.perf_counter()
@@ -1904,11 +1930,15 @@ class Client:
             )
 
         status = policy.get("status")
+        completion = upload.get("completion_context")
         transport = _RuntimeExternalPayloadTransport(
             threshold_bytes=threshold_bytes,
             max_payload_bytes=max_payload_bytes,
             request_timeout_seconds=float(request_timeout_seconds),
             status=status if isinstance(status, str) else "unknown",
+            completion_context=isinstance(completion, dict)
+            and completion.get("schema") == _PAYLOAD_COMPLETION_SCHEMA
+            and completion.get("header") == _PAYLOAD_COMPLETION_HEADER,
         )
         self._runtime_external_payload_transport_cache = transport
         self._runtime_external_payload_transport_resolved = True
@@ -1972,6 +2002,8 @@ class Client:
         worker: bool,
         transport: _RuntimeExternalPayloadTransport,
         uploaded: dict[tuple[str, str, int], RuntimeExternalPayloadReference],
+        completion: dict[str, Any] | None = None,
+        slot: tuple[str | int, ...] = (),
     ) -> Any:
         if isinstance(value, dict):
             if (
@@ -1984,6 +2016,8 @@ class Client:
                     worker=worker,
                     transport=transport,
                     uploaded=uploaded,
+                    completion=completion,
+                    slot=(*slot, "result"),
                 )
                 if "external_payload" in externalized_result:
                     normalized_command["result"] = externalized_result
@@ -2014,6 +2048,7 @@ class Client:
                         sha256=sha256,
                         worker=worker,
                         transport=transport,
+                        completion={**completion, "slot": list(slot)} if completion is not None else None,
                     )
                     uploaded[identity] = reference
                 return {"codec": codec, "external_payload": reference.to_dict()}
@@ -2024,6 +2059,8 @@ class Client:
                     worker=worker,
                     transport=transport,
                     uploaded=uploaded,
+                    completion=completion,
+                    slot=(*slot, key),
                 )
                 for key, item in value.items()
             }
@@ -2045,8 +2082,10 @@ class Client:
                     worker=worker,
                     transport=transport,
                     uploaded=uploaded,
+                    completion=completion,
+                    slot=(*slot, index),
                 )
-                for item in value
+                for index, item in enumerate(value)
             ]
         return value
 
@@ -2058,6 +2097,7 @@ class Client:
         sha256: str,
         worker: bool,
         transport: _RuntimeExternalPayloadTransport,
+        completion: dict[str, Any] | None = None,
     ) -> RuntimeExternalPayloadReference:
         headers = self._headers(worker=worker)
         headers.update({
@@ -2069,13 +2109,30 @@ class Client:
         })
 
         async def _do_request() -> httpx.Response:
+            attempt_headers = dict(headers)
             response = await self._http.request(
                 "POST",
                 f"/api{_RUNTIME_EXTERNAL_PAYLOAD_UPLOAD_PATH}",
-                headers=headers,
+                headers=attempt_headers,
                 content=data,
                 timeout=transport.request_timeout_seconds,
             )
+            if (response.status_code == 503 and completion is not None
+                    and len(response.content) <= _RUNTIME_EXTERNAL_PAYLOAD_ERROR_BODY_LIMIT):
+                try:
+                    refusal = response.json()
+                except ValueError:
+                    refusal = None
+                context = json_module.dumps(completion, separators=(",", ":"))
+                if (isinstance(refusal, dict) and refusal.get("reason") == "storage_pressure"
+                        and refusal.get("storage_state") == "draining" and len(context.encode("utf-8")) <= 4096):
+                    # Keep the same identity if the ordinary transport policy retries
+                    # an ambiguous response; never rerun application activity code.
+                    attempt_headers[_PAYLOAD_COMPLETION_HEADER] = context
+                    response = await self._http.request(
+                        "POST", f"/api{_RUNTIME_EXTERNAL_PAYLOAD_UPLOAD_PATH}", headers=attempt_headers,
+                        content=data, timeout=transport.request_timeout_seconds,
+                    )
             response.raise_for_status()
             return response
 

{durable_workflow-2.0.2 → durable_workflow-2.0.4}/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.4}/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,46 @@ 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
+    # A payload upload is content-addressed and precedes completion submission.
+    # Even a late pressure refusal can retry those same bytes. This local retry
+    # classification does not alter the original response exposed to callers.
+    if (isinstance(body, dict) and "request_admitted" not in body and exc.request.method == "POST"
+            and exc.request.url.path.endswith("/api/external-payloads/v1")):
+        body = {**body, "request_admitted": False}
+    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 +123,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 +132,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.4}/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.4/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.4
 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.4}/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.4}/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.4"
     assert release["product-train"] == project["version"]
-    assert release["registry-version"] == "2.0.2"
+    assert release["registry-version"] == "2.0.4"
     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.2 → durable_workflow-2.0.4}/tests/test_runtime_external_payload_transport.py

@@ -19,6 +19,7 @@ from durable_workflow.errors import (
     ExternalPayloadUnavailable,
     ExternalPayloadUnsupported,
     RuntimeCapabilityUnsupported,
+    ServerError,
 )
 from durable_workflow.external_storage import (
     RUNTIME_EXTERNAL_PAYLOAD_REFERENCE_SCHEMA,
@@ -174,6 +175,115 @@ class FakeRuntimePayloadServer:
         return httpx.Response(200, json=response)
 
 
+class CompletionPayloadServer(FakeRuntimePayloadServer):
+    def __init__(self, *, supported: bool = True, state: str = "draining", reject_bound: bool = False) -> None:
+        super().__init__()
+        self.supported = supported
+        self.state = state
+        self.reject_bound = reject_bound
+        self.upload_requests: list[httpx.Request] = []
+
+    def cluster_info(self) -> dict[str, Any]:
+        info = super().cluster_info()
+        if self.supported:
+            info["namespace"]["external_payload_storage"]["transport"]["upload"]["completion_context"] = {
+                "schema": "durable-workflow.v2.payload-completion-context.v1",
+                "header": "X-Durable-Workflow-Payload-Completion",
+            }
+        return info
+
+    def handler(self, request: httpx.Request) -> httpx.Response:
+        if request.method == "POST" and request.url.path == "/api/external-payloads/v1":
+            self.upload_requests.append(request)
+            if self.state != "normal" and "X-Durable-Workflow-Payload-Completion" not in request.headers:
+                return httpx.Response(503, json={"reason": "storage_pressure", "storage_state": self.state,
+                    "request_admitted": False, "retryable": True, "retry_after_seconds": 1})
+            if self.reject_bound:
+                return httpx.Response(409, json={"reason": "external_payload_completion_lease_rejected",
+                    "retryable": False, "message": "Lease rejected."})
+        return super().handler(request)
+
+
+def completion_cases() -> list[tuple[str, str, dict[str, Any], list[str | int]]]:
+    envelope = serializer.envelope("x" * 100)
+    activity = {"lease_owner": "worker", "activity_attempt_id": "attempt"}
+    cases = [
+        ("activity", "complete", {**activity, "result": envelope}, ["result"]),
+        ("activity", "fail", {**activity, "failure": {"details": envelope}}, ["failure", "details"]),
+        ("query", "complete", {"lease_owner": "worker", "query_task_attempt": 2, "result_envelope": envelope},
+            ["result_envelope"]),
+    ]
+    for kind, field in [("complete_workflow", "result"), ("schedule_activity", "arguments"),
+                        ("upsert_memo", "entries"), ("start_service_operation", "request_payload")]:
+        cases.append(("workflow", "complete", {"lease_owner": "worker", "workflow_task_attempt": 2,
+            "commands": [{"type": kind, field: envelope}]}, ["commands", 0, field]))
+    cases.append(("workflow", "complete", {"lease_owner": "worker", "workflow_task_attempt": 2,
+        "commands": [{"type": "fail_workflow", "exception": {"details": envelope}}]},
+        ["commands", 0, "exception", "details"]))
+    cases.append(("workflow", "complete", {"lease_owner": "worker", "workflow_task_attempt": 2,
+        "commands": [{"type": "record_side_effect", "result": serializer.encode("x" * 100)}]},
+        ["commands", 0, "result"]))
+    cases.append(("workflow", "complete", {"lease_owner": "worker", "workflow_task_attempt": 2,
+        "commands": [{"type": "record_side_effect", "workflow_stream": {"items": [
+            {"payload": envelope, "payload_codec": "avro"}]}}]},
+        ["commands", 0, "workflow_stream", "items", 0, "payload"]))
+    return cases
+
+
+@pytest.mark.parametrize("kind,operation,body,slot", completion_cases())
+async def test_draining_upload_retry_carries_current_completion_identity(
+    kind: str, operation: str, body: dict[str, Any], slot: list[str | int],
+) -> None:
+    server = CompletionPayloadServer()
+    async with runtime_client(server, retry_policy=TransportRetryPolicy(max_attempts=1)) as client:
+        await client._request("POST", f"/worker/{kind}-tasks/task/{operation}", worker=True, json=body)
+    first, bound = server.upload_requests
+    assert "X-Durable-Workflow-Payload-Completion" not in first.headers
+    assert json.loads(bound.headers["X-Durable-Workflow-Payload-Completion"]) == {
+        "schema": "durable-workflow.v2.payload-completion-context.v1", "kind": kind,
+        "task_id": "task", "attempt": "attempt" if kind == "activity" else 2,
+        "lease_owner": "worker", "operation": operation, "slot": slot,
+    }
+    assert first.content == bound.content
+    assert first.headers["authorization"] == bound.headers["authorization"]
+    assert first.headers["x-namespace"] == bound.headers["x-namespace"]
+    assert len(server.requests) == 1
+
+
+@pytest.mark.parametrize("supported,worker,state", [(False, True, "draining"), (True, False, "draining"),
+                                                    (True, True, "fenced")])
+async def test_completion_upload_does_not_bypass_unsupported_client_or_fenced_admission(
+    supported: bool, worker: bool, state: str,
+) -> None:
+    server = CompletionPayloadServer(supported=supported, state=state)
+    async with runtime_client(server, retry_policy=TransportRetryPolicy(max_attempts=1)) as client:
+        with pytest.raises((ServerError, ExternalPayloadError)):
+            await client._request("POST", "/worker/activity-tasks/task/complete" if worker else "/workflows",
+                worker=worker, json={"lease_owner": "worker", "activity_attempt_id": "attempt",
+                    "result" if worker else "input": serializer.envelope("x" * 100)})
+    assert len(server.upload_requests) == 1
+    assert server.requests == []
+
+
+async def test_rejected_bound_upload_does_not_loop_or_submit_completion() -> None:
+    server = CompletionPayloadServer(reject_bound=True)
+    async with runtime_client(server, retry_policy=TransportRetryPolicy(max_attempts=3)) as client:
+        with pytest.raises((ServerError, ExternalPayloadError)):
+            await client.complete_activity_task(task_id="task", activity_attempt_id="attempt", lease_owner="worker",
+                result="x" * 100)
+    assert len(server.upload_requests) == 2
+    assert server.requests == []
+
+
+async def test_normal_payload_upload_is_unchanged_when_completion_capability_exists() -> None:
+    server = CompletionPayloadServer(state="normal")
+    async with runtime_client(server) as client:
+        await client.complete_activity_task(task_id="task", activity_attempt_id="attempt", lease_owner="worker",
+            result="x" * 100)
+    assert len(server.upload_requests) == 1
+    assert "X-Durable-Workflow-Payload-Completion" not in server.upload_requests[0].headers
+
+
 def runtime_client(
     server: FakeRuntimePayloadServer,
     *,