Add workflow cancellation reason

Author: Sushisource

temporalio/worker/_workflow_instance.py

@@ -265,7 +265,7 @@ def __init__(self, det: WorkflowInstanceDetails) -> None:
         )
         self._primary_task: asyncio.Task[None] | None = None
         self._time_ns = 0
-        self._cancel_requested = False
+        self._cancel_reason: str | None = None
         self._deployment_version_for_current_task: None | (
             temporalio.bridge.proto.common.WorkerDeploymentVersion
         ) = None
@@ -595,10 +595,9 @@ def _apply(
             raise RuntimeError(f"Unrecognized job: {job.WhichOneof('variant')}")
 
     def _apply_cancel_workflow(
-        self, _job: temporalio.bridge.proto.workflow_activation.CancelWorkflow
+        self, job: temporalio.bridge.proto.workflow_activation.CancelWorkflow
     ) -> None:
-        self._cancel_requested = True
-        # TODO(cretz): Details or cancel message or whatever?
+        self._cancel_reason = job.reason
         if self._primary_task:
             # The primary task may not have started yet and we want to give the
             # workflow the ability to receive the cancellation, so we must defer
@@ -799,7 +798,6 @@ def _apply_remove_from_cache(
         self, _job: temporalio.bridge.proto.workflow_activation.RemoveFromCache
     ) -> None:
         self._deleting = True
-        self._cancel_requested = True
         # We consider eviction to be under replay so that certain code like
         # logging that avoids replaying doesn't run during eviction either
         self._is_replaying = True
@@ -1189,6 +1187,9 @@ def workflow_continue_as_new(
             )
         )
 
+    def workflow_cancellation_reason(self) -> str | None:
+        return self._cancel_reason
+
     def workflow_extern_functions(self) -> Mapping[str, Callable]:
         return self._extern_functions
 
@@ -2046,7 +2047,7 @@ async def run_child() -> Any:
                     and (t := asyncio.current_task()) is not None
                 ):
                     t.uncancel()  # type: ignore[union-attr]
-                if self._cancel_requested:
+                if self._cancel_reason is not None or self._deleting:
                     raise
 
     async def _outbound_start_nexus_operation(
@@ -2102,7 +2103,7 @@ async def operation_handle_fn() -> OutputT:
                     and (t := asyncio.current_task()) is not None
                 ):
                     t.uncancel()  # type: ignore[union-attr]
-                if self._cancel_requested:
+                if self._cancel_reason is not None or self._deleting:
                     raise
 
     #### Miscellaneous helpers ####
@@ -2588,8 +2589,9 @@ async def _run_top_level_workflow_function(self, coro: Awaitable[None]) -> None:
             # cancel later on will show the workflow as cancelled. But this is
             # a Temporal limitation in that cancellation is a state not an
             # event.
-            if self._cancel_requested and temporalio.exceptions.is_cancelled_exception(
-                err
+            if (
+                self._cancel_reason is not None
+                and temporalio.exceptions.is_cancelled_exception(err)
             ):
                 self._add_command().cancel_workflow_execution.SetInParent()
             elif self.workflow_is_failure_exception(err):

temporalio/workflow/__init__.py

@@ -59,6 +59,7 @@
     _current_update_info,
     _Runtime,
     _set_current_update_info,
+    cancellation_reason,
     current_update_info,
     deprecate_patch,
     extern_functions,
@@ -196,6 +197,7 @@
     "get_last_completion_result",
     "get_last_failure",
     "has_last_completion_result",
+    "cancellation_reason",
     "in_workflow",
     "info",
     "instance",

temporalio/workflow/_context.py

@@ -36,6 +36,7 @@
     "ParentInfo",
     "RootInfo",
     "UpdateInfo",
+    "cancellation_reason",
     "current_update_info",
     "deprecate_patch",
     "extern_functions",
@@ -295,6 +296,9 @@ def workflow_continue_as_new(
         initial_versioning_behavior: ContinueAsNewVersioningBehavior | None,
     ) -> NoReturn: ...
 
+    @abstractmethod
+    def workflow_cancellation_reason(self) -> str | None: ...
+
     @abstractmethod
     def workflow_extern_functions(self) -> Mapping[str, Callable]: ...
 
@@ -591,6 +595,27 @@ def in_workflow() -> bool:
     return _Runtime.maybe_current() is not None
 
 
+def cancellation_reason() -> str | None:
+    """Reason the workflow was cancelled, or None if no external cancellation
+    request has been received.
+
+    A non-None value (including an empty string) indicates that the workflow
+    received an explicit cancellation request from the server. This can be used
+    when catching an :py:class:`asyncio.CancelledError` to distinguish a
+    workflow-level cancel from a cancel that originated from inner asyncio task
+    cancellation.
+
+    Note, this only reflects cancellation requested via the server; it is not
+    set for cache eviction or for cancels of inner tasks/scopes.
+
+    Returns:
+        The reason string sent with the workflow cancellation request (which
+        may be empty), or ``None`` if the workflow has not been cancelled via
+        an external request.
+    """
+    return _Runtime.current().workflow_cancellation_reason()
+
+
 def memo() -> Mapping[str, Any]:
     """Current workflow's memo values, converted without type hints.
 

tests/worker/test_workflow.py

@@ -1093,6 +1093,74 @@ async def started() -> bool:
         assert (await handle.describe()).status == WorkflowExecutionStatus.CANCELED
 
 
+@workflow.defn
+class CancelReasonWorkflow:
+    def __init__(self) -> None:
+        self._started = False
+        # Reason observed when the inner task was cancelled (no external
+        # workflow cancel has happened yet at that point).
+        self._reason_inner: str | None = "unset"
+        # Reason observed in the outer CancelledError handler after the
+        # external workflow cancel has been delivered.
+        self._reason_outer: str | None = "unset"
+
+    @workflow.run
+    async def run(self) -> NoReturn:
+        self._started = True
+        task = asyncio.create_task(asyncio.sleep(1000))
+        try:
+            task.cancel()
+            await task
+        except asyncio.CancelledError:
+            self._reason_inner = workflow.cancellation_reason()
+        try:
+            await asyncio.sleep(1000)
+        except asyncio.CancelledError:
+            self._reason_outer = workflow.cancellation_reason()
+            raise
+        raise RuntimeError("unreachable")
+
+    @workflow.query
+    def started(self) -> bool:
+        return self._started
+
+    @workflow.query
+    def reason_inner(self) -> str | None:
+        return self._reason_inner
+
+    @workflow.query
+    def reason_outer(self) -> str | None:
+        return self._reason_outer
+
+
+async def test_workflow_cancellation_reason(client: Client):
+    async with new_worker(client, CancelReasonWorkflow) as worker:
+        handle = await client.start_workflow(
+            CancelReasonWorkflow.run,
+            id=f"workflow-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+        )
+
+        async def started() -> bool:
+            return await handle.query(CancelReasonWorkflow.started)
+
+        await assert_eq_eventually(True, started)
+        # Before any external cancel, reason is None even though an inner task
+        # cancel has already been observed.
+        assert await handle.query(CancelReasonWorkflow.reason_inner) is None
+
+        await handle.cancel()
+        with pytest.raises(WorkflowFailureError) as err:
+            await handle.result()
+        assert isinstance(err.value.cause, CancelledError)
+
+        # After external cancel, reason is a string (empty since the Python
+        # client does not send one) — non-None is the load-bearing distinction.
+        outer = await handle.query(CancelReasonWorkflow.reason_outer)
+        assert outer is not None
+        assert isinstance(outer, str)
+
+
 @workflow.defn
 class TrapCancelWorkflow:
     @workflow.run