Propagate Nexus serialization context to workflow updates

CHANGELOG.md

@@ -51,10 +51,10 @@ to docs, or any other relevant information.
   `TEMPORAL_WORKFLOW_TASK_DURATION_WARN_SECONDS` environment variable.
 - Nexus operation handlers now forward inbound callback links to the workflows and signals they start, instead of
   dropping them.
-- Nexus handlers now propagate `NexusSerializationContext` to the workflows and standalone
-  activities they start. These use it only for their results and failures, which are returned to
-  the Nexus caller. Standalone activities also use it for their heartbeat details, which the server
-  includes in their timeout failures.
+- Nexus handlers now propagate `NexusSerializationContext` to the workflows, workflow updates, and
+  standalone activities they start. These use it only for their results and failures, which are
+  returned to the Nexus caller. Standalone activities also use it for their heartbeat details, which
+  the server includes in their timeout failures.
   To complete such an activity asynchronously, read the context from
   `ActivityInfo.getNexusSerializationContext()`, store it with the task token, and pass it to
   `ActivityCompletionClient.withContext(NexusSerializationContext)`.

temporal-sdk/src/main/java/io/temporal/internal/client/RootWorkflowClientInvoker.java

@@ -118,6 +118,27 @@ private DataConverter workflowConverter(
     return converterFactory.forWorkflow(workflowId, runId, workflowType);
   }
 
+  /**
+   * The data converter for a workflow's or an update's result and failure. A workflow or an update
+   * started by a Nexus operation handler encodes them with the handler's Nexus serialization
+   * context, which the server returns with them. Otherwise they use the workflow context.
+   *
+   * @param propagated the context returned by the server, or null if it returned none
+   * @param workflowConverter the converter to use when the server returned no context
+   */
+  private DataConverter outcomeConverter(
+      @Nullable PropagatedSerializationContext propagated,
+      WorkflowExecution execution,
+      DataConverter workflowConverter) {
+    return propagated != null
+        ? converterFactory.forWorkflow(
+            PropagatedNexusSerializationContext.fromProto(propagated),
+            execution.getWorkflowId(),
+            execution.getRunId(),
+            null)
+        : workflowConverter;
+  }
+
   private void storeHeader(
       Header.Builder header,
       String workflowId,
@@ -475,14 +496,12 @@ private <R> R convertClosedWorkflowResult(
       GetResultInput<R> input, GetWorkflowExecutionHistoryResponse closed) {
     WorkflowExecution execution = input.getWorkflowExecution();
     DataConverter dataConverter =
-        closed.hasPropagatedNexusSerializationContext()
-            ? converterFactory.forWorkflow(
-                PropagatedNexusSerializationContext.fromProto(
-                    closed.getPropagatedNexusSerializationContext()),
-                execution.getWorkflowId(),
-                execution.getRunId(),
-                null)
-            : workflowConverter(execution);
+        outcomeConverter(
+            closed.hasPropagatedNexusSerializationContext()
+                ? closed.getPropagatedNexusSerializationContext()
+                : null,
+            execution,
+            workflowConverter(execution));
     // The long poll helpers return a response whose single event is the close event.
     HistoryEvent closeEvent = closed.getHistory().getEvents(0);
     Optional<Payloads> resultValue =
@@ -615,6 +634,17 @@ private <R> UpdateWorkflowExecutionRequest toUpdateWorkflowExecutionRequest(
             .setMeta(Meta.newBuilder().setUpdateId(input.getUpdateId()).setIdentity(identity))
             .setInput(updateInput);
 
+    // An update sent by a Nexus operation handler returns its result and failure to the Nexus
+    // caller, so they are serialized with the handler's Nexus context. It is sent with the request
+    // so the server can propagate it to the worker and to clients reading the outcome. The update's
+    // arguments keep the workflow context.
+    @Nullable
+    PropagatedSerializationContext propagatedNexusSerializationContext =
+        PropagatedNexusSerializationContext.current();
+    if (propagatedNexusSerializationContext != null) {
+      requestBuilder.setPropagatedNexusSerializationContext(propagatedNexusSerializationContext);
+    }
+
     // If this update is being issued via TemporalNexusClientImpl.startWorkflowUpdate,
     // set the fields the server needs to deliver the Nexus completion callback
     if (CurrentNexusOperationContext.isNexusContext()) {
@@ -666,15 +696,19 @@ private <R> WorkflowUpdateHandle<R> toUpdateHandle(
       UpdateWorkflowExecutionResponse result,
       DataConverter dataConverterWithWorkflowContext) {
     if (result.hasOutcome()) {
+      DataConverter dataConverter =
+          outcomeConverter(
+              result.hasPropagatedNexusSerializationContext()
+                  ? result.getPropagatedNexusSerializationContext()
+                  : null,
+              result.getUpdateRef().getWorkflowExecution(),
+              dataConverterWithWorkflowContext);
       switch (result.getOutcome().getValueCase()) {
         case SUCCESS:
           Optional<Payloads> updateResult = Optional.of(result.getOutcome().getSuccess());
           R resultValue =
               convertResultPayloads(
-                  updateResult,
-                  input.getResultClass(),
-                  input.getResultType(),
-                  dataConverterWithWorkflowContext);
+                  updateResult, input.getResultClass(), input.getResultType(), dataConverter);
           return new CompletedWorkflowUpdateHandleImpl<>(
               result.getUpdateRef().getUpdateId(),
               result.getUpdateRef().getWorkflowExecution(),
@@ -687,8 +721,7 @@ private <R> WorkflowUpdateHandle<R> toUpdateHandle(
                   result.getUpdateRef().getWorkflowExecution(),
                   result.getUpdateRef().getUpdateId(),
                   input.getUpdateName(),
-                  dataConverterWithWorkflowContext.failureToException(
-                      result.getOutcome().getFailure())));
+                  dataConverter.failureToException(result.getOutcome().getFailure())));
         default:
           throw new RuntimeException(
               "Received unexpected outcome from update request: "
@@ -748,21 +781,24 @@ public <R> PollWorkflowUpdateOutput<R> pollWorkflowUpdate(PollWorkflowUpdateInpu
         future.thenApply(
             (result) -> {
               if (result.hasOutcome()) {
+                DataConverter dataConverter =
+                    outcomeConverter(
+                        result.hasPropagatedNexusSerializationContext()
+                            ? result.getPropagatedNexusSerializationContext()
+                            : null,
+                        input.getWorkflowExecution(),
+                        dataConverterWithWorkflowContext);
                 switch (result.getOutcome().getValueCase()) {
                   case SUCCESS:
                     Optional<Payloads> updateResult = Optional.of(result.getOutcome().getSuccess());
                     return convertResultPayloads(
-                        updateResult,
-                        input.getResultClass(),
-                        input.getResultType(),
-                        dataConverterWithWorkflowContext);
+                        updateResult, input.getResultClass(), input.getResultType(), dataConverter);
                   case FAILURE:
                     throw new WorkflowUpdateException(
                         input.getWorkflowExecution(),
                         input.getUpdateId(),
                         input.getUpdateName(),
-                        dataConverterWithWorkflowContext.failureToException(
-                            result.getOutcome().getFailure()));
+                        dataConverter.failureToException(result.getOutcome().getFailure()));
                   default:
                     throw new RuntimeException(
                         "Received unexpected outcome from poll update request: "

temporal-sdk/src/main/java/io/temporal/internal/nexus/PropagatedNexusSerializationContext.java

@@ -9,10 +9,11 @@
  * Converts between {@link NexusSerializationContext} and the {@link PropagatedSerializationContext}
  * the server stores on executions started by a Nexus operation handler.
  *
- * <p>When a Nexus handler starts a workflow or a standalone activity, the SDK sends the handler's
- * serialization context with the start request. The server stores it and returns it to workers and
- * to clients reading the execution's result, so the result and failure are encoded and decoded with
- * the same context. See {@link NexusSerializationContext} for which payloads use it.
+ * <p>When a Nexus handler starts a workflow, sends a workflow update or starts a standalone
+ * activity, the SDK sends the handler's serialization context with the request. The server stores
+ * it and returns it to workers and to clients reading the execution's result, so the result and
+ * failure are encoded and decoded with the same context. See {@link NexusSerializationContext} for
+ * which payloads use it.
  */
 public final class PropagatedNexusSerializationContext {
 

temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflow.java

@@ -5,7 +5,9 @@
 import io.temporal.api.history.v1.HistoryEvent;
 import io.temporal.api.query.v1.WorkflowQuery;
 import io.temporal.internal.statemachines.UpdateProtocolCallback;
+import io.temporal.payload.context.NexusSerializationContext;
 import java.util.Optional;
+import javax.annotation.Nullable;
 
 /**
  * Manages event loop, workflow method, an abstraction level over Deterministic Runner to provide a
@@ -18,13 +20,20 @@ public interface ReplayWorkflow {
   /** Handle an external signal event. */
   void handleSignal(String signalName, Optional<Payloads> input, long eventId, Header header);
 
-  /** Handle an update workflow execution event */
+  /**
+   * Handle an update workflow execution event.
+   *
+   * @param nexusSerializationContext the serialization context of the Nexus operation handler that
+   *     sent the update, or null if it was not sent by one. When set, the update's result, failure
+   *     and rejection, which are returned to the Nexus caller, are serialized with it.
+   */
   void handleUpdate(
       String updateName,
       String updateId,
       Optional<Payloads> input,
       long eventId,
       Header header,
+      @Nullable NexusSerializationContext nexusSerializationContext,
       UpdateProtocolCallback callbacks);
 
   /**

temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowExecutor.java

@@ -17,6 +17,7 @@
 import io.temporal.internal.common.UpdateMessage;
 import io.temporal.internal.logging.LoggerTag;
 import io.temporal.internal.logging.PrefixedMdc;
+import io.temporal.internal.nexus.PropagatedNexusSerializationContext;
 import io.temporal.internal.statemachines.WorkflowStateMachines;
 import io.temporal.internal.sync.SignalHandlerInfo;
 import io.temporal.internal.sync.UpdateHandlerInfo;
@@ -194,6 +195,10 @@ public void handleWorkflowExecutionUpdated(UpdateMessage updateMessage) {
           args,
           protocolMessage.getEventId(),
           input.getHeader(),
+          update.hasPropagatedNexusSerializationContext()
+              ? PropagatedNexusSerializationContext.fromProto(
+                  update.getPropagatedNexusSerializationContext())
+              : null,
           updateMessage.getCallbacks());
     } catch (InvalidProtocolBufferException e) {
       throw new IllegalStateException("Message is not an update.");

temporal-sdk/src/main/java/io/temporal/internal/sync/SyncWorkflow.java

@@ -19,6 +19,7 @@
 import io.temporal.internal.statemachines.UpdateProtocolCallback;
 import io.temporal.internal.worker.WorkflowExecutionException;
 import io.temporal.internal.worker.WorkflowExecutorCache;
+import io.temporal.payload.context.NexusSerializationContext;
 import io.temporal.payload.context.WorkflowSerializationContext;
 import io.temporal.worker.WorkflowImplementationOptions;
 import io.temporal.workflow.UpdateInfo;
@@ -151,8 +152,11 @@ public void handleUpdate(
       Optional<Payloads> input,
       long eventId,
       Header header,
+      @Nullable NexusSerializationContext nexusSerializationContext,
       UpdateProtocolCallback callbacks) {
     final UpdateInfo updateInfo = new UpdateInfoImpl(updateName, updateId);
+    DataConverter outcomeDataConverter =
+        workflowContext.getDataConverterForUpdateOutcome(nexusSerializationContext);
     runner.executeInWorkflowThread(
         "update " + updateName,
         () -> {
@@ -165,18 +169,16 @@ public void handleUpdate(
             if (!callbacks.isReplaying()) {
               try {
                 workflowContext.setReadOnly(true);
-                workflowProc.handleValidateUpdate(updateName, updateId, input, eventId, header);
+                workflowProc.handleValidateUpdate(
+                    updateName, updateId, input, eventId, header, outcomeDataConverter);
               } catch (ReadOnlyException r) {
                 // Rethrow instead on rejecting the update to fail the WFT
                 throw r;
               } catch (WorkflowExecutionException e) {
                 callbacks.reject(e.getFailure());
                 return;
               } catch (Exception e) {
-                callbacks.reject(
-                    workflowContext
-                        .getDataConverterWithCurrentWorkflowContext()
-                        .exceptionToFailure(e));
+                callbacks.reject(outcomeDataConverter.exceptionToFailure(e));
                 return;
               } finally {
                 workflowContext.setReadOnly(false);
@@ -185,7 +187,8 @@ public void handleUpdate(
             callbacks.accept();
             try {
               Optional<Payloads> result =
-                  workflowProc.handleExecuteUpdate(updateName, updateId, input, eventId, header);
+                  workflowProc.handleExecuteUpdate(
+                      updateName, updateId, input, eventId, header, outcomeDataConverter);
               callbacks.complete(result, null);
             } catch (WorkflowExecutionException e) {
               callbacks.complete(Optional.empty(), e.getFailure());

temporal-sdk/src/main/java/io/temporal/internal/sync/SyncWorkflowContext.java

@@ -352,8 +352,14 @@ public void handleValidateUpdate(
   }
 
   public Optional<Payloads> handleExecuteUpdate(
-      String updateName, String updateId, Optional<Payloads> input, long eventId, Header header) {
-    return updateDispatcher.handleExecuteUpdate(updateName, updateId, input, eventId, header);
+      String updateName,
+      String updateId,
+      Optional<Payloads> input,
+      long eventId,
+      Header header,
+      DataConverter resultDataConverter) {
+    return updateDispatcher.handleExecuteUpdate(
+        updateName, updateId, input, eventId, header, resultDataConverter);
   }
 
   public void handleInterceptedValidateUpdate(WorkflowInboundCallsInterceptor.UpdateInput input) {
@@ -1319,8 +1325,8 @@ void setPropagatedNexusSerializationContext(
   /**
    * The data converter for the workflow result and the failure that closes the workflow, which are
    * returned to the Nexus caller: the propagated Nexus context if the workflow was started by a
-   * Nexus operation handler, otherwise the workflow context. Input, memo, signals, queries and
-   * updates always use the workflow context.
+   * Nexus operation handler, otherwise the workflow context. Input, memo, signals and queries
+   * always use the workflow context. Updates use {@link #getDataConverterForUpdateOutcome}.
    */
   DataConverter getDataConverterForResult() {
     return dataConverterWithPropagatedNexusContext != null
@@ -1657,12 +1663,17 @@ public Failure mapWorkflowExceptionToFailure(Throwable failure) {
   }
 
   /**
-   * Converts an exception thrown by an update validator or handler into the failure that rejects or
-   * fails the update. Unlike {@link #mapWorkflowExceptionToFailure}, this always uses the workflow
-   * context, because the failure belongs to the update rather than closing the workflow.
+   * The data converter for an update's result, failure and rejection, which are returned to the
+   * update's caller: the Nexus context of the Nexus operation handler that sent the update, if it
+   * was sent by one, otherwise the workflow context. Unlike {@link #getDataConverterForResult},
+   * this doesn't depend on how the workflow was started, because the outcome belongs to the update
+   * rather than closing the workflow. The update's arguments always use the workflow context.
    */
-  Failure mapUpdateExceptionToFailure(Throwable failure) {
-    return dataConverterWithCurrentWorkflowContext.exceptionToFailure(failure);
+  DataConverter getDataConverterForUpdateOutcome(
+      @Nullable NexusSerializationContext nexusSerializationContext) {
+    return nexusSerializationContext != null
+        ? dataConverter.withContext(nexusSerializationContext)
+        : dataConverterWithCurrentWorkflowContext;
   }
 
   @Nullable

temporal-sdk/src/main/java/io/temporal/internal/sync/UpdateDispatcher.java

@@ -71,8 +71,17 @@ public void handleValidateUpdate(
     }
   }
 
+  /**
+   * @param resultDataConverter the converter for the update's result, which may differ from the
+   *     workflow converter used for its arguments
+   */
   public Optional<Payloads> handleExecuteUpdate(
-      String updateName, String updateId, Optional<Payloads> input, long eventId, Header header) {
+      String updateName,
+      String updateId,
+      Optional<Payloads> input,
+      long eventId,
+      Header header,
+      DataConverter resultDataConverter) {
     WorkflowOutboundCallsInterceptor.UpdateRegistrationRequest handler =
         updateCallbacks.get(updateName);
     Object[] args;
@@ -99,7 +108,7 @@ public Optional<Payloads> handleExecuteUpdate(
               .executeUpdate(
                   new WorkflowInboundCallsInterceptor.UpdateInput(updateName, header, args))
               .getResult();
-      return dataConverterWithWorkflowContext.toPayloads(result);
+      return resultDataConverter.toPayloads(result);
     } catch (DestroyWorkflowThreadError e) {
       threadDestroyed = true;
       throw e;