Propagate Nexus serialization context to workflow updates

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.

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()
 

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 (

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

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:

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

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