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..a91b72b94 100644 --- a/temporalio/nexus/_operation_context.py +++ b/temporalio/nexus/_operation_context.py @@ -711,6 +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: temporalio.client.WorkflowUpdateStage, update_id: str | None = None, result_type: type | None = None, rpc_metadata: Mapping[str, str | bytes] = {}, @@ -718,6 +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]: + 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, 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..34c37a819 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] @@ -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"""