Use correct operation token on OPERATION_TOKEN (#2589)

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/client/WorkflowClientInternalImpl.java

@@ -17,6 +17,7 @@
 import io.temporal.common.interceptors.WorkflowClientInterceptor;
 import io.temporal.internal.WorkflowThreadMarker;
 import io.temporal.internal.client.*;
+import io.temporal.internal.client.NexusStartWorkflowResponse;
 import io.temporal.internal.client.external.GenericWorkflowClient;
 import io.temporal.internal.client.external.GenericWorkflowClientImpl;
 import io.temporal.internal.client.external.ManualActivityCompletionClientFactory;
@@ -695,12 +696,13 @@ public void deregisterWorkerFactory(WorkerFactory workerFactory) {
   }
 
   @Override
-  public WorkflowExecution startNexus(NexusStartWorkflowRequest request, Functions.Proc workflow) {
+  public NexusStartWorkflowResponse startNexus(
+      NexusStartWorkflowRequest request, Functions.Proc workflow) {
     enforceNonWorkflowThread();
     WorkflowInvocationHandler.initAsyncInvocation(InvocationType.START_NEXUS, request);
     try {
       workflow.apply();
-      return WorkflowInvocationHandler.getAsyncInvocationResult(WorkflowExecution.class);
+      return WorkflowInvocationHandler.getAsyncInvocationResult(NexusStartWorkflowResponse.class);
     } finally {
       WorkflowInvocationHandler.closeAsyncInvocation();
     }

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

@@ -0,0 +1,21 @@
+package io.temporal.internal.client;
+
+import io.temporal.api.common.v1.WorkflowExecution;
+
+public final class NexusStartWorkflowResponse {
+  private final WorkflowExecution workflowExecution;
+  private final String operationToken;
+
+  public NexusStartWorkflowResponse(WorkflowExecution workflowExecution, String operationToken) {
+    this.workflowExecution = workflowExecution;
+    this.operationToken = operationToken;
+  }
+
+  public String getOperationToken() {
+    return operationToken;
+  }
+
+  public WorkflowExecution getWorkflowExecution() {
+    return workflowExecution;
+  }
+}

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

@@ -1,6 +1,5 @@
 package io.temporal.internal.client;
 
-import io.temporal.api.common.v1.WorkflowExecution;
 import io.temporal.client.WorkflowClient;
 import io.temporal.worker.WorkerFactory;
 import io.temporal.workflow.Functions;
@@ -18,5 +17,5 @@ public interface WorkflowClientInternal {
 
   void deregisterWorkerFactory(WorkerFactory workerFactory);
 
-  WorkflowExecution startNexus(NexusStartWorkflowRequest request, Functions.Proc workflow);
+  NexusStartWorkflowResponse startNexus(NexusStartWorkflowRequest request, Functions.Proc workflow);
 }

temporal-sdk/src/main/java/io/temporal/internal/common/InternalUtils.java

@@ -1,7 +1,9 @@
 package io.temporal.internal.common;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.google.common.base.Defaults;
 import io.nexusrpc.Header;
+import io.nexusrpc.handler.HandlerException;
 import io.nexusrpc.handler.ServiceImplInstance;
 import io.temporal.api.common.v1.Callback;
 import io.temporal.api.common.v1.Link;
@@ -14,6 +16,9 @@
 import io.temporal.common.metadata.POJOWorkflowMethodMetadata;
 import io.temporal.common.metadata.WorkflowMethodType;
 import io.temporal.internal.client.NexusStartWorkflowRequest;
+import io.temporal.internal.nexus.CurrentNexusOperationContext;
+import io.temporal.internal.nexus.InternalNexusOperationContext;
+import io.temporal.internal.nexus.OperationTokenUtil;
 import java.util.*;
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
@@ -60,7 +65,7 @@ public static Object getValueOrDefault(Object value, Class<?> valueClass) {
    *     URL and headers set
    */
   @SuppressWarnings("deprecation") // Check the OPERATION_ID header for backwards compatibility
-  public static WorkflowStub createNexusBoundStub(
+  public static NexusWorkflowStarter createNexusBoundStub(
       WorkflowStub stub, NexusStartWorkflowRequest request) {
     if (!stub.getOptions().isPresent()) {
       throw new IllegalArgumentException("Options are expected to be set on the stub");
@@ -70,6 +75,19 @@ public static WorkflowStub createNexusBoundStub(
       throw new IllegalArgumentException(
           "WorkflowId is expected to be set on WorkflowOptions when used with Nexus");
     }
+    InternalNexusOperationContext nexusContext = CurrentNexusOperationContext.get();
+    // Generate the operation token for the new workflow.
+    String operationToken;
+    try {
+      operationToken =
+          OperationTokenUtil.generateWorkflowRunOperationToken(
+              options.getWorkflowId(), nexusContext.getNamespace());
+    } catch (JsonProcessingException e) {
+      // Not expected as the link is constructed by the SDK.
+      throw new HandlerException(
+          HandlerException.ErrorType.BAD_REQUEST,
+          new IllegalArgumentException("failed to generate workflow operation token", e));
+    }
     // Add the Nexus operation ID to the headers if it is not already present to support fabricating
     // a NexusOperationStarted event if the completion is received before the response to a
     // StartOperation request.
@@ -82,10 +100,10 @@ public static WorkflowStub createNexusBoundStub(
                     (a, b) -> a,
                     () -> new TreeMap<>(String.CASE_INSENSITIVE_ORDER)));
     if (!headers.containsKey(Header.OPERATION_ID)) {
-      headers.put(Header.OPERATION_ID.toLowerCase(), options.getWorkflowId());
+      headers.put(Header.OPERATION_ID.toLowerCase(), operationToken);
     }
     if (!headers.containsKey(Header.OPERATION_TOKEN)) {
-      headers.put(Header.OPERATION_TOKEN.toLowerCase(), options.getWorkflowId());
+      headers.put(Header.OPERATION_TOKEN.toLowerCase(), operationToken);
     }
     List<Link> links =
         request.getLinks() == null
@@ -134,7 +152,7 @@ public static WorkflowStub createNexusBoundStub(
             .setAttachCompletionCallbacks(true)
             .build());
 
-    return stub.newInstance(nexusWorkflowOptions.build());
+    return new NexusWorkflowStarter(stub.newInstance(nexusWorkflowOptions.build()), operationToken);
   }
 
   /** Check the method name for reserved prefixes or names. */

temporal-sdk/src/main/java/io/temporal/internal/common/NexusWorkflowStarter.java

@@ -0,0 +1,20 @@
+package io.temporal.internal.common;
+
+import io.temporal.api.common.v1.WorkflowExecution;
+import io.temporal.client.WorkflowStub;
+import io.temporal.internal.client.NexusStartWorkflowResponse;
+
+public class NexusWorkflowStarter {
+  private final WorkflowStub workflowStub;
+  private final String operationToken;
+
+  public NexusWorkflowStarter(WorkflowStub workflowStub, String operationToken) {
+    this.workflowStub = workflowStub;
+    this.operationToken = operationToken;
+  }
+
+  public NexusStartWorkflowResponse start(Object... args) {
+    WorkflowExecution workflowExecution = workflowStub.start(args);
+    return new NexusStartWorkflowResponse(workflowExecution, operationToken);
+  }
+}

temporal-sdk/src/main/java/io/temporal/nexus/WorkflowHandleInvoker.java

@@ -1,8 +1,8 @@
 package io.temporal.nexus;
 
-import io.temporal.api.common.v1.WorkflowExecution;
 import io.temporal.internal.client.NexusStartWorkflowRequest;
+import io.temporal.internal.client.NexusStartWorkflowResponse;
 
 interface WorkflowHandleInvoker {
-  WorkflowExecution invoke(NexusStartWorkflowRequest request);
+  NexusStartWorkflowResponse invoke(NexusStartWorkflowRequest request);
 }

temporal-sdk/src/main/java/io/temporal/nexus/WorkflowMethodMethodInvoker.java

@@ -1,21 +1,21 @@
 package io.temporal.nexus;
 
-import io.temporal.api.common.v1.WorkflowExecution;
 import io.temporal.internal.client.NexusStartWorkflowRequest;
+import io.temporal.internal.client.NexusStartWorkflowResponse;
 import io.temporal.internal.client.WorkflowClientInternal;
 import io.temporal.internal.nexus.CurrentNexusOperationContext;
 import io.temporal.internal.nexus.InternalNexusOperationContext;
 import io.temporal.workflow.Functions;
 
 class WorkflowMethodMethodInvoker implements WorkflowHandleInvoker {
-  private Functions.Proc workflow;
+  private final Functions.Proc workflow;
 
   public WorkflowMethodMethodInvoker(Functions.Proc workflow) {
     this.workflow = workflow;
   }
 
   @Override
-  public WorkflowExecution invoke(NexusStartWorkflowRequest request) {
+  public NexusStartWorkflowResponse invoke(NexusStartWorkflowRequest request) {
     InternalNexusOperationContext nexusCtx = CurrentNexusOperationContext.get();
     return ((WorkflowClientInternal) nexusCtx.getWorkflowClient().getInternal())
         .startNexus(request, workflow);

temporal-sdk/src/main/java/io/temporal/nexus/WorkflowRunOperationImpl.java

@@ -3,7 +3,6 @@
 import static io.temporal.internal.common.LinkConverter.workflowEventToNexusLink;
 import static io.temporal.internal.common.NexusUtil.nexusProtoLinkToLink;
 
-import com.fasterxml.jackson.core.JsonProcessingException;
 import io.nexusrpc.OperationInfo;
 import io.nexusrpc.handler.*;
 import io.nexusrpc.handler.OperationHandler;
@@ -12,6 +11,7 @@
 import io.temporal.api.enums.v1.EventType;
 import io.temporal.client.WorkflowClient;
 import io.temporal.internal.client.NexusStartWorkflowRequest;
+import io.temporal.internal.client.NexusStartWorkflowResponse;
 import io.temporal.internal.nexus.CurrentNexusOperationContext;
 import io.temporal.internal.nexus.InternalNexusOperationContext;
 import io.temporal.internal.nexus.OperationTokenUtil;
@@ -39,7 +39,9 @@ public OperationStartResult<R> start(
             nexusCtx.getTaskQueue(),
             operationStartDetails.getLinks());
 
-    WorkflowExecution workflowExec = handle.getInvoker().invoke(nexusRequest);
+    NexusStartWorkflowResponse nexusStartWorkflowResponse =
+        handle.getInvoker().invoke(nexusRequest);
+    WorkflowExecution workflowExec = nexusStartWorkflowResponse.getWorkflowExecution();
 
     // If the start workflow response returned a link use it, otherwise
     // create the link information about the new workflow and return to the caller.
@@ -59,20 +61,9 @@ public OperationStartResult<R> start(
               .build();
     }
     io.temporal.api.nexus.v1.Link nexusLink = workflowEventToNexusLink(workflowEventLink);
-    // Generate the operation token for the new workflow.
-    String operationToken;
-    try {
-      operationToken =
-          OperationTokenUtil.generateWorkflowRunOperationToken(
-              workflowExec.getWorkflowId(), nexusCtx.getNamespace());
-    } catch (JsonProcessingException e) {
-      // Not expected as the link is constructed by the SDK.
-      throw new HandlerException(
-          HandlerException.ErrorType.BAD_REQUEST,
-          new IllegalArgumentException("failed to generate workflow operation token", e));
-    }
     // Attach the link to the operation result.
-    OperationStartResult.Builder<R> result = OperationStartResult.newAsyncBuilder(operationToken);
+    OperationStartResult.Builder<R> result =
+        OperationStartResult.newAsyncBuilder(nexusStartWorkflowResponse.getOperationToken());
     if (nexusLink != null) {
       try {
         ctx.addLinks(nexusProtoLinkToLink(nexusLink));

temporal-sdk/src/main/java/io/temporal/nexus/WorkflowStubHandleInvoker.java

@@ -2,9 +2,9 @@
 
 import static io.temporal.internal.common.InternalUtils.createNexusBoundStub;
 
-import io.temporal.api.common.v1.WorkflowExecution;
 import io.temporal.client.WorkflowStub;
 import io.temporal.internal.client.NexusStartWorkflowRequest;
+import io.temporal.internal.client.NexusStartWorkflowResponse;
 
 class WorkflowStubHandleInvoker implements WorkflowHandleInvoker {
   final Object[] args;
@@ -16,7 +16,7 @@ class WorkflowStubHandleInvoker implements WorkflowHandleInvoker {
   }
 
   @Override
-  public WorkflowExecution invoke(NexusStartWorkflowRequest request) {
+  public NexusStartWorkflowResponse invoke(NexusStartWorkflowRequest request) {
     return createNexusBoundStub(stub, request).start(args);
   }
 }