Add guarded activity lease renewal client transport

src/durable_workflow/client.py

@@ -39,6 +39,7 @@
 import httpx
 
 from . import serializer
+from ._activity_lease import ActivityLeaseAuthority, current_activity_authority
 from ._cancellation_scope import CancellationScopeOpenReceipt, scope_identity
 from ._cancellation_scope_history import (
     CancellationScopeBudget,
@@ -1871,8 +1872,14 @@ async def _request(
             raise ValueError("cluster discovery credentials can only authorize /cluster/info")
 
         query_poll = worker and path == "/worker/query-tasks/poll"
+        activity_poll = worker and path == "/worker/activity-tasks/poll"
         if query_poll:
             current_query_authority.set(None)
+        if activity_poll:
+            current_activity_authority.set(None)
+        activity_authority = current_activity_authority.get()
+        if worker and activity_authority is not None and activity_authority.governs(self, self.namespace, path):
+            activity_authority.remaining()
         request_json = json
         if (
             not discovery
@@ -1904,14 +1911,20 @@ async def _request(
         plane = "worker" if worker or discovery_uses_worker_token else "control"
         status_code = "none"
         outcome = "pending"
+        request_started = time.monotonic()
 
         async def _do_request() -> httpx.Response:
+            nonlocal request_started
+            request_started = time.monotonic()
+            request_timeout = self.timeout if timeout is None else timeout
+            if worker and activity_authority is not None and activity_authority.governs(self, self.namespace, path):
+                request_timeout = min(request_timeout, activity_authority.remaining())
             resp = await self._http.request(
                 method,
                 f"/api{path}",
                 headers=self._discovery_headers() if discovery else self._headers(worker=worker),
                 json=request_json,
-                timeout=self.timeout if timeout is None else timeout,
+                timeout=request_timeout,
             )
             # Raise HTTPStatusError for 4xx/5xx so retry policy can catch it
             resp.raise_for_status()
@@ -1936,11 +1949,28 @@ async def _do_request() -> httpx.Response:
                 outcome = "ok"
                 return None
             result = resp.json()
+            if (worker and activity_authority is not None
+                    and activity_authority.governs(self, self.namespace, path)
+                    and path.endswith("/heartbeat") and isinstance(result, dict)
+                    and result.get("heartbeat_recorded") is True):
+                activity_authority.acknowledge_progress(result, request_started)
+            if activity_poll and isinstance(result, dict) and isinstance(result.get("task"), dict):
+                capabilities = result.get("server_capabilities")
+                if isinstance(capabilities, dict) and capabilities.get("activity_lease_renewal") is True:
+                    activity_authority = await ActivityLeaseAuthority.confirm_delivery(
+                        self, self.namespace, result["task"], owner=request_json["worker_id"],
+                        task_queue=request_json["task_queue"], server_time=result.get("server_time"),
+                    )
+                    current_activity_authority.set(activity_authority)
             query_authority = None
             if query_poll and isinstance(result, dict) and isinstance(result.get("task"), dict):
                 query_authority = await self._query_delivery_authority(result["task"], request_json)
                 current_query_authority.set(query_authority)
-            if not discovery and _contains_runtime_payload_envelope(result):
+            if activity_authority is not None and activity_poll and _contains_runtime_payload_envelope(result):
+                result = await activity_authority.run(
+                    lambda: self._resolve_activity_claim_payloads(result, worker=worker),
+                )
+            elif not discovery and _contains_runtime_payload_envelope(result):
                 transport = (await query_authority.run(self._runtime_external_payload_transport)
                              if query_authority is not None else await self._runtime_external_payload_transport())
                 if transport is None:
@@ -1957,6 +1987,8 @@ async def _do_request() -> httpx.Response:
                         result, worker=worker, transport=transport, path=path, request_json=request_json,
                     )
             result = self._hydrate_result_envelopes(result)
+            if activity_poll and activity_authority is not None:
+                activity_authority.remaining()
             if query_authority is not None:
                 await query_authority.confirm()
             outcome = "ok"
@@ -1976,6 +2008,12 @@ async def _do_request() -> httpx.Response:
             self.metrics.increment(CLIENT_REQUESTS, tags=tags)
             self.metrics.record(CLIENT_REQUEST_DURATION_SECONDS, time.perf_counter() - start, tags=tags)
 
+    async def _resolve_activity_claim_payloads(self, value: Any, *, worker: bool) -> Any:
+        transport = await self._runtime_external_payload_transport()
+        if transport is None:
+            raise ExternalPayloadUnsupported("activity payload reference has no authenticated runtime transport")
+        return await self._resolve_runtime_payloads(value, worker=worker, transport=transport, resolved={})
+
     async def _runtime_external_payload_transport(
         self,
     ) -> _RuntimeExternalPayloadTransport | None:
@@ -6160,6 +6198,48 @@ async def fail_activity_task(
             "POST", f"/worker/activity-tasks/{task_id}/fail", worker=True, json=body
         )
 
+    async def renew_activity_task_lease(
+        self,
+        *,
+        task_id: str,
+        activity_attempt_id: str,
+        lease_owner: str,
+        timeout: float = 5.0,
+    ) -> Any:
+        """Renew one live activity claim without recording application progress.
+
+        Requires affirmative ``activity_lease_renewal`` discovery. The caller
+        must validate the returned original identity and grant before using it.
+        This neither revives an expired claim nor extends business timeouts.
+        Discovery, retries and transport share the supplied bounded window.
+        """
+        for field, value in (
+            ("task_id", task_id), ("activity_attempt_id", activity_attempt_id), ("lease_owner", lease_owner),
+        ):
+            if not isinstance(value, str) or not value.strip():
+                raise ValueError(f"{field} must be a non-empty string")
+        if type(timeout) not in (int, float) or not math.isfinite(timeout) or not 0 < timeout <= 5:
+            raise ValueError("activity lease renewal timeout must be greater than zero and at most five seconds")
+
+        async def renew() -> Any:
+            operation = "Client.renew_activity_task_lease"
+            path = "worker_protocol.server_capabilities.activity_lease_renewal"
+            info = await self._runtime_discovery(operation=operation, required_path=path)
+            protocol = info.get("worker_protocol")
+            capabilities = protocol.get("server_capabilities") if isinstance(protocol, dict) else None
+            supported = capabilities.get("activity_lease_renewal") if isinstance(capabilities, dict) else None
+            if supported is False:
+                raise RuntimeCapabilityUnsupported(operation, path, f"This runtime does not support {path}.")
+            if supported is not True:
+                raise RuntimeDiscoveryUnavailable(operation, path, f"Runtime discovery must advertise {path}=true.")
+
+            return await self._request(
+                "POST", f"/worker/activity-tasks/{quote(task_id, safe='')}/renew-lease", worker=True,
+                json={"activity_attempt_id": activity_attempt_id, "lease_owner": lease_owner}, timeout=timeout,
+            )
+
+        return await asyncio.wait_for(renew(), timeout=timeout)
+
     async def activity_task_status(
         self,
         *,

src/durable_workflow/errors.py

@@ -36,6 +36,15 @@ def __init__(self, query_task_id: str, reason: str) -> None:
         self.reason = reason
 
 
+class ActivityLeaseAuthorityLost(DurableWorkflowError):
+    """An activity delivery lost its original live execution authority."""
+
+    def __init__(self, task_id: str, reason: str) -> None:
+        super().__init__(f"activity task {task_id!r} cannot continue: {reason}")
+        self.task_id = task_id
+        self.reason = reason
+
+
 class InvalidWorkflowDefinition(DurableWorkflowError, ValueError):
     """Workflow authoring uses an unsupported entry point or result."""