Stop default activities when a heartbeat refuses continued execution

CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Fixed
+
+- Stop default async and synchronous activities when a heartbeat explicitly
+  refuses continued execution. Preserve the Server reason separately from
+  cancellation and fence completion/failure even when user code catches the
+  refusal. The worker frees its slot after the callback exits.
+
 - Reject canonical version markers missing their change ID or supported-range
   bounds. Consistent historical duplicate ranges may still differ when each
   contains the frozen version.

README.md

@@ -142,6 +142,17 @@ independently of application heartbeats. Existing terminal cancellation remains
 available. See the [cancellation guide](docs/cooperative-cancellation-design.md)
 for immutable context, operation policies, shielded cleanup and recovery.
 
+Both `await activity.context().heartbeat()` and synchronous
+`activity.context().heartbeat_sync()` stop at an explicit Server refusal.
+They raise `ActivityLeaseAuthorityLost` with its original `.reason` for an
+elapsed deadline, closed attempt or closed run. Cancellation continues to raise
+`ActivityCancelled`. The worker submits no result or failure for an attempt
+whose heartbeat refused continuation, even if application code catches that
+exception. A heartbeat check acts only when called. Code that never heartbeats
+can continue until it returns, subject to independent lease/cancellation
+supervision; a running native thread cannot be forcibly stopped or its effects
+undone by the SDK.
+
 ## Development
 
 The installed `durable-workflow-replay-conformance` command runs an offline

docs/sdk-reference.md

@@ -966,6 +966,23 @@ compatibility manifest. A client configured with only the opposite role's token
 still fails before transport for actual worker or control operations; use
 `token` when one shared credential intentionally authorizes both planes.
 
+## Activity heartbeat refusals
+
+Use `await activity.context().heartbeat(details)` in an async activity or
+`activity.context().heartbeat_sync(details)` in a synchronous activity thread.
+Both raise `ActivityCancelled` when cancellation is requested. An explicit
+`can_continue=false` instead raises `ActivityLeaseAuthorityLost` with the
+original `.reason`, such as `activity_deadline_elapsed`, `attempt_closed` or
+`run_closed`. This is lost execution authority, not a new activity failure.
+The worker releases the slot after the callback exits and submits neither a
+result nor a failure for that refused attempt.
+
+Checks happen at the next heartbeat. Code that never heartbeats may continue
+until it returns, subject to separately negotiated lease/cancellation
+supervision. The SDK cannot forcibly stop a running native thread or undo its
+external effects. Keep side effects idempotent and heartbeat at useful progress
+boundaries.
+
 ## Metrics
 
 Pass a recorder to `Client(metrics=...)` or `Worker(metrics=...)` to collect request, poll, and task metrics. The SDK ships a no-op default, an `InMemoryMetrics` recorder for tests or custom exporter loops, and `PrometheusMetrics` for deployments that install the optional extra:

src/durable_workflow/activity.py

@@ -7,7 +7,8 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
-from .errors import ActivityCancelled
+from ._activity_lease import current_activity_authority
+from .errors import ActivityCancelled, ActivityLeaseAuthorityLost
 
 if TYPE_CHECKING:
     from .client import Client
@@ -45,6 +46,7 @@ def __init__(
         self._client = client
         self._heartbeat_callback = heartbeat_callback
         self._cancel_requested = False
+        self._cannot_continue_reason: str | None = None
         try:
             self._event_loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop()
         except RuntimeError:
@@ -61,7 +63,7 @@ def is_cancelled(self) -> bool:
         return self._cancel_requested
 
     def heartbeat_sync(self, details: dict[str, Any] | None = None) -> None:
-        """Heartbeat from a synchronous activity thread and observe cancellation.
+        """Heartbeat from a synchronous activity thread and check attempt authority.
 
         The request runs on the worker's event loop using the original attempt
         and client. Async handlers must use ``await heartbeat()`` instead.
@@ -84,7 +86,7 @@ def heartbeat_sync(self, details: dict[str, Any] | None = None) -> None:
         pending.result()
 
     async def heartbeat(self, details: dict[str, Any] | None = None) -> None:
-        """Report liveness to the server and check for a cancellation request.
+        """Report progress and check cancellation and continued attempt authority.
 
         Long-running activities should call ``heartbeat()`` periodically so the
         server can distinguish a slow-but-alive attempt from a dead worker.
@@ -94,7 +96,12 @@ async def heartbeat(self, details: dict[str, Any] | None = None) -> None:
         Raises :class:`~durable_workflow.errors.ActivityCancelled` when the
         owning workflow has requested cancellation, so the activity can exit
         cleanly at its next natural break point.
+
+        Raises :class:`~durable_workflow.errors.ActivityLeaseAuthorityLost`
+        with the original Server reason when the attempt cannot continue,
+        including an elapsed deadline, closed attempt or closed run.
         """
+        self._raise_if_cannot_continue()
         if self._heartbeat_callback is not None:
             try:
                 await self._heartbeat_callback(details)
@@ -111,6 +118,21 @@ async def heartbeat(self, details: dict[str, Any] | None = None) -> None:
         if isinstance(resp, dict) and resp.get("cancel_requested"):
             self._cancel_requested = True
             raise ActivityCancelled()
+        if isinstance(resp, dict) and resp.get("can_continue") is False:
+            reason = resp.get("reason")
+            self._cannot_continue_reason = reason if isinstance(reason, str) and reason else "activity_cannot_continue"
+            self._raise_if_cannot_continue()
+
+    def _raise_if_cannot_continue(self) -> None:
+        if self._cannot_continue_reason is not None:
+            authority = current_activity_authority.get()
+            if (authority is not None and authority.client is self._client
+                    and authority.task_id == self._info.task_id
+                    and authority.identity["activity_attempt_id"] == self._info.activity_attempt_id
+                    and authority.identity["attempt_number"] == self._info.attempt_number
+                    and authority.identity["lease_owner"] == self._info.worker_id):
+                authority.lost_reason = authority.lost_reason or self._cannot_continue_reason
+            raise ActivityLeaseAuthorityLost(self._info.task_id, self._cannot_continue_reason)
 
 
 def context() -> ActivityContext:

src/durable_workflow/client.py

@@ -6385,7 +6385,8 @@ async def heartbeat_activity_task(
         Worker-plane endpoint. Most code calls
         :meth:`~durable_workflow.ActivityContext.heartbeat` instead, which
         additionally raises :class:`~durable_workflow.errors.ActivityCancelled`
-        when the server reports the activity should stop.
+        for cancellation or :class:`~durable_workflow.errors.ActivityLeaseAuthorityLost`
+        with the original reason when execution cannot continue.
         """
         body: dict[str, Any] = {
             "activity_attempt_id": activity_attempt_id,

src/durable_workflow/worker.py

@@ -3100,7 +3100,15 @@ async def _run_activity_task_body(self, task: dict[str, Any]) -> str:
         )
         _set_context(act_ctx)
         try:
-            result = await self._invoke_remote_activity(task, activity_type, fn, args, act_ctx.info)
+            try:
+                result = await self._invoke_remote_activity(task, activity_type, fn, args, act_ctx.info)
+            finally:
+                # A caught heartbeat refusal cannot restore outcome authority,
+                # including when user code subsequently raises another error.
+                act_ctx._raise_if_cannot_continue()
+        except ActivityLeaseAuthorityLost as error:
+            log.warning("activity %s cannot continue after heartbeat: %s", task_id, error.reason)
+            return "claim_aborted"
         except _RemoteActivityExecutionAborted as error:
             log.warning("remote activity %s claim abandoned: %s", task_id, error)
             return "claim_aborted"
@@ -3134,15 +3142,7 @@ async def _run_activity_task_body(self, task: dict[str, Any]) -> str:
             )
             return "failed_non_retryable"
         except Exception as e:
-            if _is_storage_admission_error(e):
-                raise
-            log.exception("activity failed")
-            await self._report_activity_task_failure(
-                task_id, attempt_id, "activity failure", message=str(e), failure_type=type(e).__name__,
-                failure_class=_exception_class_name(e), failure_code=_exception_code(e),
-                stack_trace=traceback.format_exc(),
-            )
-            return "failed"
+            return await self._handle_activity_exception(task_id, attempt_id, e)
         finally:
             _set_context(None)
 
@@ -3162,6 +3162,17 @@ async def _run_activity_task_body(self, task: dict[str, Any]) -> str:
             return "complete_error"
         return "completed"
 
+    async def _handle_activity_exception(self, task_id: str, attempt_id: str, error: Exception) -> str:
+        if _is_storage_admission_error(error):
+            raise error
+        log.error("activity failed", exc_info=(type(error), error, error.__traceback__))
+        await self._report_activity_task_failure(
+            task_id, attempt_id, "activity failure", message=str(error), failure_type=type(error).__name__,
+            failure_class=_exception_class_name(error), failure_code=_exception_code(error),
+            stack_trace="".join(traceback.format_exception(type(error), error, error.__traceback__)),
+        )
+        return "failed"
+
     async def _decode_activity_arguments(self, task: dict[str, Any]) -> tuple[str, list[Any]] | None:
         task_id: str = task["task_id"]
         attempt_id: str = task.get("activity_attempt_id") or task.get("attempt_id", "")

tests/test_activity_heartbeat_authority.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from contextlib import suppress
+from datetime import datetime, timedelta, timezone
+from unittest.mock import AsyncMock
+
+import httpx
+import pytest
+
+from durable_workflow import activity, serializer
+from durable_workflow._activity_lease import ActivityLeaseAuthority, current_activity_authority
+from durable_workflow.activity import ActivityContext, ActivityInfo
+from durable_workflow.client import Client
+from durable_workflow.errors import ActivityLeaseAuthorityLost, NonRetryableError
+from durable_workflow.worker import Worker
+
+REASONS = ("activity_deadline_elapsed", "attempt_closed", "run_closed")
+
+
+def refusing_client(reason: str) -> AsyncMock:
+    client = AsyncMock(spec=Client)
+    client.heartbeat_activity_task.return_value = {
+        "can_continue": False, "cancel_requested": False, "reason": reason,
+    }
+    return client
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("reason", REASONS)
+@pytest.mark.parametrize("synchronous", (False, True))
+async def test_heartbeat_refusal_carries_original_reason(reason: str, synchronous: bool) -> None:
+    ctx = ActivityContext(info=ActivityInfo(
+        task_id="original-task", activity_type="heartbeat-refusal", activity_attempt_id="original-attempt",
+        attempt_number=1, task_queue="original-queue", worker_id="original-worker",
+    ), client=refusing_client(reason))
+
+    invocation = (asyncio.to_thread(ctx.heartbeat_sync, {"progress": 1}) if synchronous
+                  else ctx.heartbeat({"progress": 1}))
+    with pytest.raises(ActivityLeaseAuthorityLost) as refusal:
+        await invocation
+
+    assert refusal.value.task_id == "original-task"
+    assert refusal.value.reason == reason
+    assert not ctx.is_cancelled
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("reason", REASONS)
+@pytest.mark.parametrize("synchronous", (False, True))
+async def test_default_worker_does_not_publish_after_heartbeat_refusal(
+    reason: str, synchronous: bool,
+) -> None:
+    client = refusing_client(reason)
+    continued: list[str] = []
+
+    def sync_handler() -> str:
+        activity.context().heartbeat_sync()
+        continued.append("after refusal")
+        return "late result"
+
+    async def async_handler() -> str:
+        await activity.context().heartbeat()
+        continued.append("after refusal")
+        return "late result"
+
+    handler = activity.defn(name="heartbeat-refusal")(sync_handler if synchronous else async_handler)
+    worker = Worker(client, task_queue="original-queue", workflows=[], activities=[handler])
+    task = {
+        "task_id": "original-task", "activity_attempt_id": "original-attempt",
+        "activity_type": "heartbeat-refusal", "attempt_number": 1,
+        "arguments": serializer.encode([], codec="avro"), "payload_codec": "avro",
+    }
+
+    assert await worker._run_activity_task(task) == "claim_aborted"
+    assert continued == []
+    client.complete_activity_task.assert_not_called()
+    client.fail_activity_task.assert_not_called()
+    with pytest.raises(RuntimeError, match="outside of an activity"):
+        activity.context()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("same_original", (True, False))
+async def test_refusal_fences_only_the_matching_original_client_authority(same_original: bool) -> None:
+    paths: list[str] = []
+
+    async def respond(request: httpx.Request) -> httpx.Response:
+        paths.append(request.url.path)
+        assert request.url.path.endswith("/heartbeat")
+        return httpx.Response(200, json={
+            "heartbeat_recorded": False, "can_continue": False,
+            "cancel_requested": False, "reason": "attempt_closed",
+        })
+
+    async with Client("http://heartbeat.invalid", worker_token="synthetic-worker-token") as client:
+        await client._http.aclose()
+        client._http = httpx.AsyncClient(base_url="http://heartbeat.invalid", transport=httpx.MockTransport(respond))
+        reference = datetime.now(timezone.utc)
+        def timestamp(value: datetime) -> str:
+            return value.isoformat(timespec="microseconds").replace("+00:00", "Z")
+
+        authority = ActivityLeaseAuthority.delivery(client, client.namespace, {
+            "task_id": "original-task" if same_original else "different-task",
+            "workflow_id": "workflow-1", "run_id": "run-1", "activity_execution_id": "execution-1",
+            "activity_attempt_id": "original-attempt", "lease_owner": "original-worker",
+            "task_queue": "original-queue", "attempt_number": 1,
+            "lease_expires_at": timestamp(reference + timedelta(seconds=30)), "deadlines": {},
+        }, owner="original-worker", task_queue="original-queue", server_time=timestamp(reference),
+            request_started=time.monotonic())
+        token = current_activity_authority.set(authority)
+        try:
+            ctx = ActivityContext(info=ActivityInfo(
+                task_id="original-task", activity_type="heartbeat-refusal", activity_attempt_id="original-attempt",
+                attempt_number=1, task_queue="original-queue", worker_id="original-worker",
+            ), client=client)
+            with pytest.raises(ActivityLeaseAuthorityLost, match="attempt_closed"):
+                await ctx.heartbeat()
+            if same_original:
+                with pytest.raises(ActivityLeaseAuthorityLost, match="attempt_closed"):
+                    await client.complete_activity_task(
+                        task_id="original-task", activity_attempt_id="original-attempt",
+                        lease_owner="original-worker", result="late result",
+                    )
+                assert paths == ["/api/worker/activity-tasks/original-task/heartbeat"]
+            else:
+                assert authority.remaining() > 0
+        finally:
+            current_activity_authority.reset(token)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("late_error", (None, ValueError, NonRetryableError))
+async def test_swallowed_heartbeat_refusal_still_fences_worker_publication(late_error: type[Exception] | None) -> None:
+    client = refusing_client("attempt_closed")
+
+    @activity.defn(name="heartbeat-refusal-caught")
+    async def handler() -> str:
+        with suppress(Exception):
+            await activity.context().heartbeat()
+        if late_error is not None:
+            raise late_error("application error after closed attempt")
+        return "late result"
+
+    worker = Worker(client, task_queue="original-queue", workflows=[], activities=[handler])
+    task = {
+        "task_id": "original-task", "activity_attempt_id": "original-attempt",
+        "activity_type": "heartbeat-refusal-caught", "attempt_number": 1,
+        "arguments": serializer.encode([], codec="avro"), "payload_codec": "avro",
+    }
+
+    assert await worker._run_activity_task(task) == "claim_aborted"
+    client.complete_activity_task.assert_not_called()
+    client.fail_activity_task.assert_not_called()