NEXUS-484: Support UpdateWorkflow as a Nexus Operation (#1631)

* NEXUS-484: Add support for workflow update Nexus operations

* address comments, cleanup run_id, extend tests

* address comments 2, fix docstrings, clarify signatures and errors

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 Workflow Updates exposed as Nexus Operations. 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,10 @@ async def _start_update(
         result_type: type | None = None,
         rpc_metadata: Mapping[str, str | bytes] = {},
         rpc_timeout: timedelta | None = None,
+        # The following options are for Workflow Updates exposed as Nexus Operations. 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")
@@ -967,7 +972,7 @@ async def _start_update(
             StartWorkflowUpdateInput(
                 id=self._id,
                 run_id=self._run_id,
-                first_execution_run_id=self.first_execution_run_id,
+                first_execution_run_id=self._first_execution_run_id,
                 update_id=id,
                 update=update_name,
                 args=temporalio.common._arg_or_args(arg, args),
@@ -976,6 +981,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,6 +25,7 @@
     wait_for_worker_shutdown_sync,
 )
 from ._operation_handlers import (
+    CancelUpdateWorkflowOptions,
     CancelWorkflowRunOptions,
     TemporalOperationHandler,
 )
@@ -34,6 +35,7 @@
 __all__ = (
     "workflow_run_operation",
     "CancelWorkflowRunOptions",
+    "CancelUpdateWorkflowOptions",
     "Info",
     "LoggerAdapter",
     "NexusCallback",

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_operation_workflow_update(  # pyright: ignore[reportUnusedFunction]
+    *,
+    temporal_context: _TemporalStartOperationContext,
+    workflow_id: str,
+    update: str | Callable,
+    arg: Any = temporalio.common._arg_unset,
+    args: Sequence[Any] = [],
+    update_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 = update_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, run_id=run_id, first_execution_run_id=first_execution_run_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,
+    )

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()
+
+    async def cancel_workflow_update(
+        self,
+        ctx: TemporalCancelOperationContext,  # pyright: ignore[reportUnusedParameter]
+        options: CancelUpdateWorkflowOptions,  # pyright: ignore[reportUnusedParameter]
+    ) -> None:
+        """Cancels the Workflow Update triggered by the Nexus operation. There is no native way to cancel an accepted Update, so cancellation depends on the Update handler itself.
+        Override this method and coordinate with the specific Update handler to trigger a cancellation.
+
+
+        .. warning::
+           This API is experimental and unstable.
+        """
+        raise HandlerError(
+            """
+                There is no native way to cancel an accepted Update, so cancellation depends on the Update handler itself. 
+                Override this method and coordinate with the specific Update handler to trigger a cancellation.
+            """,
+            type=HandlerErrorType.NOT_IMPLEMENTED,
+        )

temporalio/nexus/_temporal_client.py

@@ -15,13 +15,14 @@
     overload,
 )
 
-from nexusrpc import HandlerError, HandlerErrorType
+from nexusrpc import HandlerError, HandlerErrorType, OperationError, OperationErrorState
 from nexusrpc.handler import StartOperationResultAsync, StartOperationResultSync
 from typing_extensions import Self
 
 import temporalio.common
 from temporalio.nexus._operation_context import (
     _start_nexus_backing_workflow,
+    _start_nexus_operation_workflow_update,
     _TemporalStartOperationContext,
 )
 from temporalio.types import (
@@ -33,8 +34,11 @@
     SelfType,
 )
 
+from ._token import OperationToken, OperationTokenType
+
 if TYPE_CHECKING:
     import temporalio.client
+    import temporalio.workflow
 
 
 _ResultT = TypeVar("_ResultT")
@@ -279,6 +283,91 @@ async def start_workflow(
         """
         ...
 
+    # Overload for no-param update
+    @overload
+    async def start_workflow_update(
+        self,
+        workflow_id: str,
+        update: temporalio.workflow.UpdateMethodMultiParam[[Any], ReturnType],
+        *,
+        update_id: str | 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,
+    ) -> TemporalOperationResult[ReturnType]: ...
+
+    # Overload for single-param update
+    @overload
+    async def start_workflow_update(
+        self,
+        workflow_id: str,
+        update: temporalio.workflow.UpdateMethodMultiParam[
+            [Any, ParamType], ReturnType
+        ],
+        arg: ParamType,
+        *,
+        update_id: str | 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,
+    ) -> TemporalOperationResult[ReturnType]: ...
+
+    # Overload for multi-param update
+    @overload
+    async def start_workflow_update(
+        self,
+        workflow_id: str,
+        update: temporalio.workflow.UpdateMethodMultiParam[MultiParamSpec, ReturnType],
+        *,
+        args: MultiParamSpec.args,  # type: ignore
+        update_id: str | 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,
+    ) -> TemporalOperationResult[ReturnType]: ...
+
+    # Overload for string-name update
+    @overload
+    async def start_workflow_update(
+        self,
+        workflow_id: str,
+        update: str,
+        arg: Any = temporalio.common._arg_unset,
+        *,
+        args: Sequence[Any] = [],
+        update_id: str | None = None,
+        result_type: type[ReturnType] | 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,
+    ) -> TemporalOperationResult[ReturnType]: ...
+
+    @abstractmethod
+    async def start_workflow_update(
+        self,
+        workflow_id: str,
+        update: str | Callable,
+        arg: Any = temporalio.common._arg_unset,
+        *,
+        args: Sequence[Any] = [],
+        update_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,
+    ) -> TemporalOperationResult[Any]:
+        """Start a Workflow Update-backed Nexus Operation.
+
+        .. warning::
+            This API is experimental and unstable.
+        """
+        ...
+
 
 class _TemporalNexusClient(TemporalNexusClient):  # pyright: ignore[reportUnusedClass]
     """Nexus-aware wrapper around a Temporal Client.
@@ -377,3 +466,55 @@ async def start_workflow(
             )
 
         return TemporalOperationResult.async_token(wf_handle.to_token())
+
+    async def start_workflow_update(
+        self,
+        workflow_id: str,
+        update: str | Callable,
+        arg: Any = temporalio.common._arg_unset,
+        *,
+        args: Sequence[Any] = [],
+        update_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,
+    ) -> TemporalOperationResult[Any]:
+        """Start a Workflow Update-backed Nexus Operation."""
+        if not self._temporal_context.nexus_context.callback_url:
+            raise HandlerError(
+                "callback URL is required for a workflow update Nexus operation",
+                type=HandlerErrorType.BAD_REQUEST,
+            )
+        with self._reserve_async_start():
+            update_handle = await _start_nexus_operation_workflow_update(
+                temporal_context=self._temporal_context,
+                workflow_id=workflow_id,
+                update=update,
+                arg=arg,
+                args=args,
+                update_id=update_id,
+                result_type=result_type,
+                rpc_metadata=rpc_metadata,
+                rpc_timeout=rpc_timeout,
+                run_id=run_id,
+                first_execution_run_id=first_execution_run_id,
+            )
+            # If the update has already completed, return the result synchronously
+            if update_handle._known_outcome is not None:
+                try:
+                    result = await update_handle.result()
+                except temporalio.client.WorkflowUpdateFailedError as err:
+                    raise OperationError(
+                        str(err), state=OperationErrorState.FAILED
+                    ) from err
+                return TemporalOperationResult.sync(result)
+        token = OperationToken(
+            type=OperationTokenType.UPDATE_WORKFLOW,
+            namespace=update_handle._client.namespace,
+            workflow_id=update_handle.workflow_id,
+            update_id=update_handle.id,
+            run_id=update_handle.workflow_run_id,
+        ).encode()
+        return TemporalOperationResult.async_token(token)