address comments, cleanup run_id, extend tests

Author: mavemuri

temporalio/client/_workflow.py

@@ -956,8 +956,6 @@ 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,
@@ -970,16 +968,11 @@ 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=run_id,
-                first_execution_run_id=first_execution_run_id,
+                run_id=self._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),

temporalio/nexus/_operation_context.py

@@ -743,7 +743,9 @@ async def _start_nexus_backed_workflow_update(  # pyright: ignore[reportUnusedFu
         update_id=update_id,
         run_id=run_id,
     ).encode()
-    workflow_handle = temporal_context.client.get_workflow_handle(workflow_id)
+    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,
@@ -756,6 +758,4 @@ async def _start_nexus_backed_workflow_update(  # pyright: ignore[reportUnusedFu
         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

@@ -234,22 +234,22 @@ async def cancel_workflow_run(
         )
         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.
+        """Cancels the workflow update backing the Nexus operation. Cancellation is not natively supported for update-workflow Nexus operations.
+        Inherit a TemporalOperationHandler and override this method to run cancellable workflow updates.
+
 
         .. 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.
+                Inherit a TemporalOperationHandler and override this method to run cancellable workflow updates.
             """,
             type=HandlerErrorType.NOT_IMPLEMENTED,
         )

temporalio/nexus/_temporal_client.py

@@ -25,7 +25,6 @@
     _start_nexus_backing_workflow,
     _TemporalStartOperationContext,
 )
-from temporalio.nexus._token import UpdateHandle
 from temporalio.types import (
     MethodAsyncNoParam,
     MethodAsyncSingleParam,
@@ -35,6 +34,8 @@
     SelfType,
 )
 
+from ._token import OperationToken, OperationTokenType
+
 if TYPE_CHECKING:
     import temporalio.client
     import temporalio.workflow
@@ -292,6 +293,8 @@ async def start_workflow_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
@@ -307,6 +310,8 @@ async def start_workflow_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
@@ -320,6 +325,8 @@ async def start_workflow_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
@@ -335,10 +342,10 @@ async def start_workflow_update(
         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]: ...
 
-    # draft-review: check why run_id and first_execution_run_id are not used
-    # for update workflow in python sdk
     @abstractmethod
     async def start_workflow_update(
         self,
@@ -495,7 +502,6 @@ async def start_workflow_update(
                 first_execution_run_id=first_execution_run_id,
             )
             # If the update has already completed, return the result synchronously
-            # This is in-line with the Go implementation as well
             if update_handle._known_outcome is not None:
                 try:
                     result = await update_handle.result()
@@ -504,7 +510,11 @@ async def start_workflow_update(
                         str(err), state=OperationErrorState.FAILED
                     ) from err
                 return TemporalOperationResult.sync(result)
-        nexus_handle: UpdateHandle[Any] = (
-            UpdateHandle._unsafe_from_client_workflow_update_handle(update_handle)
-        )
-        return TemporalOperationResult.async_token(nexus_handle.to_token())
+        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)

temporalio/nexus/_token.py

@@ -228,39 +228,6 @@ class UpdateHandle(Generic[OutputT]):
     # serialized token; it's only used to reject newer token versions on load.
     version: int | None = None
 
-    def _to_client_workflow_update_handle(
-        self,
-        client: temporalio.client.Client,
-        result_type: type[OutputT] | None = None,
-    ) -> temporalio.client.WorkflowUpdateHandle[Any]:
-        """Create a :py:class:`temporalio.client.WorkflowUpdateHandle` from the token."""
-        if client.namespace != self.namespace:
-            raise ValueError(
-                f"Client namespace {client.namespace} does not match "
-                f"operation token namespace {self.namespace}"
-            )
-        workflow_handle = client.get_workflow_handle(self.workflow_id)
-        return workflow_handle.get_update_handle(
-            self.update_id, result_type=result_type
-        )
-
-    @classmethod
-    def _unsafe_from_client_workflow_update_handle(
-        cls, workflow_update_handle: temporalio.client.WorkflowUpdateHandle[Any]
-    ) -> UpdateHandle[OutputT]:
-        """Create a :py:class:`UpdateHandle` from a :py:class:`temporalio.client.WorkflowUpdateHandle`.
-
-        This is a private method not intended to be used by users. It does not check
-        that the supplied client.WorkflowUpdateHandle references a workflow that has been
-        instrumented to supply the result of a Nexus operation.
-        """
-        return cls(
-            namespace=workflow_update_handle._client.namespace,
-            workflow_id=workflow_update_handle.workflow_id,
-            run_id=workflow_update_handle.workflow_run_id,
-            update_id=workflow_update_handle._id,
-        )
-
     def to_token(self) -> str:
         """Convert handle to a base64url-encoded token string."""
         return OperationToken(