some clean up

Author: Quinn-With-Two-Ns

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

@@ -24,17 +24,23 @@ public class OperationToken {
   @JsonInclude(JsonInclude.Include.NON_NULL)
   private final String activityId;
 
+  @JsonProperty("rid")
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  private final String runId;
+
   @JsonCreator
   public OperationToken(
       @JsonProperty("t") Integer type,
       @JsonProperty("ns") String namespace,
       @JsonProperty("wid") String workflowId,
       @JsonProperty("aid") String activityId,
+      @JsonProperty("rid") String runId,
       @JsonProperty("v") Integer version) {
     this.type = OperationTokenType.fromValue(type);
     this.namespace = namespace;
     this.workflowId = workflowId;
     this.activityId = activityId;
+    this.runId = runId;
     this.version = version;
   }
 
@@ -43,15 +49,26 @@ public OperationToken(OperationTokenType type, String namespace, String workflow
     this.namespace = namespace;
     this.workflowId = workflowId;
     this.activityId = null;
+    this.runId = null;
     this.version = null;
   }
 
   public OperationToken(
       OperationTokenType type, String namespace, String workflowId, String activityId) {
+    this(type, namespace, workflowId, activityId, null);
+  }
+
+  public OperationToken(
+      OperationTokenType type,
+      String namespace,
+      String workflowId,
+      String activityId,
+      String runId) {
     this.type = type;
     this.namespace = namespace;
     this.workflowId = workflowId;
     this.activityId = activityId;
+    this.runId = runId;
     this.version = null;
   }
 
@@ -74,4 +91,15 @@ public String getWorkflowId() {
   public String getActivityId() {
     return activityId;
   }
+
+  /**
+   * Returns the activity run ID embedded in the token, or {@code null} if absent.
+   *
+   * <p>Run ID is only present on activity-execution tokens that were generated AFTER the start
+   * activity RPC completed (so the run ID was known). Tokens written into the Nexus operation-token
+   * callback header are generated before that point and therefore do not carry a run ID.
+   */
+  public String getRunId() {
+    return runId;
+  }
 }

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

@@ -107,12 +107,33 @@ public static String generateWorkflowRunOperationToken(String workflowId, String
     return encoder.encodeToString(json.getBytes());
   }
 
-  /** Generate an activity execution operation token from an activity ID and namespace. */
+  /**
+   * Generate an activity execution operation token from an activity ID and namespace.
+   *
+   * <p>This overload omits the run ID. Use it when writing the token into the Nexus operation-token
+   * callback header — that token is generated before the start RPC completes, so the run ID is not
+   * yet known.
+   */
   public static String generateActivityExecutionOperationToken(String activityId, String namespace)
       throws JsonProcessingException {
+    return generateActivityExecutionOperationToken(activityId, null, namespace);
+  }
+
+  /**
+   * Generate an activity execution operation token from an activity ID, run ID, and namespace. The
+   * {@code runId} is included only when non-null.
+   *
+   * <p>This overload is used for the operation token returned to the Nexus caller from a start
+   * operation — at that point the start RPC has completed and the run ID is known. The header token
+   * written into the activity completion callback must NOT carry a run ID; use {@link
+   * #generateActivityExecutionOperationToken(String, String)} for that path.
+   */
+  public static String generateActivityExecutionOperationToken(
+      String activityId, String runId, String namespace) throws JsonProcessingException {
     String json =
         ow.writeValueAsString(
-            new OperationToken(OperationTokenType.ACTIVITY_EXECUTION, namespace, null, activityId));
+            new OperationToken(
+                OperationTokenType.ACTIVITY_EXECUTION, namespace, null, activityId, runId));
     return encoder.encodeToString(json.getBytes());
   }
 

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

@@ -2,6 +2,7 @@
 
 import io.temporal.common.Experimental;
 import java.util.Objects;
+import javax.annotation.Nullable;
 
 /**
  * Input to {@link TemporalOperationHandler#cancelActivityExecution} describing the activity
@@ -11,13 +12,27 @@
 public final class CancelActivityExecutionInput {
 
   private final String activityId;
+  private final @Nullable String runId;
 
-  public CancelActivityExecutionInput(String activityId) {
+  public CancelActivityExecutionInput(String activityId, @Nullable String runId) {
     this.activityId = Objects.requireNonNull(activityId);
+    this.runId = runId;
   }
 
   /** Returns the activity ID extracted from the operation token. */
   public String getActivityId() {
     return activityId;
   }
+
+  /**
+   * Returns the activity run ID extracted from the operation token, or {@code null} if absent.
+   *
+   * <p>Run ID is only present on operation tokens that were generated by this SDK AFTER the start
+   * activity RPC completed. Tokens originating from the activity completion callback header do not
+   * carry a run ID.
+   */
+  @Nullable
+  public String getRunId() {
+    return runId;
+  }
 }

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

@@ -1,5 +1,6 @@
 package io.temporal.nexus;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import io.nexusrpc.handler.HandlerException;
 import io.nexusrpc.handler.OperationContext;
 import io.nexusrpc.handler.OperationStartDetails;
@@ -17,6 +18,7 @@
 import io.temporal.internal.client.NexusStartWorkflowResponse;
 import io.temporal.internal.nexus.NexusStartActivityHelper;
 import io.temporal.internal.nexus.NexusStartWorkflowHelper;
+import io.temporal.internal.nexus.OperationTokenUtil;
 import io.temporal.internal.util.MethodExtractor;
 import io.temporal.workflow.Functions;
 import java.lang.reflect.Method;
@@ -257,7 +259,7 @@ public <R> TemporalOperationResult<R> startWorkflow(
   }
 
   private <R> TemporalOperationResult<R> invokeAndReturn(WorkflowHandle<R> handle) {
-    claimAsyncSlot();
+    markAsyncOperationStarted();
     try {
       NexusStartWorkflowResponse response =
           NexusStartWorkflowHelper.startWorkflowAndAttachLinks(
@@ -436,7 +438,7 @@ public <R> TemporalOperationResult<R> startActivity(
 
   private <R> TemporalOperationResult<R> startActivityImpl(
       String activityType, List<Object> args, StartActivityOptions options) {
-    claimAsyncSlot();
+    markAsyncOperationStarted();
     try {
       NexusStartActivityResponse response =
           NexusStartActivityHelper.startActivityAndAttachLinks(
@@ -473,18 +475,34 @@ private <R> TemporalOperationResult<R> startActivityImpl(
                             .build());
                 ActivityClientCallsInterceptor.StartActivityOutput out =
                     ((ActivityClientInternal) activityClient).getInvoker().startActivity(input);
-                // The invoker generated and injected the token into the callback headers before
-                // the start RPC; read it back from the output instead of regenerating it.
-                String token = out.getNexusOperationToken();
-                if (token == null) {
+                // The invoker generated and injected the runId-free token into the callback
+                // headers before the start RPC fired. The header token cannot include a run ID
+                // because the run ID isn't known until after the start RPC returns. The operation
+                // token returned to the Nexus caller can — and should — include it, so it's
+                // regenerated here from the same activity ID + the run ID the start RPC produced.
+                String headerToken = out.getNexusOperationToken();
+                if (headerToken == null) {
                   throw new HandlerException(
                       HandlerException.ErrorType.INTERNAL,
                       "invoker did not return a Nexus operation token for activity start with callback",
                       new IllegalStateException(
                           "nexusOperationToken is null on StartActivityOutput when CompletionCallback was set"));
                 }
+                String returnToken;
+                try {
+                  returnToken =
+                      OperationTokenUtil.generateActivityExecutionOperationToken(
+                          out.getActivityId(),
+                          out.getActivityRunId(),
+                          client.getOptions().getNamespace());
+                } catch (JsonProcessingException e) {
+                  throw new HandlerException(
+                      HandlerException.ErrorType.INTERNAL,
+                      "failed to generate activity operation token",
+                      e);
+                }
                 return new NexusStartActivityResponse(
-                    out.getActivityId(), out.getActivityRunId(), token);
+                    out.getActivityId(), out.getActivityRunId(), returnToken);
               });
       return TemporalOperationResult.async(response.getOperationToken());
     } catch (Throwable t) {
@@ -495,7 +513,7 @@ private <R> TemporalOperationResult<R> startActivityImpl(
     }
   }
 
-  private void claimAsyncSlot() {
+  private void markAsyncOperationStarted() {
     if (!asyncOperationStarted.compareAndSet(false, true)) {
       throw new HandlerException(
           HandlerException.ErrorType.BAD_REQUEST,

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

@@ -110,7 +110,8 @@ public final void cancel(OperationContext ctx, OperationCancelDetails details) {
         break;
       case ACTIVITY_EXECUTION:
         cancelActivityExecution(
-            cancelContext, new CancelActivityExecutionInput(token.getActivityId()));
+            cancelContext,
+            new CancelActivityExecutionInput(token.getActivityId(), token.getRunId()));
         break;
       default:
         throw new HandlerException(
@@ -154,6 +155,6 @@ protected void cancelActivityExecution(
                 .setDataConverter(wc.getOptions().getDataConverter())
                 .setIdentity(wc.getOptions().getIdentity())
                 .build());
-    ac.getHandle(input.getActivityId(), null).cancel();
+    ac.getHandle(input.getActivityId(), input.getRunId()).cancel();
   }
 }

temporal-sdk/src/test/java/io/temporal/internal/nexus/WorkflowRunTokenTest.java

@@ -100,12 +100,36 @@ public void roundTripActivityExecutionToken() throws JsonProcessingException {
     Assert.assertEquals("act-1", token.getActivityId());
     Assert.assertEquals("ns", token.getNamespace());
     Assert.assertNull(token.getWorkflowId());
+    Assert.assertNull(token.getRunId());
     Assert.assertNull(token.getVersion());
 
     // Also exercise the symmetric activityId loader.
     Assert.assertEquals("act-1", OperationTokenUtil.loadActivityIdFromOperationToken(encoded));
   }
 
+  @Test
+  public void roundTripActivityExecutionTokenWithRunId() throws JsonProcessingException {
+    String encoded =
+        OperationTokenUtil.generateActivityExecutionOperationToken("act-1", "run-1", "ns");
+    OperationToken token = OperationTokenUtil.loadOperationToken(encoded);
+    Assert.assertEquals(OperationTokenType.ACTIVITY_EXECUTION, token.getType());
+    Assert.assertEquals("act-1", token.getActivityId());
+    Assert.assertEquals("run-1", token.getRunId());
+    Assert.assertEquals("ns", token.getNamespace());
+    Assert.assertNull(token.getWorkflowId());
+    Assert.assertNull(token.getVersion());
+  }
+
+  @Test
+  public void activityExecutionTokenOmitsRunIdWhenNull() throws JsonProcessingException {
+    // The header-token path passes runId=null and must produce a payload byte-identical to the
+    // two-arg overload — the runId-aware overload is the *only* one used in code now.
+    String withRunId =
+        OperationTokenUtil.generateActivityExecutionOperationToken("act-1", null, "ns");
+    String withoutRunId = OperationTokenUtil.generateActivityExecutionOperationToken("act-1", "ns");
+    Assert.assertEquals(withoutRunId, withRunId);
+  }
+
   @Test
   public void workflowRunTokenBytesByteIdenticalSnapshot() throws JsonProcessingException {
     String encoded = OperationTokenUtil.generateWorkflowRunOperationToken("wf-1", "ns");

temporal-sdk/src/test/java/io/temporal/nexus/TemporalNexusClientImplTest.java

@@ -20,10 +20,10 @@
 import org.junit.Test;
 
 /**
- * Pure unit tests for {@link TemporalNexusClientImpl#claimAsyncSlot()} semantics. These run without
- * a Temporal server (no {@link io.temporal.testing.internal.SDKTestWorkflowRule}).
+ * Pure unit tests for {@link TemporalNexusClientImpl#markAsyncOperationStarted()} semantics. These
+ * run without a Temporal server (no {@link io.temporal.testing.internal.SDKTestWorkflowRule}).
  *
- * <p>The {@code claimAsyncSlot} guard fires as the very first statement in both {@code
+ * <p>The {@code markAsyncOperationStarted} guard fires as the very first statement in both {@code
  * startActivityImpl} and {@code invokeAndReturn}, so the second call always throws {@link
  * HandlerException}({@link HandlerException.ErrorType#BAD_REQUEST}) regardless of whether the first
  * call's downstream RPC succeeded.
@@ -80,22 +80,23 @@ public void doubleStartActivity_secondCallThrowsBadRequest() {
             .setStartToCloseTimeout(Duration.ofSeconds(10))
             .build();
 
-    // First call: claimAsyncSlot() succeeds (sets flag), RPC may throw — we don't care.
+    // First call: markAsyncOperationStarted() succeeds (sets flag), RPC may throw — we don't care.
     try {
       client.startActivity(TestActivity.class, TestActivity::doSomething, options);
     } catch (HandlerException e) {
       // If a HandlerException leaks out of the first call it must NOT be BAD_REQUEST
-      // from claimAsyncSlot — that would mean the flag was already set before setUp.
+      // from markAsyncOperationStarted — that would mean the flag was already set before setUp.
       Assert.assertNotEquals(
-          "First startActivity must not fail with BAD_REQUEST (claimAsyncSlot guard)",
+          "First startActivity must not fail with BAD_REQUEST (markAsyncOperationStarted guard)",
           HandlerException.ErrorType.BAD_REQUEST,
           e.getErrorType());
     } catch (Exception ignored) {
       // Any other exception from the RPC layer is expected; the important thing is
-      // claimAsyncSlot already ran and set asyncOperationStarted = true.
+      // markAsyncOperationStarted already ran and set asyncOperationStarted = true.
     }
 
-    // Second call: claimAsyncSlot() sees the flag and must throw BAD_REQUEST immediately.
+    // Second call: markAsyncOperationStarted() sees the flag and must throw BAD_REQUEST
+    // immediately.
     HandlerException ex =
         Assert.assertThrows(
             HandlerException.class,
@@ -125,12 +126,12 @@ public void doubleStartWorkflow_secondCallThrowsBadRequest() {
     WorkflowOptions options =
         WorkflowOptions.newBuilder().setWorkflowId("wf-1").setTaskQueue(TASK_QUEUE).build();
 
-    // First call: claimAsyncSlot() succeeds, downstream RPC may throw — we don't care.
+    // First call: markAsyncOperationStarted() succeeds, downstream RPC may throw — we don't care.
     try {
       client.startWorkflow(BlockingWorkflow.class, BlockingWorkflow::execute, "input", options);
     } catch (HandlerException e) {
       Assert.assertNotEquals(
-          "First startWorkflow must not fail with BAD_REQUEST (claimAsyncSlot guard)",
+          "First startWorkflow must not fail with BAD_REQUEST (markAsyncOperationStarted guard)",
           HandlerException.ErrorType.BAD_REQUEST,
           e.getErrorType());
     } catch (Exception ignored) {