add waitForCallback

Author: zhongkechen

README.md

@@ -21,6 +21,8 @@ Your durable function extends `DurableHandler<I, O>` and implements `handleReque
 - `ctx.stepAsync()` – Start a concurrent step  
 - `ctx.wait()` – Suspend execution without compute charges
 - `ctx.createCallback()` – Wait for external events (approvals, webhooks)
+- `ctx.waitForCallback()` – Run an isolated child context with its own checkpoint log
+- `ctx.waitForCallbackAsync()` – Start a concurrent child context
 - `ctx.invoke()` – Invoke another Lambda function and wait for the result
 - `ctx.invokeAsync()` – Start a concurrent Lambda function invocation
 - `ctx.runInChildContext()` – Run an isolated child context with its own checkpoint log

docs/core/callbacks.md

@@ -7,7 +7,7 @@ Callbacks suspend execution until an external system sends a result. Use this fo
 DurableCallbackFuture<String> callback = ctx.createCallback("approval", String.class);
 
 // Send the callback ID to an external system within a step
-ctx.step("send-notification", String.class, () -> {
+ctx.step("send-notification", String.class, stepCtx -> {
     notificationService.sendApprovalRequest(callback.callbackId(), requestDetails);
     return "notification-sent";
 });
@@ -18,6 +18,16 @@ String approvalResult = callback.get();
 
 The external system completes the callback by calling the Lambda Durable Functions API with the callback ID and result payload.
 
+#### waitForCallback() #### 
+
+`waitForCallback` simplifies callback handling by combining callback creation and submission in one operation. The SDK creates the callback, executes your submitter function with the callback ID, and waits for the result.
+
+```java
+ctx.waitForCallback("send-notification", String.class, (callbackId, stepCtx) -> {
+    notificationService.sendApprovalRequest(callbackId, requestDetails);
+})
+```
+
 #### Callback Configuration
 
 Configure timeouts and serialization to handle cases where callbacks are never completed or need custom deserialization:
@@ -30,6 +40,12 @@ var config = CallbackConfig.builder()
     .build();
 
 var callback = ctx.createCallback("approval", String.class, config);
+
+var waitForCallbackConfig = WaitForCallbackConfig.builder()
+    .callbackConfig(config)
+    .stepConfig(StepConfig.builder().retryStrategy(...).build())
+    .build();
+ctx.waitForCallback("approval", String.class, callbackId -> sendApprovalRequest(callbackId), waitForCallbackConfig);
 ```
 
 | Option | Description |

examples/src/main/java/software/amazon/lambda/durable/examples/CallbackExample.java

@@ -44,6 +44,10 @@ public String handleRequest(ApprovalRequest input, DurableContext context) {
 
         var config = CallbackConfig.builder().timeout(timeout).build();
 
+        var preapprovalCallback = context.waitForCallbackAsync("preapproval", String.class, (callbackId, ctx) -> {
+            ctx.getLogger().info("Sending callback {} to preapproval system", callbackId);
+        });
+
         var callback = context.createCallback("approval", String.class, config);
 
         // Step 2.5: Log AWS CLI command to complete the callback
@@ -57,15 +61,14 @@ public String handleRequest(ApprovalRequest input, DurableContext context) {
             return null;
         });
 
+        var preapprovalResult = preapprovalCallback.get();
+
         // Step 3: Wait for external approval (suspends execution)
         var approvalResult = callback.get();
 
         // Step 4: Process the approval
-        var result = context.step("process-approval", String.class, () -> {
-            return prepared + " - " + approvalResult;
-        });
-
-        return result;
+        return context.step(
+                "process-approval", String.class, () -> prepared + " - " + preapprovalResult + " - " + approvalResult);
     }
 }
 

examples/src/test/java/software/amazon/lambda/durable/examples/CallbackExampleTest.java

@@ -46,11 +46,19 @@ void testCallbackExampleCompletesAfterApproval() {
         var callbackId = runner.getCallbackId("approval");
         runner.completeCallback(callbackId, "\"Approved by manager\"");
 
-        // Second run - callback complete, finishes processing
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // second run - pending preapproval
+        var preapprovalCallbackId = runner.getCallbackId("preapproval-callback");
+        runner.completeCallback(preapprovalCallbackId, "\"Sent to preapprover\"");
+
+        // third run - callback complete, finishes processing
         result = runner.run(input);
 
         assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
         assertEquals(
-                "Approval request for: New laptop ($1500.0) - Approved by manager", result.getResult(String.class));
+                "Approval request for: New laptop ($1500.0) - Sent to preapprover - Approved by manager",
+                result.getResult(String.class));
     }
 }

examples/src/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java

@@ -275,10 +275,16 @@ void testErrorHandlingExample() {
     @Test
     void testCallbackExample() throws Exception {
         var runner = CloudDurableTestRunner.create(arn("callback-example"), ApprovalRequest.class, String.class);
+        var lambda = LambdaClient.create();
 
         // Start async execution
         var execution = runner.startAsync(new ApprovalRequest("Purchase order", 5000.0));
 
+        execution.pollUntil(exec -> exec.hasCallback("preapproval-callback"));
+        var preapprovalCallbackId = execution.getCallbackId("preapproval-callback");
+        lambda.sendDurableExecutionCallbackSuccess(
+                req -> req.callbackId(preapprovalCallbackId).result(SdkBytes.fromUtf8String("\"preapproved\"")));
+
         // Wait for callback to appear
         execution.pollUntil(exec -> exec.hasCallback("approval"));
 
@@ -287,7 +293,6 @@ void testCallbackExample() throws Exception {
         assertNotNull(callbackId);
 
         // Complete the callback using AWS SDK
-        var lambda = LambdaClient.create();
         lambda.sendDurableExecutionCallbackSuccess(
                 req -> req.callbackId(callbackId).result(SdkBytes.fromUtf8String("\"approved\"")));
 
@@ -297,6 +302,7 @@ void testCallbackExample() throws Exception {
 
         var finalResult = result.getResult(String.class);
         assertNotNull(finalResult);
+        assertTrue(finalResult.contains("preapproved"));
         assertTrue(finalResult.contains("Approval request for: Purchase order"));
         assertTrue(finalResult.contains("5000"));
         assertTrue(finalResult.contains("approved"));

examples/src/test/java/software/amazon/lambda/durable/examples/GenericTypesExampleTest.java

@@ -12,8 +12,7 @@ class GenericTypesExampleTest {
     @Test
     void testGenericTypesExample() {
         var handler = new GenericTypesExample();
-        var runner = LocalDurableTestRunner.create(GenericTypesExample.Input.class, handler)
-                .withSkipTime(true);
+        var runner = LocalDurableTestRunner.create(GenericTypesExample.Input.class, handler);
 
         var input = new GenericTypesExample.Input("user123");
         var result = runner.run(input);
@@ -48,8 +47,7 @@ void testGenericTypesExample() {
     @Test
     void testOperationTracking() {
         var handler = new GenericTypesExample();
-        var runner = LocalDurableTestRunner.create(GenericTypesExample.Input.class, handler)
-                .withSkipTime(true);
+        var runner = LocalDurableTestRunner.create(GenericTypesExample.Input.class, handler);
 
         var input = new GenericTypesExample.Input("user456");
         var result = runner.run(input);

sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ChildContextIntegrationTest.java

@@ -23,7 +23,7 @@ void childContextResultSurvivesReplay() {
         var childExecutionCount = new AtomicInteger(0);
 
         var runner = LocalDurableTestRunner.create(String.class, (input, ctx) -> {
-            return ctx.runInChildContext("compute", String.class, child -> {
+            return ctx.runInChildContext("compute", TypeToken.get(String.class), child -> {
                 childExecutionCount.incrementAndGet();
                 return child.step("work", String.class, () -> "result-" + input);
             });
@@ -191,7 +191,6 @@ void waitInsideChildContextSuspendsAndResumes() {
                 return child.step("after-wait", String.class, () -> "done");
             });
         });
-        runner.withSkipTime(true);
 
         var result = runner.runUntilComplete("test");
         assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
@@ -211,8 +210,6 @@ void waitInsideChildContextReturnsPendingThenCompletes() {
                 return child.step("after-wait", String.class, () -> "done");
             });
         });
-        runner.withSkipTime(false);
-
         // First run - should suspend at the wait
         var result = runner.run("test");
         assertEquals(ExecutionStatus.PENDING, result.getStatus());
@@ -246,16 +243,13 @@ void twoAsyncChildContextsBothWaitSuspendAndResume() {
             return f1.get() + "+" + f2.get();
         });
 
-        runner.withSkipTime(false);
-
         // First run - both child contexts should suspend at their waits
         // TODO: Using run() + runUntilComplete() instead of manual run/advanceTime/run due to a
         //  thread coordination race condition that causes flakiness on slow CI workers.
         var result = runner.run("test");
         assertEquals(ExecutionStatus.PENDING, result.getStatus());
 
         // Now let runUntilComplete handle the rest (with skipTime so waits auto-advance)
-        runner.withSkipTime(true);
         var finalResult = runner.runUntilComplete("test");
         assertEquals(ExecutionStatus.SUCCEEDED, finalResult.getStatus());
         assertEquals("a-done+b-done", finalResult.getResult(String.class));
@@ -285,7 +279,6 @@ void oneChildWaitsWhileOtherKeepsProcessingSuspendsAfterWorkDone() {
             });
             return busy.get() + "|" + waiting.get();
         });
-        runner.withSkipTime(false);
 
         // First run: busy child completes its work, but waiter's wait is still outstanding → PENDING
         var result = runner.run("test");

sdk-integration-tests/src/test/java/software/amazon/lambda/durable/CustomConfigIntegrationTest.java

@@ -98,19 +98,18 @@ void testCustomConfig_WithMultipleSteps() {
         var customConfig = DurableConfig.builder().withSerDes(customSerDes).build();
 
         var runner = LocalDurableTestRunner.create(
-                        String.class,
-                        (input, context) -> {
-                            var step1 = context.step("step1", String.class, () -> "Step1: " + input);
+                String.class,
+                (input, context) -> {
+                    var step1 = context.step("step1", String.class, () -> "Step1: " + input);
 
-                            // Remove wait operation to avoid complexity in this test
-                            var step2 = context.step("step2", String.class, () -> "Step2: " + input);
+                    // Remove wait operation to avoid complexity in this test
+                    var step2 = context.step("step2", String.class, () -> "Step2: " + input);
 
-                            var step3 = context.step("step3", String.class, () -> "Step3: " + input);
+                    var step3 = context.step("step3", String.class, () -> "Step3: " + input);
 
-                            return step1 + " | " + step2 + " | " + step3;
-                        },
-                        customConfig)
-                .withSkipTime(true);
+                    return step1 + " | " + step2 + " | " + step3;
+                },
+                customConfig);
 
         var result = runner.run("test");
 
@@ -137,24 +136,21 @@ void testCustomConfig_WithRetry() {
         var customConfig = DurableConfig.builder().withSerDes(customSerDes).build();
 
         var runner = LocalDurableTestRunner.create(
-                        String.class,
-                        (input, context) -> {
-                            return context.step(
-                                    "retry-step",
-                                    String.class,
-                                    () -> {
-                                        int currentAttempt = attemptCount.incrementAndGet();
-                                        // Always fail to test retry behavior (like existing RetryIntegrationTest)
-                                        throw new RuntimeException("Simulated failure attempt " + currentAttempt);
-                                    },
-                                    StepConfig.builder()
-                                            .retryStrategy(
-                                                    software.amazon.lambda.durable.retry.RetryStrategies.Presets
-                                                            .DEFAULT)
-                                            .build());
-                        },
-                        customConfig)
-                .withSkipTime(true);
+                String.class,
+                (input, context) -> {
+                    return context.step(
+                            "retry-step",
+                            String.class,
+                            () -> {
+                                int currentAttempt = attemptCount.incrementAndGet();
+                                // Always fail to test retry behavior (like existing RetryIntegrationTest)
+                                throw new RuntimeException("Simulated failure attempt " + currentAttempt);
+                            },
+                            StepConfig.builder()
+                                    .retryStrategy(software.amazon.lambda.durable.retry.RetryStrategies.Presets.DEFAULT)
+                                    .build());
+                },
+                customConfig);
 
         // First run should return PENDING (retry scheduled) - matching existing RetryIntegrationTest pattern
         var result = runner.run("test");

sdk-integration-tests/src/test/java/software/amazon/lambda/durable/IntegrationTest.java

@@ -104,7 +104,6 @@ void testFullWaitOperation() {
             var step2 = context.step("step2", String.class, () -> "Step 2 done");
             return new TestOutput(step1 + " + " + step2);
         });
-        runner.withSkipTime(true);
 
         var result = runner.runUntilComplete(new TestInput("test"));
 
@@ -167,7 +166,6 @@ void testOperationFiltering() {
             context.step("good-step", String.class, () -> "ok");
             return "done";
         });
-        runner.withSkipTime(true);
 
         var result = runner.runUntilComplete(new TestInput("test"));
 
@@ -184,9 +182,8 @@ void testWaitOperationWithManualAdvance() {
             context.wait(Duration.ofSeconds(5));
             return "done";
         });
-        runner.withSkipTime(false);
 
-        var result = runner.runUntilComplete(new TestInput("test"));
+        var result = runner.run(new TestInput("test"));
 
         assertEquals(ExecutionStatus.PENDING, result.getStatus());
         assertEquals(1, result.getSucceededOperations().size());

sdk-integration-tests/src/test/java/software/amazon/lambda/durable/StepSemanticsIntegrationTest.java

@@ -143,7 +143,6 @@ void testAtLeastOnceReExecutesAfterCheckpointFailure() {
                             .semantics(StepSemantics.AT_LEAST_ONCE_PER_RETRY)
                             .build());
         });
-        runner.withSkipTime(true);
 
         runner.run("test");
         assertEquals(1, executionCount.get());