address comments, cleanup run_id, extend tests

mavemuri · mavemuri · commit aa738a477706 · 2026-07-20T13:38:54.000-07:00
diff --git a/temporalio/client/_workflow.py b/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),
diff --git a/temporalio/nexus/__init__.py b/temporalio/nexus/__init__.py
@@ -30,7 +30,7 @@
     TemporalOperationHandler,
 )
 from ._temporal_client import TemporalNexusClient, TemporalOperationResult
-from ._token import UpdateHandle, WorkflowHandle
+from ._token import WorkflowHandle
 
 __all__ = (
     "workflow_run_operation",
@@ -51,7 +51,6 @@
     "wait_for_worker_shutdown",
     "wait_for_worker_shutdown_sync",
     "WorkflowHandle",
-    "UpdateHandle",
     "TemporalNexusClient",
     "TemporalOperationStartHandlerFunc",
     "TemporalOperationHandler",
diff --git a/temporalio/nexus/_operation_context.py b/temporalio/nexus/_operation_context.py
@@ -724,15 +724,15 @@ async def _start_nexus_backed_workflow_update(  # pyright: ignore[reportUnusedFu
     update: str | Callable,
     arg: Any = temporalio.common._arg_unset,
     args: Sequence[Any] = [],
-    id: str | None = None,
+    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 = id or temporal_context.nexus_context.request_id
+    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
@@ -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,
     )
diff --git a/temporalio/nexus/_operation_handlers.py b/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,
         )
diff --git a/temporalio/nexus/_temporal_client.py b/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
@@ -289,9 +290,11 @@ async def start_workflow_update(
         workflow_id: str,
         update: temporalio.workflow.UpdateMethodMultiParam[[Any], ReturnType],
         *,
-        id: str | None = None,
+        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
@@ -304,9 +307,11 @@ async def start_workflow_update(
         ],
         arg: ParamType,
         *,
-        id: str | None = None,
+        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
@@ -317,9 +322,11 @@ async def start_workflow_update(
         update: temporalio.workflow.UpdateMethodMultiParam[MultiParamSpec, ReturnType],
         *,
         args: MultiParamSpec.args,  # type: ignore
-        id: str | None = None,
+        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
@@ -331,14 +338,14 @@ async def start_workflow_update(
         arg: Any = temporalio.common._arg_unset,
         *,
         args: Sequence[Any] = [],
-        id: str | None = None,
+        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]: ...
 
-    # 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,
@@ -347,7 +354,7 @@ async def start_workflow_update(
         arg: Any = temporalio.common._arg_unset,
         *,
         args: Sequence[Any] = [],
-        id: str | None = None,
+        update_id: str | None = None,
         result_type: type | None = None,
         rpc_metadata: Mapping[str, str | bytes] = {},
         rpc_timeout: timedelta | None = None,
@@ -467,7 +474,7 @@ async def start_workflow_update(
         arg: Any = temporalio.common._arg_unset,
         *,
         args: Sequence[Any] = [],
-        id: str | None = None,
+        update_id: str | None = None,
         result_type: type | None = None,
         rpc_metadata: Mapping[str, str | bytes] = {},
         rpc_timeout: timedelta | None = None,
@@ -487,15 +494,14 @@ async def start_workflow_update(
                 update=update,
                 arg=arg,
                 args=args,
-                id=id,
+                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
-            # 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)
diff --git a/temporalio/nexus/_token.py b/temporalio/nexus/_token.py
@@ -211,91 +211,6 @@ def from_token(cls, token: str) -> WorkflowHandle[OutputT]:
         )
 
 
-@dataclass(frozen=True)
-class UpdateHandle(Generic[OutputT]):
-    """A handle to a workflow update that is backing a Nexus operation.
-
-    Do not instantiate this directly. Use
-    :py:func:`temporalio.nexus.TemporalNexusClient.start_workflow_update` to create a
-    handle.
-    """
-
-    namespace: str
-    workflow_id: str
-    update_id: str
-    run_id: str | None = None
-    # Version of the token. Treated as v1 if missing. This field is not included in the
-    # 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(
-            type=OperationTokenType.UPDATE_WORKFLOW,
-            namespace=self.namespace,
-            workflow_id=self.workflow_id,
-            run_id=self.run_id,
-            update_id=self.update_id,
-        ).encode()
-
-    @classmethod
-    def from_token(cls, token: str) -> UpdateHandle[OutputT]:
-        """Decodes and validates a token from its base64url-encoded string representation."""
-        op_token = OperationToken.decode(token)
-        if op_token.type != OperationTokenType.UPDATE_WORKFLOW:
-            raise TypeError(
-                f"invalid update token type: {op_token.type}, expected: {OperationTokenType.UPDATE_WORKFLOW}"
-            )
-
-        if op_token.version is not None and op_token.version != 0:
-            raise TypeError(
-                "invalid update token: 'v' field, if present, must be 0 or null/absent"
-            )
-
-        assert op_token.update_id is not None
-
-        return cls(
-            namespace=op_token.namespace,
-            workflow_id=op_token.workflow_id,
-            run_id=op_token.run_id,
-            update_id=op_token.update_id,
-            version=op_token.version,
-        )
-
-
 def _base64url_encode_no_padding(data: bytes) -> str:
     return base64.urlsafe_b64encode(data).decode("utf-8").rstrip("=")
 
diff --git a/tests/nexus/test_operation_token.py b/tests/nexus/test_operation_token.py
@@ -7,7 +7,6 @@
 from temporalio.nexus._token import (
     OperationToken,
     OperationTokenType,
-    UpdateHandle,
     WorkflowHandle,
 )
 
@@ -43,13 +42,6 @@ def test_workflow_handle_to_from_token_round_trip():
     assert WorkflowHandle[str].from_token(handle.to_token()) == handle
 
 
-def test_update_handle_to_from_token_round_trip():
-    handle = UpdateHandle[str](
-        namespace="default", workflow_id="workflow-id", update_id="update-id"
-    )
-    assert UpdateHandle[str].from_token(handle.to_token()) == handle
-
-
 @pytest.mark.parametrize(
     ("token", "expected"),
     [
diff --git a/tests/nexus/test_temporal_operation.py b/tests/nexus/test_temporal_operation.py
