feat: withRetry helper

examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java

@@ -0,0 +1,50 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples.callback;
+
+import java.time.Duration;
+import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.DurableHandler;
+import software.amazon.lambda.durable.config.WithRetryConfig;
+import software.amazon.lambda.durable.examples.types.ApprovalRequest;
+import software.amazon.lambda.durable.retry.RetryDecision;
+
+/**
+ * Example demonstrating {@code context.withRetry} with {@code context.waitForCallback}.
+ *
+ * <p>Submits an approval request to an external system via a callback. If the callback fails (e.g., the external system
+ * rejects the request), the helper retries the entire waitForCallback cycle — creating a fresh callback with a new ID
+ * each time.
+ *
+ * <p>Each attempt uses a unique callback name ({@code "approval-1"}, {@code "approval-2"}, etc.) so the execution
+ * history stays clean and replay-safe. A {@code null} name is used, so attempts are grouped under a default-named child
+ * context.
+ */
+public class RetryWaitForCallbackExample extends DurableHandler<ApprovalRequest, String> {
+
+    private static final int MAX_ATTEMPTS = 3;
+
+    @Override
+    public String handleRequest(ApprovalRequest input, DurableContext context) {
+        // Step 1: Prepare the approval request
+        var prepared = context.step(
+                "prepare",
+                String.class,
+                stepCtx -> "Approval for: " + input.description() + " ($" + input.amount() + ")");
+
+        // Step 2: waitForCallback with retry — if the external system fails, try again with a fresh callback
+        var approvalResult = context.withRetry(
+                null,
+                (attempt, ctx) -> ctx.waitForCallback(
+                        "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
+                                .info("Attempt {}: sending callback {} to approval system", attempt, callbackId)),
+                WithRetryConfig.builder()
+                        .retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS
+                                ? RetryDecision.retry(Duration.ofSeconds(2))
+                                : RetryDecision.fail())
+                        .build());
+
+        // Step 3: Process the result
+        return context.step("process-result", String.class, stepCtx -> prepared + " - Result: " + approvalResult);
+    }
+}

examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java

@@ -0,0 +1,40 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples.invoke;
+
+import java.time.Duration;
+import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.DurableHandler;
+import software.amazon.lambda.durable.config.WithRetryConfig;
+import software.amazon.lambda.durable.examples.types.GreetingRequest;
+import software.amazon.lambda.durable.retry.RetryDecision;
+
+/**
+ * Example demonstrating {@code context.withRetry} with {@code context.invoke}.
+ *
+ * <p>Retries a chained Lambda invocation up to 3 times with a fixed 2-second backoff between attempts. Each attempt
+ * uses a unique operation name ({@code "call-greeting-1"}, {@code "call-greeting-2"}, etc.) so the execution history
+ * stays clean and replay-safe.
+ *
+ * <p>A {@code null} name is used, so attempts are grouped under a default-named child context.
+ */
+public class RetryInvokeExample extends DurableHandler<GreetingRequest, String> {
+
+    private static final int MAX_ATTEMPTS = 3;
+
+    @Override
+    public String handleRequest(GreetingRequest input, DurableContext context) {
+        return context.withRetry(
+                null,
+                (attempt, ctx) -> ctx.invoke(
+                        "call-greeting-" + attempt,
+                        "simple-step-example" + input.getName() + ":$LATEST",
+                        input,
+                        String.class),
+                WithRetryConfig.builder()
+                        .retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS
+                                ? RetryDecision.retry(Duration.ofSeconds(2))
+                                : RetryDecision.fail())
+                        .build());
+    }
+}

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

@@ -710,6 +710,89 @@ void testComplexFlatMapExample() {
         assertTrue(output.contains("reason=MIN_SUCCESSFUL_REACHED"));
     }
 
+    @Test
+    void testRetryInvokeExample() {
+        var runner = CloudDurableTestRunner.create(
+                arn("retry-invoke-example"), GreetingRequest.class, String.class, lambdaClient);
+        // The handler invokes "simple-step-example" + input.getName() + ":$LATEST",
+        // so passing the functionNameSuffix as the name targets the deployed simple-step-example function
+        var result = runner.run(new GreetingRequest(functionNameSuffix));
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertNotNull(result.getResult());
+    }
+
+    @Test
+    void testRetryWaitForCallbackExampleSucceeds() {
+        var runner = CloudDurableTestRunner.create(
+                arn("retry-wait-for-callback-example"), ApprovalRequest.class, String.class, lambdaClient);
+
+        var execution = runner.startAsync(new ApprovalRequest("Server upgrade", 5000.0));
+
+        // Wait for the first callback from attempt 1
+        execution.pollUntil(exec -> exec.hasCallback("approval-1-callback"));
+        var callbackId = execution.getCallbackId("approval-1-callback");
+        assertNotNull(callbackId);
+
+        // Complete the callback on the first attempt
+        execution.completeCallback(callbackId, "\"approved\"");
+
+        var result = execution.pollUntilComplete();
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+
+        var finalResult = result.getResult();
+        assertNotNull(finalResult);
+        assertTrue(finalResult.contains("Approval for: Server upgrade"));
+        assertTrue(finalResult.contains("5000"));
+        assertTrue(finalResult.contains("approved"));
+
+        // Verify operations
+        assertNotNull(execution.getOperation("prepare"));
+        assertNotNull(execution.getOperation("process-result"));
+    }
+
+    @Test
+    void testRetryWaitForCallbackExampleRetriesAfterFailure() {
+        var runner = CloudDurableTestRunner.create(
+                arn("retry-wait-for-callback-example"), ApprovalRequest.class, String.class, lambdaClient);
+
+        var execution = runner.startAsync(new ApprovalRequest("Expensive item", 10000.0));
+
+        // Wait for the first callback from attempt 1
+        execution.pollUntil(exec -> exec.hasCallback("approval-1-callback"));
+        var callbackId1 = execution.getCallbackId("approval-1-callback");
+        assertNotNull(callbackId1);
+
+        // Fail the first attempt
+        execution.failCallback(
+                callbackId1,
+                ErrorObject.builder()
+                        .errorType("Rejected")
+                        .errorMessage("denied by reviewer")
+                        .build());
+
+        // Wait for the second callback from attempt 2 (after backoff)
+        execution.pollUntil(exec -> exec.hasCallback("approval-2-callback"));
+        var callbackId2 = execution.getCallbackId("approval-2-callback");
+        assertNotNull(callbackId2);
+
+        // Complete the second attempt
+        execution.completeCallback(callbackId2, "\"approved on retry\"");
+
+        var result = execution.pollUntilComplete();
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+
+        var finalResult = result.getResult();
+        assertNotNull(finalResult);
+        assertTrue(finalResult.contains("Approval for: Expensive item"));
+        assertTrue(finalResult.contains("10000"));
+        assertTrue(finalResult.contains("approved on retry"));
+
+        // Verify operations
+        assertNotNull(execution.getOperation("prepare"));
+        assertNotNull(execution.getOperation("process-result"));
+    }
+
     @Test
     void testWaitForConditionExample() {
         var runner = CloudDurableTestRunner.create(

examples/src/test/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExampleTest.java

@@ -0,0 +1,148 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples.callback;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.services.lambda.model.ErrorObject;
+import software.amazon.lambda.durable.examples.types.ApprovalRequest;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+
+class RetryWaitForCallbackExampleTest {
+
+    @Test
+    void succeedsOnFirstAttempt() {
+        var handler = new RetryWaitForCallbackExample();
+        var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler);
+        var input = new ApprovalRequest("New laptop", 1500.00);
+
+        // First run — prepares request, starts waitForCallback, suspends
+        var result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete the callback (waitForCallback names it "approval-1-callback" internally)
+        var callbackId = runner.getCallbackId("approval-1-callback");
+        assertNotNull(callbackId, "Callback 'approval-1-callback' should have been created");
+        runner.completeCallback(callbackId, "\"Approved by manager\"");
+
+        // Run to completion
+        result = runner.runUntilComplete(input);
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals(
+                "Approval for: New laptop ($1500.0) - Result: Approved by manager", result.getResult(String.class));
+    }
+
+    @Test
+    void retriesAfterFirstCallbackFails() {
+        var handler = new RetryWaitForCallbackExample();
+        var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler);
+        var input = new ApprovalRequest("Server upgrade", 5000.00);
+
+        // First run — prepares, starts waitForCallback attempt 1, suspends
+        var result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail the first callback
+        var callbackId1 = runner.getCallbackId("approval-1-callback");
+        assertNotNull(callbackId1);
+        runner.failCallback(
+                callbackId1,
+                ErrorObject.builder()
+                        .errorType("RejectedError")
+                        .errorMessage("Rejected by first reviewer")
+                        .build());
+
+        // Run — processes failure, hits backoff wait, suspends
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past the backoff wait
+        runner.advanceTime();
+
+        // Run — starts waitForCallback attempt 2, suspends
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete the second callback
+        var callbackId2 = runner.getCallbackId("approval-2-callback");
+        assertNotNull(callbackId2, "Callback 'approval-2-callback' should have been created after retry");
+        runner.completeCallback(callbackId2, "\"Approved on second try\"");
+
+        // Run to completion
+        result = runner.runUntilComplete(input);
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals(
+                "Approval for: Server upgrade ($5000.0) - Result: Approved on second try",
+                result.getResult(String.class));
+    }
+
+    @Test
+    void failsAfterAllRetriesExhausted() {
+        var handler = new RetryWaitForCallbackExample();
+        var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler);
+        var input = new ApprovalRequest("Expensive item", 10000.00);
+
+        // First run — starts waitForCallback attempt 1
+        var result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail callback attempt 1
+        var callbackId1 = runner.getCallbackId("approval-1-callback");
+        runner.failCallback(
+                callbackId1,
+                ErrorObject.builder()
+                        .errorType("Rejected")
+                        .errorMessage("fail 1")
+                        .build());
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff 1, run to start attempt 2
+        runner.advanceTime();
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail callback attempt 2
+        var callbackId2 = runner.getCallbackId("approval-2-callback");
+        runner.failCallback(
+                callbackId2,
+                ErrorObject.builder()
+                        .errorType("Rejected")
+                        .errorMessage("fail 2")
+                        .build());
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff 2, run to start attempt 3
+        runner.advanceTime();
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail callback attempt 3 — last attempt, retryStrategy returns fail()
+        var callbackId3 = runner.getCallbackId("approval-3-callback");
+        runner.failCallback(
+                callbackId3,
+                ErrorObject.builder()
+                        .errorType("Rejected")
+                        .errorMessage("fail 3")
+                        .build());
+        result = runner.run(input);
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+    }
+
+    @Test
+    void suspendsOnFirstRun() {
+        var handler = new RetryWaitForCallbackExample();
+        var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler);
+        var input = new ApprovalRequest("Test item", 100.00);
+
+        var result = runner.run(input);
+
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+    }
+}