Require wait_for_stage parameter for Nexus Workflow Updates (#1883)

Author: atol

CHANGELOG.md

@@ -37,6 +37,7 @@ to include examples, links to docs, or any other relevant information.
   remain available at runtime and retain their static type information.
   New code should depend on `temporalio-openai-agents` directly and import
   `temporalio.openai_agents`.
+- **Experimental**: Nexus Workflow Updates now require `wait_for_stage` to be explicitly set to `ACCEPTED`.
 
 ### Fixed
 

temporalio/nexus/_operation_context.py

@@ -711,13 +711,16 @@ async def _start_nexus_operation_workflow_update(  # pyright: ignore[reportUnuse
     update: str | Callable,
     arg: Any = temporalio.common._arg_unset,
     args: Sequence[Any] = [],
+    wait_for_stage: temporalio.client.WorkflowUpdateStage,
     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]:
+    if wait_for_stage != temporalio.client.WorkflowUpdateStage.ACCEPTED:
+        raise ValueError("Only ACCEPTED wait stage is supported")
     # 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
     workflow_handle = temporal_context.client.get_workflow_handle(
@@ -728,7 +731,7 @@ async def _start_nexus_operation_workflow_update(  # pyright: ignore[reportUnuse
             update,
             arg,
             args=args,
-            wait_for_stage=temporalio.client.WorkflowUpdateStage.ACCEPTED,  # hardcoded as nexus only supports async updates
+            wait_for_stage=wait_for_stage,
             id=update_id,
             result_type=result_type,
             rpc_metadata=rpc_metadata,

temporalio/nexus/_temporal_client.py

@@ -10,6 +10,7 @@
     Any,
     Concatenate,
     Generic,
+    Literal,
     TypeVar,
     cast,
     overload,
@@ -294,6 +295,7 @@ async def start_workflow_update(
         workflow_id: str,
         update: temporalio.workflow.UpdateMethodMultiParam[[Any], ReturnType],
         *,
+        wait_for_stage: Literal[temporalio.client.WorkflowUpdateStage.ACCEPTED],
         update_id: str | None = None,
         rpc_metadata: Mapping[str, str | bytes] = {},
         rpc_timeout: timedelta | None = None,
@@ -311,6 +313,7 @@ async def start_workflow_update(
         ],
         arg: ParamType,
         *,
+        wait_for_stage: Literal[temporalio.client.WorkflowUpdateStage.ACCEPTED],
         update_id: str | None = None,
         rpc_metadata: Mapping[str, str | bytes] = {},
         rpc_timeout: timedelta | None = None,
@@ -326,6 +329,7 @@ async def start_workflow_update(
         update: temporalio.workflow.UpdateMethodMultiParam[MultiParamSpec, ReturnType],
         *,
         args: MultiParamSpec.args,  # type: ignore
+        wait_for_stage: Literal[temporalio.client.WorkflowUpdateStage.ACCEPTED],
         update_id: str | None = None,
         rpc_metadata: Mapping[str, str | bytes] = {},
         rpc_timeout: timedelta | None = None,
@@ -342,6 +346,7 @@ async def start_workflow_update(
         arg: Any = temporalio.common._arg_unset,
         *,
         args: Sequence[Any] = [],
+        wait_for_stage: Literal[temporalio.client.WorkflowUpdateStage.ACCEPTED],
         update_id: str | None = None,
         result_type: type[ReturnType] | None = None,
         rpc_metadata: Mapping[str, str | bytes] = {},
@@ -358,6 +363,7 @@ async def start_workflow_update(
         arg: Any = temporalio.common._arg_unset,
         *,
         args: Sequence[Any] = [],
+        wait_for_stage: Literal[temporalio.client.WorkflowUpdateStage.ACCEPTED],
         update_id: str | None = None,
         result_type: type | None = None,
         rpc_metadata: Mapping[str, str | bytes] = {},
@@ -679,6 +685,7 @@ async def start_workflow_update(
         arg: Any = temporalio.common._arg_unset,
         *,
         args: Sequence[Any] = [],
+        wait_for_stage: Literal[temporalio.client.WorkflowUpdateStage.ACCEPTED],
         update_id: str | None = None,
         result_type: type | None = None,
         rpc_metadata: Mapping[str, str | bytes] = {},
@@ -699,6 +706,7 @@ async def start_workflow_update(
                 update=update,
                 arg=arg,
                 args=args,
+                wait_for_stage=wait_for_stage,
                 update_id=update_id,
                 result_type=result_type,
                 rpc_metadata=rpc_metadata,

tests/nexus/test_temporal_operation.py

@@ -2,6 +2,7 @@
 import uuid
 from dataclasses import dataclass
 from datetime import timedelta
+from typing import Any, cast
 
 import nexusrpc
 import pytest
@@ -24,6 +25,7 @@
     NexusOperationFailureError,
     WorkflowExecutionStatus,
     WorkflowFailureError,
+    WorkflowUpdateStage,
 )
 from temporalio.common import (
     NexusOperationExecutionStatus,
@@ -116,6 +118,7 @@ class TestService:
     sync_result: Operation[Input, str]
     custom_cancel: Operation[str, None]
     update_op: Operation[Input, str]
+    bad_update_stage_op: Operation[Input, str]
     query_op: Operation[str, bool]
     echo_activity: Operation[Input, str]
     error_activity: Operation[Input, None]
@@ -134,6 +137,7 @@ def __init__(self) -> None:
         self.started_custom_cancel_workflow = asyncio.Event()
         self.started_custom_cancel_activity = asyncio.Event()
         self.custom_cancel_activity_called = asyncio.Event()
+        self.bad_update_stage_error: ValueError | None = None
 
     @nexus.temporal_operation
     async def echo(
@@ -290,9 +294,30 @@ async def update_op(
             input.value,
             UpdatableWorkflow.do_update,
             input.update_value,
+            wait_for_stage=WorkflowUpdateStage.ACCEPTED,
             update_id=input.update_id,
         )
 
+    @nexus.temporal_operation
+    async def bad_update_stage_op(
+        self,
+        _ctx: nexus.TemporalStartOperationContext,
+        client: nexus.TemporalNexusClient,
+        input: Input,
+    ) -> nexus.TemporalOperationResult[str]:
+        try:
+            return await client.start_workflow_update(
+                input.value,
+                UpdatableWorkflow.do_update,
+                input.update_value,
+                # cast to bypass type checker
+                wait_for_stage=cast(Any, WorkflowUpdateStage.COMPLETED),
+                update_id=input.update_id,
+            )
+        except ValueError as err:
+            self.bad_update_stage_error = err
+            return nexus.TemporalOperationResult.sync(str(err))
+
     @nexus.temporal_operation
     async def query_op(
         self,
@@ -749,6 +774,41 @@ async def test_temporal_operation_update_workflow_delayed(
         assert expected_backward_link in handler_links
 
 
+async def test_start_workflow_update_rejects_non_accepted_wait_for_stage(
+    client: Client, env: WorkflowEnvironment
+) -> None:
+    if env.supports_time_skipping:
+        pytest.skip("Update workflow tests don't work with time-skipping server")
+    task_queue = str(uuid.uuid4())
+    endpoint_name = make_nexus_endpoint_name(task_queue)
+    await env.create_nexus_endpoint(endpoint_name, task_queue)
+    service_handler = TestServiceHandler()
+    async with Worker(
+        env.client,
+        task_queue=task_queue,
+        nexus_service_handlers=[service_handler],
+        workflows=[UpdatableWorkflow, BadUpdateStageCaller],
+    ):
+        update_workflow_id = f"updatable-workflow-{uuid.uuid4()}"
+        await client.start_workflow(
+            UpdatableWorkflow.run, id=update_workflow_id, task_queue=task_queue
+        )
+        result = await client.execute_workflow(
+            BadUpdateStageCaller.run,
+            Input(
+                value=update_workflow_id,
+                task_queue=task_queue,
+                update_value="Created",
+            ),
+            task_queue=task_queue,
+            id=f"bad-update-stage-caller-{uuid.uuid4()}",
+        )
+
+    assert isinstance(service_handler.bad_update_stage_error, ValueError)
+    assert result == str(service_handler.bad_update_stage_error)
+    assert result == "Only ACCEPTED wait stage is supported"
+
+
 async def test_temporal_operation_cancel_rejects_unknown_tokens():
     class FakeNexusTaskCancellation(OperationTaskCancellation):
         def is_cancelled(self) -> bool:
@@ -1649,6 +1709,19 @@ async def run(self, input: Input) -> str:
         return await op_handle
 
 
+@workflow.defn
+class BadUpdateStageCaller:
+    """Caller workflow for an update op that requests an unsupported update stage."""
+
+    @workflow.run
+    async def run(self, input: Input) -> str:
+        client = workflow.create_nexus_client(
+            service=TestService,
+            endpoint=make_nexus_endpoint_name(input.task_queue),
+        )
+        return await client.execute_operation(TestService.bad_update_stage_op, input)
+
+
 @workflow.defn
 class UpdatableWorkflow:
     """Workflow that accepts updates and exits when it receives a specific status"""