feat(sdk): Add subtypes for Step, Wait, Callback and Invoke operations

Added new OperationSubType values for all operation types that were
missing them (step, wait, callback and invoke). And pass these values
when creating operations in DurableContextImpl.

ChildContextOperation explicitly checks for unreachable subtypes.
Updated ChildContextOperation to check for the new subtype values.

Updated the tests for the above operations that check the subtype to
use these new subtype values.

Updated LocalMemoryExecutionClient.applyResult to preserve subtype
when updating an operation. Previously the subtype value was not
passed on to the updated operation, causing test failures.

Author: nvasiu

sdk-testing/src/main/java/software/amazon/lambda/durable/testing/local/LocalMemoryExecutionClient.java

@@ -190,6 +190,7 @@ private void applyResult(Operation op, OperationResult result) {
                     .id(op.id())
                     .name(op.name())
                     .type(op.type())
+                    .subType(op.subType())
                     .action(action)
                     .parentId(op.parentId())
                     .payload(result.result())

sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java

@@ -11,7 +11,6 @@
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import org.slf4j.LoggerFactory;
-import software.amazon.awssdk.services.lambda.model.OperationType;
 import software.amazon.lambda.durable.DurableCallbackFuture;
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.DurableContext;
@@ -146,7 +145,7 @@ public <T> DurableFuture<T> stepAsync(
 
         // Create and start step operation with TypeToken
         var operation = new StepOperation<>(
-                OperationIdentifier.of(operationId, name, OperationType.STEP), func, resultType, config, this);
+                OperationIdentifier.of(operationId, name, OperationSubType.STEP), func, resultType, config, this);
 
         operation.execute(); // Start the step (returns immediately)
 
@@ -162,7 +161,7 @@ public DurableFuture<Void> waitAsync(String name, Duration duration) {
 
         // Create and start wait operation
         var operation =
-                new WaitOperation(OperationIdentifier.of(operationId, name, OperationType.WAIT), duration, this);
+                new WaitOperation(OperationIdentifier.of(operationId, name, OperationSubType.WAIT), duration, this);
 
         operation.execute(); // Checkpoint the wait
         return operation;
@@ -187,7 +186,7 @@ public <T, U> DurableFuture<T> invokeAsync(
 
         // Create and start invoke operation
         var operation = new InvokeOperation<>(
-                OperationIdentifier.of(operationId, name, OperationType.CHAINED_INVOKE),
+                OperationIdentifier.of(operationId, name, OperationSubType.CHAINED_INVOKE),
                 functionName,
                 payload,
                 resultType,
@@ -207,7 +206,7 @@ public <T> DurableCallbackFuture<T> createCallback(String name, TypeToken<T> res
         var operationId = nextOperationId();
 
         var operation = new CallbackOperation<>(
-                OperationIdentifier.of(operationId, name, OperationType.CALLBACK), resultType, config, this);
+                OperationIdentifier.of(operationId, name, OperationSubType.CALLBACK), resultType, config, this);
         operation.execute();
 
         return operation;
@@ -248,11 +247,7 @@ private <T> DurableFuture<T> runInChildContextAsync(
         var operationId = nextOperationId();
 
         var operation = new ChildContextOperation<>(
-                OperationIdentifier.of(operationId, name, OperationType.CONTEXT, subType),
-                func,
-                resultType,
-                config,
-                this);
+                OperationIdentifier.of(operationId, name, subType), func, resultType, config, this);
 
         operation.execute();
         return operation;
@@ -277,7 +272,7 @@ public <I, O> DurableFuture<MapResult<O>> mapAsync(
         var operationId = nextOperationId();
 
         var operation = new MapOperation<>(
-                OperationIdentifier.of(operationId, name, OperationType.CONTEXT, OperationSubType.MAP),
+                OperationIdentifier.of(operationId, name, OperationSubType.MAP),
                 itemList,
                 function,
                 resultType,
@@ -293,7 +288,7 @@ public ParallelDurableFuture parallel(String name, ParallelConfig config) {
         var operationId = nextOperationId();
 
         var parallelOp = new ParallelOperation(
-                OperationIdentifier.of(operationId, name, OperationType.CONTEXT, OperationSubType.PARALLEL),
+                OperationIdentifier.of(operationId, name, OperationSubType.PARALLEL),
                 getDurableConfig().getSerDes(),
                 this,
                 config);

sdk/src/main/java/software/amazon/lambda/durable/model/OperationIdentifier.java

@@ -5,24 +5,21 @@
 import software.amazon.awssdk.services.lambda.model.OperationType;
 
 /**
- * Identifies a durable operation by its unique ID, human-readable name, type, and optional sub-type.
+ * Identifies a durable operation by its unique ID, human-readable name, type and sub-type.
  *
  * @param operationId unique sequential identifier for the operation within an execution
  * @param name human-readable name for the operation
- * @param operationType the kind of operation (STEP, WAIT, CALLBACK, etc.)
- * @param subType optional sub-type for operations that need further classification (e.g. child contexts)
+ * @param subType the operation sub-type which also determines the operation type
  */
-public record OperationIdentifier(
-        String operationId, String name, OperationType operationType, OperationSubType subType) {
+public record OperationIdentifier(String operationId, String name, OperationSubType subType) {
 
-    /** Creates an identifier without a sub-type. */
-    public static OperationIdentifier of(String operationId, String name, OperationType type) {
-        return new OperationIdentifier(operationId, name, type, null);
+    /** Returns the operation type derived from the sub-type. */
+    public OperationType operationType() {
+        return subType.getOperationType();
     }
 
-    /** Creates an identifier with a sub-type. */
-    public static OperationIdentifier of(
-            String operationId, String name, OperationType type, OperationSubType subType) {
-        return new OperationIdentifier(operationId, name, type, subType);
+    /** Creates an identifier with the given sub-type. */
+    public static OperationIdentifier of(String operationId, String name, OperationSubType subType) {
+        return new OperationIdentifier(operationId, name, subType);
     }
 }

sdk/src/main/java/software/amazon/lambda/durable/model/OperationSubType.java

@@ -2,28 +2,41 @@
 // SPDX-License-Identifier: Apache-2.0
 package software.amazon.lambda.durable.model;
 
+import software.amazon.awssdk.services.lambda.model.OperationType;
+
 /**
- * Fine-grained classification of durable operations beyond the basic operation types.
+ * Fine-grained classification of durable operations that pairs each subtype with its parent {@link OperationType}.
  *
- * <p>Used as the {@code subType} field in checkpoint updates for {@code CONTEXT} operations. Matches the
- * {@code OperationSubType} enum in the JavaScript and Python durable execution SDKs.
+ * <p>Used as the source of both the {@code type} and {@code subType} fields in checkpoint updates. Matches
+ * the {@code OperationSubType} enum in the JavaScript and Python durable execution SDKs.
  */
 public enum OperationSubType {
-    RUN_IN_CHILD_CONTEXT("RunInChildContext"),
-    MAP("Map"),
-    MAP_ITERATION("MapIteration"),
-    PARALLEL("Parallel"),
-    PARALLEL_BRANCH("ParallelBranch"),
-    WAIT_FOR_CALLBACK("WaitForCallback"),
-    WAIT_FOR_CONDITION("WaitForCondition"),
-    WITH_RETRY("WithRetry");
+    STEP(OperationType.STEP, "Step"),
+    WAIT(OperationType.WAIT, "Wait"),
+    CALLBACK(OperationType.CALLBACK, "Callback"),
+    CHAINED_INVOKE(OperationType.CHAINED_INVOKE, "ChainedInvoke"),
+    RUN_IN_CHILD_CONTEXT(OperationType.CONTEXT, "RunInChildContext"),
+    MAP(OperationType.CONTEXT, "Map"),
+    MAP_ITERATION(OperationType.CONTEXT, "MapIteration"),
+    PARALLEL(OperationType.CONTEXT, "Parallel"),
+    PARALLEL_BRANCH(OperationType.CONTEXT, "ParallelBranch"),
+    WAIT_FOR_CALLBACK(OperationType.CONTEXT, "WaitForCallback"),
+    WAIT_FOR_CONDITION(OperationType.STEP, "WaitForCondition"),
+    WITH_RETRY(OperationType.CONTEXT, "WithRetry");
 
+    private final OperationType operationType;
     private final String value;
 
-    OperationSubType(String value) {
+    OperationSubType(OperationType operationType, String value) {
+        this.operationType = operationType;
         this.value = value;
     }
 
+    /** Returns the parent {@link OperationType} for this subtype. */
+    public OperationType getOperationType() {
+        return operationType;
+    }
+
     /** Returns the wire-format string value sent in checkpoint updates. */
     public String getValue() {
         return value;

sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java

@@ -410,11 +410,11 @@ protected void sendOperationUpdate(OperationUpdate.Builder builder) {
 
     /** Sends an operation update asynchronously. */
     protected CompletableFuture<Void> sendOperationUpdateAsync(OperationUpdate.Builder builder) {
-        var updateBuilder =
-                builder.id(getOperationId()).name(getName()).type(getType()).parentId(durableContext.getParentId());
-        if (getSubType() != null) {
-            updateBuilder.subType(getSubType().getValue());
-        }
+        var updateBuilder = builder.id(getOperationId())
+                .name(getName())
+                .type(getType())
+                .subType(getSubType().getValue())
+                .parentId(durableContext.getParentId());
         var update = updateBuilder.build();
         if (replayCompletedOperation.get()) {
             // We are replaying a completed operation, so complete the completableFuture without checkpointing
@@ -444,9 +444,7 @@ protected void validateReplay(Operation checkpointed) {
                     getOperationId(), checkpointed.name(), getName())));
         }
 
-        if ((getSubType() == null && checkpointed.subType() != null)
-                || getSubType() != null
-                        && !Objects.equals(checkpointed.subType(), getSubType().getValue())) {
+        if (!Objects.equals(checkpointed.subType(), getSubType().getValue())) {
             throw terminateExecution(new NonDeterministicExecutionException(String.format(
                     "Operation subType mismatch for \"%s\". Expected \"%s\", got \"%s\"",
                     getOperationId(), checkpointed.subType(), getSubType())));

sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java

@@ -248,7 +248,8 @@ private Throwable translateException(Operation op, ErrorObject errorObject) {
             case RUN_IN_CHILD_CONTEXT, WITH_RETRY -> new ChildContextFailedException(op);
 
             // the following subtypes should not be able to reach here
-            case PARALLEL, MAP, WAIT_FOR_CONDITION -> new IllegalStateException("Unexpected sub-type: " + getSubType());
+            case PARALLEL, MAP, WAIT_FOR_CONDITION, STEP, WAIT, CALLBACK, CHAINED_INVOKE ->
+                new IllegalStateException("Unexpected sub-type: " + getSubType());
         };
     }
 

sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java

@@ -16,7 +16,6 @@
 import java.util.function.Function;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import software.amazon.awssdk.services.lambda.model.OperationType;
 import software.amazon.lambda.durable.DurableContext;
 import software.amazon.lambda.durable.TypeToken;
 import software.amazon.lambda.durable.config.NestingType;
@@ -116,7 +115,7 @@ protected <R> ChildContextOperation<R> createItem(
             SerDes serDes,
             OperationSubType branchSubType) {
         return new ChildContextOperation<>(
-                OperationIdentifier.of(operationId, name, OperationType.CONTEXT, branchSubType),
+                OperationIdentifier.of(operationId, name, branchSubType),
                 function,
                 resultType,
                 RunInChildContextConfig.builder()

sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java

@@ -8,7 +8,6 @@
 import software.amazon.awssdk.services.lambda.model.Operation;
 import software.amazon.awssdk.services.lambda.model.OperationAction;
 import software.amazon.awssdk.services.lambda.model.OperationStatus;
-import software.amazon.awssdk.services.lambda.model.OperationType;
 import software.amazon.awssdk.services.lambda.model.OperationUpdate;
 import software.amazon.awssdk.services.lambda.model.StepOptions;
 import software.amazon.lambda.durable.StepContext;
@@ -48,7 +47,7 @@ public WaitForConditionOperation(
             WaitForConditionConfig<T> config,
             DurableContextImpl durableContext) {
         super(
-                OperationIdentifier.of(operationId, name, OperationType.STEP, OperationSubType.WAIT_FOR_CONDITION),
+                OperationIdentifier.of(operationId, name, OperationSubType.WAIT_FOR_CONDITION),
                 resultTypeToken,
                 config.serDes(),
                 durableContext);

sdk/src/test/java/software/amazon/lambda/durable/ReplayValidationTest.java

@@ -21,6 +21,7 @@
 import software.amazon.lambda.durable.execution.ThreadContext;
 import software.amazon.lambda.durable.execution.ThreadType;
 import software.amazon.lambda.durable.model.DurableExecutionInput;
+import software.amazon.lambda.durable.model.OperationSubType;
 
 class ReplayValidationTest {
     private static final String EXECUTION_NAME = "exec-name";
@@ -67,6 +68,7 @@ void shouldPassValidationWhenStepTypeAndNameMatch() {
                 .id(OPERATION_ID1)
                 .name("test")
                 .type(OperationType.STEP)
+                .subType(OperationSubType.STEP.getValue())
                 .status(OperationStatus.SUCCEEDED)
                 .stepDetails(StepDetails.builder().result("\"result\"").build())
                 .build();
@@ -83,6 +85,7 @@ void shouldPassValidationWhenWaitTypeMatches() {
         var existingOp = Operation.builder()
                 .id(OPERATION_ID1)
                 .type(OperationType.WAIT)
+                .subType(OperationSubType.WAIT.getValue())
                 .status(OperationStatus.SUCCEEDED)
                 .build();
 
@@ -144,6 +147,7 @@ void shouldHandleNullNamesCorrectly() {
                 .id(OPERATION_ID1)
                 .name(null)
                 .type(OperationType.STEP)
+                .subType(OperationSubType.STEP.getValue())
                 .status(OperationStatus.SUCCEEDED)
                 .stepDetails(StepDetails.builder().result("\"result\"").build())
                 .build();

sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java

@@ -25,6 +25,7 @@
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.model.DurableExecutionInput;
 import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.model.OperationSubType;
 
 class DurableExecutionTest {
 
@@ -187,6 +188,7 @@ void testExecuteReplay() {
                 .id(OPERATION_ID1)
                 .name("step1")
                 .type(OperationType.STEP)
+                .subType(OperationSubType.STEP.getValue())
                 .status(OperationStatus.SUCCEEDED)
                 .stepDetails(StepDetails.builder().result("\"First\"").build())
                 .build();