durable-workflow 2.5.8__tar.gz → 2.5.10__tar.gz

{durable_workflow-2.5.8/src/durable_workflow.egg-info → durable_workflow-2.5.10}/PKG-INFO

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

{durable_workflow-2.5.8 → durable_workflow-2.5.10}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "durable-workflow"
-version = "2.5.8"
+version = "2.5.10"
 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.8"
-registry-version = "2.5.8"
+product-train = "2.5.10"
+registry-version = "2.5.10"
 supported-server-versions = "2.5.16"
 worker-protocol-version = "1.20"
 control-plane-version = "2"

{durable_workflow-2.5.8 → durable_workflow-2.5.10}/src/durable_workflow/client.py

@@ -28,7 +28,7 @@ import re
 import time
 import uuid
 import warnings
-from collections.abc import AsyncIterator, Callable, Mapping, Sequence
+from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
 from dataclasses import dataclass
 from importlib.metadata import PackageNotFoundError
 from importlib.metadata import version as _pkg_version
@@ -1624,6 +1624,7 @@ class Client:
             _RuntimeExternalPayloadTransport | None
         ) = None
         self._runtime_external_payload_transport_resolved = False
+        self._workflow_payload_observers: dict[str, Callable[[dict[str, Any]], Awaitable[None]]] = {}
 
     async def aclose(self) -> None:
         """Close the underlying ``httpx`` connection pool.
@@ -1631,6 +1632,7 @@ class Client:
         Equivalent to exiting the async-context-manager form of the client.
         Safe to call multiple times.
         """
+        self._workflow_payload_observers.clear()
         await self._http.aclose()
 
     async def __aenter__(self) -> Client:
@@ -1877,10 +1879,12 @@ class Client:
                         "runtime returned an external payload reference without advertising "
                         "the authenticated namespace transport"
                     )
-                result = await self._resolve_runtime_payloads(
+                result = await self._resolve_workflow_claim_payloads(
                     result,
                     worker=worker,
                     transport=transport,
+                    path=path,
+                    request_json=request_json,
                 )
             result = self._hydrate_result_envelopes(result)
             outcome = "ok"
@@ -2251,6 +2255,77 @@ class Client:
         self.external_storage_cache.put(reference, data)
         return reference
 
+    async def _resolve_workflow_claim_payloads(
+        self,
+        value: Any,
+        *,
+        worker: bool,
+        transport: _RuntimeExternalPayloadTransport,
+        path: str,
+        request_json: Any,
+    ) -> Any:
+        lease: dict[str, Any] | None = None
+        owner: Any = None
+        is_poll = worker and path == "/worker/workflow-tasks/poll"
+        history = re.fullmatch(r"/worker/workflow-tasks/([^/]+)/history", path) if worker else None
+        if isinstance(request_json, dict) and isinstance(value, dict):
+            if is_poll and isinstance(value.get("task"), dict):
+                task = value["task"]
+                if task.get("task_kind", "workflow") == "workflow":
+                    lease = dict(task)
+                    owner = request_json.get("worker_id")
+            elif history is not None:
+                lease = dict(value)
+                lease.update({
+                    "task_id": unquote(history.group(1)),
+                    "lease_owner": request_json.get("lease_owner"),
+                    "workflow_task_attempt": request_json.get("workflow_task_attempt"),
+                })
+                owner = request_json.get("lease_owner")
+        observer = self._workflow_payload_observers.get(owner) if isinstance(owner, str) else None
+        if lease is None or observer is None:
+            return await self._resolve_runtime_payloads(value, worker=worker, transport=transport)
+        if (not isinstance(lease.get("task_id"), str) or not lease["task_id"]
+                or lease.get("lease_owner", owner) != owner
+                or type(lease.get("workflow_task_attempt", 1)) is not int
+                or lease.get("workflow_task_attempt", 1) < 1
+                or lease.get("namespace", self.namespace) != self.namespace):
+            raise ExternalPayloadUnsupported("workflow payload response has an invalid task ownership fence")
+        lease.setdefault("lease_owner", owner)
+        lease.setdefault("workflow_task_attempt", 1)
+
+        async def confirm() -> None:
+            await asyncio.wait_for(observer(lease), timeout=5.0)
+
+        async def observe() -> None:
+            while True:
+                await asyncio.sleep(5.0)
+                await confirm()
+
+        # The poll body is now available, but application code cannot run until
+        # its referenced bytes arrive and the same claim is confirmed again.
+        await confirm()
+        resolution = asyncio.create_task(self._resolve_runtime_payloads(value, worker=worker, transport=transport))
+        observation = asyncio.create_task(observe())
+        try:
+            done, _ = await asyncio.wait({resolution, observation}, return_when=asyncio.FIRST_COMPLETED)
+            if observation in done:
+                await observation
+                raise ExternalPayloadUnavailable("workflow payload lease observer stopped")
+            result = await resolution
+            observation.cancel()
+            await asyncio.gather(observation, return_exceptions=True)
+            await confirm()
+            if isinstance(result, dict) and "cancellation_request" in lease:
+                target = result.get("task") if is_poll else result
+                if isinstance(target, dict):
+                    target["cancellation_request"] = lease["cancellation_request"]
+            return result
+        finally:
+            resolution.cancel()
+            observation.cancel()
+            await asyncio.gather(resolution, observation, return_exceptions=True)
+
     async def _resolve_runtime_payloads(
         self,
         value: Any,
@@ -2689,6 +2764,7 @@ class Client:
         scoped._cluster_info_lock = asyncio.Lock()
         scoped._runtime_external_payload_transport_cache = None
         scoped._runtime_external_payload_transport_resolved = False
+        scoped._workflow_payload_observers = {}
         scoped.external_storage = None
         scoped.external_storage_threshold_bytes = None
         scoped.external_storage_cache = ExternalPayloadCache()

{durable_workflow-2.5.8 → durable_workflow-2.5.10}/src/durable_workflow/worker.py

@@ -1519,6 +1519,28 @@ class Worker:
 
         raise RuntimeError("workflow yielded too many consecutive Nexus service calls")
 
+    async def _renew_workflow_payload_lease(self, task: dict[str, Any]) -> None:
+        if self._stop.is_set() or self._local_activity_shutdown.is_set():
+            raise LocalActivityExecutionAborted("worker shutdown stopped its workflow payload transfer")
+        response = await self.client.heartbeat_workflow_task(
+            task_id=task["task_id"], lease_owner=self.worker_id,
+            workflow_task_attempt=task["workflow_task_attempt"],
+        )
+        if self._stop.is_set() or self._local_activity_shutdown.is_set():
+            raise LocalActivityExecutionAborted("worker shutdown stopped its workflow payload transfer")
+        if not isinstance(response, Mapping) or any((
+            response.get("task_id") != task["task_id"],
+            response.get("lease_owner") != self.worker_id,
+            response.get("workflow_task_attempt") != task["workflow_task_attempt"],
+            response.get("renewed") is not True,
+        )):
+            raise LocalActivityExecutionAborted("workflow payload lease was not fenced and acknowledged")
+        observation = response.get("cancellation_request")
+        if observation is not None:
+            # Preserve the observation for the normal canonical delivery path.
+            # Hydrating a payload does not deliver or acknowledge cancellation.
+            self._observe_workflow_cancellation(task, observation)
+
     async def _renew_local_workflow_lease(self, task: dict[str, Any]) -> None:
         if self._cooperative_cancellation_supported and self._local_activity_shutdown.is_set():
             raise LocalActivityExecutionAborted("worker shutdown abandoned its local workflow claim")
@@ -4414,6 +4436,9 @@ class Worker:
         if self._stop.is_set():
             raise RuntimeError("worker cannot run after stop() has been called")
         self._run_started = True
+        payload_observers = getattr(self.client, "_workflow_payload_observers", None)
+        if isinstance(payload_observers, dict):
+            payload_observers[self.worker_id] = self._renew_workflow_payload_lease
 
     def worker_session(self, options: WorkerSessionOptions) -> WorkerSession:
         """Return a typed lease handle held by this worker until shutdown."""
@@ -4532,6 +4557,10 @@ class Worker:
             self._registered = False
             log.info("worker %s deregistered", self.worker_id)
         self._sticky_cache.clear()
+        payload_observers = getattr(self.client, "_workflow_payload_observers", None)
+        if (isinstance(payload_observers, dict)
+                and payload_observers.get(self.worker_id) == self._renew_workflow_payload_lease):
+            payload_observers.pop(self.worker_id)
 
     @staticmethod
     def _remaining_shutdown_time(deadline: float) -> float:

{durable_workflow-2.5.8 → durable_workflow-2.5.10}/src/durable_workflow/workflow.py

@@ -5238,7 +5238,7 @@ def _replay_state(
     pending: list[Command] = []
     advanced_cmd: Any = None
     terminal_condition_reopen_cmd: WaitCondition | None = None
-    version_decisions: dict[str, tuple[int, str]] = {}
+    version_decisions: dict[str, int] = {}
 
     def _version_error(detail: str, step: _RecordedStep | None = None) -> NonDeterministicReplayError:
         return NonDeterministicReplayError(
@@ -5256,13 +5256,12 @@ def _replay_state(
         if type(version) is not int:
             raise _version_error("Version marker must contain an integer version.", step)
         change_id = payload.get("change_id")
-        if "change_id" in payload and (not isinstance(change_id, str) or not change_id.strip()):
+        if not isinstance(change_id, str) or not change_id.strip():
             raise _version_error("Version marker must retain a non-empty change ID.", step)
-        if "min_supported" in payload or "max_supported" in payload:
-            minimum, maximum = payload.get("min_supported"), payload.get("max_supported")
-            if (type(minimum) is not int or type(maximum) is not int
-                    or not minimum <= version <= maximum):
-                raise _version_error("Recorded version marker has an invalid supported range.", step)
+        minimum, maximum = payload.get("min_supported"), payload.get("max_supported")
+        if (type(minimum) is not int or type(maximum) is not int
+                or not minimum <= version <= maximum):
+            raise _version_error("Recorded version marker has an invalid supported range.", step)
         return version
 
     def _consume_duplicate_version_markers() -> None:
@@ -5277,7 +5276,7 @@ def _replay_state(
             if _next_unconsumed_recorded_step() is not step:
                 return
             version = _recorded_version(step)
-            if version != version_decisions[change_id][0]:
+            if version != version_decisions[change_id]:
                 raise _version_error(f"Conflicting recorded decisions for change ID {change_id!r}.", step)
             # Earlier SDKs emitted a marker for every invocation. Consume these
             # aliases while retaining their physical positions for cancellation
@@ -5293,13 +5292,10 @@ def _replay_state(
         if (type(command.min_supported) is not int or type(command.max_supported) is not int
                 or command.min_supported > command.max_supported):
             raise _version_error(f"Invalid supported range for change ID {change_id!r}.")
-        family = "version" if command.result_kind == "version" else "patch"
         decision = version_decisions.get(change_id)
         step = _next_unconsumed_recorded_step()
         if decision is not None:
-            version, original_family = decision
-            if original_family != family:
-                raise _version_error(f"Change ID {change_id!r} cannot switch between version and patch helpers.", step)
+            version = decision
         elif step is not None and step.shape == "version marker":
             _assert_step_matches(command, step)
             version = _recorded_version(step)
@@ -5326,7 +5322,7 @@ def _replay_state(
                 f"Recorded version {version} for change ID {change_id!r} is outside "
                 f"the supported range {command.min_supported}..{command.max_supported}.", step,
             )
-        version_decisions.setdefault(change_id, (version, family))
+        version_decisions.setdefault(change_id, version)
         return _version_marker_result(command, version)
 
     def _parallel_leaf_kind(command: Any) -> str:

{durable_workflow-2.5.8 → durable_workflow-2.5.10/src/durable_workflow.egg-info}/PKG-INFO

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

{durable_workflow-2.5.8 → durable_workflow-2.5.10}/tests/test_patch_history.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from itertools import permutations
 from typing import Any
 
 import pytest
@@ -74,6 +75,43 @@ class PatchCancellationWorkflow:
             return [decisions, "cancelled"]
 
 
+@workflow.defn(name="tests.patch.helpers")
+class MixedPatchHelpersWorkflow:
+    def run(self, ctx, helpers):  # type: ignore[no-untyped-def]
+        decisions = []
+        for helper in helpers:
+            command = ctx.get_version("same", -1, 1) if helper == "get_version" else getattr(ctx, helper)("same")
+            decisions.append((yield command))
+        result = yield ctx.schedule_activity("old", [])
+        return [decisions, result]
+
+
+@pytest.mark.parametrize("helpers", list(permutations(["get_version", "patched", "deprecate_patch"])))
+@pytest.mark.parametrize("marker_count", [0, 1, 2])
+@pytest.mark.parametrize("completed", [False, True])
+def test_helpers_share_one_decision_on_old_history(
+    helpers: tuple[str, ...], marker_count: int, completed: bool,
+) -> None:
+    history = [version_marker(sequence) for sequence in range(1, marker_count + 1)]
+    history += old_activity(marker_count + 1, completed)
+    commands = replay(MixedPatchHelpersWorkflow, history, [helpers]).commands
+    if not completed:
+        assert commands == []
+    else:
+        assert len(commands) == 1
+        assert isinstance(commands[0], CompleteWorkflow)
+        version = 1 if marker_count else -1
+        values = {"get_version": version, "patched": version == 1, "deprecate_patch": None}
+        assert commands[0].result == [[values[helper] for helper in helpers], "recorded"]
+
+
+@pytest.mark.parametrize("helpers", list(permutations(["get_version", "patched", "deprecate_patch"])))
+def test_helpers_record_one_decision_on_fresh_history(helpers: tuple[str, ...]) -> None:
+    commands = replay(MixedPatchHelpersWorkflow, [], [helpers]).commands
+    assert [type(command) for command in commands] == [RecordVersionMarker, ScheduleActivity]
+    assert commands[0].version == 1
+
+
 @pytest.mark.parametrize("completed", [False, True])
 def test_inserted_patch_preserves_unmarked_activity(completed: bool) -> None:
     outcome = replay(PatchInsertionWorkflow, old_activity(1, completed), [])
@@ -122,7 +160,7 @@ def test_repeated_version_reuses_first_decision() -> None:
 
 
 @pytest.mark.parametrize("arguments", [[3], [1, True]])
-def test_repeated_version_rejects_incompatible_range_or_family(arguments: list[Any]) -> None:
+def test_repeated_version_rejects_incompatible_range(arguments: list[Any]) -> None:
     with pytest.raises(NonDeterministicReplayError):
         replay(VersionRangeWorkflow, [], arguments)
 
@@ -134,6 +172,41 @@ def test_conflicting_duplicate_marker_is_rejected(version: Any) -> None:
         replay(RepeatedPatchWorkflow, history, [])
 
 
+@pytest.mark.parametrize("sequence", [1, 2])
+@pytest.mark.parametrize("bounds", [
+    {}, {"min_supported": -1}, {"max_supported": 1},
+    {"min_supported": True, "max_supported": 1},
+    {"min_supported": -1, "max_supported": "1"},
+    {"min_supported": 2, "max_supported": 1},
+    {"min_supported": -1, "max_supported": 0},
+])
+def test_canonical_markers_require_valid_recorded_ranges(sequence: int, bounds: dict[str, Any]) -> None:
+    markers = [version_marker(1), version_marker(2)]
+    markers[sequence - 1]["payload"].pop("min_supported")
+    markers[sequence - 1]["payload"].pop("max_supported")
+    markers[sequence - 1]["payload"].update(bounds)
+    with pytest.raises(NonDeterministicReplayError, match="invalid supported range"):
+        replay(RepeatedPatchWorkflow, markers + old_activity(3), [])
+
+
+@pytest.mark.parametrize("identity", [{}, {"change_id": None}, {"change_id": ""}, {"change_id": True}])
+def test_canonical_marker_requires_a_change_id(identity: dict[str, Any]) -> None:
+    marker = version_marker(1)
+    marker["payload"].pop("change_id")
+    marker["payload"].update(identity)
+    with pytest.raises(NonDeterministicReplayError, match="non-empty change ID"):
+        replay(RepeatedPatchWorkflow, [marker] + old_activity(2), [])
+
+
+def test_duplicate_ranges_may_differ_around_the_same_version() -> None:
+    marker = version_marker(2)
+    marker["payload"].update(min_supported=0, max_supported=3)
+    commands = replay(RepeatedPatchWorkflow, [version_marker(1), marker] + old_activity(3), []).commands
+    assert len(commands) == 1
+    assert isinstance(commands[0], CompleteWorkflow)
+    assert commands[0].result == [True, True, "recorded"]
+
+
 @pytest.mark.parametrize("marker_count", [0, 1, 2])
 def test_patch_decisions_preserve_committed_cancellation_sequence(marker_count: int) -> None:
     history = [version_marker(sequence) for sequence in range(1, marker_count + 1)]

{durable_workflow-2.5.8 → durable_workflow-2.5.10}/tests/test_replay.py

@@ -1949,7 +1949,9 @@ class TestVersionMarker:
 
     def test_version_from_history(self) -> None:
         history = [
-            {"event_type": "VersionMarkerRecorded", "payload": {"version": 2}},
+            {"event_type": "VersionMarkerRecorded", "payload": {
+                "sequence": 1, "change_id": "change-1", "version": 2, "min_supported": 1, "max_supported": 2,
+            }},
         ]
         outcome = replay(VersionWorkflow, history, [])
         assert len(outcome.commands) == 1
@@ -1959,7 +1961,9 @@ class TestVersionMarker:
 
     def test_old_version_from_history(self) -> None:
         history = [
-            {"event_type": "VersionMarkerRecorded", "payload": {"version": 1}},
+            {"event_type": "VersionMarkerRecorded", "payload": {
+                "sequence": 1, "change_id": "change-1", "version": 1, "min_supported": 1, "max_supported": 2,
+            }},
         ]
         outcome = replay(VersionWorkflow, history, [])
         assert len(outcome.commands) == 1
@@ -1969,7 +1973,9 @@ class TestVersionMarker:
 
     def test_full_replay(self) -> None:
         history = [
-            {"event_type": "VersionMarkerRecorded", "payload": {"version": 2}},
+            {"event_type": "VersionMarkerRecorded", "payload": {
+                "sequence": 1, "change_id": "change-1", "version": 2, "min_supported": 1, "max_supported": 2,
+            }},
             {"event_type": "ActivityCompleted", "payload": {"result": _avro("done")}},
         ]
         outcome = replay(VersionWorkflow, history, [])
@@ -2006,7 +2012,9 @@ class TestVersionMarker:
 
     def test_patched_uses_recorded_marker_for_existing_patched_runs(self) -> None:
         history = [
-            {"event_type": "VersionMarkerRecorded", "payload": {"version": 1}},
+            {"event_type": "VersionMarkerRecorded", "payload": {
+                "sequence": 1, "change_id": "patch-1", "version": 1, "min_supported": -1, "max_supported": 1,
+            }},
         ]
 
         outcome = replay(PatchedWorkflow, history, [])
@@ -2017,7 +2025,9 @@ class TestVersionMarker:
 
     def test_patched_uses_legacy_default_for_existing_unpatched_runs(self) -> None:
         history = [
-            {"event_type": "VersionMarkerRecorded", "payload": {"version": -1}},
+            {"event_type": "VersionMarkerRecorded", "payload": {
+                "sequence": 1, "change_id": "patch-1", "version": -1, "min_supported": -1, "max_supported": 1,
+            }},
         ]
 
         outcome = replay(PatchedWorkflow, history, [])
@@ -2040,7 +2050,9 @@ class TestVersionMarker:
 
         replayed = replay(
             DeprecatedPatchWorkflow,
-            [{"event_type": "VersionMarkerRecorded", "payload": {"version": 1}}],
+            [{"event_type": "VersionMarkerRecorded", "payload": {
+                "sequence": 1, "change_id": "patch-1", "version": 1, "min_supported": -1, "max_supported": 1,
+            }}],
             [],
         )
         assert len(replayed.commands) == 1

{durable_workflow-2.5.8 → durable_workflow-2.5.10}/tests/test_replay_regression_corpus.py

@@ -25,7 +25,12 @@ from tests.test_golden_history_replay import (
     GoldenTimeoutWaitWorkflow,
     GoldenVersionMarkerWorkflow,
 )
-from tests.test_patch_history import PatchCancellationWorkflow, PatchInsertionWorkflow, RepeatedPatchWorkflow
+from tests.test_patch_history import (
+    MixedPatchHelpersWorkflow,
+    PatchCancellationWorkflow,
+    PatchInsertionWorkflow,
+    RepeatedPatchWorkflow,
+)
 from tests.test_pending_operation_replay import PendingOperation
 from tests.test_update_signal_condition_replay import (
     PostConditionReceiversWorkflow,
@@ -290,6 +295,7 @@ WORKFLOWS = [
     ParallelResultBindingWorkflow,
     PatchCancellationWorkflow,
     PatchInsertionWorkflow,
+    MixedPatchHelpersWorkflow,
     RepeatedPatchWorkflow,
     PendingOperation,
     PagedRecordedSideEffectsWorkflow,

{durable_workflow-2.5.8 → durable_workflow-2.5.10}/tests/test_runtime_external_payload_transport.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
+import asyncio
+import gc
 import hashlib
 import json
+import weakref
 from typing import Any
 
 import httpx
 import pytest
 
-from durable_workflow import Replayer, serializer, workflow
+from durable_workflow import Replayer, Worker, serializer, workflow
 from durable_workflow.client import Client
 from durable_workflow.errors import (
     ExternalPayloadError,
@@ -26,7 +29,7 @@ from durable_workflow.external_storage import (
     ExternalPayloadCache,
 )
 from durable_workflow.retry_policy import TransportRetryPolicy
-from durable_workflow.workflow import CompleteWorkflow, WorkflowContext
+from durable_workflow.workflow import CompleteWorkflow, LocalActivityExecutionAborted, WorkflowContext
 
 
 @workflow.defn(name="runtime-payload-replay")
@@ -309,6 +312,205 @@ def runtime_envelope(server: FakeRuntimePayloadServer, value: Any) -> dict[str,
     return {"codec": "avro", "external_payload": server.seed(blob)}
 
 
+def payload_worker(client: Client, *, capabilities: list[str] | None = None) -> Worker:
+    worker = Worker(client, task_queue="billing", worker_id="worker-1", workflows=[RuntimePayloadReplayWorkflow],
+                    capabilities=capabilities or [])
+    worker._begin_run()
+    worker._registration_done.set()
+    return worker
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("history_page", [False, True])
+async def test_worker_confirms_original_lease_before_and_after_payload_hydration(history_page: bool) -> None:
+    server = FakeRuntimePayloadServer()
+    client = runtime_client(server)
+    payload_worker(client)
+    task = {"task_id": "task-1", "lease_owner": "worker-1", "workflow_task_attempt": 3,
+            "arguments": runtime_envelope(server, ["original-input"])}
+    path = "/api/worker/workflow-tasks/task-1/history" if history_page else "/api/worker/workflow-tasks/poll"
+    server.responses[("POST", path)] = {"history_events": [task]} if history_page else {"task": task}
+    server.responses[("POST", "/api/worker/workflow-tasks/task-1/heartbeat")] = {
+        "task_id": "task-1", "lease_owner": "worker-1", "workflow_task_attempt": 3, "renewed": True,
+    }
+    original_handler = server.handler
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        if request.method == "GET" and "/external-payloads/" in request.url.path:
+            assert len([entry for entry in server.requests if entry[1].endswith("/heartbeat")]) == 1
+        return original_handler(request)
+
+    client._http = httpx.AsyncClient(base_url=client.base_url, transport=httpx.MockTransport(handler))
+    try:
+        if history_page:
+            result = await client.workflow_task_history(task_id="task-1", lease_owner="worker-1",
+                                                        workflow_task_attempt=3, next_history_page_token="page-2")
+            resolved = result["history_events"][0]
+        else:
+            result = await client.poll_workflow_task(worker_id="worker-1", task_queue="billing", timeout=0)
+            resolved = result
+        assert serializer.decode_envelope(resolved["arguments"]) == ["original-input"]
+        heartbeats = [entry for entry in server.requests if entry[1].endswith("/heartbeat")]
+        assert len(heartbeats) == 2
+        assert all(entry[2] == {"lease_owner": "worker-1", "workflow_task_attempt": 3} for entry in heartbeats)
+    finally:
+        await client.aclose()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("changed_fence", [
+    {"workflow_task_attempt": 4}, {"lease_owner": "worker-2"}, {"renewed": False},
+])
+async def test_worker_does_not_fetch_payload_when_original_claim_is_not_confirmed(
+    changed_fence: dict[str, Any],
+) -> None:
+    server = FakeRuntimePayloadServer()
+    client = runtime_client(server)
+    payload_worker(client)
+    server.responses[("POST", "/api/worker/workflow-tasks/poll")] = {"task": {
+        "task_id": "task-1", "lease_owner": "worker-1", "workflow_task_attempt": 3,
+        "arguments": runtime_envelope(server, ["never-executed"]),
+    }}
+    server.responses[("POST", "/api/worker/workflow-tasks/task-1/heartbeat")] = {
+        "task_id": "task-1", "lease_owner": "worker-1", "workflow_task_attempt": 3, "renewed": True, **changed_fence,
+    }
+    try:
+        with pytest.raises(LocalActivityExecutionAborted, match="workflow payload lease was not fenced"):
+            await client.poll_workflow_task(worker_id="worker-1", task_queue="billing", timeout=0)
+        assert server.fetch_count == 0
+    finally:
+        await client.aclose()
+
+
+@pytest.mark.asyncio
+async def test_payload_renewal_preserves_new_cancellation_for_canonical_delivery() -> None:
+    from tests.test_cooperative_cancellation import observation
+
+    server = FakeRuntimePayloadServer()
+    client = runtime_client(server)
+    worker = payload_worker(client, capabilities=["cooperative_cancellation"])
+    worker._cooperative_cancellation_supported = True
+    original = observation()
+    server.responses[("POST", "/api/worker/workflow-tasks/poll")] = {"task": {
+        "task_id": "task-1", "lease_owner": "worker-1", "workflow_task_attempt": 3,
+        "arguments": runtime_envelope(server, ["original-input"]),
+    }}
+    server.responses[("POST", "/api/worker/workflow-tasks/task-1/heartbeat")] = {
+        "task_id": "task-1", "lease_owner": "worker-1", "workflow_task_attempt": 3, "renewed": True,
+        "cancellation_request": original,
+    }
+    try:
+        result = await client.poll_workflow_task(worker_id="worker-1", task_queue="billing", timeout=0)
+        assert result["cancellation_request"] == original
+        assert "_delivered_cancellation_request_id" not in result
+        assert not any("/cancellation/" in entry[1] for entry in server.requests)
+    finally:
+        await client.aclose()
+
+
+@pytest.mark.asyncio
+async def test_stopped_worker_does_not_start_claim_payload_download() -> None:
+    server = FakeRuntimePayloadServer()
+    client = runtime_client(server)
+    worker = payload_worker(client)
+    worker._stop.set()
+    server.responses[("POST", "/api/worker/workflow-tasks/poll")] = {"task": {
+        "task_id": "task-1", "lease_owner": "worker-1", "workflow_task_attempt": 3,
+        "arguments": runtime_envelope(server, ["never-executed"]),
+    }}
+    try:
+        with pytest.raises(LocalActivityExecutionAborted, match="shutdown stopped"):
+            await client.poll_workflow_task(worker_id="worker-1", task_queue="billing", timeout=0)
+        assert server.fetch_count == 0
+        assert not any(entry[1].endswith("/heartbeat") for entry in server.requests)
+    finally:
+        await client.aclose()
+
+
+@pytest.mark.asyncio
+async def test_payload_transfer_is_cancelled_when_its_periodic_lease_observer_loses_authority(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    server = FakeRuntimePayloadServer()
+    client = runtime_client(server)
+    payload_worker(client)
+
+
+    server.responses[("POST", "/api/worker/workflow-tasks/poll")] = {"task": {
+        "task_id": "task-1", "lease_owner": "worker-1", "workflow_task_attempt": 3,
+        "arguments": runtime_envelope(server, ["never-executed"]),
+    }}
+    fetch_started = asyncio.Event()
+    fetch_cancelled = asyncio.Event()
+    observations = 0
+
+    async def heartbeat(request: httpx.Request) -> httpx.Response:
+        nonlocal observations
+        if request.url.path.endswith("/heartbeat"):
+            observations += 1
+            if observations > 1:
+                return httpx.Response(409, json={"reason": "lease_expired"})
+            return httpx.Response(200, json={"task_id": "task-1", "lease_owner": "worker-1",
+                                            "workflow_task_attempt": 3, "renewed": True})
+        if request.method == "GET" and "/external-payloads/" in request.url.path:
+            fetch_started.set()
+            try:
+                await asyncio.Event().wait()
+            finally:
+                fetch_cancelled.set()
+        return server.handler(request)
+
+    real_sleep = asyncio.sleep
+
+    async def prompt_observation(delay: float) -> None:
+        if delay == 5.0:
+            await fetch_started.wait()
+            await real_sleep(0)
+        else:
+            await real_sleep(delay)
+
+    monkeypatch.setattr("durable_workflow.client.asyncio.sleep", prompt_observation)
+    client._http = httpx.AsyncClient(base_url=client.base_url, transport=httpx.MockTransport(heartbeat))
+    try:
+        with pytest.raises(ServerError) as failure:
+            await asyncio.wait_for(client.poll_workflow_task(worker_id="worker-1", task_queue="billing", timeout=0), 2)
+        assert failure.value.reason() == "lease_expired"
+        assert fetch_cancelled.is_set()
+        assert observations == 2
+    finally:
+        await client.aclose()
+
+
+@pytest.mark.asyncio
+async def test_shared_client_does_not_retain_unstarted_or_stopped_workers() -> None:
+    client = runtime_client(FakeRuntimePayloadServer())
+    worker = Worker(client, task_queue="billing", worker_id="worker-1", workflows=[RuntimePayloadReplayWorkflow])
+    assert client._workflow_payload_observers == {}
+    reference = weakref.ref(worker)
+    worker._begin_run()
+    worker._registration_done.set()
+    assert "worker-1" in client._workflow_payload_observers
+    await worker.stop()
+    assert client._workflow_payload_observers == {}
+    del worker
+    await asyncio.sleep(0)
+    gc.collect()
+    assert reference() is None
+    await client.aclose()
+
+
+@pytest.mark.asyncio
+async def test_stopping_old_worker_preserves_a_new_worker_observer() -> None:
+    client = runtime_client(FakeRuntimePayloadServer())
+    older = payload_worker(client)
+    newer = payload_worker(client)
+    await older.stop()
+    assert client._workflow_payload_observers["worker-1"] == newer._renew_workflow_payload_lease
+    await newer.stop()
+    assert client._workflow_payload_observers == {}
+    await client.aclose()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("threshold_bytes", [1, 4096])
 async def test_nexus_arguments_offload_without_changing_recorded_arguments(threshold_bytes: int) -> None:

{durable_workflow-2.5.8 → durable_workflow-2.5.10}/tests/test_worker.py

@@ -4955,9 +4955,18 @@ class TestWorkerHeartbeats:
     async def test_heartbeat_loop_survives_transient_errors(
         self, mock_client: AsyncMock
     ) -> None:
-        mock_client.heartbeat_worker = AsyncMock(
-            side_effect=[RuntimeError("temporary"), {"acknowledged": True}]
-        )
+        recovered = asyncio.Event()
+        attempts = 0
+
+        async def heartbeat(**kwargs: object) -> dict[str, bool]:
+            nonlocal attempts
+            attempts += 1
+            if attempts == 1:
+                raise RuntimeError("temporary")
+            recovered.set()
+            return {"acknowledged": True}
+
+        mock_client.heartbeat_worker = AsyncMock(side_effect=heartbeat)
         worker = Worker(
             mock_client,
             task_queue="q1",
@@ -4967,11 +4976,13 @@ class TestWorkerHeartbeats:
             heartbeat_interval=0.02,
         )
         run_task = asyncio.create_task(worker.run())
-        await asyncio.sleep(0.15)
-        await worker.stop()
-        run_task.cancel()
-        with contextlib.suppress(asyncio.CancelledError):
-            await run_task
+        try:
+            await asyncio.wait_for(recovered.wait(), timeout=2)
+        finally:
+            await worker.stop()
+            run_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await run_task
         assert mock_client.heartbeat_worker.call_count >= 2
 
     def test_process_metrics_cpu_percent_is_instantaneous_not_lifetime(