Interrupt heartbeating activity on pause

temporalio/activity.py

@@ -34,6 +34,9 @@
     overload,
 )
 
+import temporalio.bridge
+import temporalio.bridge.proto
+import temporalio.bridge.proto.activity_task
 import temporalio.common
 import temporalio.converter
 
@@ -135,6 +138,34 @@ def _logger_details(self) -> Mapping[str, Any]:
 _current_context: contextvars.ContextVar[_Context] = contextvars.ContextVar("activity")
 
 
+@dataclass
+class _ActivityCancellationDetailsHolder:
+    details: Optional[ActivityCancellationDetails] = None
+
+
+@dataclass(frozen=True)
+class ActivityCancellationDetails:
+    """Provides the reasons for the activity's cancellation. Cancellation details are set once and do not change once set."""
+
+    not_found: bool = False
+    cancel_requested: bool = False
+    paused: bool = False
+    timed_out: bool = False
+    worker_shutdown: bool = False
+
+    @staticmethod
+    def _from_proto(
+        proto: temporalio.bridge.proto.activity_task.ActivityCancellationDetails,
+    ) -> ActivityCancellationDetails:
+        return ActivityCancellationDetails(
+            not_found=proto.is_not_found,
+            cancel_requested=proto.is_cancelled,
+            paused=proto.is_paused,
+            timed_out=proto.is_timed_out,
+            worker_shutdown=proto.is_worker_shutdown,
+        )
+
+
 @dataclass
 class _Context:
     info: Callable[[], Info]
@@ -148,6 +179,7 @@ class _Context:
         temporalio.converter.PayloadConverter,
     ]
     runtime_metric_meter: Optional[temporalio.common.MetricMeter]
+    cancellation_details: _ActivityCancellationDetailsHolder
     _logger_details: Optional[Mapping[str, Any]] = None
     _payload_converter: Optional[temporalio.converter.PayloadConverter] = None
     _metric_meter: Optional[temporalio.common.MetricMeter] = None
@@ -260,6 +292,11 @@ def info() -> Info:
     return _Context.current().info()
 
 
+def cancellation_details() -> Optional[ActivityCancellationDetails]:
+    """Cancellation details of the current activity, if any. Once set, cancellation details do not change."""
+    return _Context.current().cancellation_details.details
+
+
 def heartbeat(*details: Any) -> None:
     """Send a heartbeat for the current activity.
 

temporalio/bridge/proto/activity_task/__init__.py

@@ -1,7 +1,14 @@
-from .activity_task_pb2 import ActivityCancelReason, ActivityTask, Cancel, Start
+from .activity_task_pb2 import (
+    ActivityCancellationDetails,
+    ActivityCancelReason,
+    ActivityTask,
+    Cancel,
+    Start,
+)
 
 __all__ = [
     "ActivityCancelReason",
+    "ActivityCancellationDetails",
     "ActivityTask",
     "Cancel",
     "Start",

temporalio/bridge/proto/activity_task/activity_task_pb2.pyi

@@ -45,6 +45,8 @@ class _ActivityCancelReasonEnumTypeWrapper(
     """Activity timed out"""
     WORKER_SHUTDOWN: _ActivityCancelReason.ValueType  # 3
     """Core is shutting down and the graceful timeout has elapsed"""
+    PAUSED: _ActivityCancelReason.ValueType  # 4
+    """Activity was paused"""
 
 class ActivityCancelReason(
     _ActivityCancelReason, metaclass=_ActivityCancelReasonEnumTypeWrapper
@@ -58,6 +60,8 @@ TIMED_OUT: ActivityCancelReason.ValueType  # 2
 """Activity timed out"""
 WORKER_SHUTDOWN: ActivityCancelReason.ValueType  # 3
 """Core is shutting down and the graceful timeout has elapsed"""
+PAUSED: ActivityCancelReason.ValueType  # 4
+"""Activity was paused"""
 global___ActivityCancelReason = ActivityCancelReason
 
 class ActivityTask(google.protobuf.message.Message):
@@ -320,14 +324,66 @@ class Cancel(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
 
     REASON_FIELD_NUMBER: builtins.int
+    DETAILS_FIELD_NUMBER: builtins.int
     reason: global___ActivityCancelReason.ValueType
+    """Primary cancellation reason"""
+    @property
+    def details(self) -> global___ActivityCancellationDetails:
+        """Activity cancellation details, surfaces all cancellation reasons."""
     def __init__(
         self,
         *,
         reason: global___ActivityCancelReason.ValueType = ...,
+        details: global___ActivityCancellationDetails | None = ...,
     ) -> None: ...
+    def HasField(
+        self, field_name: typing_extensions.Literal["details", b"details"]
+    ) -> builtins.bool: ...
     def ClearField(
-        self, field_name: typing_extensions.Literal["reason", b"reason"]
+        self,
+        field_name: typing_extensions.Literal[
+            "details", b"details", "reason", b"reason"
+        ],
     ) -> None: ...
 
 global___Cancel = Cancel
+
+class ActivityCancellationDetails(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    IS_NOT_FOUND_FIELD_NUMBER: builtins.int
+    IS_CANCELLED_FIELD_NUMBER: builtins.int
+    IS_PAUSED_FIELD_NUMBER: builtins.int
+    IS_TIMED_OUT_FIELD_NUMBER: builtins.int
+    IS_WORKER_SHUTDOWN_FIELD_NUMBER: builtins.int
+    is_not_found: builtins.bool
+    is_cancelled: builtins.bool
+    is_paused: builtins.bool
+    is_timed_out: builtins.bool
+    is_worker_shutdown: builtins.bool
+    def __init__(
+        self,
+        *,
+        is_not_found: builtins.bool = ...,
+        is_cancelled: builtins.bool = ...,
+        is_paused: builtins.bool = ...,
+        is_timed_out: builtins.bool = ...,
+        is_worker_shutdown: builtins.bool = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing_extensions.Literal[
+            "is_cancelled",
+            b"is_cancelled",
+            "is_not_found",
+            b"is_not_found",
+            "is_paused",
+            b"is_paused",
+            "is_timed_out",
+            b"is_timed_out",
+            "is_worker_shutdown",
+            b"is_worker_shutdown",
+        ],
+    ) -> None: ...
+
+global___ActivityCancellationDetails = ActivityCancellationDetails

temporalio/bridge/src/client.rs

@@ -235,6 +235,9 @@ impl ClientRef {
                 "patch_schedule" => {
                     rpc_call!(retry_client, call, patch_schedule)
                 }
+                "pause_activity" => {
+                    rpc_call!(retry_client, call, pause_activity)
+                }
                 "poll_activity_task_queue" => {
                     rpc_call!(retry_client, call, poll_activity_task_queue)
                 }
@@ -325,6 +328,9 @@ impl ClientRef {
                 "trigger_workflow_rule" => {
                     rpc_call!(retry_client, call, trigger_workflow_rule)
                 }
+                "unpause_activity" => {
+                    rpc_call!(retry_client, call, unpause_activity)
+                }
                 "update_namespace" => {
                     rpc_call_on_trait!(retry_client, call, WorkflowService, update_namespace)
                 }

temporalio/client.py

@@ -6264,8 +6264,9 @@ async def heartbeat_async_activity(
                 metadata=input.rpc_metadata,
                 timeout=input.rpc_timeout,
             )
-            if resp_by_id.cancel_requested:
+            if resp_by_id.cancel_requested or resp_by_id.activity_paused:
                 raise AsyncActivityCancelledError()
+
         else:
             resp = await self._client.workflow_service.record_activity_task_heartbeat(
                 temporalio.api.workflowservice.v1.RecordActivityTaskHeartbeatRequest(
@@ -6278,7 +6279,7 @@ async def heartbeat_async_activity(
                 metadata=input.rpc_metadata,
                 timeout=input.rpc_timeout,
             )
-            if resp.cancel_requested:
+            if resp.cancel_requested or resp.activity_paused:
                 raise AsyncActivityCancelledError()
 
     async def complete_async_activity(self, input: CompleteAsyncActivityInput) -> None:

temporalio/testing/_activity.py

@@ -74,15 +74,29 @@ def __init__(self) -> None:
         self._cancelled = False
         self._worker_shutdown = False
         self._activities: Set[_Activity] = set()
+        self._cancellation_details = (
+            temporalio.activity._ActivityCancellationDetailsHolder()
+        )
 
-    def cancel(self) -> None:
+    def cancel(
+        self,
+        cancellation_details: temporalio.activity.ActivityCancellationDetails = temporalio.activity.ActivityCancellationDetails(
+            cancel_requested=True
+        ),
+    ) -> None:
         """Cancel the activity.
 
+        Args:
+            cancellation_details: details about the cancellation. These will
+                be accessible through temporalio.activity.cancellation_details()
+                in the activity after cancellation.
+
         This only has an effect on the first call.
         """
         if self._cancelled:
             return
         self._cancelled = True
+        self._cancellation_details.details = cancellation_details
         for act in self._activities:
             act.cancel()
 
@@ -154,6 +168,7 @@ def __init__(
             else self.cancel_thread_raiser.shielded,
             payload_converter_class_or_instance=env.payload_converter,
             runtime_metric_meter=env.metric_meter,
+            cancellation_details=env._cancellation_details,
         )
         self.task: Optional[asyncio.Task] = None