From 6eccfc86fd0dd752bdeb1a8d161347ee9c6fab03 Mon Sep 17 00:00:00 2001 From: Joshua Frenchwood Date: Thu, 29 Oct 2026 12:57:43 -0500 Subject: [PATCH] Propagate Nexus serialization context to workflow updates --- changelog/added/tumbling-toucan.md | 1 + temporalio/client/_impl.py | 11 +- temporalio/client/_workflow.py | 26 ++- .../converter/_serialization_context.py | 7 +- temporalio/nexus/_operation_context.py | 12 ++ temporalio/worker/_command_aware_visitor.py | 24 ++- temporalio/worker/_workflow_instance.py | 46 ++++- tests/test_serialization_context.py | 185 ++++++++++++++++++ 8 files changed, 294 insertions(+), 18 deletions(-) create mode 100644 changelog/added/tumbling-toucan.md diff --git a/changelog/added/tumbling-toucan.md b/changelog/added/tumbling-toucan.md new file mode 100644 index 000000000..baf2d8284 --- /dev/null +++ b/changelog/added/tumbling-toucan.md @@ -0,0 +1 @@ +Nexus operation handlers now propagate their `NexusSerializationContext` to the workflow updates they send, which encode their result, failure and rejection with it. Update handles decode them with the context the server returns. Requires a server that propagates the context. diff --git a/temporalio/client/_impl.py b/temporalio/client/_impl.py index d2585ef78..4e3e9b3ab 100644 --- a/temporalio/client/_impl.py +++ b/temporalio/client/_impl.py @@ -891,7 +891,7 @@ async def start_workflow_update( result_type=input.ret_type, ) if resp.HasField("outcome"): - handle._known_outcome = resp.outcome + handle._set_known_outcome(resp) if input.wait_for_stage == WorkflowUpdateStage.COMPLETED: await handle._poll_until_outcome() return handle @@ -950,6 +950,9 @@ async def _build_update_workflow_execution_request( temporalio.nexus._operation_context._apply_nexus_context_to_start_workflow_update_request( req ) + temporalio.nexus._operation_context._apply_nexus_serialization_context_to_workflow_update_request( + req + ) if input.args: req.request.input.args.payloads.extend( await data_converter.encode(input.args) @@ -1083,9 +1086,6 @@ async def _start_workflow_update_with_start( start_response = multiop_response.responses[0].start_workflow update_response = multiop_response.responses[1].update_workflow on_start(start_response) - known_outcome = ( - update_response.outcome if update_response.HasField("outcome") else None - ) if ( update_response.stage >= temporalio.api.enums.v1.UpdateWorkflowExecutionLifecycleStage.UPDATE_WORKFLOW_EXECUTION_LIFECYCLE_STAGE_ACCEPTED @@ -1097,9 +1097,10 @@ async def _start_workflow_update_with_start( id=update_req.request.meta.update_id, workflow_id=start_input.id, workflow_run_id=start_response.run_id, - known_outcome=known_outcome, result_type=update_input.ret_type, ) + if update_response.HasField("outcome"): + handle._set_known_outcome(update_response) if update_input.wait_for_stage == WorkflowUpdateStage.COMPLETED: await handle._poll_until_outcome() diff --git a/temporalio/client/_workflow.py b/temporalio/client/_workflow.py index fc2021098..8a2c423a5 100644 --- a/temporalio/client/_workflow.py +++ b/temporalio/client/_workflow.py @@ -1878,6 +1878,9 @@ def __init__( self._workflow_run_id = workflow_run_id self._result_type = result_type self._known_outcome = known_outcome + self._known_outcome_nexus_context: ( + temporalio.converter.NexusSerializationContext | None + ) = None @functools.cached_property def _data_converter(self) -> temporalio.converter.DataConverter: @@ -1888,6 +1891,18 @@ def _data_converter(self) -> temporalio.converter.DataConverter: ) ) + def _set_known_outcome( + self, + resp: temporalio.api.workflowservice.v1.UpdateWorkflowExecutionResponse + | temporalio.api.workflowservice.v1.PollWorkflowExecutionUpdateResponse, + ) -> None: + self._known_outcome = resp.outcome + # An update sent by a Nexus operation handler encodes its outcome with the handler's + # Nexus context, which the server returns with the outcome + self._known_outcome_nexus_context = temporalio.converter._serialization_context._propagated_nexus_serialization_context( + resp + ) + @property def id(self) -> str: """ID of this Update request.""" @@ -1932,14 +1947,19 @@ async def result( # Convert outcome to failure or value assert self._known_outcome + data_converter = ( + self._client.data_converter.with_context(self._known_outcome_nexus_context) + if self._known_outcome_nexus_context + else self._data_converter + ) if self._known_outcome.HasField("failure"): raise WorkflowUpdateFailedError( - await self._data_converter.decode_failure(self._known_outcome.failure), + await data_converter.decode_failure(self._known_outcome.failure), ) if not self._known_outcome.success.payloads: return None # type: ignore type_hints = [self._result_type] if self._result_type else None - results = await self._data_converter.decode( + results = await data_converter.decode( self._known_outcome.success.payloads, type_hints ) if not results: @@ -1982,7 +2002,7 @@ async def _poll_until_outcome( ) ) if res.HasField("outcome"): - self._known_outcome = res.outcome + self._set_known_outcome(res) return except RPCError as err: if ( diff --git a/temporalio/converter/_serialization_context.py b/temporalio/converter/_serialization_context.py index 73d19a5b4..7de64be04 100644 --- a/temporalio/converter/_serialization_context.py +++ b/temporalio/converter/_serialization_context.py @@ -107,9 +107,10 @@ class NexusSerializationContext(SerializationContext): receive it when decoding inputs, encoding synchronous results, and encoding failures produced while handling a Nexus task. - Workflows and standalone activities started by a handler receive this context through the - server and use it for payloads returned to the caller: results, failures, and standalone - activity heartbeat details. Clients decode these with the context the server returns. See + Workflows and standalone activities started by a handler, and workflow updates it sends, + receive this context through the server and use it for payloads returned to the caller: + results, failures, and standalone activity heartbeat details. Clients decode these with the + context the server returns. See :py:attr:`temporalio.activity.Info.nexus_serialization_context` for asynchronous completion. A standalone operation handle decodes with the context the server returns, falling back to the diff --git a/temporalio/nexus/_operation_context.py b/temporalio/nexus/_operation_context.py index bd287e4e1..c886564e4 100644 --- a/temporalio/nexus/_operation_context.py +++ b/temporalio/nexus/_operation_context.py @@ -792,6 +792,18 @@ def _apply_nexus_context_to_start_workflow_update_request( # pyright: ignore[re ) +def _apply_nexus_serialization_context_to_workflow_update_request( # pyright: ignore[reportUnusedFunction] + req: temporalio.api.workflowservice.v1.UpdateWorkflowExecutionRequest, +) -> None: + """Send the current Nexus operation's serialization context with a Workflow Update request. + + This is a no-op outside a Nexus operation context. + """ + context = _current_propagated_nexus_serialization_context() + if context is not None: + req.request.propagated_nexus_serialization_context.CopyFrom(context) + + def _apply_start_workflow_update_response_to_nexus_context( # pyright: ignore[reportUnusedFunction] resp: temporalio.api.workflowservice.v1.UpdateWorkflowExecutionResponse, ) -> None: diff --git a/temporalio/worker/_command_aware_visitor.py b/temporalio/worker/_command_aware_visitor.py index 69667079c..888c06f29 100644 --- a/temporalio/worker/_command_aware_visitor.py +++ b/temporalio/worker/_command_aware_visitor.py @@ -25,6 +25,7 @@ ScheduleNexusOperation, SignalExternalWorkflowExecution, StartChildWorkflowExecution, + UpdateResponse, WorkflowCommand, ) @@ -35,6 +36,7 @@ class CommandInfo: command_type: CommandType.ValueType command_seq: int + update_protocol_instance_id: str | None = None current_command_info: contextvars.ContextVar[CommandInfo | None] = ( @@ -123,6 +125,17 @@ async def _visit_coresdk_workflow_commands_ScheduleNexusOperation( with current_command(CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION, o.seq): await super()._visit_coresdk_workflow_commands_ScheduleNexusOperation(fs, o) + async def _visit_coresdk_workflow_commands_UpdateResponse( + self, fs: VisitorFunctions, o: UpdateResponse + ) -> None: + # Updates have no sequence number, so they're identified by protocol instance ID + with current_command( + CommandType.COMMAND_TYPE_PROTOCOL_MESSAGE, + 0, + update_protocol_instance_id=o.protocol_instance_id, + ): + await super()._visit_coresdk_workflow_commands_UpdateResponse(fs, o) + async def _visit_coresdk_workflow_commands_WorkflowCommand( self, fs: VisitorFunctions, o: WorkflowCommand ) -> None: @@ -201,11 +214,18 @@ async def _visit_coresdk_workflow_activation_ResolveNexusOperation( @contextmanager def current_command( - command_type: CommandType.ValueType, command_seq: int + command_type: CommandType.ValueType, + command_seq: int, + *, + update_protocol_instance_id: str | None = None, ) -> Iterator[None]: """Context manager for setting command info.""" token = current_command_info.set( - CommandInfo(command_type=command_type, command_seq=command_seq) + CommandInfo( + command_type=command_type, + command_seq=command_seq, + update_protocol_instance_id=update_protocol_instance_id, + ) ) try: yield diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 3e6aca853..a8ab8cc44 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -56,6 +56,7 @@ import temporalio.bridge.proto.workflow_completion import temporalio.common import temporalio.converter +import temporalio.converter._serialization_context import temporalio.exceptions import temporalio.nexus.system import temporalio.workflow @@ -310,6 +311,12 @@ def __init__(self, det: WorkflowInstanceDetails) -> None: else: self._result_payload_converter = self._workflow_context_payload_converter self._result_failure_converter = self._workflow_context_failure_converter + # Keyed by update protocol instance ID. Kept until the activation that finished the update + # has been encoded, so the codec can encode its response with the same context. + self._nexus_update_contexts: dict[ + str, temporalio.converter.NexusSerializationContext + ] = {} + self._finished_nexus_update_ids: list[str] = [] self._extern_functions = det.extern_functions self._disable_eager_activity_execution = det.disable_eager_activity_execution @@ -475,6 +482,9 @@ def activate( temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion() ) self._current_completion.successful.SetInParent() + for protocol_instance_id in self._finished_nexus_update_ids: + self._nexus_update_contexts.pop(protocol_instance_id, None) + self._finished_nexus_update_ids.clear() self._current_activation_error: Exception | None = None self._deployment_version_for_current_task = ( @@ -696,6 +706,24 @@ def _apply_do_update( # Run the validator & handler in a task. Everything, including looking up the update definition, needs to be # inside the task, since the update may not be defined until after we have started the workflow - for example # if an update is in the first WFT & is also registered dynamically at the top of workflow code. + # An update sent by a Nexus operation handler returns its result and failure to the Nexus + # caller, which decodes them with the handler's Nexus context. Its input keeps the workflow + # context. + nexus_context = temporalio.converter._serialization_context._propagated_nexus_serialization_context( + job + ) + if nexus_context: + self._nexus_update_contexts[job.protocol_instance_id] = nexus_context + result_payload_converter = self._payload_converter_with_context( + nexus_context + ) + result_failure_converter = self._failure_converter_with_context( + nexus_context + ) + else: + result_payload_converter = self._workflow_context_payload_converter + result_failure_converter = self._workflow_context_failure_converter + async def run_update() -> None: # Set the current update for the life of this task temporalio.workflow._set_current_update_info( @@ -751,9 +779,7 @@ async def run_update() -> None: # Run the handler with _inbound_update_event_group(job.id).scope(): success = await self._inbound.handle_update_handler(handler_input) - result_payloads = self._workflow_context_payload_converter.to_payloads( - [success] - ) + result_payloads = result_payload_converter.to_payloads([success]) if len(result_payloads) != 1: raise ValueError( f"Expected 1 result payload, got {len(result_payloads)}" @@ -785,9 +811,9 @@ async def run_update() -> None: job.protocol_instance_id ) command.update_response.rejected.SetInParent() - self._workflow_context_failure_converter.to_failure( + result_failure_converter.to_failure( err, - self._workflow_context_payload_converter, + result_payload_converter, command.update_response.rejected, ) else: @@ -811,6 +837,8 @@ async def run_update() -> None: raise finally: self._in_progress_updates.pop(job.id, None) + if nexus_context: + self._finished_nexus_update_ids.append(job.protocol_instance_id) self.create_task( run_update(), @@ -2501,6 +2529,14 @@ def get_serialization_context( # The result and closing failure go back to the Nexus caller return self._nexus_serialization_context + if ( + command_info.command_type + == temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_PROTOCOL_MESSAGE + and command_info.update_protocol_instance_id in self._nexus_update_contexts + ): + # The update's result and failure go back to the Nexus caller + return self._nexus_update_contexts[command_info.update_protocol_instance_id] + if ( command_info.command_type == temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK diff --git a/tests/test_serialization_context.py b/tests/test_serialization_context.py index 819ee483c..b11a13015 100644 --- a/tests/test_serialization_context.py +++ b/tests/test_serialization_context.py @@ -26,6 +26,7 @@ import temporalio.api.failure.v1 import temporalio.api.history.v1 import temporalio.api.nexus.v1 +import temporalio.api.update.v1 import temporalio.api.workflowservice.v1 from temporalio import activity, nexus, workflow from temporalio.client import ( @@ -36,8 +37,11 @@ Interceptor, NexusOperationFailureError, OutboundInterceptor, + StartWorkflowUpdateInput, WorkflowFailureError, WorkflowUpdateFailedError, + WorkflowUpdateHandle, + WorkflowUpdateStage, ) from temporalio.common import RetryPolicy from temporalio.contrib.pydantic import PydanticJSONPlainPayloadConverter @@ -2599,8 +2603,57 @@ async def nexus_started_activity(input: NexusStartedInput) -> str: return input.value +@dataclass +class NexusUpdateInput: + workflow_id: str + update_id: str + value: Literal["result", "failure", "rejection"] + + +@workflow.defn +class NexusUpdatedWorkflow: + def __init__(self) -> None: + self.finished = False + + @workflow.run + async def run(self) -> None: + await workflow.wait_condition(lambda: self.finished) + + @workflow.update + async def update(self, value: str) -> str: + if value == "failure": + raise ApplicationError( + "nexus-update-failure", "failure-detail", non_retryable=True + ) + return value + + @update.validator + def validate_update(self, value: str) -> None: + if value == "rejection": + raise ApplicationError("nexus-update-rejection", "rejection-detail") + + @workflow.signal + def finish(self) -> None: + self.finished = True + + @nexusrpc.handler.service_handler class NexusStartedExecutionHandler: + @nexus.temporal_operation + async def update_workflow( + self, + _ctx: nexus.TemporalStartOperationContext, + client: nexus.TemporalNexusClient, + input: NexusUpdateInput, + ) -> nexus.TemporalOperationResult[str]: + return await client.start_workflow_update( + input.workflow_id, + NexusUpdatedWorkflow.update, + input.value, + wait_for_stage=WorkflowUpdateStage.ACCEPTED, + update_id=input.update_id, + ) + @nexus.temporal_operation async def start_workflow( self, @@ -2924,3 +2977,135 @@ async def activity_started() -> bool: await client.get_activity_handle(activity_id, result_type=str).result() == "async-value" ) + + +class _UpdateOutcomeCapturingInterceptor(Interceptor): + def __init__(self) -> None: + super().__init__() + self.outcomes: list[temporalio.api.update.v1.Outcome] = [] + + def intercept_client(self, next: OutboundInterceptor) -> OutboundInterceptor: + return _UpdateOutcomeCapturingOutboundInterceptor(next, self) + + +class _UpdateOutcomeCapturingOutboundInterceptor(OutboundInterceptor): + def __init__( + self, next: OutboundInterceptor, parent: _UpdateOutcomeCapturingInterceptor + ) -> None: + super().__init__(next) + self._parent = parent + + async def start_workflow_update( + self, input: StartWorkflowUpdateInput + ) -> WorkflowUpdateHandle[Any]: + handle = await super().start_workflow_update(input) + if handle._known_outcome: + # Copy, since decoding the outcome's failure later modifies it in place + outcome = temporalio.api.update.v1.Outcome() + outcome.CopyFrom(handle._known_outcome) + self._parent.outcomes.append(outcome) + return handle + + +@pytest.mark.requires_local_server +@pytest.mark.parametrize("value", ["result", "failure", "rejection"]) +async def test_nexus_sent_update_uses_nexus_context( + env: WorkflowEnvironment, + value: Literal["result", "failure", "rejection"], +): + """An update sent by a Nexus handler encodes its result, failure and rejection with the + handler's context, and clients decode them with the context the server returns.""" + if env.supports_time_skipping_v1: + pytest.skip("Nexus tests don't work with the Java test server") + + task_queue = str(uuid.uuid4()) + endpoint_name = make_nexus_endpoint_name(task_queue) + context = _nexus_started_context(endpoint_name, "update_workflow") + # Rejected updates aren't stored in history, so the test reads the rejection from the + # update response the handler receives + outcome_capture = _UpdateOutcomeCapturingInterceptor() + config = _nexus_started_client(env, context).config() + config["interceptors"] = [outcome_capture] + client = Client(**config) + update_id = str(uuid.uuid4()) + + async with Worker( + client, + task_queue=task_queue, + workflows=[NexusUpdatedWorkflow], + nexus_service_handlers=[NexusStartedExecutionHandler()], + ): + await env.create_nexus_endpoint(endpoint_name, task_queue) + handle = await client.start_workflow( + NexusUpdatedWorkflow.run, + id=f"nexus-updated-workflow-{uuid.uuid4()}", + task_queue=task_queue, + ) + operation = client.create_nexus_client( + NexusStartedExecutionHandler, endpoint_name + ).execute_operation( + NexusStartedExecutionHandler.update_workflow, + NexusUpdateInput(workflow_id=handle.id, update_id=update_id, value=value), + id=str(uuid.uuid4()), + ) + if value == "result": + assert await operation == "result" + else: + with pytest.raises(NexusOperationFailureError): + await operation + + if value == "rejection": + [outcome] = outcome_capture.outcomes + [detail] = outcome.failure.application_failure_info.details.payloads + assert _marker(detail) == b"nexus" + else: + # This handle never sent the update, so it relies on the context the server returns + update_handle = handle.get_update_handle(update_id, result_type=str) + if value == "result": + assert await update_handle.result() == "result" + else: + with pytest.raises(WorkflowUpdateFailedError) as err: + await update_handle.result() + assert isinstance(err.value.cause, ApplicationError) + assert err.value.cause.details == ("failure-detail",) + + await handle.signal(NexusUpdatedWorkflow.finish) + await handle.result() + + history = await handle.fetch_history() + if value != "rejection": + [accepted] = [ + e.workflow_execution_update_accepted_event_attributes + for e in history.events + if e.HasField("workflow_execution_update_accepted_event_attributes") + ] + assert ( + accepted.accepted_request.propagated_nexus_serialization_context + == temporalio.api.nexus.v1.PropagatedSerializationContext( + endpoint=context.endpoint, + service=context.service, + operation=context.operation, + ) + ) + # The input keeps the workflow context + assert all( + _marker(p) is None + for p in accepted.accepted_request.input.args.payloads + ) + [completed] = [ + e.workflow_execution_update_completed_event_attributes + for e in history.events + if e.HasField("workflow_execution_update_completed_event_attributes") + ] + if value == "result": + [result] = completed.outcome.success.payloads + assert _marker(result) == b"nexus" + else: + failure = completed.outcome.failure + [detail] = failure.application_failure_info.details.payloads + assert _marker(detail) == b"nexus" + + # Replaying re-runs the update handler with the context from history + await Replayer( + workflows=[NexusUpdatedWorkflow], data_converter=client.data_converter + ).replay_workflow(history)