From a39d75dc3ce9849ee79bb1f290c0d02a57c5cb5c Mon Sep 17 00:00:00 2001 From: Alice Lin Date: Thu, 17 Sep 2026 09:58:35 -0700 Subject: [PATCH 1/2] Require wait_for_stage parameter for Nexus Workflow Updates --- CHANGELOG.md | 2 + temporalio/nexus/_operation_context.py | 14 ++++- temporalio/nexus/_temporal_client.py | 8 +++ tests/nexus/test_temporal_operation.py | 73 ++++++++++++++++++++++++++ 4 files changed, 96 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f016f34ff..f35d5aced 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,8 @@ to include examples, links to docs, or any other relevant information. ### :boom: Breaking Changes +- **Experimental**: Nexus Workflow Updates now require `wait_for_stage` to be explicitly set to `ACCEPTED`. + ### Fixed - Current workflow and activity payload converter accessors now return the configured converter diff --git a/temporalio/nexus/_operation_context.py b/temporalio/nexus/_operation_context.py index 327e9c51b..a5b55de3c 100644 --- a/temporalio/nexus/_operation_context.py +++ b/temporalio/nexus/_operation_context.py @@ -19,11 +19,14 @@ Any, Concatenate, Generic, + Literal, TypeVar, + cast, overload, ) import nexusrpc +from nexusrpc import HandlerError, HandlerErrorType from nexusrpc.handler import ( CancelOperationContext, OperationContext, @@ -711,6 +714,7 @@ 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: Literal[temporalio.client.WorkflowUpdateStage.ACCEPTED], update_id: str | None = None, result_type: type | None = None, rpc_metadata: Mapping[str, str | bytes] = {}, @@ -718,6 +722,14 @@ async def _start_nexus_operation_workflow_update( # pyright: ignore[reportUnuse run_id: str | None = None, first_execution_run_id: str | None = None, ) -> temporalio.client.WorkflowUpdateHandle[Any]: + # Annotations are not enforced at runtime, so validate anyway. The cast widens the + # narrowed Literal; without it the check reads as unreachable to the type checker. + if cast(Any, wait_for_stage) != temporalio.client.WorkflowUpdateStage.ACCEPTED: + raise HandlerError( + "Nexus operations only support workflow updates with " + "wait_for_stage=WorkflowUpdateStage.ACCEPTED", + type=HandlerErrorType.BAD_REQUEST, + ) # 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 +740,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, diff --git a/temporalio/nexus/_temporal_client.py b/temporalio/nexus/_temporal_client.py index 393da2fae..4fd34d939 100644 --- a/temporalio/nexus/_temporal_client.py +++ b/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, diff --git a/tests/nexus/test_temporal_operation.py b/tests/nexus/test_temporal_operation.py index 66bc2290b..976ad77c0 100644 --- a/tests/nexus/test_temporal_operation.py +++ b/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] @@ -290,6 +293,24 @@ 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]: + # Only ACCEPTED is allowed; the cast is what a handler that bypasses the + # type checker would do. + return await client.start_workflow_update( + input.value, + UpdatableWorkflow.do_update, + input.update_value, + wait_for_stage=cast(Any, WorkflowUpdateStage.COMPLETED), update_id=input.update_id, ) @@ -749,6 +770,45 @@ 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) + async with Worker( + env.client, + task_queue=task_queue, + nexus_service_handlers=[TestServiceHandler()], + 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 + ) + with pytest.raises(WorkflowFailureError) as err: + 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(err.value.cause, temporalio.exceptions.NexusOperationError) + assert isinstance(err.value.cause.cause, nexusrpc.HandlerError) + assert err.value.cause.cause.type == HandlerErrorType.BAD_REQUEST + assert ( + "Nexus operations only support workflow updates with " + "wait_for_stage=WorkflowUpdateStage.ACCEPTED" in err.value.cause.cause.message + ) + + 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""" From 69e418beb630ec02b0d84acc676a5c95aa14af05 Mon Sep 17 00:00:00 2001 From: Alice Lin Date: Fri, 18 Sep 2026 15:24:56 -0700 Subject: [PATCH 2/2] Use plain type for private function and raise ValueError --- temporalio/nexus/_operation_context.py | 15 ++----- tests/nexus/test_temporal_operation.py | 54 +++++++++++++------------- 2 files changed, 30 insertions(+), 39 deletions(-) diff --git a/temporalio/nexus/_operation_context.py b/temporalio/nexus/_operation_context.py index a5b55de3c..a91b72b94 100644 --- a/temporalio/nexus/_operation_context.py +++ b/temporalio/nexus/_operation_context.py @@ -19,14 +19,11 @@ Any, Concatenate, Generic, - Literal, TypeVar, - cast, overload, ) import nexusrpc -from nexusrpc import HandlerError, HandlerErrorType from nexusrpc.handler import ( CancelOperationContext, OperationContext, @@ -714,7 +711,7 @@ 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: Literal[temporalio.client.WorkflowUpdateStage.ACCEPTED], + wait_for_stage: temporalio.client.WorkflowUpdateStage, update_id: str | None = None, result_type: type | None = None, rpc_metadata: Mapping[str, str | bytes] = {}, @@ -722,14 +719,8 @@ async def _start_nexus_operation_workflow_update( # pyright: ignore[reportUnuse run_id: str | None = None, first_execution_run_id: str | None = None, ) -> temporalio.client.WorkflowUpdateHandle[Any]: - # Annotations are not enforced at runtime, so validate anyway. The cast widens the - # narrowed Literal; without it the check reads as unreachable to the type checker. - if cast(Any, wait_for_stage) != temporalio.client.WorkflowUpdateStage.ACCEPTED: - raise HandlerError( - "Nexus operations only support workflow updates with " - "wait_for_stage=WorkflowUpdateStage.ACCEPTED", - type=HandlerErrorType.BAD_REQUEST, - ) + 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( diff --git a/tests/nexus/test_temporal_operation.py b/tests/nexus/test_temporal_operation.py index 976ad77c0..34c37a819 100644 --- a/tests/nexus/test_temporal_operation.py +++ b/tests/nexus/test_temporal_operation.py @@ -137,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( @@ -304,15 +305,18 @@ async def bad_update_stage_op( client: nexus.TemporalNexusClient, input: Input, ) -> nexus.TemporalOperationResult[str]: - # Only ACCEPTED is allowed; the cast is what a handler that bypasses the - # type checker would do. - return await client.start_workflow_update( - input.value, - UpdatableWorkflow.do_update, - input.update_value, - wait_for_stage=cast(Any, WorkflowUpdateStage.COMPLETED), - update_id=input.update_id, - ) + 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( @@ -778,35 +782,31 @@ async def test_start_workflow_update_rejects_non_accepted_wait_for_stage( 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=[TestServiceHandler()], + 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 ) - with pytest.raises(WorkflowFailureError) as err: - await client.execute_workflow( - BadUpdateStageCaller.run, - Input( - value=update_workflow_id, - task_queue=task_queue, - update_value="Created", - ), + result = await client.execute_workflow( + BadUpdateStageCaller.run, + Input( + value=update_workflow_id, task_queue=task_queue, - id=f"bad-update-stage-caller-{uuid.uuid4()}", - ) + update_value="Created", + ), + task_queue=task_queue, + id=f"bad-update-stage-caller-{uuid.uuid4()}", + ) - assert isinstance(err.value.cause, temporalio.exceptions.NexusOperationError) - assert isinstance(err.value.cause.cause, nexusrpc.HandlerError) - assert err.value.cause.cause.type == HandlerErrorType.BAD_REQUEST - assert ( - "Nexus operations only support workflow updates with " - "wait_for_stage=WorkflowUpdateStage.ACCEPTED" in err.value.cause.cause.message - ) + 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():