NEXUS-484: Add support for workflow update Nexus operations

Author: mavemuri

temporalio/client/_impl.py

@@ -787,6 +787,11 @@ async def start_workflow_update(
             ):
                 break
 
+        # Add response link if its a Nexus operation
+        nexus_ctx = temporalio.nexus._operation_context._try_start_operation_context()
+        if nexus_ctx is not None and resp.HasField("link"):
+            nexus_ctx._add_response_link(resp.link)
+
         # Build the handle. If the user's wait stage is COMPLETED, make sure we
         # poll for result.
         handle: WorkflowUpdateHandle[Any] = WorkflowUpdateHandle(
@@ -852,6 +857,23 @@ async def _build_update_workflow_execution_request(
                 )
             ),
         )
+        # Only set Nexus fields for StartWorkflowUpdateInput, skip for UpdateWithStartUpdateWorkflowInput
+        if isinstance(input, StartWorkflowUpdateInput):
+            if input.request_id:
+                req.request.request_id = input.request_id
+            if input.links:
+                req.request.links.extend(input.links)
+            if input.callbacks:
+                req.request.completion_callbacks.extend(
+                    temporalio.api.common.v1.Callback(
+                        nexus=temporalio.api.common.v1.Callback.Nexus(
+                            url=callback.url,
+                            header=callback.headers,
+                        ),
+                        links=input.links or [],
+                    )
+                    for callback in input.callbacks
+                )
         if input.args:
             req.request.input.args.payloads.extend(
                 await data_converter.encode(input.args)

temporalio/client/_interceptor.py

@@ -322,6 +322,10 @@ class StartWorkflowUpdateInput:
     ret_type: type | None
     rpc_metadata: Mapping[str, str | bytes]
     rpc_timeout: timedelta | None
+    # The following options are for Nexus Operation-backed updates. Experimental and unstable
+    callbacks: Sequence[Callback] | None = None
+    links: Sequence[temporalio.api.common.v1.Link] | None = None
+    request_id: str | None = None
 
 
 @dataclass

temporalio/client/_workflow.py

@@ -59,6 +59,7 @@
     ReturnType,
     SelfType,
 )
+from ._callback import Callback
 from ._exceptions import (
     WorkflowContinuedAsNewError,
     WorkflowFailureError,
@@ -955,6 +956,12 @@ async def _start_update(
         result_type: type | None = None,
         rpc_metadata: Mapping[str, str | bytes] = {},
         rpc_timeout: timedelta | None = None,
+        run_id: str | None = None,
+        first_execution_run_id: str | None = None,
+        # The following options are for Nexus Operation-backed updates. Experimental and unstable
+        callbacks: Sequence[Callback] | None = None,
+        links: Sequence[temporalio.api.common.v1.Link] | None = None,
+        request_id: str | None = None,
     ) -> WorkflowUpdateHandle[Any]:
         if wait_for_stage == WorkflowUpdateStage.ADMITTED:
             raise ValueError("ADMITTED wait stage not supported")
@@ -963,11 +970,16 @@ async def _start_update(
             temporalio.workflow._UpdateDefinition.get_name_and_result_type(update)
         )
 
+        if run_id == None:
+            run_id = self._run_id
+        if first_execution_run_id == None:
+            first_execution_run_id = self.first_execution_run_id
+
         return await self._client._impl.start_workflow_update(
             StartWorkflowUpdateInput(
                 id=self._id,
-                run_id=self._run_id,
-                first_execution_run_id=self.first_execution_run_id,
+                run_id=run_id,
+                first_execution_run_id=first_execution_run_id,
                 update_id=id,
                 update=update_name,
                 args=temporalio.common._arg_or_args(arg, args),
@@ -976,6 +988,9 @@ async def _start_update(
                 rpc_metadata=rpc_metadata,
                 rpc_timeout=rpc_timeout,
                 wait_for_stage=wait_for_stage,
+                callbacks=callbacks,
+                links=links,
+                request_id=request_id,
             )
         )
 

temporalio/nexus/__init__.py

@@ -25,15 +25,17 @@
     wait_for_worker_shutdown_sync,
 )
 from ._operation_handlers import (
+    CancelUpdateWorkflowOptions,
     CancelWorkflowRunOptions,
     TemporalOperationHandler,
 )
 from ._temporal_client import TemporalNexusClient, TemporalOperationResult
-from ._token import WorkflowHandle
+from ._token import UpdateHandle, WorkflowHandle
 
 __all__ = (
     "workflow_run_operation",
     "CancelWorkflowRunOptions",
+    "CancelUpdateWorkflowOptions",
     "Info",
     "LoggerAdapter",
     "NexusCallback",
@@ -49,6 +51,7 @@
     "wait_for_worker_shutdown",
     "wait_for_worker_shutdown_sync",
     "WorkflowHandle",
+    "UpdateHandle",
     "TemporalNexusClient",
     "TemporalOperationStartHandlerFunc",
     "TemporalOperationHandler",

temporalio/nexus/_operation_context.py

@@ -715,3 +715,47 @@ async def _start_nexus_backing_workflow(
         )
 
     return WorkflowHandle[ReturnType]._unsafe_from_client_workflow_handle(wf_handle)
+
+
+async def _start_nexus_backed_workflow_update(  # pyright: ignore[reportUnusedFunction]
+    *,
+    temporal_context: _TemporalStartOperationContext,
+    workflow_id: str,
+    update: str | Callable,
+    arg: Any = temporalio.common._arg_unset,
+    args: Sequence[Any] = [],
+    id: str | None = None,
+    result_type: type | None = None,
+    rpc_metadata: Mapping[str, str | bytes] = {},
+    rpc_timeout: timedelta | None = None,
+    run_id: str | None = None,
+    first_execution_run_id: str | None = None,
+) -> temporalio.client.WorkflowUpdateHandle[Any]:
+    # Default update ID to the Nexus request ID for retry-safety (matches sdk-go).
+    update_id = id or temporal_context.nexus_context.request_id
+    # This token is different from the actual token returned to the caller
+    # because return token will have the run_id that is unknowable before
+    # making the call. If run_id is passed, then it will be the same
+    token = OperationToken(
+        type=OperationTokenType.UPDATE_WORKFLOW,
+        namespace=temporal_context.client.namespace,
+        workflow_id=workflow_id,
+        update_id=update_id,
+        run_id=run_id,
+    ).encode()
+    workflow_handle = temporal_context.client.get_workflow_handle(workflow_id)
+    return await workflow_handle._start_update(
+        update,
+        arg,
+        args=args,
+        wait_for_stage=temporalio.client.WorkflowUpdateStage.ACCEPTED,  # hardcoded as nexus only supports async updates
+        id=update_id,
+        result_type=result_type,
+        rpc_metadata=rpc_metadata,
+        rpc_timeout=rpc_timeout,
+        callbacks=temporal_context._get_callbacks(token),
+        links=temporal_context._get_request_links(),
+        request_id=temporal_context.nexus_context.request_id,
+        run_id=run_id,
+        first_execution_run_id=first_execution_run_id,
+    )

temporalio/nexus/_operation_handlers.py

@@ -138,6 +138,25 @@ class CancelWorkflowRunOptions:
     """The ID of the workflow to cancel."""
 
 
+@dataclass(frozen=True)
+class CancelUpdateWorkflowOptions:
+    """Options for cancelling the workflow update backing a Nexus operation.
+
+    These options are built by :py:class:`TemporalOperationHandler` and passed to
+    :py:meth:`TemporalOperationHandler.cancel_workflow_update`.
+
+    .. warning::
+       This API is experimental and unstable.
+    """
+
+    workflow_id: str
+    """The ID of the workflow where the update is running."""
+    update_id: str
+    """The ID of the update to cancel."""
+    run_id: str
+    """The workflow runID that accepted the update"""
+
+
 class TemporalOperationHandler(OperationHandler[InputT, OutputT], ABC):
     """Operation handler for Nexus operations that interact with Temporal.
     Implementations override the start_operation method.
@@ -190,6 +209,15 @@ async def cancel(self, ctx: CancelOperationContext, token: str) -> None:
                     workflow_id=operation_token.workflow_id
                 )
                 await self.cancel_workflow_run(cancel_ctx, options)
+            case OperationTokenType.UPDATE_WORKFLOW:
+                assert operation_token.update_id is not None
+                assert operation_token.run_id is not None
+                cancel_options = CancelUpdateWorkflowOptions(
+                    workflow_id=operation_token.workflow_id,
+                    update_id=operation_token.update_id,
+                    run_id=operation_token.run_id,
+                )
+                await self.cancel_workflow_update(cancel_ctx, cancel_options)
 
     async def cancel_workflow_run(
         self,
@@ -205,3 +233,23 @@ async def cancel_workflow_run(
             options.workflow_id
         )
         await workflow_handle.cancel()
+
+    # draft-review: maybe just move it inline, no need for a function just to error out
+    # check after review in case theres some other way to override/supply custom cancels
+    async def cancel_workflow_update(
+        self,
+        ctx: TemporalCancelOperationContext,  # pyright: ignore[reportUnusedParameter]
+        options: CancelUpdateWorkflowOptions,  # pyright: ignore[reportUnusedParameter]
+    ) -> None:
+        """Cancels the workflow update backing the Nexus operation.
+
+        .. warning::
+           This API is experimental and unstable.
+        """
+        raise HandlerError(
+            """
+                Cancellation is not natively supported for update-workflow Nexus operations.
+                Override a TemporalOperationHandler and implement this method to run cancellable workflow updates.
+            """,
+            type=HandlerErrorType.NOT_IMPLEMENTED,
+        )