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

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

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

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

{durable_workflow-2.5.8 → durable_workflow-2.5.9}/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.9}/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.9/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.9
 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.9}/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.9}/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(