From 2fe64a7b1cf8cee26c064725327eae646346ee4c Mon Sep 17 00:00:00 2001 From: hsilan Date: Wed, 22 Apr 2026 15:04:25 -0700 Subject: [PATCH 01/19] feat: first attempt at retryableOperation --- .../durable/config/RetryOperationConfig.java | 103 +++++ .../durable/util/RetryOperationHelper.java | 135 ++++++ .../durable/util/RetryableOperation.java | 26 ++ .../config/RetryOperationConfigTest.java | 80 ++++ .../util/RetryOperationHelperTest.java | 420 ++++++++++++++++++ .../durable/util/RetryableOperationTest.java | 52 +++ 6 files changed, 816 insertions(+) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java new file mode 100644 index 000000000..696d1d253 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java @@ -0,0 +1,103 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.config; + +import software.amazon.lambda.durable.retry.RetryStrategy; + +/** + * Configuration for {@link software.amazon.lambda.durable.util.RetryOperationHelper#retryOperation}. + * + *

Uses the same {@link RetryStrategy} shape that developers already know from {@link StepConfig}, so there are zero + * new retry concepts to learn. + */ +public class RetryOperationConfig { + private final RetryStrategy retryStrategy; + private final boolean wrapInChildContext; + + private RetryOperationConfig(Builder builder) { + this.retryStrategy = builder.retryStrategy; + this.wrapInChildContext = builder.wrapInChildContext; + } + + /** + * Returns the retry strategy. Same type as {@link StepConfig#retryStrategy()}. + * + * @return the retry strategy, never null + */ + public RetryStrategy retryStrategy() { + return retryStrategy; + } + + /** + * Whether to wrap the retry loop in {@code runInChildContext} so all attempts are grouped under a single named + * operation in execution history. Only applies when a name is provided to the named form of {@code retryOperation}. + * Defaults to {@code true}. + * + * @return true if child-context wrapping is enabled + */ + public boolean wrapInChildContext() { + return wrapInChildContext; + } + + /** + * Creates a new builder for {@code RetryOperationConfig}. + * + * @return a new builder instance + */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for creating {@link RetryOperationConfig} instances. */ + public static class Builder { + private RetryStrategy retryStrategy; + private boolean wrapInChildContext = true; + + private Builder() {} + + /** + * Sets the retry strategy. Required. + * + *

Reuses the exact same {@link RetryStrategy} interface from {@link StepConfig}. All existing factory + * methods ({@link software.amazon.lambda.durable.retry.RetryStrategies#exponentialBackoff}, + * {@link software.amazon.lambda.durable.retry.RetryStrategies#fixedDelay}, presets, and custom lambdas) work + * without modification. + * + * @param retryStrategy the retry strategy to use + * @return this builder for method chaining + */ + public Builder retryStrategy(RetryStrategy retryStrategy) { + this.retryStrategy = retryStrategy; + return this; + } + + /** + * Controls whether the retry loop is wrapped in a child context. Only meaningful for the named form of + * {@code retryOperation}. Defaults to {@code true}. + * + *

When {@code true}, all attempts and backoff waits are grouped under a single named operation in execution + * history, providing a cleaner view and isolated operation ID space. Set to {@code false} to flatten attempts + * into the parent context. + * + * @param wrapInChildContext whether to wrap in a child context + * @return this builder for method chaining + */ + public Builder wrapInChildContext(boolean wrapInChildContext) { + this.wrapInChildContext = wrapInChildContext; + return this; + } + + /** + * Builds the {@link RetryOperationConfig} instance. + * + * @return a new config with the configured options + * @throws IllegalArgumentException if retryStrategy is not set + */ + public RetryOperationConfig build() { + if (retryStrategy == null) { + throw new IllegalArgumentException("retryStrategy is required"); + } + return new RetryOperationConfig(this); + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java new file mode 100644 index 000000000..c63f0aad0 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java @@ -0,0 +1,135 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.util; + +import java.time.Duration; +import java.util.Objects; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.RetryOperationConfig; +import software.amazon.lambda.durable.retry.RetryDecision; + +/** + * Replay-safe retry loop for any durable operation. + * + *

Provides the same retry-with-backoff pattern that {@code context.step()} has built in, but for operations that + * cannot live inside a step ({@code waitForCallback}, {@code invoke}, {@code waitForCondition}, etc.). + * + *

Every side-effect in the loop is a durable operation, so the loop is replay-safe by construction. On replay, + * completed operations return cached results instantly and the loop fast-forwards to the current attempt. + * + *

Usage — callback retry

+ * + * 
{@code
+ * var result = RetryOperationHelper.retryOperation(
+ *     context,
+ *     "approval",
+ *     (ctx, attempt) -> ctx.waitForCallback(
+ *         "approval-" + attempt,
+ *         String.class,
+ *         (callbackId, stepCtx) -> sendApprovalEmail(approverEmail, callbackId)
+ *     ),
+ *     RetryOperationConfig.builder()
+ *         .retryStrategy(RetryStrategies.exponentialBackoff(
+ *             3, Duration.ofSeconds(2), Duration.ofSeconds(30), 2.0, JitterStrategy.FULL))
+ *         .build()
+ * );
+ * }
+ * + *

Usage — invoke retry (anonymous form)

+ * + * 
{@code
+ * var result = RetryOperationHelper.retryOperation(
+ *     context,
+ *     (ctx, attempt) -> ctx.invoke(
+ *         "charge-" + attempt, paymentFnArn, new ChargeRequest(orderId), String.class),
+ *     RetryOperationConfig.builder()
+ *         .retryStrategy((err, att) -> att < 3
+ *             ? RetryDecision.retry(Duration.ofSeconds(1))
+ *             : RetryDecision.fail())
+ *         .build()
+ * );
+ * }
+ */ +public final class RetryOperationHelper { + + private static final Duration DEFAULT_BACKOFF_DELAY = Duration.ofSeconds(1); + private static final String BACKOFF_SUFFIX = "-backoff-"; + private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-"; + + private RetryOperationHelper() { + // utility class + } + + /** + * Named form — wraps the retry loop in {@code runInChildContext} by default so all attempts are grouped under a + * single named operation in execution history. + * + *

The child-context wrapping can be disabled via + * {@link RetryOperationConfig.Builder#wrapInChildContext(boolean)}. + * + * @param the result type + * @param context the durable context + * @param name operation name (used for child context and backoff wait names) + * @param operation the retryable operation — receives the context and 1-based attempt number + * @param config retry configuration including the retry strategy + * @return the operation result + */ + @SuppressWarnings("unchecked") + public static T retryOperation( + DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(name, "name cannot be null"); + Objects.requireNonNull(operation, "operation cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + + if (config.wrapInChildContext()) { + return (T) context.runInChildContext( + name, new TypeToken() {}, childCtx -> executeRetryLoop(childCtx, name, operation, config)); + } + return executeRetryLoop(context, name, operation, config); + } + + /** + * Anonymous form — runs the retry loop directly in the caller's context. No child-context wrapping is applied + * regardless of the {@code wrapInChildContext} config setting. + * + * @param the result type + * @param context the durable context + * @param operation the retryable operation — receives the context and 1-based attempt number + * @param config retry configuration including the retry strategy + * @return the operation result + */ + public static T retryOperation( + DurableContext context, RetryableOperation operation, RetryOperationConfig config) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(operation, "operation cannot be null"); + Objects.requireNonNull(config, "config cannot be null"); + + return executeRetryLoop(context, null, operation, config); + } + + /** + * Core retry loop. Replay-safe because every side-effect is a durable operation: the user's operation calls durable + * primitives, and backoff uses {@code context.wait()}. + */ + private static T executeRetryLoop( + DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) { + var attempt = 1; + while (true) { + try { + return operation.execute(context, attempt); + } catch (Exception e) { + RetryDecision decision = config.retryStrategy().makeRetryDecision(e, attempt); + if (!decision.shouldRetry()) { + throw e; + } + + var delay = decision.delay().isZero() ? DEFAULT_BACKOFF_DELAY : decision.delay(); + var waitName = name != null ? name + BACKOFF_SUFFIX + attempt : ANONYMOUS_BACKOFF_PREFIX + attempt; + context.wait(waitName, delay); + attempt++; + } + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java new file mode 100644 index 000000000..d6b2496c1 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java @@ -0,0 +1,26 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.util; + +import software.amazon.lambda.durable.DurableContext; + +/** + * A durable operation that can be retried end-to-end by {@link RetryOperationHelper}. + * + *

Receives the durable context and the 1-based attempt number so callers can generate unique operation names per + * attempt (e.g., {@code "approval-" + attempt}). + * + * @param the result type + */ +@FunctionalInterface +public interface RetryableOperation { + + /** + * Executes the durable operation. + * + * @param context the durable context to use for durable operations + * @param attempt the current attempt number (1-based: first attempt is 1) + * @return the operation result + */ + T execute(DurableContext context, int attempt); +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java new file mode 100644 index 000000000..726da5b61 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java @@ -0,0 +1,80 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.config; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.retry.RetryDecision; +import software.amazon.lambda.durable.retry.RetryStrategies; + +class RetryOperationConfigTest { + + @Test + void builderWithRetryStrategy() { + var strategy = RetryStrategies.Presets.DEFAULT; + + var config = RetryOperationConfig.builder().retryStrategy(strategy).build(); + + assertEquals(strategy, config.retryStrategy()); + } + + @Test + void builderWithoutRetryStrategy_shouldThrow() { + var exception = assertThrows(IllegalArgumentException.class, () -> RetryOperationConfig.builder() + .build()); + + assertEquals("retryStrategy is required", exception.getMessage()); + } + + @Test + void wrapInChildContext_defaultsToTrue() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build(); + + assertTrue(config.wrapInChildContext()); + } + + @Test + void wrapInChildContext_canBeSetToFalse() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .wrapInChildContext(false) + .build(); + + assertFalse(config.wrapInChildContext()); + } + + @Test + void wrapInChildContext_canBeSetToTrue() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .wrapInChildContext(true) + .build(); + + assertTrue(config.wrapInChildContext()); + } + + @Test + void builderChaining() { + var strategy = RetryStrategies.Presets.DEFAULT; + + var config = RetryOperationConfig.builder() + .retryStrategy(strategy) + .wrapInChildContext(false) + .build(); + + assertEquals(strategy, config.retryStrategy()); + assertFalse(config.wrapInChildContext()); + } + + @Test + void builderWithCustomLambdaRetryStrategy() { + var config = RetryOperationConfig.builder() + .retryStrategy((error, attempt) -> RetryDecision.fail()) + .build(); + + assertNotNull(config.retryStrategy()); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java new file mode 100644 index 000000000..9558f3c83 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java @@ -0,0 +1,420 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.util; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.*; +import static org.mockito.Mockito.*; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.function.Function; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.RetryOperationConfig; +import software.amazon.lambda.durable.retry.RetryDecision; +import software.amazon.lambda.durable.retry.RetryStrategies; + +class RetryOperationHelperTest { + + private DurableContext context; + + @BeforeEach + void setUp() { + context = mock(DurableContext.class); + } + + // --- Named form tests --- + + @Nested + class NamedForm { + + @Test + void successOnFirstAttempt_wrapsInChildContext() { + // runInChildContext should be called; delegate to the function immediately + when(context.runInChildContext(eq("my-op"), any(TypeToken.class), any())) + .thenAnswer(invocation -> { + Function func = invocation.getArgument(2); + return func.apply(context); + }); + + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build(); + + var result = RetryOperationHelper.retryOperation(context, "my-op", (ctx, attempt) -> "success", config); + + assertEquals("success", result); + verify(context).runInChildContext(eq("my-op"), any(TypeToken.class), any()); + } + + @Test + void successOnFirstAttempt_noChildContext_whenDisabled() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .wrapInChildContext(false) + .build(); + + var result = RetryOperationHelper.retryOperation(context, "my-op", (ctx, attempt) -> "direct", config); + + assertEquals("direct", result); + verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any()); + } + + @Test + void retriesWithBackoffWaits_namedForm() { + // Disable child context wrapping so we can directly verify wait calls + var callCount = new int[] {0}; + var config = RetryOperationConfig.builder() + .retryStrategy((error, attempt) -> + attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail()) + .wrapInChildContext(false) + .build(); + + var result = RetryOperationHelper.retryOperation( + context, + "my-op", + (ctx, attempt) -> { + callCount[0]++; + if (attempt < 3) { + throw new RuntimeException("fail-" + attempt); + } + return "success-on-3"; + }, + config); + + assertEquals("success-on-3", result); + assertEquals(3, callCount[0]); + verify(context).wait("my-op-backoff-1", Duration.ofSeconds(5)); + verify(context).wait("my-op-backoff-2", Duration.ofSeconds(5)); + verify(context, times(2)).wait(anyString(), any(Duration.class)); + } + + @Test + void rethrowsWhenRetryStrategyReturnsFail() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .wrapInChildContext(false) + .build(); + + var exception = assertThrows( + RuntimeException.class, + () -> RetryOperationHelper.retryOperation( + context, + "my-op", + (ctx, attempt) -> { + throw new RuntimeException("terminal"); + }, + config)); + + assertEquals("terminal", exception.getMessage()); + verify(context, never()).wait(anyString(), any(Duration.class)); + } + + @Test + void usesDefaultDelayWhenRetryDecisionDelayIsZero() { + var config = RetryOperationConfig.builder() + .retryStrategy( + (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail()) + .wrapInChildContext(false) + .build(); + + var callCount = new int[] {0}; + var result = RetryOperationHelper.retryOperation( + context, + "my-op", + (ctx, attempt) -> { + callCount[0]++; + if (attempt == 1) { + throw new RuntimeException("fail"); + } + return "ok"; + }, + config); + + assertEquals("ok", result); + // Zero delay should be replaced with 1-second default + verify(context).wait("my-op-backoff-1", Duration.ofSeconds(1)); + } + + @Test + void nullContext_shouldThrow() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build(); + + assertThrows( + NullPointerException.class, + () -> RetryOperationHelper.retryOperation(null, "name", (ctx, a) -> "x", config)); + } + + @Test + void nullName_shouldThrow() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build(); + + assertThrows( + NullPointerException.class, + () -> RetryOperationHelper.retryOperation(context, null, (ctx, a) -> "x", config)); + } + + @Test + void nullOperation_shouldThrow() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build(); + + assertThrows( + NullPointerException.class, + () -> RetryOperationHelper.retryOperation(context, "name", null, config)); + } + + @Test + void nullConfig_shouldThrow() { + assertThrows( + NullPointerException.class, + () -> RetryOperationHelper.retryOperation(context, "name", (ctx, a) -> "x", null)); + } + } + + // --- Anonymous form tests --- + + @Nested + class AnonymousForm { + + @Test + void successOnFirstAttempt() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build(); + + var result = RetryOperationHelper.retryOperation(context, (ctx, attempt) -> "anonymous-success", config); + + assertEquals("anonymous-success", result); + verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any()); + } + + @Test + void retriesWithAnonymousBackoffNames() { + var config = RetryOperationConfig.builder() + .retryStrategy((error, attempt) -> + attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail()) + .build(); + + var result = RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> { + if (attempt < 3) { + throw new RuntimeException("fail"); + } + return "done"; + }, + config); + + assertEquals("done", result); + verify(context).wait("retry-backoff-1", Duration.ofSeconds(2)); + verify(context).wait("retry-backoff-2", Duration.ofSeconds(2)); + } + + @Test + void neverWrapsInChildContext_evenWhenConfigSaysTrue() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .wrapInChildContext(true) // should be ignored for anonymous form + .build(); + + RetryOperationHelper.retryOperation(context, (ctx, attempt) -> "result", config); + + verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any()); + } + + @Test + void rethrowsOriginalException() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build(); + + var original = new IllegalStateException("original error"); + var thrown = assertThrows( + IllegalStateException.class, + () -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> { + throw original; + }, + config)); + + assertSame(original, thrown); + } + + @Test + void nullContext_shouldThrow() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build(); + + assertThrows( + NullPointerException.class, + () -> RetryOperationHelper.retryOperation(null, (ctx, a) -> "x", config)); + } + + @Test + void nullOperation_shouldThrow() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build(); + + assertThrows( + NullPointerException.class, + () -> RetryOperationHelper.retryOperation(context, (RetryableOperation) null, config)); + } + + @Test + void nullConfig_shouldThrow() { + assertThrows( + NullPointerException.class, + () -> RetryOperationHelper.retryOperation(context, (ctx, a) -> "x", null)); + } + } + + // --- Retry behavior tests --- + + @Nested + class RetryBehavior { + + @Test + void passesCorrectAttemptNumberToOperation() { + var attempts = new ArrayList(); + var config = RetryOperationConfig.builder() + .retryStrategy((error, attempt) -> + attempt < 4 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail()) + .wrapInChildContext(false) + .build(); + + RetryOperationHelper.retryOperation( + context, + "track", + (ctx, attempt) -> { + attempts.add(attempt); + if (attempt < 4) { + throw new RuntimeException("not yet"); + } + return "done"; + }, + config); + + assertEquals(4, attempts.size()); + assertEquals(1, attempts.get(0)); + assertEquals(2, attempts.get(1)); + assertEquals(3, attempts.get(2)); + assertEquals(4, attempts.get(3)); + } + + @Test + void passesErrorToRetryStrategy() { + var errors = new ArrayList(); + var config = RetryOperationConfig.builder() + .retryStrategy((error, attempt) -> { + errors.add(error); + return attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail(); + }) + .build(); + + assertThrows( + RuntimeException.class, + () -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> { + throw new RuntimeException("error-" + attempt); + }, + config)); + + assertEquals(3, errors.size()); + assertEquals("error-1", errors.get(0).getMessage()); + assertEquals("error-2", errors.get(1).getMessage()); + assertEquals("error-3", errors.get(2).getMessage()); + } + + @Test + void respectsCustomDelayFromRetryDecision() { + var config = RetryOperationConfig.builder() + .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(attempt * 10L))) + .build(); + + RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> { + if (attempt <= 2) { + throw new RuntimeException("fail"); + } + return "ok"; + }, + config); + + verify(context).wait("retry-backoff-1", Duration.ofSeconds(10)); + verify(context).wait("retry-backoff-2", Duration.ofSeconds(20)); + } + + @Test + void passesContextToOperation() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build(); + + RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> { + assertSame(context, ctx); + return "verified"; + }, + config); + } + + @Test + void namedFormPassesChildContextToOperation_whenWrapped() { + var childContext = mock(DurableContext.class); + when(context.runInChildContext(eq("wrapped"), any(TypeToken.class), any())) + .thenAnswer(invocation -> { + Function func = invocation.getArgument(2); + return func.apply(childContext); + }); + + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .wrapInChildContext(true) + .build(); + + RetryOperationHelper.retryOperation( + context, + "wrapped", + (ctx, attempt) -> { + assertSame(childContext, ctx); + return "verified"; + }, + config); + } + + @Test + void rethrowsLastExceptionWhenAllRetriesExhausted() { + var config = RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1))) + .build(); + + var thrown = assertThrows( + RuntimeException.class, + () -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> { + throw new RuntimeException("attempt-" + attempt); + }, + config)); + + // The last attempt's exception is rethrown + assertEquals("attempt-3", thrown.getMessage()); + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java new file mode 100644 index 000000000..f5893a37c --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java @@ -0,0 +1,52 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.util; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.DurableContext; + +class RetryableOperationTest { + + @Test + void canBeImplementedAsLambda() { + RetryableOperation operation = (ctx, attempt) -> "result-" + attempt; + var context = mock(DurableContext.class); + + assertEquals("result-1", operation.execute(context, 1)); + assertEquals("result-3", operation.execute(context, 3)); + } + + @Test + void receivesContextAndAttempt() { + var context = mock(DurableContext.class); + RetryableOperation operation = (ctx, attempt) -> { + assertSame(context, ctx); + return "attempt-" + attempt; + }; + + assertEquals("attempt-1", operation.execute(context, 1)); + assertEquals("attempt-2", operation.execute(context, 2)); + } + + @Test + void canThrowExceptions() { + RetryableOperation operation = (ctx, attempt) -> { + throw new RuntimeException("failed on attempt " + attempt); + }; + var context = mock(DurableContext.class); + + var exception = assertThrows(RuntimeException.class, () -> operation.execute(context, 1)); + assertEquals("failed on attempt 1", exception.getMessage()); + } + + @Test + void canReturnNull() { + RetryableOperation operation = (ctx, attempt) -> null; + var context = mock(DurableContext.class); + + assertNull(operation.execute(context, 1)); + } +} From 7fd1628be2a9b8bd10ad2a55b6a737b5c3c1137d Mon Sep 17 00:00:00 2001 From: hsilan Date: Wed, 22 Apr 2026 15:48:47 -0700 Subject: [PATCH 02/19] feat: add example tests and integration tests for new RetryableOperation util --- .../callback/RetryWaitForCallbackExample.java | 50 ++++ .../examples/invoke/RetryInvokeExample.java | 41 +++ .../RetryWaitForCallbackExampleTest.java | 148 ++++++++++ .../invoke/RetryInvokeExampleTest.java | 130 +++++++++ .../durable/RetryInvokeIntegrationTest.java | 198 ++++++++++++++ .../RetryWaitForCallbackIntegrationTest.java | 253 ++++++++++++++++++ .../durable/util/RetryOperationHelper.java | 8 + .../util/RetryOperationHelperTest.java | 42 +++ 8 files changed, 870 insertions(+) create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExampleTest.java create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExampleTest.java create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java new file mode 100644 index 000000000..e16ca8f43 --- /dev/null +++ b/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.RetryOperationConfig; +import software.amazon.lambda.durable.examples.types.ApprovalRequest; +import software.amazon.lambda.durable.retry.RetryDecision; +import software.amazon.lambda.durable.util.RetryOperationHelper; + +/** + * Example demonstrating {@link RetryOperationHelper} with {@code context.waitForCallback}. + * + *

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. + * + *

Each attempt uses a unique callback name ({@code "approval-1"}, {@code "approval-2"}, etc.) so the execution + * history stays clean and replay-safe. The anonymous form is used, so attempts run directly in the caller's context. + */ +public class RetryWaitForCallbackExample extends DurableHandler { + + 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 = RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> ctx.waitForCallback( + "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger() + .info("Attempt {}: sending callback {} to approval system", attempt, callbackId)), + RetryOperationConfig.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); + } +} diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java new file mode 100644 index 000000000..774c7c681 --- /dev/null +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java @@ -0,0 +1,41 @@ +// 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.RetryOperationConfig; +import software.amazon.lambda.durable.examples.types.GreetingRequest; +import software.amazon.lambda.durable.retry.RetryDecision; +import software.amazon.lambda.durable.util.RetryOperationHelper; + +/** + * Example demonstrating {@link RetryOperationHelper} with {@code context.invoke}. + * + *

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. + * + *

The anonymous form is used, so attempts run directly in the caller's context without child-context wrapping. + */ +public class RetryInvokeExample extends DurableHandler { + + private static final int MAX_ATTEMPTS = 3; + + @Override + public String handleRequest(GreetingRequest input, DurableContext context) { + return RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> ctx.invoke( + "call-greeting-" + attempt, + "simple-step-example" + input.getName() + ":$LATEST", + input, + String.class), + RetryOperationConfig.builder() + .retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS + ? RetryDecision.retry(Duration.ofSeconds(2)) + : RetryDecision.fail()) + .build()); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExampleTest.java new file mode 100644 index 000000000..ea5e8a4a8 --- /dev/null +++ b/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()); + } +} diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExampleTest.java new file mode 100644 index 000000000..7d8e66104 --- /dev/null +++ b/examples/src/test/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExampleTest.java @@ -0,0 +1,130 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.examples.invoke; + +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.GreetingRequest; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class RetryInvokeExampleTest { + + @Test + void succeedsOnFirstAttempt() { + var handler = new RetryInvokeExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + var input = new GreetingRequest("world"); + + // First run — starts the invoke, suspends waiting for result + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Complete the first invoke attempt + runner.completeChainedInvoke("call-greeting-1", "\"hello world\""); + result = runner.run(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("hello world", result.getResult(String.class)); + } + + @Test + void retriesAfterFirstAttemptFails() { + var handler = new RetryInvokeExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + var input = new GreetingRequest("world"); + + // First run — starts invoke attempt 1 + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail the first invoke attempt + runner.failChainedInvoke( + "call-greeting-1", + ErrorObject.builder() + .errorType("TransientError") + .errorMessage("Service unavailable") + .build()); + + // Second run — processes the failure, does backoff wait, starts invoke attempt 2 + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past the backoff wait + runner.advanceTime(); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Complete the second invoke attempt + runner.completeChainedInvoke("call-greeting-2", "\"hello on retry\""); + result = runner.run(input); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("hello on retry", result.getResult(String.class)); + } + + @Test + void failsAfterAllRetriesExhausted() { + var handler = new RetryInvokeExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + var input = new GreetingRequest("world"); + + // First run — starts invoke attempt 1 + var result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail attempt 1 + runner.failChainedInvoke( + "call-greeting-1", + ErrorObject.builder() + .errorType("TransientError") + .errorMessage("fail 1") + .build()); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past backoff wait 1 + runner.advanceTime(); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail attempt 2 + runner.failChainedInvoke( + "call-greeting-2", + ErrorObject.builder() + .errorType("TransientError") + .errorMessage("fail 2") + .build()); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past backoff wait 2 + runner.advanceTime(); + result = runner.run(input); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail attempt 3 — this is the last attempt, retryStrategy returns fail() + runner.failChainedInvoke( + "call-greeting-3", + ErrorObject.builder() + .errorType("TransientError") + .errorMessage("fail 3") + .build()); + result = runner.run(input); + + assertEquals(ExecutionStatus.FAILED, result.getStatus()); + } + + @Test + void suspendsOnFirstRun() { + var handler = new RetryInvokeExample(); + var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler); + var input = new GreetingRequest("test"); + + var result = runner.run(input); + + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + } +} diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java new file mode 100644 index 000000000..1d7185b95 --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java @@ -0,0 +1,198 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.*; + +import java.time.Duration; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.lambda.durable.config.RetryOperationConfig; +import software.amazon.lambda.durable.exception.InvokeFailedException; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.retry.RetryDecision; +import software.amazon.lambda.durable.retry.RetryStrategies; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; +import software.amazon.lambda.durable.util.RetryOperationHelper; + +class RetryInvokeIntegrationTest { + + @Test + void invokeSucceedsOnFirstAttempt() { + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class), + RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2))) + .build())); + + var result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + runner.completeChainedInvoke("invoke-1", "\"success\""); + result = runner.run("test"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("success", result.getResult(String.class)); + } + + @Test + void invokeRetriesAfterFailure() { + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class), + RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2))) + .build())); + + // First run — invoke attempt 1 starts + var result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail attempt 1 + runner.failChainedInvoke( + "invoke-1", + ErrorObject.builder() + .errorType("TransientError") + .errorMessage("service unavailable") + .build()); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past backoff wait + runner.advanceTime(); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Complete attempt 2 + runner.completeChainedInvoke("invoke-2", "\"recovered\""); + result = runner.run("test"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("recovered", result.getResult(String.class)); + } + + @Test + void invokeFailsAfterAllRetriesExhausted() { + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class), + RetryOperationConfig.builder() + .retryStrategy((error, attempt) -> + attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail()) + .build())); + + // Attempt 1 + var result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + runner.failChainedInvoke( + "invoke-1", ErrorObject.builder().errorMessage("fail 1").build()); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past backoff + runner.advanceTime(); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Attempt 2 — last attempt + runner.failChainedInvoke( + "invoke-2", ErrorObject.builder().errorMessage("fail 2").build()); + result = runner.run("test"); + + assertEquals(ExecutionStatus.FAILED, result.getStatus()); + } + + @Test + void invokeRetryWithCustomBackoffDelay() { + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class), + RetryOperationConfig.builder() + .retryStrategy((error, attempt) -> attempt < 3 + ? RetryDecision.retry(Duration.ofSeconds(attempt * 5L)) + : RetryDecision.fail()) + .build())); + + // Attempt 1 fails + var result = runner.run("test"); + runner.failChainedInvoke( + "invoke-1", ErrorObject.builder().errorMessage("fail").build()); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past first backoff (5s) + runner.advanceTime(); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Attempt 2 succeeds + runner.completeChainedInvoke("invoke-2", "\"ok\""); + result = runner.run("test"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("ok", result.getResult(String.class)); + } + + @Test + void invokeRetryWithStepsBeforeAndAfter() { + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var prefix = context.step("prepare", String.class, stepCtx -> "prepared"); + + var invokeResult = RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class), + RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1))) + .build()); + + return context.step("finalize", String.class, stepCtx -> prefix + " -> " + invokeResult + " -> done"); + }); + + // First run — prepare step completes, invoke starts + var result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Complete invoke + runner.completeChainedInvoke("invoke-1", "\"invoked\""); + result = runner.runUntilComplete("test"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("prepared -> invoked -> done", result.getResult(String.class)); + } + + @Test + void invokeRetryPreservesOriginalExceptionType() { + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + try { + return RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class), + RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build()); + } catch (InvokeFailedException e) { + assertEquals("invoke failed", e.getMessage()); + throw e; + } + }); + + var result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + runner.failChainedInvoke( + "invoke-1", ErrorObject.builder().errorMessage("invoke failed").build()); + result = runner.run("test"); + + assertEquals(ExecutionStatus.FAILED, result.getStatus()); + } +} diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java new file mode 100644 index 000000000..f7806d577 --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java @@ -0,0 +1,253 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable; + +import static org.junit.jupiter.api.Assertions.*; + +import java.time.Duration; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.lambda.model.ErrorObject; +import software.amazon.lambda.durable.config.RetryOperationConfig; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.retry.RetryDecision; +import software.amazon.lambda.durable.retry.RetryStrategies; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; +import software.amazon.lambda.durable.util.RetryOperationHelper; + +class RetryWaitForCallbackIntegrationTest { + + @Test + void waitForCallbackSucceedsOnFirstAttempt() { + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> ctx.waitForCallback( + "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger() + .info("Submitting callback {}", callbackId)), + RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2))) + .build())); + + var result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // waitForCallback("approval-1", ...) creates "approval-1-callback" internally + var callbackId = runner.getCallbackId("approval-1-callback"); + assertNotNull(callbackId); + runner.completeCallback(callbackId, "\"approved\""); + + result = runner.runUntilComplete("test"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("approved", result.getResult(String.class)); + } + + @Test + void waitForCallbackRetriesAfterFailure() { + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> ctx.waitForCallback( + "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger() + .info("Attempt {} callback {}", attempt, callbackId)), + RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2))) + .build())); + + // First run — starts waitForCallback attempt 1 + var result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Fail attempt 1 + var callbackId1 = runner.getCallbackId("approval-1-callback"); + assertNotNull(callbackId1); + runner.failCallback( + callbackId1, + ErrorObject.builder() + .errorType("Rejected") + .errorMessage("denied by reviewer") + .build()); + + // Run — processes failure, hits backoff wait + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past backoff + runner.advanceTime(); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Complete attempt 2 + var callbackId2 = runner.getCallbackId("approval-2-callback"); + assertNotNull(callbackId2); + runner.completeCallback(callbackId2, "\"approved on retry\""); + + result = runner.runUntilComplete("test"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("approved on retry", result.getResult(String.class)); + } + + @Test + void waitForCallbackFailsAfterAllRetriesExhausted() { + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> + ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}), + RetryOperationConfig.builder() + .retryStrategy((error, attempt) -> + attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail()) + .build())); + + // Attempt 1 + var result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + runner.failCallback( + runner.getCallbackId("approval-1-callback"), + ErrorObject.builder().errorMessage("fail 1").build()); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past backoff + runner.advanceTime(); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Attempt 2 — last attempt + runner.failCallback( + runner.getCallbackId("approval-2-callback"), + ErrorObject.builder().errorMessage("fail 2").build()); + result = runner.run("test"); + + assertEquals(ExecutionStatus.FAILED, result.getStatus()); + } + + @Test + void waitForCallbackRetryWithStepsBeforeAndAfter() { + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { + var prefix = context.step("prepare", String.class, stepCtx -> "prepared"); + + var callbackResult = RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> + ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}), + RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1))) + .build()); + + return context.step("finalize", String.class, stepCtx -> prefix + " -> " + callbackResult + " -> done"); + }); + + // First run — prepare completes, waitForCallback starts + var result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Complete callback + var callbackId = runner.getCallbackId("approval-1-callback"); + runner.completeCallback(callbackId, "\"approved\""); + + result = runner.runUntilComplete("test"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("prepared -> approved -> done", result.getResult(String.class)); + } + + @Test + void waitForCallbackRetryMultipleFailuresThenSuccess() { + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> + ctx.waitForCallback("cb-" + attempt, String.class, (callbackId, stepCtx) -> {}), + RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.fixedDelay(4, Duration.ofSeconds(1))) + .build())); + + // Attempt 1 — fail + var result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + runner.failCallback( + runner.getCallbackId("cb-1-callback"), + ErrorObject.builder().errorMessage("fail 1").build()); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + runner.advanceTime(); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Attempt 2 — fail + runner.failCallback( + runner.getCallbackId("cb-2-callback"), + ErrorObject.builder().errorMessage("fail 2").build()); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + runner.advanceTime(); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Attempt 3 — succeed + runner.completeCallback(runner.getCallbackId("cb-3-callback"), "\"third time's the charm\""); + result = runner.runUntilComplete("test"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("third time's the charm", result.getResult(String.class)); + } + + @Test + void waitForCallbackRetryWithSubmitterLogic() { + // Verify the submitter runs on each retry attempt + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> + ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> { + // Submitter runs each attempt — in a real scenario this would + // send the callbackId to an external system + stepCtx.getLogger().info("Attempt {} submitting {}", attempt, callbackId); + }), + RetryOperationConfig.builder() + .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2))) + .build())); + + // Attempt 1 + var result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Verify submitter step was created for attempt 1 + var submitterOp = runner.getOperation("approval-1-submitter"); + assertNotNull(submitterOp, "Submitter step should exist for attempt 1"); + + // Fail attempt 1 + runner.failCallback( + runner.getCallbackId("approval-1-callback"), + ErrorObject.builder().errorMessage("rejected").build()); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Advance past backoff, start attempt 2 + runner.advanceTime(); + result = runner.run("test"); + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + // Verify submitter step was created for attempt 2 + var submitterOp2 = runner.getOperation("approval-2-submitter"); + assertNotNull(submitterOp2, "Submitter step should exist for attempt 2"); + + // Complete attempt 2 + runner.completeCallback(runner.getCallbackId("approval-2-callback"), "\"approved\""); + result = runner.runUntilComplete("test"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("approved", result.getResult(String.class)); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java index c63f0aad0..f4eaf8c28 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java @@ -7,6 +7,8 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.RetryOperationConfig; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.retry.RetryDecision; /** @@ -112,6 +114,9 @@ public static T retryOperation( /** * Core retry loop. Replay-safe because every side-effect is a durable operation: the user's operation calls durable * primitives, and backoff uses {@code context.wait()}. + * + *

{@link SuspendExecutionException} and {@link UnrecoverableDurableExecutionException} are never retried — they + * are internal SDK control flow signals that must propagate immediately. */ private static T executeRetryLoop( DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) { @@ -119,6 +124,9 @@ private static T executeRetryLoop( while (true) { try { return operation.execute(context, attempt); + } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) { + // Internal SDK control flow — never retry, always propagate + throw e; } catch (Exception e) { RetryDecision decision = config.retryStrategy().makeRetryDecision(e, attempt); if (!decision.shouldRetry()) { diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java index 9558f3c83..265fa968c 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java @@ -15,6 +15,8 @@ import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.RetryOperationConfig; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; +import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.retry.RetryDecision; import software.amazon.lambda.durable.retry.RetryStrategies; @@ -416,5 +418,45 @@ void rethrowsLastExceptionWhenAllRetriesExhausted() { // The last attempt's exception is rethrown assertEquals("attempt-3", thrown.getMessage()); } + + @Test + void propagatesSuspendExecutionExceptionWithoutRetrying() { + var config = RetryOperationConfig.builder() + .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1))) + .build(); + + assertThrows( + SuspendExecutionException.class, + () -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> { + throw new SuspendExecutionException(); + }, + config)); + + // Should never reach the wait — SuspendExecutionException propagates immediately + verify(context, never()).wait(anyString(), any(Duration.class)); + } + + @Test + void propagatesUnrecoverableDurableExecutionExceptionWithoutRetrying() { + var config = RetryOperationConfig.builder() + .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1))) + .build(); + + assertThrows( + UnrecoverableDurableExecutionException.class, + () -> RetryOperationHelper.retryOperation( + context, + (ctx, attempt) -> { + throw new UnrecoverableDurableExecutionException( + software.amazon.awssdk.services.lambda.model.ErrorObject.builder() + .errorMessage("unrecoverable") + .build()); + }, + config)); + + verify(context, never()).wait(anyString(), any(Duration.class)); + } } } From 30e29687f9344d84cf4d9fb0d16370d4f0f7c788 Mon Sep 17 00:00:00 2001 From: hsilan Date: Mon, 27 Apr 2026 10:48:48 -0700 Subject: [PATCH 03/19] chore: rename to withRetry from RetryableOperation --- .../lambda/durable/util/RetryOperationHelper.java | 7 +++---- .../util/{RetryableOperation.java => WithRetry.java} | 2 +- .../lambda/durable/util/RetryOperationHelperTest.java | 2 +- ...{RetryableOperationTest.java => WithRetryTest.java} | 10 +++++----- 4 files changed, 10 insertions(+), 11 deletions(-) rename sdk/src/main/java/software/amazon/lambda/durable/util/{RetryableOperation.java => WithRetry.java} (95%) rename sdk/src/test/java/software/amazon/lambda/durable/util/{RetryableOperationTest.java => WithRetryTest.java} (81%) diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java index f4eaf8c28..965af8fee 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java @@ -79,7 +79,7 @@ private RetryOperationHelper() { */ @SuppressWarnings("unchecked") public static T retryOperation( - DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) { + DurableContext context, String name, WithRetry operation, RetryOperationConfig config) { Objects.requireNonNull(context, "context cannot be null"); Objects.requireNonNull(name, "name cannot be null"); Objects.requireNonNull(operation, "operation cannot be null"); @@ -102,8 +102,7 @@ public static T retryOperation( * @param config retry configuration including the retry strategy * @return the operation result */ - public static T retryOperation( - DurableContext context, RetryableOperation operation, RetryOperationConfig config) { + public static T retryOperation(DurableContext context, WithRetry operation, RetryOperationConfig config) { Objects.requireNonNull(context, "context cannot be null"); Objects.requireNonNull(operation, "operation cannot be null"); Objects.requireNonNull(config, "config cannot be null"); @@ -119,7 +118,7 @@ public static T retryOperation( * are internal SDK control flow signals that must propagate immediately. */ private static T executeRetryLoop( - DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) { + DurableContext context, String name, WithRetry operation, RetryOperationConfig config) { var attempt = 1; while (true) { try { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetry.java similarity index 95% rename from sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java rename to sdk/src/main/java/software/amazon/lambda/durable/util/WithRetry.java index d6b2496c1..5b6784919 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetry.java @@ -13,7 +13,7 @@ * @param the result type */ @FunctionalInterface -public interface RetryableOperation { +public interface WithRetry { /** * Executes the durable operation. diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java index 265fa968c..692e5c8d2 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java @@ -272,7 +272,7 @@ void nullOperation_shouldThrow() { assertThrows( NullPointerException.class, - () -> RetryOperationHelper.retryOperation(context, (RetryableOperation) null, config)); + () -> RetryOperationHelper.retryOperation(context, (WithRetry) null, config)); } @Test diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryTest.java similarity index 81% rename from sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java rename to sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryTest.java index f5893a37c..0b1069e8e 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryTest.java @@ -8,11 +8,11 @@ import org.junit.jupiter.api.Test; import software.amazon.lambda.durable.DurableContext; -class RetryableOperationTest { +class WithRetryTest { @Test void canBeImplementedAsLambda() { - RetryableOperation operation = (ctx, attempt) -> "result-" + attempt; + WithRetry operation = (ctx, attempt) -> "result-" + attempt; var context = mock(DurableContext.class); assertEquals("result-1", operation.execute(context, 1)); @@ -22,7 +22,7 @@ void canBeImplementedAsLambda() { @Test void receivesContextAndAttempt() { var context = mock(DurableContext.class); - RetryableOperation operation = (ctx, attempt) -> { + WithRetry operation = (ctx, attempt) -> { assertSame(context, ctx); return "attempt-" + attempt; }; @@ -33,7 +33,7 @@ void receivesContextAndAttempt() { @Test void canThrowExceptions() { - RetryableOperation operation = (ctx, attempt) -> { + WithRetry operation = (ctx, attempt) -> { throw new RuntimeException("failed on attempt " + attempt); }; var context = mock(DurableContext.class); @@ -44,7 +44,7 @@ void canThrowExceptions() { @Test void canReturnNull() { - RetryableOperation operation = (ctx, attempt) -> null; + WithRetry operation = (ctx, attempt) -> null; var context = mock(DurableContext.class); assertNull(operation.execute(context, 1)); From e12daee92346419ad44d239429669e09227a3ddc Mon Sep 17 00:00:00 2001 From: hsilan Date: Mon, 27 Apr 2026 10:56:54 -0700 Subject: [PATCH 04/19] chore: rename to WithRetry from RetryOperation --- .../callback/RetryWaitForCallbackExample.java | 10 +- .../examples/invoke/RetryInvokeExample.java | 10 +- .../durable/RetryInvokeIntegrationTest.java | 28 +++--- .../RetryWaitForCallbackIntegrationTest.java | 28 +++--- ...rationConfig.java => WithRetryConfig.java} | 16 +-- .../amazon/lambda/durable/util/WithRetry.java | 2 +- ...rationHelper.java => WithRetryHelper.java} | 23 +++-- ...nfigTest.java => WithRetryConfigTest.java} | 18 ++-- ...lperTest.java => WithRetryHelperTest.java} | 99 +++++++++---------- 9 files changed, 115 insertions(+), 119 deletions(-) rename sdk/src/main/java/software/amazon/lambda/durable/config/{RetryOperationConfig.java => WithRetryConfig.java} (89%) rename sdk/src/main/java/software/amazon/lambda/durable/util/{RetryOperationHelper.java => WithRetryHelper.java} (90%) rename sdk/src/test/java/software/amazon/lambda/durable/config/{RetryOperationConfigTest.java => WithRetryConfigTest.java} (78%) rename sdk/src/test/java/software/amazon/lambda/durable/util/{RetryOperationHelperTest.java => WithRetryHelperTest.java} (81%) diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java index e16ca8f43..85b2ffcfb 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java @@ -5,13 +5,13 @@ import java.time.Duration; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; -import software.amazon.lambda.durable.config.RetryOperationConfig; +import software.amazon.lambda.durable.config.WithRetryConfig; import software.amazon.lambda.durable.examples.types.ApprovalRequest; import software.amazon.lambda.durable.retry.RetryDecision; -import software.amazon.lambda.durable.util.RetryOperationHelper; +import software.amazon.lambda.durable.util.WithRetryHelper; /** - * Example demonstrating {@link RetryOperationHelper} with {@code context.waitForCallback}. + * Example demonstrating {@link WithRetryHelper} with {@code context.waitForCallback}. * *

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 @@ -33,12 +33,12 @@ public String handleRequest(ApprovalRequest input, DurableContext context) { 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 = RetryOperationHelper.retryOperation( + var approvalResult = WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.waitForCallback( "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger() .info("Attempt {}: sending callback {} to approval system", attempt, callbackId)), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail()) diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java index 774c7c681..e1058226f 100644 --- a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java +++ b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java @@ -5,13 +5,13 @@ import java.time.Duration; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.DurableHandler; -import software.amazon.lambda.durable.config.RetryOperationConfig; +import software.amazon.lambda.durable.config.WithRetryConfig; import software.amazon.lambda.durable.examples.types.GreetingRequest; import software.amazon.lambda.durable.retry.RetryDecision; -import software.amazon.lambda.durable.util.RetryOperationHelper; +import software.amazon.lambda.durable.util.WithRetryHelper; /** - * Example demonstrating {@link RetryOperationHelper} with {@code context.invoke}. + * Example demonstrating {@link WithRetryHelper} with {@code context.invoke}. * *

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 @@ -25,14 +25,14 @@ public class RetryInvokeExample extends DurableHandler @Override public String handleRequest(GreetingRequest input, DurableContext context) { - return RetryOperationHelper.retryOperation( + return WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.invoke( "call-greeting-" + attempt, "simple-step-example" + input.getName() + ":$LATEST", input, String.class), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail()) diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java index 1d7185b95..268b279be 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java @@ -7,13 +7,13 @@ import java.time.Duration; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.ErrorObject; -import software.amazon.lambda.durable.config.RetryOperationConfig; +import software.amazon.lambda.durable.config.WithRetryConfig; import software.amazon.lambda.durable.exception.InvokeFailedException; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.retry.RetryDecision; import software.amazon.lambda.durable.retry.RetryStrategies; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; -import software.amazon.lambda.durable.util.RetryOperationHelper; +import software.amazon.lambda.durable.util.WithRetryHelper; class RetryInvokeIntegrationTest { @@ -21,10 +21,10 @@ class RetryInvokeIntegrationTest { void invokeSucceedsOnFirstAttempt() { var runner = LocalDurableTestRunner.create( String.class, - (input, context) -> RetryOperationHelper.retryOperation( + (input, context) -> WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2))) .build())); @@ -42,10 +42,10 @@ void invokeSucceedsOnFirstAttempt() { void invokeRetriesAfterFailure() { var runner = LocalDurableTestRunner.create( String.class, - (input, context) -> RetryOperationHelper.retryOperation( + (input, context) -> WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2))) .build())); @@ -80,10 +80,10 @@ void invokeRetriesAfterFailure() { void invokeFailsAfterAllRetriesExhausted() { var runner = LocalDurableTestRunner.create( String.class, - (input, context) -> RetryOperationHelper.retryOperation( + (input, context) -> WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy((error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail()) .build())); @@ -114,10 +114,10 @@ void invokeFailsAfterAllRetriesExhausted() { void invokeRetryWithCustomBackoffDelay() { var runner = LocalDurableTestRunner.create( String.class, - (input, context) -> RetryOperationHelper.retryOperation( + (input, context) -> WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy((error, attempt) -> attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(attempt * 5L)) : RetryDecision.fail()) @@ -148,10 +148,10 @@ void invokeRetryWithStepsBeforeAndAfter() { var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { var prefix = context.step("prepare", String.class, stepCtx -> "prepared"); - var invokeResult = RetryOperationHelper.retryOperation( + var invokeResult = WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1))) .build()); @@ -174,10 +174,10 @@ void invokeRetryWithStepsBeforeAndAfter() { void invokeRetryPreservesOriginalExceptionType() { var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { try { - return RetryOperationHelper.retryOperation( + return WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy(RetryStrategies.Presets.NO_RETRY) .build()); } catch (InvokeFailedException e) { diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java index f7806d577..d7bec449b 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java @@ -7,12 +7,12 @@ import java.time.Duration; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.ErrorObject; -import software.amazon.lambda.durable.config.RetryOperationConfig; +import software.amazon.lambda.durable.config.WithRetryConfig; import software.amazon.lambda.durable.model.ExecutionStatus; import software.amazon.lambda.durable.retry.RetryDecision; import software.amazon.lambda.durable.retry.RetryStrategies; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; -import software.amazon.lambda.durable.util.RetryOperationHelper; +import software.amazon.lambda.durable.util.WithRetryHelper; class RetryWaitForCallbackIntegrationTest { @@ -20,12 +20,12 @@ class RetryWaitForCallbackIntegrationTest { void waitForCallbackSucceedsOnFirstAttempt() { var runner = LocalDurableTestRunner.create( String.class, - (input, context) -> RetryOperationHelper.retryOperation( + (input, context) -> WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.waitForCallback( "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger() .info("Submitting callback {}", callbackId)), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2))) .build())); @@ -47,12 +47,12 @@ void waitForCallbackSucceedsOnFirstAttempt() { void waitForCallbackRetriesAfterFailure() { var runner = LocalDurableTestRunner.create( String.class, - (input, context) -> RetryOperationHelper.retryOperation( + (input, context) -> WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.waitForCallback( "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger() .info("Attempt {} callback {}", attempt, callbackId)), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2))) .build())); @@ -94,11 +94,11 @@ void waitForCallbackRetriesAfterFailure() { void waitForCallbackFailsAfterAllRetriesExhausted() { var runner = LocalDurableTestRunner.create( String.class, - (input, context) -> RetryOperationHelper.retryOperation( + (input, context) -> WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy((error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail()) .build())); @@ -132,11 +132,11 @@ void waitForCallbackRetryWithStepsBeforeAndAfter() { var runner = LocalDurableTestRunner.create(String.class, (input, context) -> { var prefix = context.step("prepare", String.class, stepCtx -> "prepared"); - var callbackResult = RetryOperationHelper.retryOperation( + var callbackResult = WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1))) .build()); @@ -161,11 +161,11 @@ void waitForCallbackRetryWithStepsBeforeAndAfter() { void waitForCallbackRetryMultipleFailuresThenSuccess() { var runner = LocalDurableTestRunner.create( String.class, - (input, context) -> RetryOperationHelper.retryOperation( + (input, context) -> WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.waitForCallback("cb-" + attempt, String.class, (callbackId, stepCtx) -> {}), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy(RetryStrategies.fixedDelay(4, Duration.ofSeconds(1))) .build())); @@ -207,7 +207,7 @@ void waitForCallbackRetryWithSubmitterLogic() { // Verify the submitter runs on each retry attempt var runner = LocalDurableTestRunner.create( String.class, - (input, context) -> RetryOperationHelper.retryOperation( + (input, context) -> WithRetryHelper.retryOperation( context, (ctx, attempt) -> ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> { @@ -215,7 +215,7 @@ void waitForCallbackRetryWithSubmitterLogic() { // send the callbackId to an external system stepCtx.getLogger().info("Attempt {} submitting {}", attempt, callbackId); }), - RetryOperationConfig.builder() + WithRetryConfig.builder() .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2))) .build())); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java similarity index 89% rename from sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java rename to sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java index 696d1d253..0c7bf1d34 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java @@ -5,16 +5,16 @@ import software.amazon.lambda.durable.retry.RetryStrategy; /** - * Configuration for {@link software.amazon.lambda.durable.util.RetryOperationHelper#retryOperation}. + * Configuration for {@link software.amazon.lambda.durable.util.WithRetryHelper#retryOperation}. * *

Uses the same {@link RetryStrategy} shape that developers already know from {@link StepConfig}, so there are zero * new retry concepts to learn. */ -public class RetryOperationConfig { +public class WithRetryConfig { private final RetryStrategy retryStrategy; private final boolean wrapInChildContext; - private RetryOperationConfig(Builder builder) { + private WithRetryConfig(Builder builder) { this.retryStrategy = builder.retryStrategy; this.wrapInChildContext = builder.wrapInChildContext; } @@ -40,7 +40,7 @@ public boolean wrapInChildContext() { } /** - * Creates a new builder for {@code RetryOperationConfig}. + * Creates a new builder for {@code WithRetryConfig}. * * @return a new builder instance */ @@ -48,7 +48,7 @@ public static Builder builder() { return new Builder(); } - /** Builder for creating {@link RetryOperationConfig} instances. */ + /** Builder for creating {@link WithRetryConfig} instances. */ public static class Builder { private RetryStrategy retryStrategy; private boolean wrapInChildContext = true; @@ -88,16 +88,16 @@ public Builder wrapInChildContext(boolean wrapInChildContext) { } /** - * Builds the {@link RetryOperationConfig} instance. + * Builds the {@link WithRetryConfig} instance. * * @return a new config with the configured options * @throws IllegalArgumentException if retryStrategy is not set */ - public RetryOperationConfig build() { + public WithRetryConfig build() { if (retryStrategy == null) { throw new IllegalArgumentException("retryStrategy is required"); } - return new RetryOperationConfig(this); + return new WithRetryConfig(this); } } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetry.java b/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetry.java index 5b6784919..87f7be964 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetry.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetry.java @@ -5,7 +5,7 @@ import software.amazon.lambda.durable.DurableContext; /** - * A durable operation that can be retried end-to-end by {@link RetryOperationHelper}. + * A durable operation that can be retried end-to-end by {@link WithRetryHelper}. * *

Receives the durable context and the 1-based attempt number so callers can generate unique operation names per * attempt (e.g., {@code "approval-" + attempt}). diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java similarity index 90% rename from sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java rename to sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java index 965af8fee..938a63741 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java @@ -6,7 +6,7 @@ import java.util.Objects; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.TypeToken; -import software.amazon.lambda.durable.config.RetryOperationConfig; +import software.amazon.lambda.durable.config.WithRetryConfig; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.retry.RetryDecision; @@ -23,7 +23,7 @@ *

Usage — callback retry

* * 
{@code
- * var result = RetryOperationHelper.retryOperation(
+ * var result = WithRetryHelper.retryOperation(
  *     context,
  *     "approval",
  *     (ctx, attempt) -> ctx.waitForCallback(
@@ -31,7 +31,7 @@
  *         String.class,
  *         (callbackId, stepCtx) -> sendApprovalEmail(approverEmail, callbackId)
  *     ),
- *     RetryOperationConfig.builder()
+ *     WithRetryConfig.builder()
  *         .retryStrategy(RetryStrategies.exponentialBackoff(
  *             3, Duration.ofSeconds(2), Duration.ofSeconds(30), 2.0, JitterStrategy.FULL))
  *         .build()
@@ -41,11 +41,11 @@
  * 
Usage — invoke retry (anonymous form)

  *
  * 
{@code
- * var result = RetryOperationHelper.retryOperation(
+ * var result = WithRetryHelper.retryOperation(
  *     context,
  *     (ctx, attempt) -> ctx.invoke(
  *         "charge-" + attempt, paymentFnArn, new ChargeRequest(orderId), String.class),
- *     RetryOperationConfig.builder()
+ *     WithRetryConfig.builder()
  *         .retryStrategy((err, att) -> att < 3
  *             ? RetryDecision.retry(Duration.ofSeconds(1))
  *             : RetryDecision.fail())
@@ -53,13 +53,13 @@
  * );
  * }

  */
-public final class RetryOperationHelper {
+public final class WithRetryHelper {
 
     private static final Duration DEFAULT_BACKOFF_DELAY = Duration.ofSeconds(1);
     private static final String BACKOFF_SUFFIX = "-backoff-";
     private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-";
 
-    private RetryOperationHelper() {
+    private WithRetryHelper() {
         // utility class
     }
 
@@ -67,8 +67,7 @@ private RetryOperationHelper() {
      * Named form — wraps the retry loop in {@code runInChildContext} by default so all attempts are grouped under a
      * single named operation in execution history.
      *
-     * 
The child-context wrapping can be disabled via
-     * {@link RetryOperationConfig.Builder#wrapInChildContext(boolean)}.
+     * 
The child-context wrapping can be disabled via {@link WithRetryConfig.Builder#wrapInChildContext(boolean)}.
      *
      * @param  the result type
      * @param context the durable context
@@ -79,7 +78,7 @@ private RetryOperationHelper() {
      */
     @SuppressWarnings("unchecked")
     public static  T retryOperation(
-            DurableContext context, String name, WithRetry operation, RetryOperationConfig config) {
+            DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
         Objects.requireNonNull(context, "context cannot be null");
         Objects.requireNonNull(name, "name cannot be null");
         Objects.requireNonNull(operation, "operation cannot be null");
@@ -102,7 +101,7 @@ public static  T retryOperation(
      * @param config retry configuration including the retry strategy
      * @return the operation result
      */
-    public static  T retryOperation(DurableContext context, WithRetry operation, RetryOperationConfig config) {
+    public static  T retryOperation(DurableContext context, WithRetry operation, WithRetryConfig config) {
         Objects.requireNonNull(context, "context cannot be null");
         Objects.requireNonNull(operation, "operation cannot be null");
         Objects.requireNonNull(config, "config cannot be null");
@@ -118,7 +117,7 @@ public static  T retryOperation(DurableContext context, WithRetry operatio
      * are internal SDK control flow signals that must propagate immediately.
      */
     private static  T executeRetryLoop(
-            DurableContext context, String name, WithRetry operation, RetryOperationConfig config) {
+            DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
         var attempt = 1;
         while (true) {
             try {
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/WithRetryConfigTest.java
similarity index 78%
rename from sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
rename to sdk/src/test/java/software/amazon/lambda/durable/config/WithRetryConfigTest.java
index 726da5b61..26357d312 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/config/WithRetryConfigTest.java
@@ -8,28 +8,28 @@
 import software.amazon.lambda.durable.retry.RetryDecision;
 import software.amazon.lambda.durable.retry.RetryStrategies;
 
-class RetryOperationConfigTest {
+class WithRetryConfigTest {
 
     @Test
     void builderWithRetryStrategy() {
         var strategy = RetryStrategies.Presets.DEFAULT;
 
-        var config = RetryOperationConfig.builder().retryStrategy(strategy).build();
+        var config = WithRetryConfig.builder().retryStrategy(strategy).build();
 
         assertEquals(strategy, config.retryStrategy());
     }
 
     @Test
     void builderWithoutRetryStrategy_shouldThrow() {
-        var exception = assertThrows(IllegalArgumentException.class, () -> RetryOperationConfig.builder()
-                .build());
+        var exception = assertThrows(
+                IllegalArgumentException.class, () -> WithRetryConfig.builder().build());
 
         assertEquals("retryStrategy is required", exception.getMessage());
     }
 
     @Test
     void wrapInChildContext_defaultsToTrue() {
-        var config = RetryOperationConfig.builder()
+        var config = WithRetryConfig.builder()
                 .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                 .build();
 
@@ -38,7 +38,7 @@ void wrapInChildContext_defaultsToTrue() {
 
     @Test
     void wrapInChildContext_canBeSetToFalse() {
-        var config = RetryOperationConfig.builder()
+        var config = WithRetryConfig.builder()
                 .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                 .wrapInChildContext(false)
                 .build();
@@ -48,7 +48,7 @@ void wrapInChildContext_canBeSetToFalse() {
 
     @Test
     void wrapInChildContext_canBeSetToTrue() {
-        var config = RetryOperationConfig.builder()
+        var config = WithRetryConfig.builder()
                 .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                 .wrapInChildContext(true)
                 .build();
@@ -60,7 +60,7 @@ void wrapInChildContext_canBeSetToTrue() {
     void builderChaining() {
         var strategy = RetryStrategies.Presets.DEFAULT;
 
-        var config = RetryOperationConfig.builder()
+        var config = WithRetryConfig.builder()
                 .retryStrategy(strategy)
                 .wrapInChildContext(false)
                 .build();
@@ -71,7 +71,7 @@ void builderChaining() {
 
     @Test
     void builderWithCustomLambdaRetryStrategy() {
-        var config = RetryOperationConfig.builder()
+        var config = WithRetryConfig.builder()
                 .retryStrategy((error, attempt) -> RetryDecision.fail())
                 .build();
 
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
similarity index 81%
rename from sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
rename to sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
index 692e5c8d2..994015be6 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
@@ -14,13 +14,13 @@
 import org.junit.jupiter.api.Test;
 import software.amazon.lambda.durable.DurableContext;
 import software.amazon.lambda.durable.TypeToken;
-import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.config.WithRetryConfig;
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.execution.SuspendExecutionException;
 import software.amazon.lambda.durable.retry.RetryDecision;
 import software.amazon.lambda.durable.retry.RetryStrategies;
 
-class RetryOperationHelperTest {
+class WithRetryHelperTest {
 
     private DurableContext context;
 
@@ -43,11 +43,11 @@ void successOnFirstAttempt_wrapsInChildContext() {
                         return func.apply(context);
                     });
 
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var result = RetryOperationHelper.retryOperation(context, "my-op", (ctx, attempt) -> "success", config);
+            var result = WithRetryHelper.retryOperation(context, "my-op", (ctx, attempt) -> "success", config);
 
             assertEquals("success", result);
             verify(context).runInChildContext(eq("my-op"), any(TypeToken.class), any());
@@ -55,12 +55,12 @@ void successOnFirstAttempt_wrapsInChildContext() {
 
         @Test
         void successOnFirstAttempt_noChildContext_whenDisabled() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .wrapInChildContext(false)
                     .build();
 
-            var result = RetryOperationHelper.retryOperation(context, "my-op", (ctx, attempt) -> "direct", config);
+            var result = WithRetryHelper.retryOperation(context, "my-op", (ctx, attempt) -> "direct", config);
 
             assertEquals("direct", result);
             verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
@@ -70,13 +70,13 @@ void successOnFirstAttempt_noChildContext_whenDisabled() {
         void retriesWithBackoffWaits_namedForm() {
             // Disable child context wrapping so we can directly verify wait calls
             var callCount = new int[] {0};
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) ->
                             attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
                     .wrapInChildContext(false)
                     .build();
 
-            var result = RetryOperationHelper.retryOperation(
+            var result = WithRetryHelper.retryOperation(
                     context,
                     "my-op",
                     (ctx, attempt) -> {
@@ -97,14 +97,14 @@ void retriesWithBackoffWaits_namedForm() {
 
         @Test
         void rethrowsWhenRetryStrategyReturnsFail() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .wrapInChildContext(false)
                     .build();
 
             var exception = assertThrows(
                     RuntimeException.class,
-                    () -> RetryOperationHelper.retryOperation(
+                    () -> WithRetryHelper.retryOperation(
                             context,
                             "my-op",
                             (ctx, attempt) -> {
@@ -118,14 +118,14 @@ void rethrowsWhenRetryStrategyReturnsFail() {
 
         @Test
         void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(
                             (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail())
                     .wrapInChildContext(false)
                     .build();
 
             var callCount = new int[] {0};
-            var result = RetryOperationHelper.retryOperation(
+            var result = WithRetryHelper.retryOperation(
                     context,
                     "my-op",
                     (ctx, attempt) -> {
@@ -144,42 +144,41 @@ void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
 
         @Test
         void nullContext_shouldThrow() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
             assertThrows(
                     NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(null, "name", (ctx, a) -> "x", config));
+                    () -> WithRetryHelper.retryOperation(null, "name", (ctx, a) -> "x", config));
         }
 
         @Test
         void nullName_shouldThrow() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
             assertThrows(
                     NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, null, (ctx, a) -> "x", config));
+                    () -> WithRetryHelper.retryOperation(context, null, (ctx, a) -> "x", config));
         }
 
         @Test
         void nullOperation_shouldThrow() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
             assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, "name", null, config));
+                    NullPointerException.class, () -> WithRetryHelper.retryOperation(context, "name", null, config));
         }
 
         @Test
         void nullConfig_shouldThrow() {
             assertThrows(
                     NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, "name", (ctx, a) -> "x", null));
+                    () -> WithRetryHelper.retryOperation(context, "name", (ctx, a) -> "x", null));
         }
     }
 
@@ -190,11 +189,11 @@ class AnonymousForm {
 
         @Test
         void successOnFirstAttempt() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var result = RetryOperationHelper.retryOperation(context, (ctx, attempt) -> "anonymous-success", config);
+            var result = WithRetryHelper.retryOperation(context, (ctx, attempt) -> "anonymous-success", config);
 
             assertEquals("anonymous-success", result);
             verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
@@ -202,12 +201,12 @@ void successOnFirstAttempt() {
 
         @Test
         void retriesWithAnonymousBackoffNames() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) ->
                             attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
                     .build();
 
-            var result = RetryOperationHelper.retryOperation(
+            var result = WithRetryHelper.retryOperation(
                     context,
                     (ctx, attempt) -> {
                         if (attempt < 3) {
@@ -224,26 +223,26 @@ void retriesWithAnonymousBackoffNames() {
 
         @Test
         void neverWrapsInChildContext_evenWhenConfigSaysTrue() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .wrapInChildContext(true) // should be ignored for anonymous form
                     .build();
 
-            RetryOperationHelper.retryOperation(context, (ctx, attempt) -> "result", config);
+            WithRetryHelper.retryOperation(context, (ctx, attempt) -> "result", config);
 
             verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
         }
 
         @Test
         void rethrowsOriginalException() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
             var original = new IllegalStateException("original error");
             var thrown = assertThrows(
                     IllegalStateException.class,
-                    () -> RetryOperationHelper.retryOperation(
+                    () -> WithRetryHelper.retryOperation(
                             context,
                             (ctx, attempt) -> {
                                 throw original;
@@ -255,31 +254,29 @@ void rethrowsOriginalException() {
 
         @Test
         void nullContext_shouldThrow() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
             assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(null, (ctx, a) -> "x", config));
+                    NullPointerException.class, () -> WithRetryHelper.retryOperation(null, (ctx, a) -> "x", config));
         }
 
         @Test
         void nullOperation_shouldThrow() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
             assertThrows(
                     NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, (WithRetry) null, config));
+                    () -> WithRetryHelper.retryOperation(context, (WithRetry) null, config));
         }
 
         @Test
         void nullConfig_shouldThrow() {
             assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, (ctx, a) -> "x", null));
+                    NullPointerException.class, () -> WithRetryHelper.retryOperation(context, (ctx, a) -> "x", null));
         }
     }
 
@@ -291,13 +288,13 @@ class RetryBehavior {
         @Test
         void passesCorrectAttemptNumberToOperation() {
             var attempts = new ArrayList();
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) ->
                             attempt < 4 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
                     .wrapInChildContext(false)
                     .build();
 
-            RetryOperationHelper.retryOperation(
+            WithRetryHelper.retryOperation(
                     context,
                     "track",
                     (ctx, attempt) -> {
@@ -319,7 +316,7 @@ void passesCorrectAttemptNumberToOperation() {
         @Test
         void passesErrorToRetryStrategy() {
             var errors = new ArrayList();
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) -> {
                         errors.add(error);
                         return attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail();
@@ -328,7 +325,7 @@ void passesErrorToRetryStrategy() {
 
             assertThrows(
                     RuntimeException.class,
-                    () -> RetryOperationHelper.retryOperation(
+                    () -> WithRetryHelper.retryOperation(
                             context,
                             (ctx, attempt) -> {
                                 throw new RuntimeException("error-" + attempt);
@@ -343,11 +340,11 @@ void passesErrorToRetryStrategy() {
 
         @Test
         void respectsCustomDelayFromRetryDecision() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(attempt * 10L)))
                     .build();
 
-            RetryOperationHelper.retryOperation(
+            WithRetryHelper.retryOperation(
                     context,
                     (ctx, attempt) -> {
                         if (attempt <= 2) {
@@ -363,11 +360,11 @@ void respectsCustomDelayFromRetryDecision() {
 
         @Test
         void passesContextToOperation() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            RetryOperationHelper.retryOperation(
+            WithRetryHelper.retryOperation(
                     context,
                     (ctx, attempt) -> {
                         assertSame(context, ctx);
@@ -385,12 +382,12 @@ void namedFormPassesChildContextToOperation_whenWrapped() {
                         return func.apply(childContext);
                     });
 
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .wrapInChildContext(true)
                     .build();
 
-            RetryOperationHelper.retryOperation(
+            WithRetryHelper.retryOperation(
                     context,
                     "wrapped",
                     (ctx, attempt) -> {
@@ -402,13 +399,13 @@ void namedFormPassesChildContextToOperation_whenWrapped() {
 
         @Test
         void rethrowsLastExceptionWhenAllRetriesExhausted() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
                     .build();
 
             var thrown = assertThrows(
                     RuntimeException.class,
-                    () -> RetryOperationHelper.retryOperation(
+                    () -> WithRetryHelper.retryOperation(
                             context,
                             (ctx, attempt) -> {
                                 throw new RuntimeException("attempt-" + attempt);
@@ -421,13 +418,13 @@ void rethrowsLastExceptionWhenAllRetriesExhausted() {
 
         @Test
         void propagatesSuspendExecutionExceptionWithoutRetrying() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1)))
                     .build();
 
             assertThrows(
                     SuspendExecutionException.class,
-                    () -> RetryOperationHelper.retryOperation(
+                    () -> WithRetryHelper.retryOperation(
                             context,
                             (ctx, attempt) -> {
                                 throw new SuspendExecutionException();
@@ -440,13 +437,13 @@ void propagatesSuspendExecutionExceptionWithoutRetrying() {
 
         @Test
         void propagatesUnrecoverableDurableExecutionExceptionWithoutRetrying() {
-            var config = RetryOperationConfig.builder()
+            var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1)))
                     .build();
 
             assertThrows(
                     UnrecoverableDurableExecutionException.class,
-                    () -> RetryOperationHelper.retryOperation(
+                    () -> WithRetryHelper.retryOperation(
                             context,
                             (ctx, attempt) -> {
                                 throw new UnrecoverableDurableExecutionException(

From 9ff33b8355d90713fd2a760eeccee039aa64b760 Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Mon, 27 Apr 2026 11:28:44 -0700
Subject: [PATCH 05/19] chore: add some edge case tests for WithRetryHelperTest

---
 .../durable/util/WithRetryHelperTest.java     | 114 ++++++++++++++++++
 1 file changed, 114 insertions(+)

diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
index 994015be6..51b502d4c 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
@@ -15,6 +15,7 @@
 import software.amazon.lambda.durable.DurableContext;
 import software.amazon.lambda.durable.TypeToken;
 import software.amazon.lambda.durable.config.WithRetryConfig;
+import software.amazon.lambda.durable.exception.SerDesException;
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.execution.SuspendExecutionException;
 import software.amazon.lambda.durable.retry.RetryDecision;
@@ -180,6 +181,19 @@ void nullConfig_shouldThrow() {
                     NullPointerException.class,
                     () -> WithRetryHelper.retryOperation(context, "name", (ctx, a) -> "x", null));
         }
+
+        @Test
+        void operationReturnsNull() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(false)
+                    .build();
+
+            var result = WithRetryHelper.retryOperation(
+                    context, "my-op", (WithRetry) (ctx, attempt) -> null, config);
+
+            assertNull(result);
+        }
     }
 
     // --- Anonymous form tests ---
@@ -278,6 +292,38 @@ void nullConfig_shouldThrow() {
             assertThrows(
                     NullPointerException.class, () -> WithRetryHelper.retryOperation(context, (ctx, a) -> "x", null));
         }
+
+        @Test
+        void operationReturnsNull() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            var result = WithRetryHelper.retryOperation(context, (WithRetry) (ctx, attempt) -> null, config);
+
+            assertNull(result);
+        }
+
+        @Test
+        void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(
+                            (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail())
+                    .build();
+
+            var result = WithRetryHelper.retryOperation(
+                    context,
+                    (ctx, attempt) -> {
+                        if (attempt == 1) {
+                            throw new RuntimeException("fail");
+                        }
+                        return "ok";
+                    },
+                    config);
+
+            assertEquals("ok", result);
+            verify(context).wait("retry-backoff-1", Duration.ofSeconds(1));
+        }
     }
 
     // --- Retry behavior tests ---
@@ -455,5 +501,73 @@ void propagatesUnrecoverableDurableExecutionExceptionWithoutRetrying() {
 
             verify(context, never()).wait(anyString(), any(Duration.class));
         }
+
+        @Test
+        void propagatesSuspendExecutionExceptionOnLaterAttempt() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1)))
+                    .build();
+
+            assertThrows(
+                    SuspendExecutionException.class,
+                    () -> WithRetryHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                if (attempt == 1) {
+                                    throw new RuntimeException("transient");
+                                }
+                                // Second attempt triggers suspend — must propagate, not retry
+                                throw new SuspendExecutionException();
+                            },
+                            config));
+
+            // First attempt retried (one backoff wait), second attempt suspended immediately
+            verify(context, times(1)).wait(anyString(), any(Duration.class));
+        }
+
+        @Test
+        void propagatesUnrecoverableDurableExecutionExceptionOnLaterAttempt() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1)))
+                    .build();
+
+            assertThrows(
+                    UnrecoverableDurableExecutionException.class,
+                    () -> WithRetryHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                if (attempt == 1) {
+                                    throw new RuntimeException("transient");
+                                }
+                                throw new UnrecoverableDurableExecutionException(
+                                        software.amazon.awssdk.services.lambda.model.ErrorObject.builder()
+                                                .errorMessage("unrecoverable on attempt 2")
+                                                .build());
+                            },
+                            config));
+
+            verify(context, times(1)).wait(anyString(), any(Duration.class));
+        }
+
+        @Test
+        void preservesCheckedExceptionSubclassType() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            var original = new SerDesException("deserialization failed", new RuntimeException("bad json"));
+
+            var thrown = assertThrows(
+                    SerDesException.class,
+                    () -> WithRetryHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                throw original;
+                            },
+                            config));
+
+            assertSame(original, thrown);
+            assertEquals("deserialization failed", thrown.getMessage());
+        }
     }
 }

From b102747168374bb11f745330cb891b63fc275614 Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Mon, 27 Apr 2026 11:33:41 -0700
Subject: [PATCH 06/19] chore: rename retryOperation to withRetry

---
 .../callback/RetryWaitForCallbackExample.java |  2 +-
 .../examples/invoke/RetryInvokeExample.java   |  2 +-
 .../durable/RetryInvokeIntegrationTest.java   | 12 ++--
 .../RetryWaitForCallbackIntegrationTest.java  | 12 ++--
 .../durable/config/WithRetryConfig.java       |  6 +-
 .../lambda/durable/util/WithRetryHelper.java  |  9 ++-
 .../durable/util/WithRetryHelperTest.java     | 66 +++++++++----------
 7 files changed, 52 insertions(+), 57 deletions(-)

diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
index 85b2ffcfb..4f7afbaef 100644
--- a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
+++ b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
@@ -33,7 +33,7 @@ public String handleRequest(ApprovalRequest input, DurableContext context) {
                 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 = WithRetryHelper.retryOperation(
+        var approvalResult = WithRetryHelper.withRetry(
                 context,
                 (ctx, attempt) -> ctx.waitForCallback(
                         "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
index e1058226f..a764f7c36 100644
--- a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
+++ b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
@@ -25,7 +25,7 @@ public class RetryInvokeExample extends DurableHandler
 
     @Override
     public String handleRequest(GreetingRequest input, DurableContext context) {
-        return WithRetryHelper.retryOperation(
+        return WithRetryHelper.withRetry(
                 context,
                 (ctx, attempt) -> ctx.invoke(
                         "call-greeting-" + attempt,
diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
index 268b279be..f7ed4408c 100644
--- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
+++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
@@ -21,7 +21,7 @@ class RetryInvokeIntegrationTest {
     void invokeSucceedsOnFirstAttempt() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.retryOperation(
+                (input, context) -> WithRetryHelper.withRetry(
                         context,
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
@@ -42,7 +42,7 @@ void invokeSucceedsOnFirstAttempt() {
     void invokeRetriesAfterFailure() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.retryOperation(
+                (input, context) -> WithRetryHelper.withRetry(
                         context,
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
@@ -80,7 +80,7 @@ void invokeRetriesAfterFailure() {
     void invokeFailsAfterAllRetriesExhausted() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.retryOperation(
+                (input, context) -> WithRetryHelper.withRetry(
                         context,
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
@@ -114,7 +114,7 @@ void invokeFailsAfterAllRetriesExhausted() {
     void invokeRetryWithCustomBackoffDelay() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.retryOperation(
+                (input, context) -> WithRetryHelper.withRetry(
                         context,
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
@@ -148,7 +148,7 @@ void invokeRetryWithStepsBeforeAndAfter() {
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             var prefix = context.step("prepare", String.class, stepCtx -> "prepared");
 
-            var invokeResult = WithRetryHelper.retryOperation(
+            var invokeResult = WithRetryHelper.withRetry(
                     context,
                     (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                     WithRetryConfig.builder()
@@ -174,7 +174,7 @@ void invokeRetryWithStepsBeforeAndAfter() {
     void invokeRetryPreservesOriginalExceptionType() {
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             try {
-                return WithRetryHelper.retryOperation(
+                return WithRetryHelper.withRetry(
                         context,
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
index d7bec449b..2f08315ff 100644
--- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
+++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
@@ -20,7 +20,7 @@ class RetryWaitForCallbackIntegrationTest {
     void waitForCallbackSucceedsOnFirstAttempt() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.retryOperation(
+                (input, context) -> WithRetryHelper.withRetry(
                         context,
                         (ctx, attempt) -> ctx.waitForCallback(
                                 "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
@@ -47,7 +47,7 @@ void waitForCallbackSucceedsOnFirstAttempt() {
     void waitForCallbackRetriesAfterFailure() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.retryOperation(
+                (input, context) -> WithRetryHelper.withRetry(
                         context,
                         (ctx, attempt) -> ctx.waitForCallback(
                                 "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
@@ -94,7 +94,7 @@ void waitForCallbackRetriesAfterFailure() {
     void waitForCallbackFailsAfterAllRetriesExhausted() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.retryOperation(
+                (input, context) -> WithRetryHelper.withRetry(
                         context,
                         (ctx, attempt) ->
                                 ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
@@ -132,7 +132,7 @@ void waitForCallbackRetryWithStepsBeforeAndAfter() {
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             var prefix = context.step("prepare", String.class, stepCtx -> "prepared");
 
-            var callbackResult = WithRetryHelper.retryOperation(
+            var callbackResult = WithRetryHelper.withRetry(
                     context,
                     (ctx, attempt) ->
                             ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
@@ -161,7 +161,7 @@ void waitForCallbackRetryWithStepsBeforeAndAfter() {
     void waitForCallbackRetryMultipleFailuresThenSuccess() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.retryOperation(
+                (input, context) -> WithRetryHelper.withRetry(
                         context,
                         (ctx, attempt) ->
                                 ctx.waitForCallback("cb-" + attempt, String.class, (callbackId, stepCtx) -> {}),
@@ -207,7 +207,7 @@ void waitForCallbackRetryWithSubmitterLogic() {
         // Verify the submitter runs on each retry attempt
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.retryOperation(
+                (input, context) -> WithRetryHelper.withRetry(
                         context,
                         (ctx, attempt) ->
                                 ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java
index 0c7bf1d34..624b7df32 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java
@@ -5,7 +5,7 @@
 import software.amazon.lambda.durable.retry.RetryStrategy;
 
 /**
- * Configuration for {@link software.amazon.lambda.durable.util.WithRetryHelper#retryOperation}.
+ * Configuration for {@link software.amazon.lambda.durable.util.WithRetryHelper#withRetry}.
  *
  * 
Uses the same {@link RetryStrategy} shape that developers already know from {@link StepConfig}, so there are zero
  * new retry concepts to learn.
@@ -30,7 +30,7 @@ public RetryStrategy retryStrategy() {
 
     /**
      * Whether to wrap the retry loop in {@code runInChildContext} so all attempts are grouped under a single named
-     * operation in execution history. Only applies when a name is provided to the named form of {@code retryOperation}.
+     * operation in execution history. Only applies when a name is provided to the named form of {@code withRetry}.
      * Defaults to {@code true}.
      *
      * @return true if child-context wrapping is enabled
@@ -73,7 +73,7 @@ public Builder retryStrategy(RetryStrategy retryStrategy) {
 
         /**
          * Controls whether the retry loop is wrapped in a child context. Only meaningful for the named form of
-         * {@code retryOperation}. Defaults to {@code true}.
+         * {@code withRetry}. Defaults to {@code true}.
          *
          * 
When {@code true}, all attempts and backoff waits are grouped under a single named operation in execution
          * history, providing a cleaner view and isolated operation ID space. Set to {@code false} to flatten attempts
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java
index 938a63741..dbbb0a05d 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java
@@ -23,7 +23,7 @@
  * 
Usage — callback retry

  *
  * 
{@code
- * var result = WithRetryHelper.retryOperation(
+ * var result = WithRetryHelper.withRetry(
  *     context,
  *     "approval",
  *     (ctx, attempt) -> ctx.waitForCallback(
@@ -41,7 +41,7 @@
  * 
Usage — invoke retry (anonymous form)

  *
  * 
{@code
- * var result = WithRetryHelper.retryOperation(
+ * var result = WithRetryHelper.withRetry(
  *     context,
  *     (ctx, attempt) -> ctx.invoke(
  *         "charge-" + attempt, paymentFnArn, new ChargeRequest(orderId), String.class),
@@ -77,8 +77,7 @@ private WithRetryHelper() {
      * @return the operation result
      */
     @SuppressWarnings("unchecked")
-    public static  T retryOperation(
-            DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
+    public static  T withRetry(DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
         Objects.requireNonNull(context, "context cannot be null");
         Objects.requireNonNull(name, "name cannot be null");
         Objects.requireNonNull(operation, "operation cannot be null");
@@ -101,7 +100,7 @@ public static  T retryOperation(
      * @param config retry configuration including the retry strategy
      * @return the operation result
      */
-    public static  T retryOperation(DurableContext context, WithRetry operation, WithRetryConfig config) {
+    public static  T withRetry(DurableContext context, WithRetry operation, WithRetryConfig config) {
         Objects.requireNonNull(context, "context cannot be null");
         Objects.requireNonNull(operation, "operation cannot be null");
         Objects.requireNonNull(config, "config cannot be null");
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
index 51b502d4c..c6cd81fec 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
@@ -48,7 +48,7 @@ void successOnFirstAttempt_wrapsInChildContext() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var result = WithRetryHelper.retryOperation(context, "my-op", (ctx, attempt) -> "success", config);
+            var result = WithRetryHelper.withRetry(context, "my-op", (ctx, attempt) -> "success", config);
 
             assertEquals("success", result);
             verify(context).runInChildContext(eq("my-op"), any(TypeToken.class), any());
@@ -61,7 +61,7 @@ void successOnFirstAttempt_noChildContext_whenDisabled() {
                     .wrapInChildContext(false)
                     .build();
 
-            var result = WithRetryHelper.retryOperation(context, "my-op", (ctx, attempt) -> "direct", config);
+            var result = WithRetryHelper.withRetry(context, "my-op", (ctx, attempt) -> "direct", config);
 
             assertEquals("direct", result);
             verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
@@ -77,7 +77,7 @@ void retriesWithBackoffWaits_namedForm() {
                     .wrapInChildContext(false)
                     .build();
 
-            var result = WithRetryHelper.retryOperation(
+            var result = WithRetryHelper.withRetry(
                     context,
                     "my-op",
                     (ctx, attempt) -> {
@@ -105,7 +105,7 @@ void rethrowsWhenRetryStrategyReturnsFail() {
 
             var exception = assertThrows(
                     RuntimeException.class,
-                    () -> WithRetryHelper.retryOperation(
+                    () -> WithRetryHelper.withRetry(
                             context,
                             "my-op",
                             (ctx, attempt) -> {
@@ -126,7 +126,7 @@ void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
                     .build();
 
             var callCount = new int[] {0};
-            var result = WithRetryHelper.retryOperation(
+            var result = WithRetryHelper.withRetry(
                     context,
                     "my-op",
                     (ctx, attempt) -> {
@@ -150,8 +150,7 @@ void nullContext_shouldThrow() {
                     .build();
 
             assertThrows(
-                    NullPointerException.class,
-                    () -> WithRetryHelper.retryOperation(null, "name", (ctx, a) -> "x", config));
+                    NullPointerException.class, () -> WithRetryHelper.withRetry(null, "name", (ctx, a) -> "x", config));
         }
 
         @Test
@@ -162,7 +161,7 @@ void nullName_shouldThrow() {
 
             assertThrows(
                     NullPointerException.class,
-                    () -> WithRetryHelper.retryOperation(context, null, (ctx, a) -> "x", config));
+                    () -> WithRetryHelper.withRetry(context, null, (ctx, a) -> "x", config));
         }
 
         @Test
@@ -171,15 +170,14 @@ void nullOperation_shouldThrow() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            assertThrows(
-                    NullPointerException.class, () -> WithRetryHelper.retryOperation(context, "name", null, config));
+            assertThrows(NullPointerException.class, () -> WithRetryHelper.withRetry(context, "name", null, config));
         }
 
         @Test
         void nullConfig_shouldThrow() {
             assertThrows(
                     NullPointerException.class,
-                    () -> WithRetryHelper.retryOperation(context, "name", (ctx, a) -> "x", null));
+                    () -> WithRetryHelper.withRetry(context, "name", (ctx, a) -> "x", null));
         }
 
         @Test
@@ -189,8 +187,8 @@ void operationReturnsNull() {
                     .wrapInChildContext(false)
                     .build();
 
-            var result = WithRetryHelper.retryOperation(
-                    context, "my-op", (WithRetry) (ctx, attempt) -> null, config);
+            var result =
+                    WithRetryHelper.withRetry(context, "my-op", (WithRetry) (ctx, attempt) -> null, config);
 
             assertNull(result);
         }
@@ -207,7 +205,7 @@ void successOnFirstAttempt() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var result = WithRetryHelper.retryOperation(context, (ctx, attempt) -> "anonymous-success", config);
+            var result = WithRetryHelper.withRetry(context, (ctx, attempt) -> "anonymous-success", config);
 
             assertEquals("anonymous-success", result);
             verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
@@ -220,7 +218,7 @@ void retriesWithAnonymousBackoffNames() {
                             attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
                     .build();
 
-            var result = WithRetryHelper.retryOperation(
+            var result = WithRetryHelper.withRetry(
                     context,
                     (ctx, attempt) -> {
                         if (attempt < 3) {
@@ -242,7 +240,7 @@ void neverWrapsInChildContext_evenWhenConfigSaysTrue() {
                     .wrapInChildContext(true) // should be ignored for anonymous form
                     .build();
 
-            WithRetryHelper.retryOperation(context, (ctx, attempt) -> "result", config);
+            WithRetryHelper.withRetry(context, (ctx, attempt) -> "result", config);
 
             verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
         }
@@ -256,7 +254,7 @@ void rethrowsOriginalException() {
             var original = new IllegalStateException("original error");
             var thrown = assertThrows(
                     IllegalStateException.class,
-                    () -> WithRetryHelper.retryOperation(
+                    () -> WithRetryHelper.withRetry(
                             context,
                             (ctx, attempt) -> {
                                 throw original;
@@ -272,8 +270,7 @@ void nullContext_shouldThrow() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            assertThrows(
-                    NullPointerException.class, () -> WithRetryHelper.retryOperation(null, (ctx, a) -> "x", config));
+            assertThrows(NullPointerException.class, () -> WithRetryHelper.withRetry(null, (ctx, a) -> "x", config));
         }
 
         @Test
@@ -284,13 +281,12 @@ void nullOperation_shouldThrow() {
 
             assertThrows(
                     NullPointerException.class,
-                    () -> WithRetryHelper.retryOperation(context, (WithRetry) null, config));
+                    () -> WithRetryHelper.withRetry(context, (WithRetry) null, config));
         }
 
         @Test
         void nullConfig_shouldThrow() {
-            assertThrows(
-                    NullPointerException.class, () -> WithRetryHelper.retryOperation(context, (ctx, a) -> "x", null));
+            assertThrows(NullPointerException.class, () -> WithRetryHelper.withRetry(context, (ctx, a) -> "x", null));
         }
 
         @Test
@@ -299,7 +295,7 @@ void operationReturnsNull() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var result = WithRetryHelper.retryOperation(context, (WithRetry) (ctx, attempt) -> null, config);
+            var result = WithRetryHelper.withRetry(context, (WithRetry) (ctx, attempt) -> null, config);
 
             assertNull(result);
         }
@@ -311,7 +307,7 @@ void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
                             (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail())
                     .build();
 
-            var result = WithRetryHelper.retryOperation(
+            var result = WithRetryHelper.withRetry(
                     context,
                     (ctx, attempt) -> {
                         if (attempt == 1) {
@@ -340,7 +336,7 @@ void passesCorrectAttemptNumberToOperation() {
                     .wrapInChildContext(false)
                     .build();
 
-            WithRetryHelper.retryOperation(
+            WithRetryHelper.withRetry(
                     context,
                     "track",
                     (ctx, attempt) -> {
@@ -371,7 +367,7 @@ void passesErrorToRetryStrategy() {
 
             assertThrows(
                     RuntimeException.class,
-                    () -> WithRetryHelper.retryOperation(
+                    () -> WithRetryHelper.withRetry(
                             context,
                             (ctx, attempt) -> {
                                 throw new RuntimeException("error-" + attempt);
@@ -390,7 +386,7 @@ void respectsCustomDelayFromRetryDecision() {
                     .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(attempt * 10L)))
                     .build();
 
-            WithRetryHelper.retryOperation(
+            WithRetryHelper.withRetry(
                     context,
                     (ctx, attempt) -> {
                         if (attempt <= 2) {
@@ -410,7 +406,7 @@ void passesContextToOperation() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            WithRetryHelper.retryOperation(
+            WithRetryHelper.withRetry(
                     context,
                     (ctx, attempt) -> {
                         assertSame(context, ctx);
@@ -433,7 +429,7 @@ void namedFormPassesChildContextToOperation_whenWrapped() {
                     .wrapInChildContext(true)
                     .build();
 
-            WithRetryHelper.retryOperation(
+            WithRetryHelper.withRetry(
                     context,
                     "wrapped",
                     (ctx, attempt) -> {
@@ -451,7 +447,7 @@ void rethrowsLastExceptionWhenAllRetriesExhausted() {
 
             var thrown = assertThrows(
                     RuntimeException.class,
-                    () -> WithRetryHelper.retryOperation(
+                    () -> WithRetryHelper.withRetry(
                             context,
                             (ctx, attempt) -> {
                                 throw new RuntimeException("attempt-" + attempt);
@@ -470,7 +466,7 @@ void propagatesSuspendExecutionExceptionWithoutRetrying() {
 
             assertThrows(
                     SuspendExecutionException.class,
-                    () -> WithRetryHelper.retryOperation(
+                    () -> WithRetryHelper.withRetry(
                             context,
                             (ctx, attempt) -> {
                                 throw new SuspendExecutionException();
@@ -489,7 +485,7 @@ void propagatesUnrecoverableDurableExecutionExceptionWithoutRetrying() {
 
             assertThrows(
                     UnrecoverableDurableExecutionException.class,
-                    () -> WithRetryHelper.retryOperation(
+                    () -> WithRetryHelper.withRetry(
                             context,
                             (ctx, attempt) -> {
                                 throw new UnrecoverableDurableExecutionException(
@@ -510,7 +506,7 @@ void propagatesSuspendExecutionExceptionOnLaterAttempt() {
 
             assertThrows(
                     SuspendExecutionException.class,
-                    () -> WithRetryHelper.retryOperation(
+                    () -> WithRetryHelper.withRetry(
                             context,
                             (ctx, attempt) -> {
                                 if (attempt == 1) {
@@ -533,7 +529,7 @@ void propagatesUnrecoverableDurableExecutionExceptionOnLaterAttempt() {
 
             assertThrows(
                     UnrecoverableDurableExecutionException.class,
-                    () -> WithRetryHelper.retryOperation(
+                    () -> WithRetryHelper.withRetry(
                             context,
                             (ctx, attempt) -> {
                                 if (attempt == 1) {
@@ -559,7 +555,7 @@ void preservesCheckedExceptionSubclassType() {
 
             var thrown = assertThrows(
                     SerDesException.class,
-                    () -> WithRetryHelper.retryOperation(
+                    () -> WithRetryHelper.withRetry(
                             context,
                             (ctx, attempt) -> {
                                 throw original;

From d901827d2ceb1732d9d4f1265b41abede43b42e6 Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Wed, 29 Apr 2026 14:54:07 -0700
Subject: [PATCH 07/19] chore: added WithRetryAsync functions

---
 .../lambda/durable/util/WithRetryHelper.java  |  79 +++++-
 .../durable/util/WithRetryHelperTest.java     | 229 +++++++++++++++++-
 2 files changed, 289 insertions(+), 19 deletions(-)

diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java
index dbbb0a05d..237cd9530 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java
@@ -5,6 +5,7 @@
 import java.time.Duration;
 import java.util.Objects;
 import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.DurableFuture;
 import software.amazon.lambda.durable.TypeToken;
 import software.amazon.lambda.durable.config.WithRetryConfig;
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
@@ -52,6 +53,24 @@
  *         .build()
  * );
  * }

+ *
+ * 
Usage — async form returning DurableFuture

+ *
+ * 
{@code
+ * DurableFuture future = WithRetryHelper.withRetryAsync(
+ *     context,
+ *     "approval",
+ *     (ctx, attempt) -> ctx.waitForCallback(
+ *         "approval-" + attempt, String.class,
+ *         (callbackId, stepCtx) -> sendApprovalEmail(approverEmail, callbackId)
+ *     ),
+ *     WithRetryConfig.builder()
+ *         .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
+ *         .build()
+ * );
+ * // ... do other work ...
+ * var result = future.get();
+ * }

  */
 public final class WithRetryHelper {
 
@@ -64,48 +83,88 @@ private WithRetryHelper() {
     }
 
     /**
-     * Named form — wraps the retry loop in {@code runInChildContext} by default so all attempts are grouped under a
-     * single named operation in execution history.
+     * Named async form — wraps the retry loop in {@code runInChildContextAsync} by default so all attempts are grouped
+     * under a single named operation in execution history, and returns a {@link DurableFuture} that can be composed or
+     * blocked on.
      *
      * 
The child-context wrapping can be disabled via {@link WithRetryConfig.Builder#wrapInChildContext(boolean)}.
+     * When disabled, the retry loop executes immediately and the returned future is already completed.
      *
      * @param  the result type
      * @param context the durable context
      * @param name operation name (used for child context and backoff wait names)
      * @param operation the retryable operation — receives the context and 1-based attempt number
      * @param config retry configuration including the retry strategy
-     * @return the operation result
+     * @return a future representing the operation result
      */
     @SuppressWarnings("unchecked")
-    public static  T withRetry(DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
+    public static  DurableFuture withRetryAsync(
+            DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
         Objects.requireNonNull(context, "context cannot be null");
         Objects.requireNonNull(name, "name cannot be null");
         Objects.requireNonNull(operation, "operation cannot be null");
         Objects.requireNonNull(config, "config cannot be null");
 
         if (config.wrapInChildContext()) {
-            return (T) context.runInChildContext(
+            return (DurableFuture) context.runInChildContextAsync(
                     name, new TypeToken() {}, childCtx -> executeRetryLoop(childCtx, name, operation, config));
         }
-        return executeRetryLoop(context, name, operation, config);
+        return new CompletedDurableFuture<>(executeRetryLoop(context, name, operation, config));
     }
 
     /**
-     * Anonymous form — runs the retry loop directly in the caller's context. No child-context wrapping is applied
-     * regardless of the {@code wrapInChildContext} config setting.
+     * Named sync form — wraps the retry loop in {@code runInChildContext} by default so all attempts are grouped under
+     * a single named operation in execution history, and blocks until the result is available.
+     *
+     * 
Equivalent to {@code withRetryAsync(context, name, operation, config).get()}.
      *
      * @param  the result type
      * @param context the durable context
+     * @param name operation name (used for child context and backoff wait names)
      * @param operation the retryable operation — receives the context and 1-based attempt number
      * @param config retry configuration including the retry strategy
      * @return the operation result
      */
-    public static  T withRetry(DurableContext context, WithRetry operation, WithRetryConfig config) {
+    public static  T withRetry(DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
+        return withRetryAsync(context, name, operation, config).get();
+    }
+
+    /**
+     * Anonymous async form — runs the retry loop directly in the caller's context and returns a {@link DurableFuture}.
+     * No child-context wrapping is applied regardless of the {@code wrapInChildContext} config setting.
+     *
+     * 
Because the anonymous form executes the retry loop inline (no child context), the returned future is always
+     * already completed.
+     *
+     * @param  the result type
+     * @param context the durable context
+     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param config retry configuration including the retry strategy
+     * @return a future representing the operation result
+     */
+    public static  DurableFuture withRetryAsync(
+            DurableContext context, WithRetry operation, WithRetryConfig config) {
         Objects.requireNonNull(context, "context cannot be null");
         Objects.requireNonNull(operation, "operation cannot be null");
         Objects.requireNonNull(config, "config cannot be null");
 
-        return executeRetryLoop(context, null, operation, config);
+        return new CompletedDurableFuture<>(executeRetryLoop(context, null, operation, config));
+    }
+
+    /**
+     * Anonymous sync form — runs the retry loop directly in the caller's context. No child-context wrapping is applied
+     * regardless of the {@code wrapInChildContext} config setting.
+     *
+     * 
Equivalent to {@code withRetryAsync(context, operation, config).get()}.
+     *
+     * @param  the result type
+     * @param context the durable context
+     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param config retry configuration including the retry strategy
+     * @return the operation result
+     */
+    public static  T withRetry(DurableContext context, WithRetry operation, WithRetryConfig config) {
+        return withRetryAsync(context, operation, config).get();
     }
 
     /**
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
index c6cd81fec..31b917b23 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
@@ -13,6 +13,7 @@
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.DurableFuture;
 import software.amazon.lambda.durable.TypeToken;
 import software.amazon.lambda.durable.config.WithRetryConfig;
 import software.amazon.lambda.durable.exception.SerDesException;
@@ -37,11 +38,11 @@ class NamedForm {
 
         @Test
         void successOnFirstAttempt_wrapsInChildContext() {
-            // runInChildContext should be called; delegate to the function immediately
-            when(context.runInChildContext(eq("my-op"), any(TypeToken.class), any()))
+            // runInChildContextAsync should be called; delegate to the function immediately
+            when(context.runInChildContextAsync(eq("my-op"), any(TypeToken.class), any()))
                     .thenAnswer(invocation -> {
                         Function func = invocation.getArgument(2);
-                        return func.apply(context);
+                        return new CompletedDurableFuture<>(func.apply(context));
                     });
 
             var config = WithRetryConfig.builder()
@@ -51,7 +52,7 @@ void successOnFirstAttempt_wrapsInChildContext() {
             var result = WithRetryHelper.withRetry(context, "my-op", (ctx, attempt) -> "success", config);
 
             assertEquals("success", result);
-            verify(context).runInChildContext(eq("my-op"), any(TypeToken.class), any());
+            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
         }
 
         @Test
@@ -64,7 +65,7 @@ void successOnFirstAttempt_noChildContext_whenDisabled() {
             var result = WithRetryHelper.withRetry(context, "my-op", (ctx, attempt) -> "direct", config);
 
             assertEquals("direct", result);
-            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+            verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
         }
 
         @Test
@@ -208,7 +209,7 @@ void successOnFirstAttempt() {
             var result = WithRetryHelper.withRetry(context, (ctx, attempt) -> "anonymous-success", config);
 
             assertEquals("anonymous-success", result);
-            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+            verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
         }
 
         @Test
@@ -242,7 +243,7 @@ void neverWrapsInChildContext_evenWhenConfigSaysTrue() {
 
             WithRetryHelper.withRetry(context, (ctx, attempt) -> "result", config);
 
-            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+            verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
         }
 
         @Test
@@ -322,6 +323,216 @@ void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
         }
     }
 
+    // --- Async form tests ---
+
+    @Nested
+    class AsyncForm {
+
+        @Test
+        void namedAsyncReturnsDurableFuture() {
+            when(context.runInChildContextAsync(eq("my-op"), any(TypeToken.class), any()))
+                    .thenAnswer(invocation -> {
+                        Function func = invocation.getArgument(2);
+                        return new CompletedDurableFuture<>(func.apply(context));
+                    });
+
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            DurableFuture future =
+                    WithRetryHelper.withRetryAsync(context, "my-op", (ctx, attempt) -> "async-result", config);
+
+            assertNotNull(future);
+            assertEquals("async-result", future.get());
+        }
+
+        @Test
+        void namedAsyncWithoutChildContext_returnsCompletedFuture() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(false)
+                    .build();
+
+            DurableFuture future =
+                    WithRetryHelper.withRetryAsync(context, "my-op", (ctx, attempt) -> "direct-async", config);
+
+            assertNotNull(future);
+            assertInstanceOf(CompletedDurableFuture.class, future);
+            assertEquals("direct-async", future.get());
+            verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
+        }
+
+        @Test
+        void anonymousAsyncReturnsDurableFuture() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            DurableFuture future =
+                    WithRetryHelper.withRetryAsync(context, (ctx, attempt) -> "anon-async", config);
+
+            assertNotNull(future);
+            assertInstanceOf(CompletedDurableFuture.class, future);
+            assertEquals("anon-async", future.get());
+        }
+
+        @Test
+        void anonymousAsyncNeverWrapsInChildContext() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(true)
+                    .build();
+
+            WithRetryHelper.withRetryAsync(context, (ctx, attempt) -> "result", config);
+
+            verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
+        }
+
+        @Test
+        void namedAsyncRetriesWithBackoff() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy((error, attempt) ->
+                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
+                    .wrapInChildContext(false)
+                    .build();
+
+            var callCount = new int[] {0};
+            DurableFuture future = WithRetryHelper.withRetryAsync(
+                    context,
+                    "my-op",
+                    (ctx, attempt) -> {
+                        callCount[0]++;
+                        if (attempt < 3) {
+                            throw new RuntimeException("fail-" + attempt);
+                        }
+                        return "success-on-3";
+                    },
+                    config);
+
+            assertEquals("success-on-3", future.get());
+            assertEquals(3, callCount[0]);
+            verify(context).wait("my-op-backoff-1", Duration.ofSeconds(5));
+            verify(context).wait("my-op-backoff-2", Duration.ofSeconds(5));
+        }
+
+        @Test
+        void anonymousAsyncRetriesWithBackoff() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy((error, attempt) ->
+                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
+                    .build();
+
+            DurableFuture future = WithRetryHelper.withRetryAsync(
+                    context,
+                    (ctx, attempt) -> {
+                        if (attempt < 3) {
+                            throw new RuntimeException("fail");
+                        }
+                        return "done";
+                    },
+                    config);
+
+            assertEquals("done", future.get());
+            verify(context).wait("retry-backoff-1", Duration.ofSeconds(2));
+            verify(context).wait("retry-backoff-2", Duration.ofSeconds(2));
+        }
+
+        @Test
+        void namedAsyncNullContext_shouldThrow() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> WithRetryHelper.withRetryAsync(null, "name", (ctx, a) -> "x", config));
+        }
+
+        @Test
+        void namedAsyncNullName_shouldThrow() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> WithRetryHelper.withRetryAsync(context, null, (ctx, a) -> "x", config));
+        }
+
+        @Test
+        void namedAsyncNullOperation_shouldThrow() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class, () -> WithRetryHelper.withRetryAsync(context, "name", null, config));
+        }
+
+        @Test
+        void namedAsyncNullConfig_shouldThrow() {
+            assertThrows(
+                    NullPointerException.class,
+                    () -> WithRetryHelper.withRetryAsync(context, "name", (ctx, a) -> "x", null));
+        }
+
+        @Test
+        void anonymousAsyncNullContext_shouldThrow() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class, () -> WithRetryHelper.withRetryAsync(null, (ctx, a) -> "x", config));
+        }
+
+        @Test
+        void anonymousAsyncNullOperation_shouldThrow() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> WithRetryHelper.withRetryAsync(context, (WithRetry) null, config));
+        }
+
+        @Test
+        void anonymousAsyncNullConfig_shouldThrow() {
+            assertThrows(
+                    NullPointerException.class, () -> WithRetryHelper.withRetryAsync(context, (ctx, a) -> "x", null));
+        }
+
+        @Test
+        void asyncOperationReturnsNull() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(false)
+                    .build();
+
+            DurableFuture future = WithRetryHelper.withRetryAsync(
+                    context, "my-op", (WithRetry) (ctx, attempt) -> null, config);
+
+            assertNull(future.get());
+        }
+
+        @Test
+        void withRetryDelegatesToWithRetryAsync() {
+            // Verify that withRetry (sync) produces the same result as withRetryAsync().get()
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(false)
+                    .build();
+
+            var syncResult = WithRetryHelper.withRetry(context, "op", (ctx, attempt) -> "value", config);
+            var asyncResult = WithRetryHelper.withRetryAsync(context, "op", (ctx, attempt) -> "value", config)
+                    .get();
+
+            assertEquals(syncResult, asyncResult);
+        }
+    }
+
     // --- Retry behavior tests ---
 
     @Nested
@@ -418,10 +629,10 @@ void passesContextToOperation() {
         @Test
         void namedFormPassesChildContextToOperation_whenWrapped() {
             var childContext = mock(DurableContext.class);
-            when(context.runInChildContext(eq("wrapped"), any(TypeToken.class), any()))
+            when(context.runInChildContextAsync(eq("wrapped"), any(TypeToken.class), any()))
                     .thenAnswer(invocation -> {
                         Function func = invocation.getArgument(2);
-                        return func.apply(childContext);
+                        return new CompletedDurableFuture<>(func.apply(childContext));
                     });
 
             var config = WithRetryConfig.builder()

From 23ed3df4f27af2584cd1e26e6376cfcbbe722437 Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Wed, 22 Apr 2026 15:04:25 -0700
Subject: [PATCH 08/19] feat: first attempt at retryableOperation

---
 .../durable/config/RetryOperationConfig.java  | 103 +++++
 .../durable/util/RetryOperationHelper.java    | 135 ++++++
 .../durable/util/RetryableOperation.java      |  26 ++
 .../config/RetryOperationConfigTest.java      |  80 ++++
 .../util/RetryOperationHelperTest.java        | 420 ++++++++++++++++++
 .../durable/util/RetryableOperationTest.java  |  52 +++
 6 files changed, 816 insertions(+)
 create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java
 create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
 create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java
 create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
 create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
 create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java

diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java
new file mode 100644
index 000000000..696d1d253
--- /dev/null
+++ b/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java
@@ -0,0 +1,103 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.config;
+
+import software.amazon.lambda.durable.retry.RetryStrategy;
+
+/**
+ * Configuration for {@link software.amazon.lambda.durable.util.RetryOperationHelper#retryOperation}.
+ *
+ * 
Uses the same {@link RetryStrategy} shape that developers already know from {@link StepConfig}, so there are zero
+ * new retry concepts to learn.
+ */
+public class RetryOperationConfig {
+    private final RetryStrategy retryStrategy;
+    private final boolean wrapInChildContext;
+
+    private RetryOperationConfig(Builder builder) {
+        this.retryStrategy = builder.retryStrategy;
+        this.wrapInChildContext = builder.wrapInChildContext;
+    }
+
+    /**
+     * Returns the retry strategy. Same type as {@link StepConfig#retryStrategy()}.
+     *
+     * @return the retry strategy, never null
+     */
+    public RetryStrategy retryStrategy() {
+        return retryStrategy;
+    }
+
+    /**
+     * Whether to wrap the retry loop in {@code runInChildContext} so all attempts are grouped under a single named
+     * operation in execution history. Only applies when a name is provided to the named form of {@code retryOperation}.
+     * Defaults to {@code true}.
+     *
+     * @return true if child-context wrapping is enabled
+     */
+    public boolean wrapInChildContext() {
+        return wrapInChildContext;
+    }
+
+    /**
+     * Creates a new builder for {@code RetryOperationConfig}.
+     *
+     * @return a new builder instance
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /** Builder for creating {@link RetryOperationConfig} instances. */
+    public static class Builder {
+        private RetryStrategy retryStrategy;
+        private boolean wrapInChildContext = true;
+
+        private Builder() {}
+
+        /**
+         * Sets the retry strategy. Required.
+         *
+         * 
Reuses the exact same {@link RetryStrategy} interface from {@link StepConfig}. All existing factory
+         * methods ({@link software.amazon.lambda.durable.retry.RetryStrategies#exponentialBackoff},
+         * {@link software.amazon.lambda.durable.retry.RetryStrategies#fixedDelay}, presets, and custom lambdas) work
+         * without modification.
+         *
+         * @param retryStrategy the retry strategy to use
+         * @return this builder for method chaining
+         */
+        public Builder retryStrategy(RetryStrategy retryStrategy) {
+            this.retryStrategy = retryStrategy;
+            return this;
+        }
+
+        /**
+         * Controls whether the retry loop is wrapped in a child context. Only meaningful for the named form of
+         * {@code retryOperation}. Defaults to {@code true}.
+         *
+         * 
When {@code true}, all attempts and backoff waits are grouped under a single named operation in execution
+         * history, providing a cleaner view and isolated operation ID space. Set to {@code false} to flatten attempts
+         * into the parent context.
+         *
+         * @param wrapInChildContext whether to wrap in a child context
+         * @return this builder for method chaining
+         */
+        public Builder wrapInChildContext(boolean wrapInChildContext) {
+            this.wrapInChildContext = wrapInChildContext;
+            return this;
+        }
+
+        /**
+         * Builds the {@link RetryOperationConfig} instance.
+         *
+         * @return a new config with the configured options
+         * @throws IllegalArgumentException if retryStrategy is not set
+         */
+        public RetryOperationConfig build() {
+            if (retryStrategy == null) {
+                throw new IllegalArgumentException("retryStrategy is required");
+            }
+            return new RetryOperationConfig(this);
+        }
+    }
+}
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
new file mode 100644
index 000000000..c63f0aad0
--- /dev/null
+++ b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
@@ -0,0 +1,135 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.util;
+
+import java.time.Duration;
+import java.util.Objects;
+import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.TypeToken;
+import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.retry.RetryDecision;
+
+/**
+ * Replay-safe retry loop for any durable operation.
+ *
+ * 
Provides the same retry-with-backoff pattern that {@code context.step()} has built in, but for operations that
+ * cannot live inside a step ({@code waitForCallback}, {@code invoke}, {@code waitForCondition}, etc.).
+ *
+ * 
Every side-effect in the loop is a durable operation, so the loop is replay-safe by construction. On replay,
+ * completed operations return cached results instantly and the loop fast-forwards to the current attempt.
+ *
+ * 
Usage — callback retry

+ *
+ * 
{@code
+ * var result = RetryOperationHelper.retryOperation(
+ *     context,
+ *     "approval",
+ *     (ctx, attempt) -> ctx.waitForCallback(
+ *         "approval-" + attempt,
+ *         String.class,
+ *         (callbackId, stepCtx) -> sendApprovalEmail(approverEmail, callbackId)
+ *     ),
+ *     RetryOperationConfig.builder()
+ *         .retryStrategy(RetryStrategies.exponentialBackoff(
+ *             3, Duration.ofSeconds(2), Duration.ofSeconds(30), 2.0, JitterStrategy.FULL))
+ *         .build()
+ * );
+ * }

+ *
+ * 
Usage — invoke retry (anonymous form)

+ *
+ * 
{@code
+ * var result = RetryOperationHelper.retryOperation(
+ *     context,
+ *     (ctx, attempt) -> ctx.invoke(
+ *         "charge-" + attempt, paymentFnArn, new ChargeRequest(orderId), String.class),
+ *     RetryOperationConfig.builder()
+ *         .retryStrategy((err, att) -> att < 3
+ *             ? RetryDecision.retry(Duration.ofSeconds(1))
+ *             : RetryDecision.fail())
+ *         .build()
+ * );
+ * }

+ */
+public final class RetryOperationHelper {
+
+    private static final Duration DEFAULT_BACKOFF_DELAY = Duration.ofSeconds(1);
+    private static final String BACKOFF_SUFFIX = "-backoff-";
+    private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-";
+
+    private RetryOperationHelper() {
+        // utility class
+    }
+
+    /**
+     * Named form — wraps the retry loop in {@code runInChildContext} by default so all attempts are grouped under a
+     * single named operation in execution history.
+     *
+     * 
The child-context wrapping can be disabled via
+     * {@link RetryOperationConfig.Builder#wrapInChildContext(boolean)}.
+     *
+     * @param  the result type
+     * @param context the durable context
+     * @param name operation name (used for child context and backoff wait names)
+     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param config retry configuration including the retry strategy
+     * @return the operation result
+     */
+    @SuppressWarnings("unchecked")
+    public static  T retryOperation(
+            DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) {
+        Objects.requireNonNull(context, "context cannot be null");
+        Objects.requireNonNull(name, "name cannot be null");
+        Objects.requireNonNull(operation, "operation cannot be null");
+        Objects.requireNonNull(config, "config cannot be null");
+
+        if (config.wrapInChildContext()) {
+            return (T) context.runInChildContext(
+                    name, new TypeToken() {}, childCtx -> executeRetryLoop(childCtx, name, operation, config));
+        }
+        return executeRetryLoop(context, name, operation, config);
+    }
+
+    /**
+     * Anonymous form — runs the retry loop directly in the caller's context. No child-context wrapping is applied
+     * regardless of the {@code wrapInChildContext} config setting.
+     *
+     * @param  the result type
+     * @param context the durable context
+     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param config retry configuration including the retry strategy
+     * @return the operation result
+     */
+    public static  T retryOperation(
+            DurableContext context, RetryableOperation operation, RetryOperationConfig config) {
+        Objects.requireNonNull(context, "context cannot be null");
+        Objects.requireNonNull(operation, "operation cannot be null");
+        Objects.requireNonNull(config, "config cannot be null");
+
+        return executeRetryLoop(context, null, operation, config);
+    }
+
+    /**
+     * Core retry loop. Replay-safe because every side-effect is a durable operation: the user's operation calls durable
+     * primitives, and backoff uses {@code context.wait()}.
+     */
+    private static  T executeRetryLoop(
+            DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) {
+        var attempt = 1;
+        while (true) {
+            try {
+                return operation.execute(context, attempt);
+            } catch (Exception e) {
+                RetryDecision decision = config.retryStrategy().makeRetryDecision(e, attempt);
+                if (!decision.shouldRetry()) {
+                    throw e;
+                }
+
+                var delay = decision.delay().isZero() ? DEFAULT_BACKOFF_DELAY : decision.delay();
+                var waitName = name != null ? name + BACKOFF_SUFFIX + attempt : ANONYMOUS_BACKOFF_PREFIX + attempt;
+                context.wait(waitName, delay);
+                attempt++;
+            }
+        }
+    }
+}
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java
new file mode 100644
index 000000000..d6b2496c1
--- /dev/null
+++ b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java
@@ -0,0 +1,26 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.util;
+
+import software.amazon.lambda.durable.DurableContext;
+
+/**
+ * A durable operation that can be retried end-to-end by {@link RetryOperationHelper}.
+ *
+ * 
Receives the durable context and the 1-based attempt number so callers can generate unique operation names per
+ * attempt (e.g., {@code "approval-" + attempt}).
+ *
+ * @param  the result type
+ */
+@FunctionalInterface
+public interface RetryableOperation {
+
+    /**
+     * Executes the durable operation.
+     *
+     * @param context the durable context to use for durable operations
+     * @param attempt the current attempt number (1-based: first attempt is 1)
+     * @return the operation result
+     */
+    T execute(DurableContext context, int attempt);
+}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
new file mode 100644
index 000000000..726da5b61
--- /dev/null
+++ b/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
@@ -0,0 +1,80 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.config;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.lambda.durable.retry.RetryDecision;
+import software.amazon.lambda.durable.retry.RetryStrategies;
+
+class RetryOperationConfigTest {
+
+    @Test
+    void builderWithRetryStrategy() {
+        var strategy = RetryStrategies.Presets.DEFAULT;
+
+        var config = RetryOperationConfig.builder().retryStrategy(strategy).build();
+
+        assertEquals(strategy, config.retryStrategy());
+    }
+
+    @Test
+    void builderWithoutRetryStrategy_shouldThrow() {
+        var exception = assertThrows(IllegalArgumentException.class, () -> RetryOperationConfig.builder()
+                .build());
+
+        assertEquals("retryStrategy is required", exception.getMessage());
+    }
+
+    @Test
+    void wrapInChildContext_defaultsToTrue() {
+        var config = RetryOperationConfig.builder()
+                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                .build();
+
+        assertTrue(config.wrapInChildContext());
+    }
+
+    @Test
+    void wrapInChildContext_canBeSetToFalse() {
+        var config = RetryOperationConfig.builder()
+                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                .wrapInChildContext(false)
+                .build();
+
+        assertFalse(config.wrapInChildContext());
+    }
+
+    @Test
+    void wrapInChildContext_canBeSetToTrue() {
+        var config = RetryOperationConfig.builder()
+                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                .wrapInChildContext(true)
+                .build();
+
+        assertTrue(config.wrapInChildContext());
+    }
+
+    @Test
+    void builderChaining() {
+        var strategy = RetryStrategies.Presets.DEFAULT;
+
+        var config = RetryOperationConfig.builder()
+                .retryStrategy(strategy)
+                .wrapInChildContext(false)
+                .build();
+
+        assertEquals(strategy, config.retryStrategy());
+        assertFalse(config.wrapInChildContext());
+    }
+
+    @Test
+    void builderWithCustomLambdaRetryStrategy() {
+        var config = RetryOperationConfig.builder()
+                .retryStrategy((error, attempt) -> RetryDecision.fail())
+                .build();
+
+        assertNotNull(config.retryStrategy());
+    }
+}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
new file mode 100644
index 000000000..9558f3c83
--- /dev/null
+++ b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
@@ -0,0 +1,420 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.util;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.ArgumentMatchers.*;
+import static org.mockito.Mockito.*;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.function.Function;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.TypeToken;
+import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.retry.RetryDecision;
+import software.amazon.lambda.durable.retry.RetryStrategies;
+
+class RetryOperationHelperTest {
+
+    private DurableContext context;
+
+    @BeforeEach
+    void setUp() {
+        context = mock(DurableContext.class);
+    }
+
+    // --- Named form tests ---
+
+    @Nested
+    class NamedForm {
+
+        @Test
+        void successOnFirstAttempt_wrapsInChildContext() {
+            // runInChildContext should be called; delegate to the function immediately
+            when(context.runInChildContext(eq("my-op"), any(TypeToken.class), any()))
+                    .thenAnswer(invocation -> {
+                        Function func = invocation.getArgument(2);
+                        return func.apply(context);
+                    });
+
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            var result = RetryOperationHelper.retryOperation(context, "my-op", (ctx, attempt) -> "success", config);
+
+            assertEquals("success", result);
+            verify(context).runInChildContext(eq("my-op"), any(TypeToken.class), any());
+        }
+
+        @Test
+        void successOnFirstAttempt_noChildContext_whenDisabled() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(false)
+                    .build();
+
+            var result = RetryOperationHelper.retryOperation(context, "my-op", (ctx, attempt) -> "direct", config);
+
+            assertEquals("direct", result);
+            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+        }
+
+        @Test
+        void retriesWithBackoffWaits_namedForm() {
+            // Disable child context wrapping so we can directly verify wait calls
+            var callCount = new int[] {0};
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) ->
+                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
+                    .wrapInChildContext(false)
+                    .build();
+
+            var result = RetryOperationHelper.retryOperation(
+                    context,
+                    "my-op",
+                    (ctx, attempt) -> {
+                        callCount[0]++;
+                        if (attempt < 3) {
+                            throw new RuntimeException("fail-" + attempt);
+                        }
+                        return "success-on-3";
+                    },
+                    config);
+
+            assertEquals("success-on-3", result);
+            assertEquals(3, callCount[0]);
+            verify(context).wait("my-op-backoff-1", Duration.ofSeconds(5));
+            verify(context).wait("my-op-backoff-2", Duration.ofSeconds(5));
+            verify(context, times(2)).wait(anyString(), any(Duration.class));
+        }
+
+        @Test
+        void rethrowsWhenRetryStrategyReturnsFail() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(false)
+                    .build();
+
+            var exception = assertThrows(
+                    RuntimeException.class,
+                    () -> RetryOperationHelper.retryOperation(
+                            context,
+                            "my-op",
+                            (ctx, attempt) -> {
+                                throw new RuntimeException("terminal");
+                            },
+                            config));
+
+            assertEquals("terminal", exception.getMessage());
+            verify(context, never()).wait(anyString(), any(Duration.class));
+        }
+
+        @Test
+        void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(
+                            (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail())
+                    .wrapInChildContext(false)
+                    .build();
+
+            var callCount = new int[] {0};
+            var result = RetryOperationHelper.retryOperation(
+                    context,
+                    "my-op",
+                    (ctx, attempt) -> {
+                        callCount[0]++;
+                        if (attempt == 1) {
+                            throw new RuntimeException("fail");
+                        }
+                        return "ok";
+                    },
+                    config);
+
+            assertEquals("ok", result);
+            // Zero delay should be replaced with 1-second default
+            verify(context).wait("my-op-backoff-1", Duration.ofSeconds(1));
+        }
+
+        @Test
+        void nullContext_shouldThrow() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(null, "name", (ctx, a) -> "x", config));
+        }
+
+        @Test
+        void nullName_shouldThrow() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(context, null, (ctx, a) -> "x", config));
+        }
+
+        @Test
+        void nullOperation_shouldThrow() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(context, "name", null, config));
+        }
+
+        @Test
+        void nullConfig_shouldThrow() {
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(context, "name", (ctx, a) -> "x", null));
+        }
+    }
+
+    // --- Anonymous form tests ---
+
+    @Nested
+    class AnonymousForm {
+
+        @Test
+        void successOnFirstAttempt() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            var result = RetryOperationHelper.retryOperation(context, (ctx, attempt) -> "anonymous-success", config);
+
+            assertEquals("anonymous-success", result);
+            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+        }
+
+        @Test
+        void retriesWithAnonymousBackoffNames() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) ->
+                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
+                    .build();
+
+            var result = RetryOperationHelper.retryOperation(
+                    context,
+                    (ctx, attempt) -> {
+                        if (attempt < 3) {
+                            throw new RuntimeException("fail");
+                        }
+                        return "done";
+                    },
+                    config);
+
+            assertEquals("done", result);
+            verify(context).wait("retry-backoff-1", Duration.ofSeconds(2));
+            verify(context).wait("retry-backoff-2", Duration.ofSeconds(2));
+        }
+
+        @Test
+        void neverWrapsInChildContext_evenWhenConfigSaysTrue() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(true) // should be ignored for anonymous form
+                    .build();
+
+            RetryOperationHelper.retryOperation(context, (ctx, attempt) -> "result", config);
+
+            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+        }
+
+        @Test
+        void rethrowsOriginalException() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            var original = new IllegalStateException("original error");
+            var thrown = assertThrows(
+                    IllegalStateException.class,
+                    () -> RetryOperationHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                throw original;
+                            },
+                            config));
+
+            assertSame(original, thrown);
+        }
+
+        @Test
+        void nullContext_shouldThrow() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(null, (ctx, a) -> "x", config));
+        }
+
+        @Test
+        void nullOperation_shouldThrow() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(context, (RetryableOperation) null, config));
+        }
+
+        @Test
+        void nullConfig_shouldThrow() {
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(context, (ctx, a) -> "x", null));
+        }
+    }
+
+    // --- Retry behavior tests ---
+
+    @Nested
+    class RetryBehavior {
+
+        @Test
+        void passesCorrectAttemptNumberToOperation() {
+            var attempts = new ArrayList();
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) ->
+                            attempt < 4 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
+                    .wrapInChildContext(false)
+                    .build();
+
+            RetryOperationHelper.retryOperation(
+                    context,
+                    "track",
+                    (ctx, attempt) -> {
+                        attempts.add(attempt);
+                        if (attempt < 4) {
+                            throw new RuntimeException("not yet");
+                        }
+                        return "done";
+                    },
+                    config);
+
+            assertEquals(4, attempts.size());
+            assertEquals(1, attempts.get(0));
+            assertEquals(2, attempts.get(1));
+            assertEquals(3, attempts.get(2));
+            assertEquals(4, attempts.get(3));
+        }
+
+        @Test
+        void passesErrorToRetryStrategy() {
+            var errors = new ArrayList();
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) -> {
+                        errors.add(error);
+                        return attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail();
+                    })
+                    .build();
+
+            assertThrows(
+                    RuntimeException.class,
+                    () -> RetryOperationHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                throw new RuntimeException("error-" + attempt);
+                            },
+                            config));
+
+            assertEquals(3, errors.size());
+            assertEquals("error-1", errors.get(0).getMessage());
+            assertEquals("error-2", errors.get(1).getMessage());
+            assertEquals("error-3", errors.get(2).getMessage());
+        }
+
+        @Test
+        void respectsCustomDelayFromRetryDecision() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(attempt * 10L)))
+                    .build();
+
+            RetryOperationHelper.retryOperation(
+                    context,
+                    (ctx, attempt) -> {
+                        if (attempt <= 2) {
+                            throw new RuntimeException("fail");
+                        }
+                        return "ok";
+                    },
+                    config);
+
+            verify(context).wait("retry-backoff-1", Duration.ofSeconds(10));
+            verify(context).wait("retry-backoff-2", Duration.ofSeconds(20));
+        }
+
+        @Test
+        void passesContextToOperation() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            RetryOperationHelper.retryOperation(
+                    context,
+                    (ctx, attempt) -> {
+                        assertSame(context, ctx);
+                        return "verified";
+                    },
+                    config);
+        }
+
+        @Test
+        void namedFormPassesChildContextToOperation_whenWrapped() {
+            var childContext = mock(DurableContext.class);
+            when(context.runInChildContext(eq("wrapped"), any(TypeToken.class), any()))
+                    .thenAnswer(invocation -> {
+                        Function func = invocation.getArgument(2);
+                        return func.apply(childContext);
+                    });
+
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(true)
+                    .build();
+
+            RetryOperationHelper.retryOperation(
+                    context,
+                    "wrapped",
+                    (ctx, attempt) -> {
+                        assertSame(childContext, ctx);
+                        return "verified";
+                    },
+                    config);
+        }
+
+        @Test
+        void rethrowsLastExceptionWhenAllRetriesExhausted() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
+                    .build();
+
+            var thrown = assertThrows(
+                    RuntimeException.class,
+                    () -> RetryOperationHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                throw new RuntimeException("attempt-" + attempt);
+                            },
+                            config));
+
+            // The last attempt's exception is rethrown
+            assertEquals("attempt-3", thrown.getMessage());
+        }
+    }
+}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java
new file mode 100644
index 000000000..f5893a37c
--- /dev/null
+++ b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java
@@ -0,0 +1,52 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.util;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.lambda.durable.DurableContext;
+
+class RetryableOperationTest {
+
+    @Test
+    void canBeImplementedAsLambda() {
+        RetryableOperation operation = (ctx, attempt) -> "result-" + attempt;
+        var context = mock(DurableContext.class);
+
+        assertEquals("result-1", operation.execute(context, 1));
+        assertEquals("result-3", operation.execute(context, 3));
+    }
+
+    @Test
+    void receivesContextAndAttempt() {
+        var context = mock(DurableContext.class);
+        RetryableOperation operation = (ctx, attempt) -> {
+            assertSame(context, ctx);
+            return "attempt-" + attempt;
+        };
+
+        assertEquals("attempt-1", operation.execute(context, 1));
+        assertEquals("attempt-2", operation.execute(context, 2));
+    }
+
+    @Test
+    void canThrowExceptions() {
+        RetryableOperation operation = (ctx, attempt) -> {
+            throw new RuntimeException("failed on attempt " + attempt);
+        };
+        var context = mock(DurableContext.class);
+
+        var exception = assertThrows(RuntimeException.class, () -> operation.execute(context, 1));
+        assertEquals("failed on attempt 1", exception.getMessage());
+    }
+
+    @Test
+    void canReturnNull() {
+        RetryableOperation operation = (ctx, attempt) -> null;
+        var context = mock(DurableContext.class);
+
+        assertNull(operation.execute(context, 1));
+    }
+}

From ac7697d0c0bbbb15933f70c9598f189100db9840 Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Wed, 22 Apr 2026 15:48:47 -0700
Subject: [PATCH 09/19] feat: add example tests and integration tests for new
 RetryableOperation util

---
 .../callback/RetryWaitForCallbackExample.java |  50 ++++
 .../examples/invoke/RetryInvokeExample.java   |  41 +++
 .../RetryWaitForCallbackExampleTest.java      | 148 ++++++++++
 .../invoke/RetryInvokeExampleTest.java        | 130 +++++++++
 .../durable/RetryInvokeIntegrationTest.java   | 198 ++++++++++++++
 .../RetryWaitForCallbackIntegrationTest.java  | 253 ++++++++++++++++++
 .../durable/util/RetryOperationHelper.java    |   8 +
 .../util/RetryOperationHelperTest.java        |  42 +++
 8 files changed, 870 insertions(+)
 create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
 create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
 create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExampleTest.java
 create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExampleTest.java
 create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
 create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java

diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
new file mode 100644
index 000000000..e16ca8f43
--- /dev/null
+++ b/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.RetryOperationConfig;
+import software.amazon.lambda.durable.examples.types.ApprovalRequest;
+import software.amazon.lambda.durable.retry.RetryDecision;
+import software.amazon.lambda.durable.util.RetryOperationHelper;
+
+/**
+ * Example demonstrating {@link RetryOperationHelper} with {@code context.waitForCallback}.
+ *
+ * 
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.
+ *
+ * 
Each attempt uses a unique callback name ({@code "approval-1"}, {@code "approval-2"}, etc.) so the execution
+ * history stays clean and replay-safe. The anonymous form is used, so attempts run directly in the caller's context.
+ */
+public class RetryWaitForCallbackExample extends DurableHandler {
+
+    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 = RetryOperationHelper.retryOperation(
+                context,
+                (ctx, attempt) -> ctx.waitForCallback(
+                        "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
+                                .info("Attempt {}: sending callback {} to approval system", attempt, callbackId)),
+                RetryOperationConfig.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);
+    }
+}
diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
new file mode 100644
index 000000000..774c7c681
--- /dev/null
+++ b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
@@ -0,0 +1,41 @@
+// 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.RetryOperationConfig;
+import software.amazon.lambda.durable.examples.types.GreetingRequest;
+import software.amazon.lambda.durable.retry.RetryDecision;
+import software.amazon.lambda.durable.util.RetryOperationHelper;
+
+/**
+ * Example demonstrating {@link RetryOperationHelper} with {@code context.invoke}.
+ *
+ * 
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.
+ *
+ * 
The anonymous form is used, so attempts run directly in the caller's context without child-context wrapping.
+ */
+public class RetryInvokeExample extends DurableHandler {
+
+    private static final int MAX_ATTEMPTS = 3;
+
+    @Override
+    public String handleRequest(GreetingRequest input, DurableContext context) {
+        return RetryOperationHelper.retryOperation(
+                context,
+                (ctx, attempt) -> ctx.invoke(
+                        "call-greeting-" + attempt,
+                        "simple-step-example" + input.getName() + ":$LATEST",
+                        input,
+                        String.class),
+                RetryOperationConfig.builder()
+                        .retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS
+                                ? RetryDecision.retry(Duration.ofSeconds(2))
+                                : RetryDecision.fail())
+                        .build());
+    }
+}
diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExampleTest.java
new file mode 100644
index 000000000..ea5e8a4a8
--- /dev/null
+++ b/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());
+    }
+}
diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExampleTest.java
new file mode 100644
index 000000000..7d8e66104
--- /dev/null
+++ b/examples/src/test/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExampleTest.java
@@ -0,0 +1,130 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples.invoke;
+
+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.GreetingRequest;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+
+class RetryInvokeExampleTest {
+
+    @Test
+    void succeedsOnFirstAttempt() {
+        var handler = new RetryInvokeExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+        var input = new GreetingRequest("world");
+
+        // First run — starts the invoke, suspends waiting for result
+        var result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete the first invoke attempt
+        runner.completeChainedInvoke("call-greeting-1", "\"hello world\"");
+        result = runner.run(input);
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("hello world", result.getResult(String.class));
+    }
+
+    @Test
+    void retriesAfterFirstAttemptFails() {
+        var handler = new RetryInvokeExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+        var input = new GreetingRequest("world");
+
+        // First run — starts invoke attempt 1
+        var result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail the first invoke attempt
+        runner.failChainedInvoke(
+                "call-greeting-1",
+                ErrorObject.builder()
+                        .errorType("TransientError")
+                        .errorMessage("Service unavailable")
+                        .build());
+
+        // Second run — processes the failure, does backoff wait, starts invoke attempt 2
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past the backoff wait
+        runner.advanceTime();
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete the second invoke attempt
+        runner.completeChainedInvoke("call-greeting-2", "\"hello on retry\"");
+        result = runner.run(input);
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("hello on retry", result.getResult(String.class));
+    }
+
+    @Test
+    void failsAfterAllRetriesExhausted() {
+        var handler = new RetryInvokeExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+        var input = new GreetingRequest("world");
+
+        // First run — starts invoke attempt 1
+        var result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail attempt 1
+        runner.failChainedInvoke(
+                "call-greeting-1",
+                ErrorObject.builder()
+                        .errorType("TransientError")
+                        .errorMessage("fail 1")
+                        .build());
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff wait 1
+        runner.advanceTime();
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail attempt 2
+        runner.failChainedInvoke(
+                "call-greeting-2",
+                ErrorObject.builder()
+                        .errorType("TransientError")
+                        .errorMessage("fail 2")
+                        .build());
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff wait 2
+        runner.advanceTime();
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail attempt 3 — this is the last attempt, retryStrategy returns fail()
+        runner.failChainedInvoke(
+                "call-greeting-3",
+                ErrorObject.builder()
+                        .errorType("TransientError")
+                        .errorMessage("fail 3")
+                        .build());
+        result = runner.run(input);
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+    }
+
+    @Test
+    void suspendsOnFirstRun() {
+        var handler = new RetryInvokeExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+        var input = new GreetingRequest("test");
+
+        var result = runner.run(input);
+
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+    }
+}
diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
new file mode 100644
index 000000000..1d7185b95
--- /dev/null
+++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
@@ -0,0 +1,198 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import java.time.Duration;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.services.lambda.model.ErrorObject;
+import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.exception.InvokeFailedException;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.retry.RetryDecision;
+import software.amazon.lambda.durable.retry.RetryStrategies;
+import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+import software.amazon.lambda.durable.util.RetryOperationHelper;
+
+class RetryInvokeIntegrationTest {
+
+    @Test
+    void invokeSucceedsOnFirstAttempt() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
+                                .build()));
+
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.completeChainedInvoke("invoke-1", "\"success\"");
+        result = runner.run("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("success", result.getResult(String.class));
+    }
+
+    @Test
+    void invokeRetriesAfterFailure() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
+                                .build()));
+
+        // First run — invoke attempt 1 starts
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail attempt 1
+        runner.failChainedInvoke(
+                "invoke-1",
+                ErrorObject.builder()
+                        .errorType("TransientError")
+                        .errorMessage("service unavailable")
+                        .build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff wait
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete attempt 2
+        runner.completeChainedInvoke("invoke-2", "\"recovered\"");
+        result = runner.run("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("recovered", result.getResult(String.class));
+    }
+
+    @Test
+    void invokeFailsAfterAllRetriesExhausted() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        RetryOperationConfig.builder()
+                                .retryStrategy((error, attempt) ->
+                                        attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
+                                .build()));
+
+        // Attempt 1
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.failChainedInvoke(
+                "invoke-1", ErrorObject.builder().errorMessage("fail 1").build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Attempt 2 — last attempt
+        runner.failChainedInvoke(
+                "invoke-2", ErrorObject.builder().errorMessage("fail 2").build());
+        result = runner.run("test");
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+    }
+
+    @Test
+    void invokeRetryWithCustomBackoffDelay() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        RetryOperationConfig.builder()
+                                .retryStrategy((error, attempt) -> attempt < 3
+                                        ? RetryDecision.retry(Duration.ofSeconds(attempt * 5L))
+                                        : RetryDecision.fail())
+                                .build()));
+
+        // Attempt 1 fails
+        var result = runner.run("test");
+        runner.failChainedInvoke(
+                "invoke-1", ErrorObject.builder().errorMessage("fail").build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past first backoff (5s)
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Attempt 2 succeeds
+        runner.completeChainedInvoke("invoke-2", "\"ok\"");
+        result = runner.run("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("ok", result.getResult(String.class));
+    }
+
+    @Test
+    void invokeRetryWithStepsBeforeAndAfter() {
+        var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
+            var prefix = context.step("prepare", String.class, stepCtx -> "prepared");
+
+            var invokeResult = RetryOperationHelper.retryOperation(
+                    context,
+                    (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                    RetryOperationConfig.builder()
+                            .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
+                            .build());
+
+            return context.step("finalize", String.class, stepCtx -> prefix + " -> " + invokeResult + " -> done");
+        });
+
+        // First run — prepare step completes, invoke starts
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete invoke
+        runner.completeChainedInvoke("invoke-1", "\"invoked\"");
+        result = runner.runUntilComplete("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("prepared -> invoked -> done", result.getResult(String.class));
+    }
+
+    @Test
+    void invokeRetryPreservesOriginalExceptionType() {
+        var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
+            try {
+                return RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                                .build());
+            } catch (InvokeFailedException e) {
+                assertEquals("invoke failed", e.getMessage());
+                throw e;
+            }
+        });
+
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.failChainedInvoke(
+                "invoke-1", ErrorObject.builder().errorMessage("invoke failed").build());
+        result = runner.run("test");
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+    }
+}
diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
new file mode 100644
index 000000000..f7806d577
--- /dev/null
+++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
@@ -0,0 +1,253 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import java.time.Duration;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.services.lambda.model.ErrorObject;
+import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.retry.RetryDecision;
+import software.amazon.lambda.durable.retry.RetryStrategies;
+import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+import software.amazon.lambda.durable.util.RetryOperationHelper;
+
+class RetryWaitForCallbackIntegrationTest {
+
+    @Test
+    void waitForCallbackSucceedsOnFirstAttempt() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.waitForCallback(
+                                "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
+                                        .info("Submitting callback {}", callbackId)),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
+                                .build()));
+
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // waitForCallback("approval-1", ...) creates "approval-1-callback" internally
+        var callbackId = runner.getCallbackId("approval-1-callback");
+        assertNotNull(callbackId);
+        runner.completeCallback(callbackId, "\"approved\"");
+
+        result = runner.runUntilComplete("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("approved", result.getResult(String.class));
+    }
+
+    @Test
+    void waitForCallbackRetriesAfterFailure() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.waitForCallback(
+                                "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
+                                        .info("Attempt {} callback {}", attempt, callbackId)),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
+                                .build()));
+
+        // First run — starts waitForCallback attempt 1
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail attempt 1
+        var callbackId1 = runner.getCallbackId("approval-1-callback");
+        assertNotNull(callbackId1);
+        runner.failCallback(
+                callbackId1,
+                ErrorObject.builder()
+                        .errorType("Rejected")
+                        .errorMessage("denied by reviewer")
+                        .build());
+
+        // Run — processes failure, hits backoff wait
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete attempt 2
+        var callbackId2 = runner.getCallbackId("approval-2-callback");
+        assertNotNull(callbackId2);
+        runner.completeCallback(callbackId2, "\"approved on retry\"");
+
+        result = runner.runUntilComplete("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("approved on retry", result.getResult(String.class));
+    }
+
+    @Test
+    void waitForCallbackFailsAfterAllRetriesExhausted() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) ->
+                                ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
+                        RetryOperationConfig.builder()
+                                .retryStrategy((error, attempt) ->
+                                        attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
+                                .build()));
+
+        // Attempt 1
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.failCallback(
+                runner.getCallbackId("approval-1-callback"),
+                ErrorObject.builder().errorMessage("fail 1").build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Attempt 2 — last attempt
+        runner.failCallback(
+                runner.getCallbackId("approval-2-callback"),
+                ErrorObject.builder().errorMessage("fail 2").build());
+        result = runner.run("test");
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+    }
+
+    @Test
+    void waitForCallbackRetryWithStepsBeforeAndAfter() {
+        var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
+            var prefix = context.step("prepare", String.class, stepCtx -> "prepared");
+
+            var callbackResult = RetryOperationHelper.retryOperation(
+                    context,
+                    (ctx, attempt) ->
+                            ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
+                    RetryOperationConfig.builder()
+                            .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
+                            .build());
+
+            return context.step("finalize", String.class, stepCtx -> prefix + " -> " + callbackResult + " -> done");
+        });
+
+        // First run — prepare completes, waitForCallback starts
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete callback
+        var callbackId = runner.getCallbackId("approval-1-callback");
+        runner.completeCallback(callbackId, "\"approved\"");
+
+        result = runner.runUntilComplete("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("prepared -> approved -> done", result.getResult(String.class));
+    }
+
+    @Test
+    void waitForCallbackRetryMultipleFailuresThenSuccess() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) ->
+                                ctx.waitForCallback("cb-" + attempt, String.class, (callbackId, stepCtx) -> {}),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.fixedDelay(4, Duration.ofSeconds(1)))
+                                .build()));
+
+        // Attempt 1 — fail
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.failCallback(
+                runner.getCallbackId("cb-1-callback"),
+                ErrorObject.builder().errorMessage("fail 1").build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Attempt 2 — fail
+        runner.failCallback(
+                runner.getCallbackId("cb-2-callback"),
+                ErrorObject.builder().errorMessage("fail 2").build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Attempt 3 — succeed
+        runner.completeCallback(runner.getCallbackId("cb-3-callback"), "\"third time's the charm\"");
+        result = runner.runUntilComplete("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("third time's the charm", result.getResult(String.class));
+    }
+
+    @Test
+    void waitForCallbackRetryWithSubmitterLogic() {
+        // Verify the submitter runs on each retry attempt
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) ->
+                                ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {
+                                    // Submitter runs each attempt — in a real scenario this would
+                                    // send the callbackId to an external system
+                                    stepCtx.getLogger().info("Attempt {} submitting {}", attempt, callbackId);
+                                }),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
+                                .build()));
+
+        // Attempt 1
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Verify submitter step was created for attempt 1
+        var submitterOp = runner.getOperation("approval-1-submitter");
+        assertNotNull(submitterOp, "Submitter step should exist for attempt 1");
+
+        // Fail attempt 1
+        runner.failCallback(
+                runner.getCallbackId("approval-1-callback"),
+                ErrorObject.builder().errorMessage("rejected").build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff, start attempt 2
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Verify submitter step was created for attempt 2
+        var submitterOp2 = runner.getOperation("approval-2-submitter");
+        assertNotNull(submitterOp2, "Submitter step should exist for attempt 2");
+
+        // Complete attempt 2
+        runner.completeCallback(runner.getCallbackId("approval-2-callback"), "\"approved\"");
+        result = runner.runUntilComplete("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("approved", result.getResult(String.class));
+    }
+}
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
index c63f0aad0..f4eaf8c28 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
@@ -7,6 +7,8 @@
 import software.amazon.lambda.durable.DurableContext;
 import software.amazon.lambda.durable.TypeToken;
 import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
+import software.amazon.lambda.durable.execution.SuspendExecutionException;
 import software.amazon.lambda.durable.retry.RetryDecision;
 
 /**
@@ -112,6 +114,9 @@ public static  T retryOperation(
     /**
      * Core retry loop. Replay-safe because every side-effect is a durable operation: the user's operation calls durable
      * primitives, and backoff uses {@code context.wait()}.
+     *
+     * 
{@link SuspendExecutionException} and {@link UnrecoverableDurableExecutionException} are never retried — they
+     * are internal SDK control flow signals that must propagate immediately.
      */
     private static  T executeRetryLoop(
             DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) {
@@ -119,6 +124,9 @@ private static  T executeRetryLoop(
         while (true) {
             try {
                 return operation.execute(context, attempt);
+            } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) {
+                // Internal SDK control flow — never retry, always propagate
+                throw e;
             } catch (Exception e) {
                 RetryDecision decision = config.retryStrategy().makeRetryDecision(e, attempt);
                 if (!decision.shouldRetry()) {
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
index 9558f3c83..265fa968c 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
@@ -15,6 +15,8 @@
 import software.amazon.lambda.durable.DurableContext;
 import software.amazon.lambda.durable.TypeToken;
 import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
+import software.amazon.lambda.durable.execution.SuspendExecutionException;
 import software.amazon.lambda.durable.retry.RetryDecision;
 import software.amazon.lambda.durable.retry.RetryStrategies;
 
@@ -416,5 +418,45 @@ void rethrowsLastExceptionWhenAllRetriesExhausted() {
             // The last attempt's exception is rethrown
             assertEquals("attempt-3", thrown.getMessage());
         }
+
+        @Test
+        void propagatesSuspendExecutionExceptionWithoutRetrying() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1)))
+                    .build();
+
+            assertThrows(
+                    SuspendExecutionException.class,
+                    () -> RetryOperationHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                throw new SuspendExecutionException();
+                            },
+                            config));
+
+            // Should never reach the wait — SuspendExecutionException propagates immediately
+            verify(context, never()).wait(anyString(), any(Duration.class));
+        }
+
+        @Test
+        void propagatesUnrecoverableDurableExecutionExceptionWithoutRetrying() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1)))
+                    .build();
+
+            assertThrows(
+                    UnrecoverableDurableExecutionException.class,
+                    () -> RetryOperationHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                throw new UnrecoverableDurableExecutionException(
+                                        software.amazon.awssdk.services.lambda.model.ErrorObject.builder()
+                                                .errorMessage("unrecoverable")
+                                                .build());
+                            },
+                            config));
+
+            verify(context, never()).wait(anyString(), any(Duration.class));
+        }
     }
 }

From 344f828c8994742807b2336e4aa543a2eaba6791 Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Wed, 29 Apr 2026 15:44:50 -0700
Subject: [PATCH 10/19] chore: set default retry strategy when missing config

---
 .../lambda/durable/config/WithRetryConfig.java     | 14 ++++++--------
 .../lambda/durable/config/WithRetryConfigTest.java |  7 +++----
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java
index 624b7df32..8872206e7 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java
@@ -2,13 +2,14 @@
 // SPDX-License-Identifier: Apache-2.0
 package software.amazon.lambda.durable.config;
 
+import software.amazon.lambda.durable.retry.RetryStrategies;
 import software.amazon.lambda.durable.retry.RetryStrategy;
 
 /**
  * Configuration for {@link software.amazon.lambda.durable.util.WithRetryHelper#withRetry}.
  *
  * 
Uses the same {@link RetryStrategy} shape that developers already know from {@link StepConfig}, so there are zero
- * new retry concepts to learn.
+ * new retry concepts to learn. If no retry strategy is specified, {@link RetryStrategies.Presets#DEFAULT} is used.
  */
 public class WithRetryConfig {
     private final RetryStrategy retryStrategy;
@@ -20,12 +21,13 @@ private WithRetryConfig(Builder builder) {
     }
 
     /**
-     * Returns the retry strategy. Same type as {@link StepConfig#retryStrategy()}.
+     * Returns the retry strategy, or the default strategy if not specified. Same type as
+     * {@link StepConfig#retryStrategy()}.
      *
      * @return the retry strategy, never null
      */
     public RetryStrategy retryStrategy() {
-        return retryStrategy;
+        return retryStrategy != null ? retryStrategy : RetryStrategies.Presets.DEFAULT;
     }
 
     /**
@@ -56,7 +58,7 @@ public static class Builder {
         private Builder() {}
 
         /**
-         * Sets the retry strategy. Required.
+         * Sets the retry strategy. Optional — defaults to {@link RetryStrategies.Presets#DEFAULT} if not set.
          *
          * 
Reuses the exact same {@link RetryStrategy} interface from {@link StepConfig}. All existing factory
          * methods ({@link software.amazon.lambda.durable.retry.RetryStrategies#exponentialBackoff},
@@ -91,12 +93,8 @@ public Builder wrapInChildContext(boolean wrapInChildContext) {
          * Builds the {@link WithRetryConfig} instance.
          *
          * @return a new config with the configured options
-         * @throws IllegalArgumentException if retryStrategy is not set
          */
         public WithRetryConfig build() {
-            if (retryStrategy == null) {
-                throw new IllegalArgumentException("retryStrategy is required");
-            }
             return new WithRetryConfig(this);
         }
     }
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/WithRetryConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/WithRetryConfigTest.java
index 26357d312..c70c1b256 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/config/WithRetryConfigTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/config/WithRetryConfigTest.java
@@ -20,11 +20,10 @@ void builderWithRetryStrategy() {
     }
 
     @Test
-    void builderWithoutRetryStrategy_shouldThrow() {
-        var exception = assertThrows(
-                IllegalArgumentException.class, () -> WithRetryConfig.builder().build());
+    void builderWithoutRetryStrategy_usesDefault() {
+        var config = WithRetryConfig.builder().build();
 
-        assertEquals("retryStrategy is required", exception.getMessage());
+        assertEquals(RetryStrategies.Presets.DEFAULT, config.retryStrategy());
     }
 
     @Test

From eae41bd02ad2d35e8366ac93e3962e72d27e8aac Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Wed, 29 Apr 2026 17:04:05 -0700
Subject: [PATCH 11/19] feat: refactor WithRetry to be implemented by consumers
 of DurableContext

---
 .../callback/RetryWaitForCallbackExample.java |   9 +-
 .../examples/invoke/RetryInvokeExample.java   |   8 +-
 .../durable/RetryInvokeIntegrationTest.java   |  19 +-
 .../RetryWaitForCallbackIntegrationTest.java  |  19 +-
 .../amazon/lambda/durable/DurableContext.java |  66 ++
 .../durable/config/RetryOperationConfig.java  | 103 ---
 .../durable/config/WithRetryConfig.java       |  28 +-
 .../durable/context/DurableContextImpl.java   |  94 +++
 .../durable/{util => model}/WithRetry.java    |   5 +-
 .../durable/util/RetryOperationHelper.java    | 143 -----
 .../durable/util/RetryableOperation.java      |  26 -
 .../lambda/durable/util/WithRetryHelper.java  | 199 ------
 .../config/RetryOperationConfigTest.java      |  80 ---
 .../durable/config/WithRetryConfigTest.java   |  49 +-
 .../DurableContextWithRetryTest.java}         | 589 ++++++++++++------
 .../{util => model}/WithRetryTest.java        |   2 +-
 .../util/RetryOperationHelperTest.java        | 462 --------------
 .../durable/util/RetryableOperationTest.java  |  52 --
 18 files changed, 627 insertions(+), 1326 deletions(-)
 delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java
 rename sdk/src/main/java/software/amazon/lambda/durable/{util => model}/WithRetry.java (81%)
 delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
 delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java
 delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java
 delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
 rename sdk/src/test/java/software/amazon/lambda/durable/{util/WithRetryHelperTest.java => context/DurableContextWithRetryTest.java} (56%)
 rename sdk/src/test/java/software/amazon/lambda/durable/{util => model}/WithRetryTest.java (97%)
 delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
 delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java

diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
index 4f7afbaef..953366050 100644
--- a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
+++ b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
@@ -8,17 +8,17 @@
 import software.amazon.lambda.durable.config.WithRetryConfig;
 import software.amazon.lambda.durable.examples.types.ApprovalRequest;
 import software.amazon.lambda.durable.retry.RetryDecision;
-import software.amazon.lambda.durable.util.WithRetryHelper;
 
 /**
- * Example demonstrating {@link WithRetryHelper} with {@code context.waitForCallback}.
+ * Example demonstrating {@code context.withRetry} with {@code context.waitForCallback}.
  *
  * 
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.
  *
  * 
Each attempt uses a unique callback name ({@code "approval-1"}, {@code "approval-2"}, etc.) so the execution
- * history stays clean and replay-safe. The anonymous form is used, so attempts run directly in the caller's context.
+ * history stays clean and replay-safe. The anonymous form is used, so attempts are grouped under a default-named child
+ * context.
  */
 public class RetryWaitForCallbackExample extends DurableHandler {
 
@@ -33,8 +33,7 @@ public String handleRequest(ApprovalRequest input, DurableContext context) {
                 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 = WithRetryHelper.withRetry(
-                context,
+        var approvalResult = context.withRetry(
                 (ctx, attempt) -> ctx.waitForCallback(
                         "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
                                 .info("Attempt {}: sending callback {} to approval system", attempt, callbackId)),
diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
index a764f7c36..7850b2d7b 100644
--- a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
+++ b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
@@ -8,16 +8,15 @@
 import software.amazon.lambda.durable.config.WithRetryConfig;
 import software.amazon.lambda.durable.examples.types.GreetingRequest;
 import software.amazon.lambda.durable.retry.RetryDecision;
-import software.amazon.lambda.durable.util.WithRetryHelper;
 
 /**
- * Example demonstrating {@link WithRetryHelper} with {@code context.invoke}.
+ * Example demonstrating {@code context.withRetry} with {@code context.invoke}.
  *
  * 
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.
  *
- * 
The anonymous form is used, so attempts run directly in the caller's context without child-context wrapping.
+ * 
The anonymous form is used, so attempts are grouped under a default-named child context.
  */
 public class RetryInvokeExample extends DurableHandler {
 
@@ -25,8 +24,7 @@ public class RetryInvokeExample extends DurableHandler
 
     @Override
     public String handleRequest(GreetingRequest input, DurableContext context) {
-        return WithRetryHelper.withRetry(
-                context,
+        return context.withRetry(
                 (ctx, attempt) -> ctx.invoke(
                         "call-greeting-" + attempt,
                         "simple-step-example" + input.getName() + ":$LATEST",
diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
index f7ed4408c..fd206831d 100644
--- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
+++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
@@ -13,7 +13,6 @@
 import software.amazon.lambda.durable.retry.RetryDecision;
 import software.amazon.lambda.durable.retry.RetryStrategies;
 import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
-import software.amazon.lambda.durable.util.WithRetryHelper;
 
 class RetryInvokeIntegrationTest {
 
@@ -21,8 +20,7 @@ class RetryInvokeIntegrationTest {
     void invokeSucceedsOnFirstAttempt() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.withRetry(
-                        context,
+                (input, context) -> context.withRetry(
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
@@ -42,8 +40,7 @@ void invokeSucceedsOnFirstAttempt() {
     void invokeRetriesAfterFailure() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.withRetry(
-                        context,
+                (input, context) -> context.withRetry(
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
@@ -80,8 +77,7 @@ void invokeRetriesAfterFailure() {
     void invokeFailsAfterAllRetriesExhausted() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.withRetry(
-                        context,
+                (input, context) -> context.withRetry(
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy((error, attempt) ->
@@ -114,8 +110,7 @@ void invokeFailsAfterAllRetriesExhausted() {
     void invokeRetryWithCustomBackoffDelay() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.withRetry(
-                        context,
+                (input, context) -> context.withRetry(
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy((error, attempt) -> attempt < 3
@@ -148,8 +143,7 @@ void invokeRetryWithStepsBeforeAndAfter() {
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             var prefix = context.step("prepare", String.class, stepCtx -> "prepared");
 
-            var invokeResult = WithRetryHelper.withRetry(
-                    context,
+            var invokeResult = context.withRetry(
                     (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                     WithRetryConfig.builder()
                             .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
@@ -174,8 +168,7 @@ void invokeRetryWithStepsBeforeAndAfter() {
     void invokeRetryPreservesOriginalExceptionType() {
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             try {
-                return WithRetryHelper.withRetry(
-                        context,
+                return context.withRetry(
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy(RetryStrategies.Presets.NO_RETRY)
diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
index 2f08315ff..1cd7b7e50 100644
--- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
+++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
@@ -12,7 +12,6 @@
 import software.amazon.lambda.durable.retry.RetryDecision;
 import software.amazon.lambda.durable.retry.RetryStrategies;
 import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
-import software.amazon.lambda.durable.util.WithRetryHelper;
 
 class RetryWaitForCallbackIntegrationTest {
 
@@ -20,8 +19,7 @@ class RetryWaitForCallbackIntegrationTest {
     void waitForCallbackSucceedsOnFirstAttempt() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.withRetry(
-                        context,
+                (input, context) -> context.withRetry(
                         (ctx, attempt) -> ctx.waitForCallback(
                                 "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
                                         .info("Submitting callback {}", callbackId)),
@@ -47,8 +45,7 @@ void waitForCallbackSucceedsOnFirstAttempt() {
     void waitForCallbackRetriesAfterFailure() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.withRetry(
-                        context,
+                (input, context) -> context.withRetry(
                         (ctx, attempt) -> ctx.waitForCallback(
                                 "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
                                         .info("Attempt {} callback {}", attempt, callbackId)),
@@ -94,8 +91,7 @@ void waitForCallbackRetriesAfterFailure() {
     void waitForCallbackFailsAfterAllRetriesExhausted() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.withRetry(
-                        context,
+                (input, context) -> context.withRetry(
                         (ctx, attempt) ->
                                 ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
                         WithRetryConfig.builder()
@@ -132,8 +128,7 @@ void waitForCallbackRetryWithStepsBeforeAndAfter() {
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             var prefix = context.step("prepare", String.class, stepCtx -> "prepared");
 
-            var callbackResult = WithRetryHelper.withRetry(
-                    context,
+            var callbackResult = context.withRetry(
                     (ctx, attempt) ->
                             ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
                     WithRetryConfig.builder()
@@ -161,8 +156,7 @@ void waitForCallbackRetryWithStepsBeforeAndAfter() {
     void waitForCallbackRetryMultipleFailuresThenSuccess() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.withRetry(
-                        context,
+                (input, context) -> context.withRetry(
                         (ctx, attempt) ->
                                 ctx.waitForCallback("cb-" + attempt, String.class, (callbackId, stepCtx) -> {}),
                         WithRetryConfig.builder()
@@ -207,8 +201,7 @@ void waitForCallbackRetryWithSubmitterLogic() {
         // Verify the submitter runs on each retry attempt
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> WithRetryHelper.withRetry(
-                        context,
+                (input, context) -> context.withRetry(
                         (ctx, attempt) ->
                                 ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {
                                     // Submitter runs each attempt — in a real scenario this would
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
index 615ce70e8..9518e5060 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
@@ -16,9 +16,11 @@
 import software.amazon.lambda.durable.config.StepConfig;
 import software.amazon.lambda.durable.config.WaitForCallbackConfig;
 import software.amazon.lambda.durable.config.WaitForConditionConfig;
+import software.amazon.lambda.durable.config.WithRetryConfig;
 import software.amazon.lambda.durable.context.BaseContext;
 import software.amazon.lambda.durable.model.MapResult;
 import software.amazon.lambda.durable.model.WaitForConditionResult;
+import software.amazon.lambda.durable.model.WithRetry;
 
 public interface DurableContext extends BaseContext {
     /**
@@ -727,6 +729,70 @@  DurableFuture waitForConditionAsync(
             BiFunction> checkFunc,
             WaitForConditionConfig config);
 
+    // =============== withRetry ================
+
+    /**
+     * Replay-safe retry loop for any durable operation (named form, sync).
+     *
+     * 
Provides the same retry-with-backoff pattern that {@code step()} has built in, but for operations that cannot
+     * live inside a step ({@code waitForCallback}, {@code invoke}, {@code waitForCondition}, etc.).
+     *
+     * 
Every side-effect in the loop is a durable operation, so the loop is replay-safe by construction. On replay,
+     * completed operations return cached results instantly and the loop fast-forwards to the current attempt.
+     *
+     * 
By default, the retry loop runs directly on the caller's context. If
+     * {@link WithRetryConfig#wrapInChildContext()} is enabled, the loop is wrapped in a child context so all attempts
+     * are grouped under a single named operation in execution history.
+     *
+     * @param  the result type
+     * @param name operation name (used for backoff wait names, and as the child context name when wrapping)
+     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param config retry configuration including the retry strategy and child context wrapping
+     * @return the operation result
+     */
+     T withRetry(String name, WithRetry operation, WithRetryConfig config);
+
+    /**
+     * Replay-safe retry loop for any durable operation (anonymous form, sync).
+     *
+     * 
By default, the retry loop runs directly on the caller's context. If
+     * {@link WithRetryConfig#wrapInChildContext()} is enabled, the loop is wrapped in a child context with a default
+     * name so all attempts are grouped under a single operation in execution history.
+     *
+     * @param  the result type
+     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param config retry configuration including the retry strategy and child context wrapping
+     * @return the operation result
+     */
+     T withRetry(WithRetry operation, WithRetryConfig config);
+
+    /**
+     * Replay-safe retry loop for any durable operation (named form, async).
+     *
+     * 
Wraps the retry loop in {@code runInChildContextAsync} so all attempts are grouped under a single named
+     * operation in execution history, and returns a {@link DurableFuture} that can be composed or blocked on.
+     *
+     * @param  the result type
+     * @param name operation name (used for child context and backoff wait names)
+     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param config retry configuration including the retry strategy
+     * @return a future representing the operation result
+     */
+     DurableFuture withRetryAsync(String name, WithRetry operation, WithRetryConfig config);
+
+    /**
+     * Replay-safe retry loop for any durable operation (anonymous form, async).
+     *
+     * 
Wraps the retry loop in {@code runInChildContextAsync} with a default name so all attempts are grouped under a
+     * single operation in execution history, and returns a {@link DurableFuture} that can be composed or blocked on.
+     *
+     * @param  the result type
+     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param config retry configuration including the retry strategy
+     * @return a future representing the operation result
+     */
+     DurableFuture withRetryAsync(WithRetry operation, WithRetryConfig config);
+
     /**
      * Function applied to each item in a map operation.
      *
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java
deleted file mode 100644
index 696d1d253..000000000
--- a/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java
+++ /dev/null
@@ -1,103 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.config;
-
-import software.amazon.lambda.durable.retry.RetryStrategy;
-
-/**
- * Configuration for {@link software.amazon.lambda.durable.util.RetryOperationHelper#retryOperation}.
- *
- * 
Uses the same {@link RetryStrategy} shape that developers already know from {@link StepConfig}, so there are zero
- * new retry concepts to learn.
- */
-public class RetryOperationConfig {
-    private final RetryStrategy retryStrategy;
-    private final boolean wrapInChildContext;
-
-    private RetryOperationConfig(Builder builder) {
-        this.retryStrategy = builder.retryStrategy;
-        this.wrapInChildContext = builder.wrapInChildContext;
-    }
-
-    /**
-     * Returns the retry strategy. Same type as {@link StepConfig#retryStrategy()}.
-     *
-     * @return the retry strategy, never null
-     */
-    public RetryStrategy retryStrategy() {
-        return retryStrategy;
-    }
-
-    /**
-     * Whether to wrap the retry loop in {@code runInChildContext} so all attempts are grouped under a single named
-     * operation in execution history. Only applies when a name is provided to the named form of {@code retryOperation}.
-     * Defaults to {@code true}.
-     *
-     * @return true if child-context wrapping is enabled
-     */
-    public boolean wrapInChildContext() {
-        return wrapInChildContext;
-    }
-
-    /**
-     * Creates a new builder for {@code RetryOperationConfig}.
-     *
-     * @return a new builder instance
-     */
-    public static Builder builder() {
-        return new Builder();
-    }
-
-    /** Builder for creating {@link RetryOperationConfig} instances. */
-    public static class Builder {
-        private RetryStrategy retryStrategy;
-        private boolean wrapInChildContext = true;
-
-        private Builder() {}
-
-        /**
-         * Sets the retry strategy. Required.
-         *
-         * 
Reuses the exact same {@link RetryStrategy} interface from {@link StepConfig}. All existing factory
-         * methods ({@link software.amazon.lambda.durable.retry.RetryStrategies#exponentialBackoff},
-         * {@link software.amazon.lambda.durable.retry.RetryStrategies#fixedDelay}, presets, and custom lambdas) work
-         * without modification.
-         *
-         * @param retryStrategy the retry strategy to use
-         * @return this builder for method chaining
-         */
-        public Builder retryStrategy(RetryStrategy retryStrategy) {
-            this.retryStrategy = retryStrategy;
-            return this;
-        }
-
-        /**
-         * Controls whether the retry loop is wrapped in a child context. Only meaningful for the named form of
-         * {@code retryOperation}. Defaults to {@code true}.
-         *
-         * 
When {@code true}, all attempts and backoff waits are grouped under a single named operation in execution
-         * history, providing a cleaner view and isolated operation ID space. Set to {@code false} to flatten attempts
-         * into the parent context.
-         *
-         * @param wrapInChildContext whether to wrap in a child context
-         * @return this builder for method chaining
-         */
-        public Builder wrapInChildContext(boolean wrapInChildContext) {
-            this.wrapInChildContext = wrapInChildContext;
-            return this;
-        }
-
-        /**
-         * Builds the {@link RetryOperationConfig} instance.
-         *
-         * @return a new config with the configured options
-         * @throws IllegalArgumentException if retryStrategy is not set
-         */
-        public RetryOperationConfig build() {
-            if (retryStrategy == null) {
-                throw new IllegalArgumentException("retryStrategy is required");
-            }
-            return new RetryOperationConfig(this);
-        }
-    }
-}
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java
index 8872206e7..23ff43830 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/config/WithRetryConfig.java
@@ -6,7 +6,7 @@
 import software.amazon.lambda.durable.retry.RetryStrategy;
 
 /**
- * Configuration for {@link software.amazon.lambda.durable.util.WithRetryHelper#withRetry}.
+ * Configuration for {@link software.amazon.lambda.durable.DurableContext#withRetry}.
  *
  * 
Uses the same {@link RetryStrategy} shape that developers already know from {@link StepConfig}, so there are zero
  * new retry concepts to learn. If no retry strategy is specified, {@link RetryStrategies.Presets#DEFAULT} is used.
@@ -31,11 +31,15 @@ public RetryStrategy retryStrategy() {
     }
 
     /**
-     * Whether to wrap the retry loop in {@code runInChildContext} so all attempts are grouped under a single named
-     * operation in execution history. Only applies when a name is provided to the named form of {@code withRetry}.
-     * Defaults to {@code true}.
+     * Returns whether the sync {@code withRetry} should wrap the retry loop in a child context.
      *
-     * @return true if child-context wrapping is enabled
+     * 
When {@code true}, the sync form behaves like the async form — all retry attempts are grouped under a single
+     * named child context in execution history. When {@code false} (the default), the retry loop runs directly on the
+     * caller's context.
+     *
+     * 
This setting has no effect on the async {@code withRetryAsync} methods, which always wrap in a child context.
+     *
+     * @return {@code true} if the sync retry loop should be wrapped in a child context
      */
     public boolean wrapInChildContext() {
         return wrapInChildContext;
@@ -53,7 +57,7 @@ public static Builder builder() {
     /** Builder for creating {@link WithRetryConfig} instances. */
     public static class Builder {
         private RetryStrategy retryStrategy;
-        private boolean wrapInChildContext = true;
+        private boolean wrapInChildContext;
 
         private Builder() {}
 
@@ -74,14 +78,14 @@ public Builder retryStrategy(RetryStrategy retryStrategy) {
         }
 
         /**
-         * Controls whether the retry loop is wrapped in a child context. Only meaningful for the named form of
-         * {@code withRetry}. Defaults to {@code true}.
+         * Sets whether the sync {@code withRetry} should wrap the retry loop in a child context. Optional — defaults to
+         * {@code false}.
          *
-         * 
When {@code true}, all attempts and backoff waits are grouped under a single named operation in execution
-         * history, providing a cleaner view and isolated operation ID space. Set to {@code false} to flatten attempts
-         * into the parent context.
+         * 
When enabled, the sync form groups all retry attempts under a single named child context in execution
+         * history, matching the behavior of the async form. This is useful when you want operation isolation but don't
+         * need a {@link software.amazon.lambda.durable.DurableFuture}.
          *
-         * @param wrapInChildContext whether to wrap in a child context
+         * @param wrapInChildContext {@code true} to wrap in a child context
          * @return this builder for method chaining
          */
         public Builder wrapInChildContext(boolean wrapInChildContext) {
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
index 86b39f0eb..82fa29bce 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
@@ -27,14 +27,18 @@
 import software.amazon.lambda.durable.config.StepConfig;
 import software.amazon.lambda.durable.config.WaitForCallbackConfig;
 import software.amazon.lambda.durable.config.WaitForConditionConfig;
+import software.amazon.lambda.durable.config.WithRetryConfig;
+import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.execution.ExecutionManager;
 import software.amazon.lambda.durable.execution.OperationIdGenerator;
+import software.amazon.lambda.durable.execution.SuspendExecutionException;
 import software.amazon.lambda.durable.execution.ThreadType;
 import software.amazon.lambda.durable.logging.DurableLogger;
 import software.amazon.lambda.durable.model.MapResult;
 import software.amazon.lambda.durable.model.OperationIdentifier;
 import software.amazon.lambda.durable.model.OperationSubType;
 import software.amazon.lambda.durable.model.WaitForConditionResult;
+import software.amazon.lambda.durable.model.WithRetry;
 import software.amazon.lambda.durable.operation.CallbackOperation;
 import software.amazon.lambda.durable.operation.ChildContextOperation;
 import software.amazon.lambda.durable.operation.InvokeOperation;
@@ -43,6 +47,7 @@
 import software.amazon.lambda.durable.operation.StepOperation;
 import software.amazon.lambda.durable.operation.WaitForConditionOperation;
 import software.amazon.lambda.durable.operation.WaitOperation;
+import software.amazon.lambda.durable.retry.RetryDecision;
 import software.amazon.lambda.durable.util.ParameterValidator;
 
 /**
@@ -366,6 +371,95 @@ public  DurableFuture waitForConditionAsync(
         return operation;
     }
 
+    // =============== withRetry ================
+
+    private static final Duration DEFAULT_BACKOFF_DELAY = Duration.ofSeconds(1);
+    private static final String BACKOFF_SUFFIX = "-backoff-";
+    private static final String ANONYMOUS_CHILD_CONTEXT_NAME = "retry";
+    private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-";
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public  T withRetry(String name, WithRetry operation, WithRetryConfig config) {
+        Objects.requireNonNull(name, "name cannot be null");
+        Objects.requireNonNull(operation, "operation cannot be null");
+        Objects.requireNonNull(config, "config cannot be null");
+
+        if (config.wrapInChildContext()) {
+            return (T) runInChildContext(
+                    name, new TypeToken() {}, childCtx -> executeRetryLoop(childCtx, name, operation, config));
+        }
+        return executeRetryLoop(this, name, operation, config);
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public  T withRetry(WithRetry operation, WithRetryConfig config) {
+        Objects.requireNonNull(operation, "operation cannot be null");
+        Objects.requireNonNull(config, "config cannot be null");
+
+        if (config.wrapInChildContext()) {
+            return (T) runInChildContext(
+                    ANONYMOUS_CHILD_CONTEXT_NAME,
+                    new TypeToken() {},
+                    childCtx -> executeRetryLoop(childCtx, null, operation, config));
+        }
+        return executeRetryLoop(this, null, operation, config);
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public  DurableFuture withRetryAsync(String name, WithRetry operation, WithRetryConfig config) {
+        Objects.requireNonNull(name, "name cannot be null");
+        Objects.requireNonNull(operation, "operation cannot be null");
+        Objects.requireNonNull(config, "config cannot be null");
+
+        return (DurableFuture) runInChildContextAsync(
+                name, new TypeToken() {}, childCtx -> executeRetryLoop(childCtx, name, operation, config));
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public  DurableFuture withRetryAsync(WithRetry operation, WithRetryConfig config) {
+        Objects.requireNonNull(operation, "operation cannot be null");
+        Objects.requireNonNull(config, "config cannot be null");
+
+        return (DurableFuture) runInChildContextAsync(
+                ANONYMOUS_CHILD_CONTEXT_NAME,
+                new TypeToken() {},
+                childCtx -> executeRetryLoop(childCtx, null, operation, config));
+    }
+
+    /**
+     * Core retry loop. Replay-safe because every side-effect is a durable operation: the user's operation calls durable
+     * primitives, and backoff uses {@code context.wait()}.
+     *
+     * 
{@link SuspendExecutionException} and {@link UnrecoverableDurableExecutionException} are never retried — they
+     * are internal SDK control flow signals that must propagate immediately.
+     */
+    private static  T executeRetryLoop(
+            DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
+        var attempt = 1;
+        while (true) {
+            try {
+                return operation.execute(context, attempt);
+            } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) {
+                // Internal SDK control flow — never retry, always propagate
+                throw e;
+            } catch (Exception e) {
+                RetryDecision decision = config.retryStrategy().makeRetryDecision(e, attempt);
+                if (!decision.shouldRetry()) {
+                    throw e;
+                }
+
+                var delay = decision.delay().isZero() ? DEFAULT_BACKOFF_DELAY : decision.delay();
+                var waitName = name != null ? name + BACKOFF_SUFFIX + attempt : ANONYMOUS_BACKOFF_PREFIX + attempt;
+                context.wait(waitName, delay);
+                attempt++;
+            }
+        }
+    }
+
     // =============== accessors ================
     @Override
     public DurableLogger getLogger() {
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetry.java b/sdk/src/main/java/software/amazon/lambda/durable/model/WithRetry.java
similarity index 81%
rename from sdk/src/main/java/software/amazon/lambda/durable/util/WithRetry.java
rename to sdk/src/main/java/software/amazon/lambda/durable/model/WithRetry.java
index 87f7be964..3a5675eaa 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetry.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/model/WithRetry.java
@@ -1,11 +1,12 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.util;
+package software.amazon.lambda.durable.model;
 
 import software.amazon.lambda.durable.DurableContext;
 
 /**
- * A durable operation that can be retried end-to-end by {@link WithRetryHelper}.
+ * A durable operation that can be retried end-to-end by
+ * {@link software.amazon.lambda.durable.DurableContext#withRetry}.
  *
  * 
Receives the durable context and the 1-based attempt number so callers can generate unique operation names per
  * attempt (e.g., {@code "approval-" + attempt}).
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
deleted file mode 100644
index f4eaf8c28..000000000
--- a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
+++ /dev/null
@@ -1,143 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.util;
-
-import java.time.Duration;
-import java.util.Objects;
-import software.amazon.lambda.durable.DurableContext;
-import software.amazon.lambda.durable.TypeToken;
-import software.amazon.lambda.durable.config.RetryOperationConfig;
-import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
-import software.amazon.lambda.durable.execution.SuspendExecutionException;
-import software.amazon.lambda.durable.retry.RetryDecision;
-
-/**
- * Replay-safe retry loop for any durable operation.
- *
- * 
Provides the same retry-with-backoff pattern that {@code context.step()} has built in, but for operations that
- * cannot live inside a step ({@code waitForCallback}, {@code invoke}, {@code waitForCondition}, etc.).
- *
- * 
Every side-effect in the loop is a durable operation, so the loop is replay-safe by construction. On replay,
- * completed operations return cached results instantly and the loop fast-forwards to the current attempt.
- *
- * 
Usage — callback retry

- *
- * 
{@code
- * var result = RetryOperationHelper.retryOperation(
- *     context,
- *     "approval",
- *     (ctx, attempt) -> ctx.waitForCallback(
- *         "approval-" + attempt,
- *         String.class,
- *         (callbackId, stepCtx) -> sendApprovalEmail(approverEmail, callbackId)
- *     ),
- *     RetryOperationConfig.builder()
- *         .retryStrategy(RetryStrategies.exponentialBackoff(
- *             3, Duration.ofSeconds(2), Duration.ofSeconds(30), 2.0, JitterStrategy.FULL))
- *         .build()
- * );
- * }

- *
- * 
Usage — invoke retry (anonymous form)

- *
- * 
{@code
- * var result = RetryOperationHelper.retryOperation(
- *     context,
- *     (ctx, attempt) -> ctx.invoke(
- *         "charge-" + attempt, paymentFnArn, new ChargeRequest(orderId), String.class),
- *     RetryOperationConfig.builder()
- *         .retryStrategy((err, att) -> att < 3
- *             ? RetryDecision.retry(Duration.ofSeconds(1))
- *             : RetryDecision.fail())
- *         .build()
- * );
- * }

- */
-public final class RetryOperationHelper {
-
-    private static final Duration DEFAULT_BACKOFF_DELAY = Duration.ofSeconds(1);
-    private static final String BACKOFF_SUFFIX = "-backoff-";
-    private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-";
-
-    private RetryOperationHelper() {
-        // utility class
-    }
-
-    /**
-     * Named form — wraps the retry loop in {@code runInChildContext} by default so all attempts are grouped under a
-     * single named operation in execution history.
-     *
-     * 
The child-context wrapping can be disabled via
-     * {@link RetryOperationConfig.Builder#wrapInChildContext(boolean)}.
-     *
-     * @param  the result type
-     * @param context the durable context
-     * @param name operation name (used for child context and backoff wait names)
-     * @param operation the retryable operation — receives the context and 1-based attempt number
-     * @param config retry configuration including the retry strategy
-     * @return the operation result
-     */
-    @SuppressWarnings("unchecked")
-    public static  T retryOperation(
-            DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) {
-        Objects.requireNonNull(context, "context cannot be null");
-        Objects.requireNonNull(name, "name cannot be null");
-        Objects.requireNonNull(operation, "operation cannot be null");
-        Objects.requireNonNull(config, "config cannot be null");
-
-        if (config.wrapInChildContext()) {
-            return (T) context.runInChildContext(
-                    name, new TypeToken() {}, childCtx -> executeRetryLoop(childCtx, name, operation, config));
-        }
-        return executeRetryLoop(context, name, operation, config);
-    }
-
-    /**
-     * Anonymous form — runs the retry loop directly in the caller's context. No child-context wrapping is applied
-     * regardless of the {@code wrapInChildContext} config setting.
-     *
-     * @param  the result type
-     * @param context the durable context
-     * @param operation the retryable operation — receives the context and 1-based attempt number
-     * @param config retry configuration including the retry strategy
-     * @return the operation result
-     */
-    public static  T retryOperation(
-            DurableContext context, RetryableOperation operation, RetryOperationConfig config) {
-        Objects.requireNonNull(context, "context cannot be null");
-        Objects.requireNonNull(operation, "operation cannot be null");
-        Objects.requireNonNull(config, "config cannot be null");
-
-        return executeRetryLoop(context, null, operation, config);
-    }
-
-    /**
-     * Core retry loop. Replay-safe because every side-effect is a durable operation: the user's operation calls durable
-     * primitives, and backoff uses {@code context.wait()}.
-     *
-     * 
{@link SuspendExecutionException} and {@link UnrecoverableDurableExecutionException} are never retried — they
-     * are internal SDK control flow signals that must propagate immediately.
-     */
-    private static  T executeRetryLoop(
-            DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) {
-        var attempt = 1;
-        while (true) {
-            try {
-                return operation.execute(context, attempt);
-            } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) {
-                // Internal SDK control flow — never retry, always propagate
-                throw e;
-            } catch (Exception e) {
-                RetryDecision decision = config.retryStrategy().makeRetryDecision(e, attempt);
-                if (!decision.shouldRetry()) {
-                    throw e;
-                }
-
-                var delay = decision.delay().isZero() ? DEFAULT_BACKOFF_DELAY : decision.delay();
-                var waitName = name != null ? name + BACKOFF_SUFFIX + attempt : ANONYMOUS_BACKOFF_PREFIX + attempt;
-                context.wait(waitName, delay);
-                attempt++;
-            }
-        }
-    }
-}
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java
deleted file mode 100644
index d6b2496c1..000000000
--- a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java
+++ /dev/null
@@ -1,26 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.util;
-
-import software.amazon.lambda.durable.DurableContext;
-
-/**
- * A durable operation that can be retried end-to-end by {@link RetryOperationHelper}.
- *
- * 
Receives the durable context and the 1-based attempt number so callers can generate unique operation names per
- * attempt (e.g., {@code "approval-" + attempt}).
- *
- * @param  the result type
- */
-@FunctionalInterface
-public interface RetryableOperation {
-
-    /**
-     * Executes the durable operation.
-     *
-     * @param context the durable context to use for durable operations
-     * @param attempt the current attempt number (1-based: first attempt is 1)
-     * @return the operation result
-     */
-    T execute(DurableContext context, int attempt);
-}
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java
deleted file mode 100644
index 237cd9530..000000000
--- a/sdk/src/main/java/software/amazon/lambda/durable/util/WithRetryHelper.java
+++ /dev/null
@@ -1,199 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.util;
-
-import java.time.Duration;
-import java.util.Objects;
-import software.amazon.lambda.durable.DurableContext;
-import software.amazon.lambda.durable.DurableFuture;
-import software.amazon.lambda.durable.TypeToken;
-import software.amazon.lambda.durable.config.WithRetryConfig;
-import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
-import software.amazon.lambda.durable.execution.SuspendExecutionException;
-import software.amazon.lambda.durable.retry.RetryDecision;
-
-/**
- * Replay-safe retry loop for any durable operation.
- *
- * 
Provides the same retry-with-backoff pattern that {@code context.step()} has built in, but for operations that
- * cannot live inside a step ({@code waitForCallback}, {@code invoke}, {@code waitForCondition}, etc.).
- *
- * 
Every side-effect in the loop is a durable operation, so the loop is replay-safe by construction. On replay,
- * completed operations return cached results instantly and the loop fast-forwards to the current attempt.
- *
- * 
Usage — callback retry

- *
- * 
{@code
- * var result = WithRetryHelper.withRetry(
- *     context,
- *     "approval",
- *     (ctx, attempt) -> ctx.waitForCallback(
- *         "approval-" + attempt,
- *         String.class,
- *         (callbackId, stepCtx) -> sendApprovalEmail(approverEmail, callbackId)
- *     ),
- *     WithRetryConfig.builder()
- *         .retryStrategy(RetryStrategies.exponentialBackoff(
- *             3, Duration.ofSeconds(2), Duration.ofSeconds(30), 2.0, JitterStrategy.FULL))
- *         .build()
- * );
- * }

- *
- * 
Usage — invoke retry (anonymous form)

- *
- * 
{@code
- * var result = WithRetryHelper.withRetry(
- *     context,
- *     (ctx, attempt) -> ctx.invoke(
- *         "charge-" + attempt, paymentFnArn, new ChargeRequest(orderId), String.class),
- *     WithRetryConfig.builder()
- *         .retryStrategy((err, att) -> att < 3
- *             ? RetryDecision.retry(Duration.ofSeconds(1))
- *             : RetryDecision.fail())
- *         .build()
- * );
- * }

- *
- * 
Usage — async form returning DurableFuture

- *
- * 
{@code
- * DurableFuture future = WithRetryHelper.withRetryAsync(
- *     context,
- *     "approval",
- *     (ctx, attempt) -> ctx.waitForCallback(
- *         "approval-" + attempt, String.class,
- *         (callbackId, stepCtx) -> sendApprovalEmail(approverEmail, callbackId)
- *     ),
- *     WithRetryConfig.builder()
- *         .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
- *         .build()
- * );
- * // ... do other work ...
- * var result = future.get();
- * }

- */
-public final class WithRetryHelper {
-
-    private static final Duration DEFAULT_BACKOFF_DELAY = Duration.ofSeconds(1);
-    private static final String BACKOFF_SUFFIX = "-backoff-";
-    private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-";
-
-    private WithRetryHelper() {
-        // utility class
-    }
-
-    /**
-     * Named async form — wraps the retry loop in {@code runInChildContextAsync} by default so all attempts are grouped
-     * under a single named operation in execution history, and returns a {@link DurableFuture} that can be composed or
-     * blocked on.
-     *
-     * 
The child-context wrapping can be disabled via {@link WithRetryConfig.Builder#wrapInChildContext(boolean)}.
-     * When disabled, the retry loop executes immediately and the returned future is already completed.
-     *
-     * @param  the result type
-     * @param context the durable context
-     * @param name operation name (used for child context and backoff wait names)
-     * @param operation the retryable operation — receives the context and 1-based attempt number
-     * @param config retry configuration including the retry strategy
-     * @return a future representing the operation result
-     */
-    @SuppressWarnings("unchecked")
-    public static  DurableFuture withRetryAsync(
-            DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
-        Objects.requireNonNull(context, "context cannot be null");
-        Objects.requireNonNull(name, "name cannot be null");
-        Objects.requireNonNull(operation, "operation cannot be null");
-        Objects.requireNonNull(config, "config cannot be null");
-
-        if (config.wrapInChildContext()) {
-            return (DurableFuture) context.runInChildContextAsync(
-                    name, new TypeToken() {}, childCtx -> executeRetryLoop(childCtx, name, operation, config));
-        }
-        return new CompletedDurableFuture<>(executeRetryLoop(context, name, operation, config));
-    }
-
-    /**
-     * Named sync form — wraps the retry loop in {@code runInChildContext} by default so all attempts are grouped under
-     * a single named operation in execution history, and blocks until the result is available.
-     *
-     * 
Equivalent to {@code withRetryAsync(context, name, operation, config).get()}.
-     *
-     * @param  the result type
-     * @param context the durable context
-     * @param name operation name (used for child context and backoff wait names)
-     * @param operation the retryable operation — receives the context and 1-based attempt number
-     * @param config retry configuration including the retry strategy
-     * @return the operation result
-     */
-    public static  T withRetry(DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
-        return withRetryAsync(context, name, operation, config).get();
-    }
-
-    /**
-     * Anonymous async form — runs the retry loop directly in the caller's context and returns a {@link DurableFuture}.
-     * No child-context wrapping is applied regardless of the {@code wrapInChildContext} config setting.
-     *
-     * 
Because the anonymous form executes the retry loop inline (no child context), the returned future is always
-     * already completed.
-     *
-     * @param  the result type
-     * @param context the durable context
-     * @param operation the retryable operation — receives the context and 1-based attempt number
-     * @param config retry configuration including the retry strategy
-     * @return a future representing the operation result
-     */
-    public static  DurableFuture withRetryAsync(
-            DurableContext context, WithRetry operation, WithRetryConfig config) {
-        Objects.requireNonNull(context, "context cannot be null");
-        Objects.requireNonNull(operation, "operation cannot be null");
-        Objects.requireNonNull(config, "config cannot be null");
-
-        return new CompletedDurableFuture<>(executeRetryLoop(context, null, operation, config));
-    }
-
-    /**
-     * Anonymous sync form — runs the retry loop directly in the caller's context. No child-context wrapping is applied
-     * regardless of the {@code wrapInChildContext} config setting.
-     *
-     * 
Equivalent to {@code withRetryAsync(context, operation, config).get()}.
-     *
-     * @param  the result type
-     * @param context the durable context
-     * @param operation the retryable operation — receives the context and 1-based attempt number
-     * @param config retry configuration including the retry strategy
-     * @return the operation result
-     */
-    public static  T withRetry(DurableContext context, WithRetry operation, WithRetryConfig config) {
-        return withRetryAsync(context, operation, config).get();
-    }
-
-    /**
-     * Core retry loop. Replay-safe because every side-effect is a durable operation: the user's operation calls durable
-     * primitives, and backoff uses {@code context.wait()}.
-     *
-     * 
{@link SuspendExecutionException} and {@link UnrecoverableDurableExecutionException} are never retried — they
-     * are internal SDK control flow signals that must propagate immediately.
-     */
-    private static  T executeRetryLoop(
-            DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
-        var attempt = 1;
-        while (true) {
-            try {
-                return operation.execute(context, attempt);
-            } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) {
-                // Internal SDK control flow — never retry, always propagate
-                throw e;
-            } catch (Exception e) {
-                RetryDecision decision = config.retryStrategy().makeRetryDecision(e, attempt);
-                if (!decision.shouldRetry()) {
-                    throw e;
-                }
-
-                var delay = decision.delay().isZero() ? DEFAULT_BACKOFF_DELAY : decision.delay();
-                var waitName = name != null ? name + BACKOFF_SUFFIX + attempt : ANONYMOUS_BACKOFF_PREFIX + attempt;
-                context.wait(waitName, delay);
-                attempt++;
-            }
-        }
-    }
-}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
deleted file mode 100644
index 726da5b61..000000000
--- a/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
+++ /dev/null
@@ -1,80 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.config;
-
-import static org.junit.jupiter.api.Assertions.*;
-
-import org.junit.jupiter.api.Test;
-import software.amazon.lambda.durable.retry.RetryDecision;
-import software.amazon.lambda.durable.retry.RetryStrategies;
-
-class RetryOperationConfigTest {
-
-    @Test
-    void builderWithRetryStrategy() {
-        var strategy = RetryStrategies.Presets.DEFAULT;
-
-        var config = RetryOperationConfig.builder().retryStrategy(strategy).build();
-
-        assertEquals(strategy, config.retryStrategy());
-    }
-
-    @Test
-    void builderWithoutRetryStrategy_shouldThrow() {
-        var exception = assertThrows(IllegalArgumentException.class, () -> RetryOperationConfig.builder()
-                .build());
-
-        assertEquals("retryStrategy is required", exception.getMessage());
-    }
-
-    @Test
-    void wrapInChildContext_defaultsToTrue() {
-        var config = RetryOperationConfig.builder()
-                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                .build();
-
-        assertTrue(config.wrapInChildContext());
-    }
-
-    @Test
-    void wrapInChildContext_canBeSetToFalse() {
-        var config = RetryOperationConfig.builder()
-                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                .wrapInChildContext(false)
-                .build();
-
-        assertFalse(config.wrapInChildContext());
-    }
-
-    @Test
-    void wrapInChildContext_canBeSetToTrue() {
-        var config = RetryOperationConfig.builder()
-                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                .wrapInChildContext(true)
-                .build();
-
-        assertTrue(config.wrapInChildContext());
-    }
-
-    @Test
-    void builderChaining() {
-        var strategy = RetryStrategies.Presets.DEFAULT;
-
-        var config = RetryOperationConfig.builder()
-                .retryStrategy(strategy)
-                .wrapInChildContext(false)
-                .build();
-
-        assertEquals(strategy, config.retryStrategy());
-        assertFalse(config.wrapInChildContext());
-    }
-
-    @Test
-    void builderWithCustomLambdaRetryStrategy() {
-        var config = RetryOperationConfig.builder()
-                .retryStrategy((error, attempt) -> RetryDecision.fail())
-                .build();
-
-        assertNotNull(config.retryStrategy());
-    }
-}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/WithRetryConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/WithRetryConfigTest.java
index c70c1b256..b1c02db9a 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/config/WithRetryConfigTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/config/WithRetryConfigTest.java
@@ -27,53 +27,54 @@ void builderWithoutRetryStrategy_usesDefault() {
     }
 
     @Test
-    void wrapInChildContext_defaultsToTrue() {
-        var config = WithRetryConfig.builder()
-                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                .build();
+    void builderChaining() {
+        var strategy = RetryStrategies.Presets.DEFAULT;
 
-        assertTrue(config.wrapInChildContext());
+        var config = WithRetryConfig.builder().retryStrategy(strategy).build();
+
+        assertEquals(strategy, config.retryStrategy());
     }
 
     @Test
-    void wrapInChildContext_canBeSetToFalse() {
+    void builderWithCustomLambdaRetryStrategy() {
         var config = WithRetryConfig.builder()
-                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                .wrapInChildContext(false)
+                .retryStrategy((error, attempt) -> RetryDecision.fail())
                 .build();
 
+        assertNotNull(config.retryStrategy());
+    }
+
+    @Test
+    void wrapInChildContext_defaultsFalse() {
+        var config = WithRetryConfig.builder().build();
+
         assertFalse(config.wrapInChildContext());
     }
 
     @Test
-    void wrapInChildContext_canBeSetToTrue() {
-        var config = WithRetryConfig.builder()
-                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                .wrapInChildContext(true)
-                .build();
+    void wrapInChildContext_canBeEnabled() {
+        var config = WithRetryConfig.builder().wrapInChildContext(true).build();
 
         assertTrue(config.wrapInChildContext());
     }
 
     @Test
-    void builderChaining() {
-        var strategy = RetryStrategies.Presets.DEFAULT;
+    void wrapInChildContext_canBeExplicitlyDisabled() {
+        var config = WithRetryConfig.builder().wrapInChildContext(false).build();
 
-        var config = WithRetryConfig.builder()
-                .retryStrategy(strategy)
-                .wrapInChildContext(false)
-                .build();
-
-        assertEquals(strategy, config.retryStrategy());
         assertFalse(config.wrapInChildContext());
     }
 
     @Test
-    void builderWithCustomLambdaRetryStrategy() {
+    void builderChaining_withWrapInChildContext() {
+        var strategy = RetryStrategies.Presets.DEFAULT;
+
         var config = WithRetryConfig.builder()
-                .retryStrategy((error, attempt) -> RetryDecision.fail())
+                .retryStrategy(strategy)
+                .wrapInChildContext(true)
                 .build();
 
-        assertNotNull(config.retryStrategy());
+        assertEquals(strategy, config.retryStrategy());
+        assertTrue(config.wrapInChildContext());
     }
 }
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
similarity index 56%
rename from sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
rename to sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
index 31b917b23..8a3e1bade 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryHelperTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
@@ -1,6 +1,6 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.util;
+package software.amazon.lambda.durable.context;
 
 import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.ArgumentMatchers.*;
@@ -8,6 +8,7 @@
 
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Objects;
 import java.util.function.Function;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Nested;
@@ -19,16 +20,139 @@
 import software.amazon.lambda.durable.exception.SerDesException;
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.execution.SuspendExecutionException;
+import software.amazon.lambda.durable.model.WithRetry;
 import software.amazon.lambda.durable.retry.RetryDecision;
 import software.amazon.lambda.durable.retry.RetryStrategies;
+import software.amazon.lambda.durable.util.CompletedDurableFuture;
 
-class WithRetryHelperTest {
+@SuppressWarnings("unchecked")
+class DurableContextWithRetryTest {
 
     private DurableContext context;
+    private DurableContext childContext;
 
     @BeforeEach
     void setUp() {
         context = mock(DurableContext.class);
+        childContext = mock(DurableContext.class);
+        stubWithRetryMethods(context);
+        stubWithRetryMethods(childContext);
+    }
+
+    /**
+     * Stubs the withRetry/withRetryAsync methods on a mock DurableContext so they execute the real retry loop logic.
+     * This is needed because withRetry/withRetryAsync are abstract interface methods, and mocks return null by default.
+     */
+    private void stubWithRetryMethods(DurableContext mock) {
+        // Named sync form
+        when(mock.withRetry(any(), nullable(WithRetry.class), nullable(WithRetryConfig.class)))
+                .thenAnswer(invocation -> {
+                    String name = invocation.getArgument(0);
+                    WithRetry operation = invocation.getArgument(1);
+                    WithRetryConfig config = invocation.getArgument(2);
+                    Objects.requireNonNull(name, "name cannot be null");
+                    Objects.requireNonNull(operation, "operation cannot be null");
+                    Objects.requireNonNull(config, "config cannot be null");
+                    if (config.wrapInChildContext()) {
+                        return mock.runInChildContext(
+                                name,
+                                new TypeToken() {},
+                                childCtx -> executeRetryLoop(childCtx, name, operation, config));
+                    }
+                    return executeRetryLoop(mock, name, operation, config);
+                });
+
+        // Anonymous sync form
+        when(mock.withRetry(nullable(WithRetry.class), nullable(WithRetryConfig.class)))
+                .thenAnswer(invocation -> {
+                    WithRetry operation = invocation.getArgument(0);
+                    WithRetryConfig config = invocation.getArgument(1);
+                    Objects.requireNonNull(operation, "operation cannot be null");
+                    Objects.requireNonNull(config, "config cannot be null");
+                    if (config.wrapInChildContext()) {
+                        return mock.runInChildContext(
+                                "retry",
+                                new TypeToken() {},
+                                childCtx -> executeRetryLoop(childCtx, null, operation, config));
+                    }
+                    return executeRetryLoop(mock, null, operation, config);
+                });
+
+        // Named async form
+        when(mock.withRetryAsync(any(), nullable(WithRetry.class), nullable(WithRetryConfig.class)))
+                .thenAnswer(invocation -> {
+                    String name = invocation.getArgument(0);
+                    WithRetry operation = invocation.getArgument(1);
+                    WithRetryConfig config = invocation.getArgument(2);
+                    Objects.requireNonNull(name, "name cannot be null");
+                    Objects.requireNonNull(operation, "operation cannot be null");
+                    Objects.requireNonNull(config, "config cannot be null");
+                    return mock.runInChildContextAsync(
+                            name,
+                            new TypeToken() {},
+                            childCtx -> executeRetryLoop(childCtx, name, operation, config));
+                });
+
+        // Anonymous async form
+        when(mock.withRetryAsync(nullable(WithRetry.class), nullable(WithRetryConfig.class)))
+                .thenAnswer(invocation -> {
+                    WithRetry operation = invocation.getArgument(0);
+                    WithRetryConfig config = invocation.getArgument(1);
+                    Objects.requireNonNull(operation, "operation cannot be null");
+                    Objects.requireNonNull(config, "config cannot be null");
+                    return mock.runInChildContextAsync(
+                            "retry",
+                            new TypeToken() {},
+                            childCtx -> executeRetryLoop(childCtx, null, operation, config));
+                });
+    }
+
+    /** Replicates the retry loop logic from DurableContextImpl for test stubbing. */
+    private static  T executeRetryLoop(
+            DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
+        var attempt = 1;
+        while (true) {
+            try {
+                return operation.execute(context, attempt);
+            } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) {
+                throw e;
+            } catch (Exception e) {
+                var decision = config.retryStrategy().makeRetryDecision(e, attempt);
+                if (!decision.shouldRetry()) {
+                    throw e;
+                }
+                var delay = decision.delay().isZero() ? Duration.ofSeconds(1) : decision.delay();
+                var waitName = name != null ? name + "-backoff-" + attempt : "retry-backoff-" + attempt;
+                context.wait(waitName, delay);
+                attempt++;
+            }
+        }
+    }
+
+    /** Stubs runInChildContextAsync to immediately execute the function with childContext. */
+    private void stubChildContext(String name) {
+        when(context.runInChildContextAsync(eq(name), any(TypeToken.class), any()))
+                .thenAnswer(invocation -> {
+                    Function func = invocation.getArgument(2);
+                    return new CompletedDurableFuture<>(func.apply(childContext));
+                });
+    }
+
+    /** Stubs runInChildContextAsync for any name to immediately execute the function with childContext. */
+    private void stubChildContextAnyName() {
+        when(context.runInChildContextAsync(anyString(), any(TypeToken.class), any()))
+                .thenAnswer(invocation -> {
+                    Function func = invocation.getArgument(2);
+                    return new CompletedDurableFuture<>(func.apply(childContext));
+                });
+    }
+
+    /** Stubs runInChildContext (sync) to immediately execute the function with childContext. */
+    private void stubChildContextSync(String name) {
+        when(context.runInChildContext(eq(name), any(TypeToken.class), any())).thenAnswer(invocation -> {
+            Function func = invocation.getArgument(2);
+            return func.apply(childContext);
+        });
     }
 
     // --- Named form tests ---
@@ -37,49 +161,37 @@ void setUp() {
     class NamedForm {
 
         @Test
-        void successOnFirstAttempt_wrapsInChildContext() {
-            // runInChildContextAsync should be called; delegate to the function immediately
-            when(context.runInChildContextAsync(eq("my-op"), any(TypeToken.class), any()))
-                    .thenAnswer(invocation -> {
-                        Function func = invocation.getArgument(2);
-                        return new CompletedDurableFuture<>(func.apply(context));
-                    });
-
+        void successOnFirstAttempt_runsDirectlyOnContext() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var result = WithRetryHelper.withRetry(context, "my-op", (ctx, attempt) -> "success", config);
+            var result = context.withRetry("my-op", (ctx, attempt) -> "success", config);
 
             assertEquals("success", result);
-            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
+            verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
         }
 
         @Test
-        void successOnFirstAttempt_noChildContext_whenDisabled() {
+        void doesNotWrapInChildContext() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(false)
                     .build();
 
-            var result = WithRetryHelper.withRetry(context, "my-op", (ctx, attempt) -> "direct", config);
+            context.withRetry("my-op", (ctx, attempt) -> "result", config);
 
-            assertEquals("direct", result);
             verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
         }
 
         @Test
         void retriesWithBackoffWaits_namedForm() {
-            // Disable child context wrapping so we can directly verify wait calls
             var callCount = new int[] {0};
             var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) ->
                             attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
-                    .wrapInChildContext(false)
                     .build();
 
-            var result = WithRetryHelper.withRetry(
-                    context,
+            var result = context.withRetry(
                     "my-op",
                     (ctx, attempt) -> {
                         callCount[0]++;
@@ -92,6 +204,7 @@ void retriesWithBackoffWaits_namedForm() {
 
             assertEquals("success-on-3", result);
             assertEquals(3, callCount[0]);
+            // Backoff waits happen on the caller's context directly
             verify(context).wait("my-op-backoff-1", Duration.ofSeconds(5));
             verify(context).wait("my-op-backoff-2", Duration.ofSeconds(5));
             verify(context, times(2)).wait(anyString(), any(Duration.class));
@@ -101,13 +214,11 @@ void retriesWithBackoffWaits_namedForm() {
         void rethrowsWhenRetryStrategyReturnsFail() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(false)
                     .build();
 
             var exception = assertThrows(
                     RuntimeException.class,
-                    () -> WithRetryHelper.withRetry(
-                            context,
+                    () -> context.withRetry(
                             "my-op",
                             (ctx, attempt) -> {
                                 throw new RuntimeException("terminal");
@@ -123,12 +234,10 @@ void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(
                             (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail())
-                    .wrapInChildContext(false)
                     .build();
 
             var callCount = new int[] {0};
-            var result = WithRetryHelper.withRetry(
-                    context,
+            var result = context.withRetry(
                     "my-op",
                     (ctx, attempt) -> {
                         callCount[0]++;
@@ -144,25 +253,13 @@ void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
             verify(context).wait("my-op-backoff-1", Duration.ofSeconds(1));
         }
 
-        @Test
-        void nullContext_shouldThrow() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class, () -> WithRetryHelper.withRetry(null, "name", (ctx, a) -> "x", config));
-        }
-
         @Test
         void nullName_shouldThrow() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            assertThrows(
-                    NullPointerException.class,
-                    () -> WithRetryHelper.withRetry(context, null, (ctx, a) -> "x", config));
+            assertThrows(NullPointerException.class, () -> context.withRetry(null, (ctx, a) -> "x", config));
         }
 
         @Test
@@ -171,28 +268,39 @@ void nullOperation_shouldThrow() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            assertThrows(NullPointerException.class, () -> WithRetryHelper.withRetry(context, "name", null, config));
+            assertThrows(NullPointerException.class, () -> context.withRetry("name", null, config));
         }
 
         @Test
         void nullConfig_shouldThrow() {
-            assertThrows(
-                    NullPointerException.class,
-                    () -> WithRetryHelper.withRetry(context, "name", (ctx, a) -> "x", null));
+            assertThrows(NullPointerException.class, () -> context.withRetry("name", (ctx, a) -> "x", null));
         }
 
         @Test
         void operationReturnsNull() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(false)
                     .build();
 
-            var result =
-                    WithRetryHelper.withRetry(context, "my-op", (WithRetry) (ctx, attempt) -> null, config);
+            var result = context.withRetry("my-op", (WithRetry) (ctx, attempt) -> null, config);
 
             assertNull(result);
         }
+
+        @Test
+        void passesCallerContextToOperation() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            context.withRetry(
+                    "my-op",
+                    (ctx, attempt) -> {
+                        assertSame(context, ctx);
+                        return "verified";
+                    },
+                    config);
+        }
     }
 
     // --- Anonymous form tests ---
@@ -206,9 +314,19 @@ void successOnFirstAttempt() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var result = WithRetryHelper.withRetry(context, (ctx, attempt) -> "anonymous-success", config);
+            var result = context.withRetry((ctx, attempt) -> "anonymous-success", config);
 
             assertEquals("anonymous-success", result);
+        }
+
+        @Test
+        void doesNotWrapInChildContext() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            context.withRetry((ctx, attempt) -> "result", config);
+
             verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
         }
 
@@ -219,8 +337,7 @@ void retriesWithAnonymousBackoffNames() {
                             attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
                     .build();
 
-            var result = WithRetryHelper.withRetry(
-                    context,
+            var result = context.withRetry(
                     (ctx, attempt) -> {
                         if (attempt < 3) {
                             throw new RuntimeException("fail");
@@ -234,18 +351,6 @@ void retriesWithAnonymousBackoffNames() {
             verify(context).wait("retry-backoff-2", Duration.ofSeconds(2));
         }
 
-        @Test
-        void neverWrapsInChildContext_evenWhenConfigSaysTrue() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(true) // should be ignored for anonymous form
-                    .build();
-
-            WithRetryHelper.withRetry(context, (ctx, attempt) -> "result", config);
-
-            verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
-        }
-
         @Test
         void rethrowsOriginalException() {
             var config = WithRetryConfig.builder()
@@ -255,8 +360,7 @@ void rethrowsOriginalException() {
             var original = new IllegalStateException("original error");
             var thrown = assertThrows(
                     IllegalStateException.class,
-                    () -> WithRetryHelper.withRetry(
-                            context,
+                    () -> context.withRetry(
                             (ctx, attempt) -> {
                                 throw original;
                             },
@@ -265,29 +369,18 @@ void rethrowsOriginalException() {
             assertSame(original, thrown);
         }
 
-        @Test
-        void nullContext_shouldThrow() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(NullPointerException.class, () -> WithRetryHelper.withRetry(null, (ctx, a) -> "x", config));
-        }
-
         @Test
         void nullOperation_shouldThrow() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            assertThrows(
-                    NullPointerException.class,
-                    () -> WithRetryHelper.withRetry(context, (WithRetry) null, config));
+            assertThrows(NullPointerException.class, () -> context.withRetry((WithRetry) null, config));
         }
 
         @Test
         void nullConfig_shouldThrow() {
-            assertThrows(NullPointerException.class, () -> WithRetryHelper.withRetry(context, (ctx, a) -> "x", null));
+            assertThrows(NullPointerException.class, () -> context.withRetry((ctx, a) -> "x", null));
         }
 
         @Test
@@ -296,7 +389,7 @@ void operationReturnsNull() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var result = WithRetryHelper.withRetry(context, (WithRetry) (ctx, attempt) -> null, config);
+            var result = context.withRetry((WithRetry) (ctx, attempt) -> null, config);
 
             assertNull(result);
         }
@@ -308,8 +401,7 @@ void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
                             (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail())
                     .build();
 
-            var result = WithRetryHelper.withRetry(
-                    context,
+            var result = context.withRetry(
                     (ctx, attempt) -> {
                         if (attempt == 1) {
                             throw new RuntimeException("fail");
@@ -321,6 +413,20 @@ void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
             assertEquals("ok", result);
             verify(context).wait("retry-backoff-1", Duration.ofSeconds(1));
         }
+
+        @Test
+        void passesCallerContextToOperation() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            context.withRetry(
+                    (ctx, attempt) -> {
+                        assertSame(context, ctx);
+                        return "verified";
+                    },
+                    config);
+        }
     }
 
     // --- Async form tests ---
@@ -330,76 +436,69 @@ class AsyncForm {
 
         @Test
         void namedAsyncReturnsDurableFuture() {
-            when(context.runInChildContextAsync(eq("my-op"), any(TypeToken.class), any()))
-                    .thenAnswer(invocation -> {
-                        Function func = invocation.getArgument(2);
-                        return new CompletedDurableFuture<>(func.apply(context));
-                    });
+            stubChildContext("my-op");
 
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            DurableFuture future =
-                    WithRetryHelper.withRetryAsync(context, "my-op", (ctx, attempt) -> "async-result", config);
+            DurableFuture future = context.withRetryAsync("my-op", (ctx, attempt) -> "async-result", config);
 
             assertNotNull(future);
             assertEquals("async-result", future.get());
         }
 
         @Test
-        void namedAsyncWithoutChildContext_returnsCompletedFuture() {
+        void namedAsyncAlwaysWrapsInChildContext() {
+            stubChildContext("my-op");
+
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(false)
                     .build();
 
-            DurableFuture future =
-                    WithRetryHelper.withRetryAsync(context, "my-op", (ctx, attempt) -> "direct-async", config);
+            context.withRetryAsync("my-op", (ctx, attempt) -> "result", config);
 
-            assertNotNull(future);
-            assertInstanceOf(CompletedDurableFuture.class, future);
-            assertEquals("direct-async", future.get());
-            verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
+            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
         }
 
         @Test
         void anonymousAsyncReturnsDurableFuture() {
+            stubChildContext("retry");
+
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            DurableFuture future =
-                    WithRetryHelper.withRetryAsync(context, (ctx, attempt) -> "anon-async", config);
+            DurableFuture future = context.withRetryAsync((ctx, attempt) -> "anon-async", config);
 
             assertNotNull(future);
-            assertInstanceOf(CompletedDurableFuture.class, future);
             assertEquals("anon-async", future.get());
         }
 
         @Test
-        void anonymousAsyncNeverWrapsInChildContext() {
+        void anonymousAsyncWrapsInChildContext() {
+            stubChildContext("retry");
+
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(true)
                     .build();
 
-            WithRetryHelper.withRetryAsync(context, (ctx, attempt) -> "result", config);
+            context.withRetryAsync((ctx, attempt) -> "result", config);
 
-            verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
+            verify(context).runInChildContextAsync(eq("retry"), any(TypeToken.class), any());
         }
 
         @Test
         void namedAsyncRetriesWithBackoff() {
+            stubChildContext("my-op");
+
             var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) ->
                             attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
-                    .wrapInChildContext(false)
                     .build();
 
             var callCount = new int[] {0};
-            DurableFuture future = WithRetryHelper.withRetryAsync(
-                    context,
+            DurableFuture future = context.withRetryAsync(
                     "my-op",
                     (ctx, attempt) -> {
                         callCount[0]++;
@@ -412,19 +511,20 @@ void namedAsyncRetriesWithBackoff() {
 
             assertEquals("success-on-3", future.get());
             assertEquals(3, callCount[0]);
-            verify(context).wait("my-op-backoff-1", Duration.ofSeconds(5));
-            verify(context).wait("my-op-backoff-2", Duration.ofSeconds(5));
+            verify(childContext).wait("my-op-backoff-1", Duration.ofSeconds(5));
+            verify(childContext).wait("my-op-backoff-2", Duration.ofSeconds(5));
         }
 
         @Test
         void anonymousAsyncRetriesWithBackoff() {
+            stubChildContext("retry");
+
             var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) ->
                             attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
                     .build();
 
-            DurableFuture future = WithRetryHelper.withRetryAsync(
-                    context,
+            DurableFuture future = context.withRetryAsync(
                     (ctx, attempt) -> {
                         if (attempt < 3) {
                             throw new RuntimeException("fail");
@@ -434,19 +534,8 @@ void anonymousAsyncRetriesWithBackoff() {
                     config);
 
             assertEquals("done", future.get());
-            verify(context).wait("retry-backoff-1", Duration.ofSeconds(2));
-            verify(context).wait("retry-backoff-2", Duration.ofSeconds(2));
-        }
-
-        @Test
-        void namedAsyncNullContext_shouldThrow() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class,
-                    () -> WithRetryHelper.withRetryAsync(null, "name", (ctx, a) -> "x", config));
+            verify(childContext).wait("retry-backoff-1", Duration.ofSeconds(2));
+            verify(childContext).wait("retry-backoff-2", Duration.ofSeconds(2));
         }
 
         @Test
@@ -455,9 +544,7 @@ void namedAsyncNullName_shouldThrow() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            assertThrows(
-                    NullPointerException.class,
-                    () -> WithRetryHelper.withRetryAsync(context, null, (ctx, a) -> "x", config));
+            assertThrows(NullPointerException.class, () -> context.withRetryAsync(null, (ctx, a) -> "x", config));
         }
 
         @Test
@@ -466,25 +553,12 @@ void namedAsyncNullOperation_shouldThrow() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            assertThrows(
-                    NullPointerException.class, () -> WithRetryHelper.withRetryAsync(context, "name", null, config));
+            assertThrows(NullPointerException.class, () -> context.withRetryAsync("name", null, config));
         }
 
         @Test
         void namedAsyncNullConfig_shouldThrow() {
-            assertThrows(
-                    NullPointerException.class,
-                    () -> WithRetryHelper.withRetryAsync(context, "name", (ctx, a) -> "x", null));
-        }
-
-        @Test
-        void anonymousAsyncNullContext_shouldThrow() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class, () -> WithRetryHelper.withRetryAsync(null, (ctx, a) -> "x", config));
+            assertThrows(NullPointerException.class, () -> context.withRetryAsync("name", (ctx, a) -> "x", null));
         }
 
         @Test
@@ -493,40 +567,38 @@ void anonymousAsyncNullOperation_shouldThrow() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            assertThrows(
-                    NullPointerException.class,
-                    () -> WithRetryHelper.withRetryAsync(context, (WithRetry) null, config));
+            assertThrows(NullPointerException.class, () -> context.withRetryAsync((WithRetry) null, config));
         }
 
         @Test
         void anonymousAsyncNullConfig_shouldThrow() {
-            assertThrows(
-                    NullPointerException.class, () -> WithRetryHelper.withRetryAsync(context, (ctx, a) -> "x", null));
+            assertThrows(NullPointerException.class, () -> context.withRetryAsync((ctx, a) -> "x", null));
         }
 
         @Test
         void asyncOperationReturnsNull() {
+            stubChildContext("my-op");
+
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(false)
                     .build();
 
-            DurableFuture future = WithRetryHelper.withRetryAsync(
-                    context, "my-op", (WithRetry) (ctx, attempt) -> null, config);
+            DurableFuture future =
+                    context.withRetryAsync("my-op", (WithRetry) (ctx, attempt) -> null, config);
 
             assertNull(future.get());
         }
 
         @Test
-        void withRetryDelegatesToWithRetryAsync() {
-            // Verify that withRetry (sync) produces the same result as withRetryAsync().get()
+        void syncAndAsyncProduceSameResult() {
+            stubChildContextAnyName();
+
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(false)
                     .build();
 
-            var syncResult = WithRetryHelper.withRetry(context, "op", (ctx, attempt) -> "value", config);
-            var asyncResult = WithRetryHelper.withRetryAsync(context, "op", (ctx, attempt) -> "value", config)
+            var syncResult = context.withRetry("op", (ctx, attempt) -> "value", config);
+            var asyncResult = context.withRetryAsync("op", (ctx, attempt) -> "value", config)
                     .get();
 
             assertEquals(syncResult, asyncResult);
@@ -544,11 +616,34 @@ void passesCorrectAttemptNumberToOperation() {
             var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) ->
                             attempt < 4 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
-                    .wrapInChildContext(false)
                     .build();
 
-            WithRetryHelper.withRetry(
-                    context,
+            context.withRetry(
+                    (ctx, attempt) -> {
+                        attempts.add(attempt);
+                        if (attempt < 4) {
+                            throw new RuntimeException("not yet");
+                        }
+                        return "done";
+                    },
+                    config);
+
+            assertEquals(4, attempts.size());
+            assertEquals(1, attempts.get(0));
+            assertEquals(2, attempts.get(1));
+            assertEquals(3, attempts.get(2));
+            assertEquals(4, attempts.get(3));
+        }
+
+        @Test
+        void passesCorrectAttemptNumberToOperation_namedForm() {
+            var attempts = new ArrayList();
+            var config = WithRetryConfig.builder()
+                    .retryStrategy((error, attempt) ->
+                            attempt < 4 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
+                    .build();
+
+            context.withRetry(
                     "track",
                     (ctx, attempt) -> {
                         attempts.add(attempt);
@@ -578,8 +673,7 @@ void passesErrorToRetryStrategy() {
 
             assertThrows(
                     RuntimeException.class,
-                    () -> WithRetryHelper.withRetry(
-                            context,
+                    () -> context.withRetry(
                             (ctx, attempt) -> {
                                 throw new RuntimeException("error-" + attempt);
                             },
@@ -597,8 +691,7 @@ void respectsCustomDelayFromRetryDecision() {
                     .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(attempt * 10L)))
                     .build();
 
-            WithRetryHelper.withRetry(
-                    context,
+            context.withRetry(
                     (ctx, attempt) -> {
                         if (attempt <= 2) {
                             throw new RuntimeException("fail");
@@ -612,13 +705,12 @@ void respectsCustomDelayFromRetryDecision() {
         }
 
         @Test
-        void passesContextToOperation() {
+        void anonymousFormPassesCallerContextToOperation() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            WithRetryHelper.withRetry(
-                    context,
+            context.withRetry(
                     (ctx, attempt) -> {
                         assertSame(context, ctx);
                         return "verified";
@@ -627,24 +719,15 @@ void passesContextToOperation() {
         }
 
         @Test
-        void namedFormPassesChildContextToOperation_whenWrapped() {
-            var childContext = mock(DurableContext.class);
-            when(context.runInChildContextAsync(eq("wrapped"), any(TypeToken.class), any()))
-                    .thenAnswer(invocation -> {
-                        Function func = invocation.getArgument(2);
-                        return new CompletedDurableFuture<>(func.apply(childContext));
-                    });
-
+        void namedFormPassesCallerContextToOperation() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(true)
                     .build();
 
-            WithRetryHelper.withRetry(
-                    context,
+            context.withRetry(
                     "wrapped",
                     (ctx, attempt) -> {
-                        assertSame(childContext, ctx);
+                        assertSame(context, ctx);
                         return "verified";
                     },
                     config);
@@ -658,8 +741,7 @@ void rethrowsLastExceptionWhenAllRetriesExhausted() {
 
             var thrown = assertThrows(
                     RuntimeException.class,
-                    () -> WithRetryHelper.withRetry(
-                            context,
+                    () -> context.withRetry(
                             (ctx, attempt) -> {
                                 throw new RuntimeException("attempt-" + attempt);
                             },
@@ -677,8 +759,7 @@ void propagatesSuspendExecutionExceptionWithoutRetrying() {
 
             assertThrows(
                     SuspendExecutionException.class,
-                    () -> WithRetryHelper.withRetry(
-                            context,
+                    () -> context.withRetry(
                             (ctx, attempt) -> {
                                 throw new SuspendExecutionException();
                             },
@@ -696,8 +777,7 @@ void propagatesUnrecoverableDurableExecutionExceptionWithoutRetrying() {
 
             assertThrows(
                     UnrecoverableDurableExecutionException.class,
-                    () -> WithRetryHelper.withRetry(
-                            context,
+                    () -> context.withRetry(
                             (ctx, attempt) -> {
                                 throw new UnrecoverableDurableExecutionException(
                                         software.amazon.awssdk.services.lambda.model.ErrorObject.builder()
@@ -717,8 +797,7 @@ void propagatesSuspendExecutionExceptionOnLaterAttempt() {
 
             assertThrows(
                     SuspendExecutionException.class,
-                    () -> WithRetryHelper.withRetry(
-                            context,
+                    () -> context.withRetry(
                             (ctx, attempt) -> {
                                 if (attempt == 1) {
                                     throw new RuntimeException("transient");
@@ -740,8 +819,7 @@ void propagatesUnrecoverableDurableExecutionExceptionOnLaterAttempt() {
 
             assertThrows(
                     UnrecoverableDurableExecutionException.class,
-                    () -> WithRetryHelper.withRetry(
-                            context,
+                    () -> context.withRetry(
                             (ctx, attempt) -> {
                                 if (attempt == 1) {
                                     throw new RuntimeException("transient");
@@ -766,8 +844,7 @@ void preservesCheckedExceptionSubclassType() {
 
             var thrown = assertThrows(
                     SerDesException.class,
-                    () -> WithRetryHelper.withRetry(
-                            context,
+                    () -> context.withRetry(
                             (ctx, attempt) -> {
                                 throw original;
                             },
@@ -777,4 +854,144 @@ void preservesCheckedExceptionSubclassType() {
             assertEquals("deserialization failed", thrown.getMessage());
         }
     }
+
+    // --- WrapInChildContext tests ---
+
+    @Nested
+    class WrapInChildContext {
+
+        @Test
+        void namedForm_wrapsInChildContextWhenEnabled() {
+            stubChildContextSync("my-op");
+
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(true)
+                    .build();
+
+            var result = context.withRetry("my-op", (ctx, attempt) -> "wrapped", config);
+
+            assertEquals("wrapped", result);
+            verify(context).runInChildContext(eq("my-op"), any(TypeToken.class), any());
+        }
+
+        @Test
+        void namedForm_doesNotWrapWhenDisabled() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(false)
+                    .build();
+
+            var result = context.withRetry("my-op", (ctx, attempt) -> "direct", config);
+
+            assertEquals("direct", result);
+            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+        }
+
+        @Test
+        void namedForm_doesNotWrapByDefault() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            context.withRetry("my-op", (ctx, attempt) -> "direct", config);
+
+            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+        }
+
+        @Test
+        void anonymousForm_wrapsInChildContextWhenEnabled() {
+            stubChildContextSync("retry");
+
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(true)
+                    .build();
+
+            var result = context.withRetry((ctx, attempt) -> "wrapped", config);
+
+            assertEquals("wrapped", result);
+            verify(context).runInChildContext(eq("retry"), any(TypeToken.class), any());
+        }
+
+        @Test
+        void anonymousForm_doesNotWrapWhenDisabled() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(false)
+                    .build();
+
+            var result = context.withRetry((ctx, attempt) -> "direct", config);
+
+            assertEquals("direct", result);
+            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+        }
+
+        @Test
+        void namedForm_passesChildContextToOperationWhenWrapped() {
+            stubChildContextSync("my-op");
+
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(true)
+                    .build();
+
+            context.withRetry(
+                    "my-op",
+                    (ctx, attempt) -> {
+                        assertSame(childContext, ctx);
+                        return "verified";
+                    },
+                    config);
+        }
+
+        @Test
+        void namedForm_retriesWithBackoffOnChildContextWhenWrapped() {
+            stubChildContextSync("my-op");
+
+            var config = WithRetryConfig.builder()
+                    .retryStrategy((error, attempt) ->
+                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
+                    .wrapInChildContext(true)
+                    .build();
+
+            var result = context.withRetry(
+                    "my-op",
+                    (ctx, attempt) -> {
+                        if (attempt < 3) {
+                            throw new RuntimeException("fail-" + attempt);
+                        }
+                        return "success-on-3";
+                    },
+                    config);
+
+            assertEquals("success-on-3", result);
+            verify(childContext).wait("my-op-backoff-1", Duration.ofSeconds(5));
+            verify(childContext).wait("my-op-backoff-2", Duration.ofSeconds(5));
+        }
+
+        @Test
+        void anonymousForm_retriesWithBackoffOnChildContextWhenWrapped() {
+            stubChildContextSync("retry");
+
+            var config = WithRetryConfig.builder()
+                    .retryStrategy((error, attempt) ->
+                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
+                    .wrapInChildContext(true)
+                    .build();
+
+            var result = context.withRetry(
+                    (ctx, attempt) -> {
+                        if (attempt < 3) {
+                            throw new RuntimeException("fail");
+                        }
+                        return "done";
+                    },
+                    config);
+
+            assertEquals("done", result);
+            verify(childContext).wait("retry-backoff-1", Duration.ofSeconds(2));
+            verify(childContext).wait("retry-backoff-2", Duration.ofSeconds(2));
+        }
+    }
 }
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryTest.java b/sdk/src/test/java/software/amazon/lambda/durable/model/WithRetryTest.java
similarity index 97%
rename from sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryTest.java
rename to sdk/src/test/java/software/amazon/lambda/durable/model/WithRetryTest.java
index 0b1069e8e..db6c83953 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/util/WithRetryTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/model/WithRetryTest.java
@@ -1,6 +1,6 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.util;
+package software.amazon.lambda.durable.model;
 
 import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.Mockito.*;
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
deleted file mode 100644
index 265fa968c..000000000
--- a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
+++ /dev/null
@@ -1,462 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.util;
-
-import static org.junit.jupiter.api.Assertions.*;
-import static org.mockito.ArgumentMatchers.*;
-import static org.mockito.Mockito.*;
-
-import java.time.Duration;
-import java.util.ArrayList;
-import java.util.function.Function;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Nested;
-import org.junit.jupiter.api.Test;
-import software.amazon.lambda.durable.DurableContext;
-import software.amazon.lambda.durable.TypeToken;
-import software.amazon.lambda.durable.config.RetryOperationConfig;
-import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
-import software.amazon.lambda.durable.execution.SuspendExecutionException;
-import software.amazon.lambda.durable.retry.RetryDecision;
-import software.amazon.lambda.durable.retry.RetryStrategies;
-
-class RetryOperationHelperTest {
-
-    private DurableContext context;
-
-    @BeforeEach
-    void setUp() {
-        context = mock(DurableContext.class);
-    }
-
-    // --- Named form tests ---
-
-    @Nested
-    class NamedForm {
-
-        @Test
-        void successOnFirstAttempt_wrapsInChildContext() {
-            // runInChildContext should be called; delegate to the function immediately
-            when(context.runInChildContext(eq("my-op"), any(TypeToken.class), any()))
-                    .thenAnswer(invocation -> {
-                        Function func = invocation.getArgument(2);
-                        return func.apply(context);
-                    });
-
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var result = RetryOperationHelper.retryOperation(context, "my-op", (ctx, attempt) -> "success", config);
-
-            assertEquals("success", result);
-            verify(context).runInChildContext(eq("my-op"), any(TypeToken.class), any());
-        }
-
-        @Test
-        void successOnFirstAttempt_noChildContext_whenDisabled() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(false)
-                    .build();
-
-            var result = RetryOperationHelper.retryOperation(context, "my-op", (ctx, attempt) -> "direct", config);
-
-            assertEquals("direct", result);
-            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
-        }
-
-        @Test
-        void retriesWithBackoffWaits_namedForm() {
-            // Disable child context wrapping so we can directly verify wait calls
-            var callCount = new int[] {0};
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) ->
-                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
-                    .wrapInChildContext(false)
-                    .build();
-
-            var result = RetryOperationHelper.retryOperation(
-                    context,
-                    "my-op",
-                    (ctx, attempt) -> {
-                        callCount[0]++;
-                        if (attempt < 3) {
-                            throw new RuntimeException("fail-" + attempt);
-                        }
-                        return "success-on-3";
-                    },
-                    config);
-
-            assertEquals("success-on-3", result);
-            assertEquals(3, callCount[0]);
-            verify(context).wait("my-op-backoff-1", Duration.ofSeconds(5));
-            verify(context).wait("my-op-backoff-2", Duration.ofSeconds(5));
-            verify(context, times(2)).wait(anyString(), any(Duration.class));
-        }
-
-        @Test
-        void rethrowsWhenRetryStrategyReturnsFail() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(false)
-                    .build();
-
-            var exception = assertThrows(
-                    RuntimeException.class,
-                    () -> RetryOperationHelper.retryOperation(
-                            context,
-                            "my-op",
-                            (ctx, attempt) -> {
-                                throw new RuntimeException("terminal");
-                            },
-                            config));
-
-            assertEquals("terminal", exception.getMessage());
-            verify(context, never()).wait(anyString(), any(Duration.class));
-        }
-
-        @Test
-        void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(
-                            (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail())
-                    .wrapInChildContext(false)
-                    .build();
-
-            var callCount = new int[] {0};
-            var result = RetryOperationHelper.retryOperation(
-                    context,
-                    "my-op",
-                    (ctx, attempt) -> {
-                        callCount[0]++;
-                        if (attempt == 1) {
-                            throw new RuntimeException("fail");
-                        }
-                        return "ok";
-                    },
-                    config);
-
-            assertEquals("ok", result);
-            // Zero delay should be replaced with 1-second default
-            verify(context).wait("my-op-backoff-1", Duration.ofSeconds(1));
-        }
-
-        @Test
-        void nullContext_shouldThrow() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(null, "name", (ctx, a) -> "x", config));
-        }
-
-        @Test
-        void nullName_shouldThrow() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, null, (ctx, a) -> "x", config));
-        }
-
-        @Test
-        void nullOperation_shouldThrow() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, "name", null, config));
-        }
-
-        @Test
-        void nullConfig_shouldThrow() {
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, "name", (ctx, a) -> "x", null));
-        }
-    }
-
-    // --- Anonymous form tests ---
-
-    @Nested
-    class AnonymousForm {
-
-        @Test
-        void successOnFirstAttempt() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var result = RetryOperationHelper.retryOperation(context, (ctx, attempt) -> "anonymous-success", config);
-
-            assertEquals("anonymous-success", result);
-            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
-        }
-
-        @Test
-        void retriesWithAnonymousBackoffNames() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) ->
-                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
-                    .build();
-
-            var result = RetryOperationHelper.retryOperation(
-                    context,
-                    (ctx, attempt) -> {
-                        if (attempt < 3) {
-                            throw new RuntimeException("fail");
-                        }
-                        return "done";
-                    },
-                    config);
-
-            assertEquals("done", result);
-            verify(context).wait("retry-backoff-1", Duration.ofSeconds(2));
-            verify(context).wait("retry-backoff-2", Duration.ofSeconds(2));
-        }
-
-        @Test
-        void neverWrapsInChildContext_evenWhenConfigSaysTrue() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(true) // should be ignored for anonymous form
-                    .build();
-
-            RetryOperationHelper.retryOperation(context, (ctx, attempt) -> "result", config);
-
-            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
-        }
-
-        @Test
-        void rethrowsOriginalException() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var original = new IllegalStateException("original error");
-            var thrown = assertThrows(
-                    IllegalStateException.class,
-                    () -> RetryOperationHelper.retryOperation(
-                            context,
-                            (ctx, attempt) -> {
-                                throw original;
-                            },
-                            config));
-
-            assertSame(original, thrown);
-        }
-
-        @Test
-        void nullContext_shouldThrow() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(null, (ctx, a) -> "x", config));
-        }
-
-        @Test
-        void nullOperation_shouldThrow() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, (RetryableOperation) null, config));
-        }
-
-        @Test
-        void nullConfig_shouldThrow() {
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, (ctx, a) -> "x", null));
-        }
-    }
-
-    // --- Retry behavior tests ---
-
-    @Nested
-    class RetryBehavior {
-
-        @Test
-        void passesCorrectAttemptNumberToOperation() {
-            var attempts = new ArrayList();
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) ->
-                            attempt < 4 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
-                    .wrapInChildContext(false)
-                    .build();
-
-            RetryOperationHelper.retryOperation(
-                    context,
-                    "track",
-                    (ctx, attempt) -> {
-                        attempts.add(attempt);
-                        if (attempt < 4) {
-                            throw new RuntimeException("not yet");
-                        }
-                        return "done";
-                    },
-                    config);
-
-            assertEquals(4, attempts.size());
-            assertEquals(1, attempts.get(0));
-            assertEquals(2, attempts.get(1));
-            assertEquals(3, attempts.get(2));
-            assertEquals(4, attempts.get(3));
-        }
-
-        @Test
-        void passesErrorToRetryStrategy() {
-            var errors = new ArrayList();
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) -> {
-                        errors.add(error);
-                        return attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail();
-                    })
-                    .build();
-
-            assertThrows(
-                    RuntimeException.class,
-                    () -> RetryOperationHelper.retryOperation(
-                            context,
-                            (ctx, attempt) -> {
-                                throw new RuntimeException("error-" + attempt);
-                            },
-                            config));
-
-            assertEquals(3, errors.size());
-            assertEquals("error-1", errors.get(0).getMessage());
-            assertEquals("error-2", errors.get(1).getMessage());
-            assertEquals("error-3", errors.get(2).getMessage());
-        }
-
-        @Test
-        void respectsCustomDelayFromRetryDecision() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(attempt * 10L)))
-                    .build();
-
-            RetryOperationHelper.retryOperation(
-                    context,
-                    (ctx, attempt) -> {
-                        if (attempt <= 2) {
-                            throw new RuntimeException("fail");
-                        }
-                        return "ok";
-                    },
-                    config);
-
-            verify(context).wait("retry-backoff-1", Duration.ofSeconds(10));
-            verify(context).wait("retry-backoff-2", Duration.ofSeconds(20));
-        }
-
-        @Test
-        void passesContextToOperation() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            RetryOperationHelper.retryOperation(
-                    context,
-                    (ctx, attempt) -> {
-                        assertSame(context, ctx);
-                        return "verified";
-                    },
-                    config);
-        }
-
-        @Test
-        void namedFormPassesChildContextToOperation_whenWrapped() {
-            var childContext = mock(DurableContext.class);
-            when(context.runInChildContext(eq("wrapped"), any(TypeToken.class), any()))
-                    .thenAnswer(invocation -> {
-                        Function func = invocation.getArgument(2);
-                        return func.apply(childContext);
-                    });
-
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(true)
-                    .build();
-
-            RetryOperationHelper.retryOperation(
-                    context,
-                    "wrapped",
-                    (ctx, attempt) -> {
-                        assertSame(childContext, ctx);
-                        return "verified";
-                    },
-                    config);
-        }
-
-        @Test
-        void rethrowsLastExceptionWhenAllRetriesExhausted() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
-                    .build();
-
-            var thrown = assertThrows(
-                    RuntimeException.class,
-                    () -> RetryOperationHelper.retryOperation(
-                            context,
-                            (ctx, attempt) -> {
-                                throw new RuntimeException("attempt-" + attempt);
-                            },
-                            config));
-
-            // The last attempt's exception is rethrown
-            assertEquals("attempt-3", thrown.getMessage());
-        }
-
-        @Test
-        void propagatesSuspendExecutionExceptionWithoutRetrying() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1)))
-                    .build();
-
-            assertThrows(
-                    SuspendExecutionException.class,
-                    () -> RetryOperationHelper.retryOperation(
-                            context,
-                            (ctx, attempt) -> {
-                                throw new SuspendExecutionException();
-                            },
-                            config));
-
-            // Should never reach the wait — SuspendExecutionException propagates immediately
-            verify(context, never()).wait(anyString(), any(Duration.class));
-        }
-
-        @Test
-        void propagatesUnrecoverableDurableExecutionExceptionWithoutRetrying() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1)))
-                    .build();
-
-            assertThrows(
-                    UnrecoverableDurableExecutionException.class,
-                    () -> RetryOperationHelper.retryOperation(
-                            context,
-                            (ctx, attempt) -> {
-                                throw new UnrecoverableDurableExecutionException(
-                                        software.amazon.awssdk.services.lambda.model.ErrorObject.builder()
-                                                .errorMessage("unrecoverable")
-                                                .build());
-                            },
-                            config));
-
-            verify(context, never()).wait(anyString(), any(Duration.class));
-        }
-    }
-}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java
deleted file mode 100644
index f5893a37c..000000000
--- a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java
+++ /dev/null
@@ -1,52 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.util;
-
-import static org.junit.jupiter.api.Assertions.*;
-import static org.mockito.Mockito.*;
-
-import org.junit.jupiter.api.Test;
-import software.amazon.lambda.durable.DurableContext;
-
-class RetryableOperationTest {
-
-    @Test
-    void canBeImplementedAsLambda() {
-        RetryableOperation operation = (ctx, attempt) -> "result-" + attempt;
-        var context = mock(DurableContext.class);
-
-        assertEquals("result-1", operation.execute(context, 1));
-        assertEquals("result-3", operation.execute(context, 3));
-    }
-
-    @Test
-    void receivesContextAndAttempt() {
-        var context = mock(DurableContext.class);
-        RetryableOperation operation = (ctx, attempt) -> {
-            assertSame(context, ctx);
-            return "attempt-" + attempt;
-        };
-
-        assertEquals("attempt-1", operation.execute(context, 1));
-        assertEquals("attempt-2", operation.execute(context, 2));
-    }
-
-    @Test
-    void canThrowExceptions() {
-        RetryableOperation operation = (ctx, attempt) -> {
-            throw new RuntimeException("failed on attempt " + attempt);
-        };
-        var context = mock(DurableContext.class);
-
-        var exception = assertThrows(RuntimeException.class, () -> operation.execute(context, 1));
-        assertEquals("failed on attempt 1", exception.getMessage());
-    }
-
-    @Test
-    void canReturnNull() {
-        RetryableOperation operation = (ctx, attempt) -> null;
-        var context = mock(DurableContext.class);
-
-        assertNull(operation.execute(context, 1));
-    }
-}

From 4badc897d1be94faf42ee92c4dccdf4047841beb Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Wed, 22 Apr 2026 15:04:25 -0700
Subject: [PATCH 12/19] feat: first attempt at retryableOperation

---
 .../durable/config/RetryOperationConfig.java  | 103 +++++
 .../durable/util/RetryOperationHelper.java    | 135 ++++++
 .../durable/util/RetryableOperation.java      |  26 ++
 .../config/RetryOperationConfigTest.java      |  80 ++++
 .../util/RetryOperationHelperTest.java        | 420 ++++++++++++++++++
 .../durable/util/RetryableOperationTest.java  |  52 +++
 6 files changed, 816 insertions(+)
 create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java
 create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
 create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java
 create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
 create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
 create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java

diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java
new file mode 100644
index 000000000..696d1d253
--- /dev/null
+++ b/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java
@@ -0,0 +1,103 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.config;
+
+import software.amazon.lambda.durable.retry.RetryStrategy;
+
+/**
+ * Configuration for {@link software.amazon.lambda.durable.util.RetryOperationHelper#retryOperation}.
+ *
+ * 
Uses the same {@link RetryStrategy} shape that developers already know from {@link StepConfig}, so there are zero
+ * new retry concepts to learn.
+ */
+public class RetryOperationConfig {
+    private final RetryStrategy retryStrategy;
+    private final boolean wrapInChildContext;
+
+    private RetryOperationConfig(Builder builder) {
+        this.retryStrategy = builder.retryStrategy;
+        this.wrapInChildContext = builder.wrapInChildContext;
+    }
+
+    /**
+     * Returns the retry strategy. Same type as {@link StepConfig#retryStrategy()}.
+     *
+     * @return the retry strategy, never null
+     */
+    public RetryStrategy retryStrategy() {
+        return retryStrategy;
+    }
+
+    /**
+     * Whether to wrap the retry loop in {@code runInChildContext} so all attempts are grouped under a single named
+     * operation in execution history. Only applies when a name is provided to the named form of {@code retryOperation}.
+     * Defaults to {@code true}.
+     *
+     * @return true if child-context wrapping is enabled
+     */
+    public boolean wrapInChildContext() {
+        return wrapInChildContext;
+    }
+
+    /**
+     * Creates a new builder for {@code RetryOperationConfig}.
+     *
+     * @return a new builder instance
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /** Builder for creating {@link RetryOperationConfig} instances. */
+    public static class Builder {
+        private RetryStrategy retryStrategy;
+        private boolean wrapInChildContext = true;
+
+        private Builder() {}
+
+        /**
+         * Sets the retry strategy. Required.
+         *
+         * 
Reuses the exact same {@link RetryStrategy} interface from {@link StepConfig}. All existing factory
+         * methods ({@link software.amazon.lambda.durable.retry.RetryStrategies#exponentialBackoff},
+         * {@link software.amazon.lambda.durable.retry.RetryStrategies#fixedDelay}, presets, and custom lambdas) work
+         * without modification.
+         *
+         * @param retryStrategy the retry strategy to use
+         * @return this builder for method chaining
+         */
+        public Builder retryStrategy(RetryStrategy retryStrategy) {
+            this.retryStrategy = retryStrategy;
+            return this;
+        }
+
+        /**
+         * Controls whether the retry loop is wrapped in a child context. Only meaningful for the named form of
+         * {@code retryOperation}. Defaults to {@code true}.
+         *
+         * 
When {@code true}, all attempts and backoff waits are grouped under a single named operation in execution
+         * history, providing a cleaner view and isolated operation ID space. Set to {@code false} to flatten attempts
+         * into the parent context.
+         *
+         * @param wrapInChildContext whether to wrap in a child context
+         * @return this builder for method chaining
+         */
+        public Builder wrapInChildContext(boolean wrapInChildContext) {
+            this.wrapInChildContext = wrapInChildContext;
+            return this;
+        }
+
+        /**
+         * Builds the {@link RetryOperationConfig} instance.
+         *
+         * @return a new config with the configured options
+         * @throws IllegalArgumentException if retryStrategy is not set
+         */
+        public RetryOperationConfig build() {
+            if (retryStrategy == null) {
+                throw new IllegalArgumentException("retryStrategy is required");
+            }
+            return new RetryOperationConfig(this);
+        }
+    }
+}
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
new file mode 100644
index 000000000..c63f0aad0
--- /dev/null
+++ b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
@@ -0,0 +1,135 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.util;
+
+import java.time.Duration;
+import java.util.Objects;
+import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.TypeToken;
+import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.retry.RetryDecision;
+
+/**
+ * Replay-safe retry loop for any durable operation.
+ *
+ * 
Provides the same retry-with-backoff pattern that {@code context.step()} has built in, but for operations that
+ * cannot live inside a step ({@code waitForCallback}, {@code invoke}, {@code waitForCondition}, etc.).
+ *
+ * 
Every side-effect in the loop is a durable operation, so the loop is replay-safe by construction. On replay,
+ * completed operations return cached results instantly and the loop fast-forwards to the current attempt.
+ *
+ * 
Usage — callback retry

+ *
+ * 
{@code
+ * var result = RetryOperationHelper.retryOperation(
+ *     context,
+ *     "approval",
+ *     (ctx, attempt) -> ctx.waitForCallback(
+ *         "approval-" + attempt,
+ *         String.class,
+ *         (callbackId, stepCtx) -> sendApprovalEmail(approverEmail, callbackId)
+ *     ),
+ *     RetryOperationConfig.builder()
+ *         .retryStrategy(RetryStrategies.exponentialBackoff(
+ *             3, Duration.ofSeconds(2), Duration.ofSeconds(30), 2.0, JitterStrategy.FULL))
+ *         .build()
+ * );
+ * }

+ *
+ * 
Usage — invoke retry (anonymous form)

+ *
+ * 
{@code
+ * var result = RetryOperationHelper.retryOperation(
+ *     context,
+ *     (ctx, attempt) -> ctx.invoke(
+ *         "charge-" + attempt, paymentFnArn, new ChargeRequest(orderId), String.class),
+ *     RetryOperationConfig.builder()
+ *         .retryStrategy((err, att) -> att < 3
+ *             ? RetryDecision.retry(Duration.ofSeconds(1))
+ *             : RetryDecision.fail())
+ *         .build()
+ * );
+ * }

+ */
+public final class RetryOperationHelper {
+
+    private static final Duration DEFAULT_BACKOFF_DELAY = Duration.ofSeconds(1);
+    private static final String BACKOFF_SUFFIX = "-backoff-";
+    private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-";
+
+    private RetryOperationHelper() {
+        // utility class
+    }
+
+    /**
+     * Named form — wraps the retry loop in {@code runInChildContext} by default so all attempts are grouped under a
+     * single named operation in execution history.
+     *
+     * 
The child-context wrapping can be disabled via
+     * {@link RetryOperationConfig.Builder#wrapInChildContext(boolean)}.
+     *
+     * @param  the result type
+     * @param context the durable context
+     * @param name operation name (used for child context and backoff wait names)
+     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param config retry configuration including the retry strategy
+     * @return the operation result
+     */
+    @SuppressWarnings("unchecked")
+    public static  T retryOperation(
+            DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) {
+        Objects.requireNonNull(context, "context cannot be null");
+        Objects.requireNonNull(name, "name cannot be null");
+        Objects.requireNonNull(operation, "operation cannot be null");
+        Objects.requireNonNull(config, "config cannot be null");
+
+        if (config.wrapInChildContext()) {
+            return (T) context.runInChildContext(
+                    name, new TypeToken() {}, childCtx -> executeRetryLoop(childCtx, name, operation, config));
+        }
+        return executeRetryLoop(context, name, operation, config);
+    }
+
+    /**
+     * Anonymous form — runs the retry loop directly in the caller's context. No child-context wrapping is applied
+     * regardless of the {@code wrapInChildContext} config setting.
+     *
+     * @param  the result type
+     * @param context the durable context
+     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param config retry configuration including the retry strategy
+     * @return the operation result
+     */
+    public static  T retryOperation(
+            DurableContext context, RetryableOperation operation, RetryOperationConfig config) {
+        Objects.requireNonNull(context, "context cannot be null");
+        Objects.requireNonNull(operation, "operation cannot be null");
+        Objects.requireNonNull(config, "config cannot be null");
+
+        return executeRetryLoop(context, null, operation, config);
+    }
+
+    /**
+     * Core retry loop. Replay-safe because every side-effect is a durable operation: the user's operation calls durable
+     * primitives, and backoff uses {@code context.wait()}.
+     */
+    private static  T executeRetryLoop(
+            DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) {
+        var attempt = 1;
+        while (true) {
+            try {
+                return operation.execute(context, attempt);
+            } catch (Exception e) {
+                RetryDecision decision = config.retryStrategy().makeRetryDecision(e, attempt);
+                if (!decision.shouldRetry()) {
+                    throw e;
+                }
+
+                var delay = decision.delay().isZero() ? DEFAULT_BACKOFF_DELAY : decision.delay();
+                var waitName = name != null ? name + BACKOFF_SUFFIX + attempt : ANONYMOUS_BACKOFF_PREFIX + attempt;
+                context.wait(waitName, delay);
+                attempt++;
+            }
+        }
+    }
+}
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java
new file mode 100644
index 000000000..d6b2496c1
--- /dev/null
+++ b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java
@@ -0,0 +1,26 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.util;
+
+import software.amazon.lambda.durable.DurableContext;
+
+/**
+ * A durable operation that can be retried end-to-end by {@link RetryOperationHelper}.
+ *
+ * 
Receives the durable context and the 1-based attempt number so callers can generate unique operation names per
+ * attempt (e.g., {@code "approval-" + attempt}).
+ *
+ * @param  the result type
+ */
+@FunctionalInterface
+public interface RetryableOperation {
+
+    /**
+     * Executes the durable operation.
+     *
+     * @param context the durable context to use for durable operations
+     * @param attempt the current attempt number (1-based: first attempt is 1)
+     * @return the operation result
+     */
+    T execute(DurableContext context, int attempt);
+}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
new file mode 100644
index 000000000..726da5b61
--- /dev/null
+++ b/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
@@ -0,0 +1,80 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.config;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.lambda.durable.retry.RetryDecision;
+import software.amazon.lambda.durable.retry.RetryStrategies;
+
+class RetryOperationConfigTest {
+
+    @Test
+    void builderWithRetryStrategy() {
+        var strategy = RetryStrategies.Presets.DEFAULT;
+
+        var config = RetryOperationConfig.builder().retryStrategy(strategy).build();
+
+        assertEquals(strategy, config.retryStrategy());
+    }
+
+    @Test
+    void builderWithoutRetryStrategy_shouldThrow() {
+        var exception = assertThrows(IllegalArgumentException.class, () -> RetryOperationConfig.builder()
+                .build());
+
+        assertEquals("retryStrategy is required", exception.getMessage());
+    }
+
+    @Test
+    void wrapInChildContext_defaultsToTrue() {
+        var config = RetryOperationConfig.builder()
+                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                .build();
+
+        assertTrue(config.wrapInChildContext());
+    }
+
+    @Test
+    void wrapInChildContext_canBeSetToFalse() {
+        var config = RetryOperationConfig.builder()
+                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                .wrapInChildContext(false)
+                .build();
+
+        assertFalse(config.wrapInChildContext());
+    }
+
+    @Test
+    void wrapInChildContext_canBeSetToTrue() {
+        var config = RetryOperationConfig.builder()
+                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                .wrapInChildContext(true)
+                .build();
+
+        assertTrue(config.wrapInChildContext());
+    }
+
+    @Test
+    void builderChaining() {
+        var strategy = RetryStrategies.Presets.DEFAULT;
+
+        var config = RetryOperationConfig.builder()
+                .retryStrategy(strategy)
+                .wrapInChildContext(false)
+                .build();
+
+        assertEquals(strategy, config.retryStrategy());
+        assertFalse(config.wrapInChildContext());
+    }
+
+    @Test
+    void builderWithCustomLambdaRetryStrategy() {
+        var config = RetryOperationConfig.builder()
+                .retryStrategy((error, attempt) -> RetryDecision.fail())
+                .build();
+
+        assertNotNull(config.retryStrategy());
+    }
+}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
new file mode 100644
index 000000000..9558f3c83
--- /dev/null
+++ b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
@@ -0,0 +1,420 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.util;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.ArgumentMatchers.*;
+import static org.mockito.Mockito.*;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.function.Function;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.TypeToken;
+import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.retry.RetryDecision;
+import software.amazon.lambda.durable.retry.RetryStrategies;
+
+class RetryOperationHelperTest {
+
+    private DurableContext context;
+
+    @BeforeEach
+    void setUp() {
+        context = mock(DurableContext.class);
+    }
+
+    // --- Named form tests ---
+
+    @Nested
+    class NamedForm {
+
+        @Test
+        void successOnFirstAttempt_wrapsInChildContext() {
+            // runInChildContext should be called; delegate to the function immediately
+            when(context.runInChildContext(eq("my-op"), any(TypeToken.class), any()))
+                    .thenAnswer(invocation -> {
+                        Function func = invocation.getArgument(2);
+                        return func.apply(context);
+                    });
+
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            var result = RetryOperationHelper.retryOperation(context, "my-op", (ctx, attempt) -> "success", config);
+
+            assertEquals("success", result);
+            verify(context).runInChildContext(eq("my-op"), any(TypeToken.class), any());
+        }
+
+        @Test
+        void successOnFirstAttempt_noChildContext_whenDisabled() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(false)
+                    .build();
+
+            var result = RetryOperationHelper.retryOperation(context, "my-op", (ctx, attempt) -> "direct", config);
+
+            assertEquals("direct", result);
+            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+        }
+
+        @Test
+        void retriesWithBackoffWaits_namedForm() {
+            // Disable child context wrapping so we can directly verify wait calls
+            var callCount = new int[] {0};
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) ->
+                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
+                    .wrapInChildContext(false)
+                    .build();
+
+            var result = RetryOperationHelper.retryOperation(
+                    context,
+                    "my-op",
+                    (ctx, attempt) -> {
+                        callCount[0]++;
+                        if (attempt < 3) {
+                            throw new RuntimeException("fail-" + attempt);
+                        }
+                        return "success-on-3";
+                    },
+                    config);
+
+            assertEquals("success-on-3", result);
+            assertEquals(3, callCount[0]);
+            verify(context).wait("my-op-backoff-1", Duration.ofSeconds(5));
+            verify(context).wait("my-op-backoff-2", Duration.ofSeconds(5));
+            verify(context, times(2)).wait(anyString(), any(Duration.class));
+        }
+
+        @Test
+        void rethrowsWhenRetryStrategyReturnsFail() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(false)
+                    .build();
+
+            var exception = assertThrows(
+                    RuntimeException.class,
+                    () -> RetryOperationHelper.retryOperation(
+                            context,
+                            "my-op",
+                            (ctx, attempt) -> {
+                                throw new RuntimeException("terminal");
+                            },
+                            config));
+
+            assertEquals("terminal", exception.getMessage());
+            verify(context, never()).wait(anyString(), any(Duration.class));
+        }
+
+        @Test
+        void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(
+                            (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail())
+                    .wrapInChildContext(false)
+                    .build();
+
+            var callCount = new int[] {0};
+            var result = RetryOperationHelper.retryOperation(
+                    context,
+                    "my-op",
+                    (ctx, attempt) -> {
+                        callCount[0]++;
+                        if (attempt == 1) {
+                            throw new RuntimeException("fail");
+                        }
+                        return "ok";
+                    },
+                    config);
+
+            assertEquals("ok", result);
+            // Zero delay should be replaced with 1-second default
+            verify(context).wait("my-op-backoff-1", Duration.ofSeconds(1));
+        }
+
+        @Test
+        void nullContext_shouldThrow() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(null, "name", (ctx, a) -> "x", config));
+        }
+
+        @Test
+        void nullName_shouldThrow() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(context, null, (ctx, a) -> "x", config));
+        }
+
+        @Test
+        void nullOperation_shouldThrow() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(context, "name", null, config));
+        }
+
+        @Test
+        void nullConfig_shouldThrow() {
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(context, "name", (ctx, a) -> "x", null));
+        }
+    }
+
+    // --- Anonymous form tests ---
+
+    @Nested
+    class AnonymousForm {
+
+        @Test
+        void successOnFirstAttempt() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            var result = RetryOperationHelper.retryOperation(context, (ctx, attempt) -> "anonymous-success", config);
+
+            assertEquals("anonymous-success", result);
+            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+        }
+
+        @Test
+        void retriesWithAnonymousBackoffNames() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) ->
+                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
+                    .build();
+
+            var result = RetryOperationHelper.retryOperation(
+                    context,
+                    (ctx, attempt) -> {
+                        if (attempt < 3) {
+                            throw new RuntimeException("fail");
+                        }
+                        return "done";
+                    },
+                    config);
+
+            assertEquals("done", result);
+            verify(context).wait("retry-backoff-1", Duration.ofSeconds(2));
+            verify(context).wait("retry-backoff-2", Duration.ofSeconds(2));
+        }
+
+        @Test
+        void neverWrapsInChildContext_evenWhenConfigSaysTrue() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(true) // should be ignored for anonymous form
+                    .build();
+
+            RetryOperationHelper.retryOperation(context, (ctx, attempt) -> "result", config);
+
+            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+        }
+
+        @Test
+        void rethrowsOriginalException() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            var original = new IllegalStateException("original error");
+            var thrown = assertThrows(
+                    IllegalStateException.class,
+                    () -> RetryOperationHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                throw original;
+                            },
+                            config));
+
+            assertSame(original, thrown);
+        }
+
+        @Test
+        void nullContext_shouldThrow() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(null, (ctx, a) -> "x", config));
+        }
+
+        @Test
+        void nullOperation_shouldThrow() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(context, (RetryableOperation) null, config));
+        }
+
+        @Test
+        void nullConfig_shouldThrow() {
+            assertThrows(
+                    NullPointerException.class,
+                    () -> RetryOperationHelper.retryOperation(context, (ctx, a) -> "x", null));
+        }
+    }
+
+    // --- Retry behavior tests ---
+
+    @Nested
+    class RetryBehavior {
+
+        @Test
+        void passesCorrectAttemptNumberToOperation() {
+            var attempts = new ArrayList();
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) ->
+                            attempt < 4 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
+                    .wrapInChildContext(false)
+                    .build();
+
+            RetryOperationHelper.retryOperation(
+                    context,
+                    "track",
+                    (ctx, attempt) -> {
+                        attempts.add(attempt);
+                        if (attempt < 4) {
+                            throw new RuntimeException("not yet");
+                        }
+                        return "done";
+                    },
+                    config);
+
+            assertEquals(4, attempts.size());
+            assertEquals(1, attempts.get(0));
+            assertEquals(2, attempts.get(1));
+            assertEquals(3, attempts.get(2));
+            assertEquals(4, attempts.get(3));
+        }
+
+        @Test
+        void passesErrorToRetryStrategy() {
+            var errors = new ArrayList();
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) -> {
+                        errors.add(error);
+                        return attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail();
+                    })
+                    .build();
+
+            assertThrows(
+                    RuntimeException.class,
+                    () -> RetryOperationHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                throw new RuntimeException("error-" + attempt);
+                            },
+                            config));
+
+            assertEquals(3, errors.size());
+            assertEquals("error-1", errors.get(0).getMessage());
+            assertEquals("error-2", errors.get(1).getMessage());
+            assertEquals("error-3", errors.get(2).getMessage());
+        }
+
+        @Test
+        void respectsCustomDelayFromRetryDecision() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(attempt * 10L)))
+                    .build();
+
+            RetryOperationHelper.retryOperation(
+                    context,
+                    (ctx, attempt) -> {
+                        if (attempt <= 2) {
+                            throw new RuntimeException("fail");
+                        }
+                        return "ok";
+                    },
+                    config);
+
+            verify(context).wait("retry-backoff-1", Duration.ofSeconds(10));
+            verify(context).wait("retry-backoff-2", Duration.ofSeconds(20));
+        }
+
+        @Test
+        void passesContextToOperation() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            RetryOperationHelper.retryOperation(
+                    context,
+                    (ctx, attempt) -> {
+                        assertSame(context, ctx);
+                        return "verified";
+                    },
+                    config);
+        }
+
+        @Test
+        void namedFormPassesChildContextToOperation_whenWrapped() {
+            var childContext = mock(DurableContext.class);
+            when(context.runInChildContext(eq("wrapped"), any(TypeToken.class), any()))
+                    .thenAnswer(invocation -> {
+                        Function func = invocation.getArgument(2);
+                        return func.apply(childContext);
+                    });
+
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .wrapInChildContext(true)
+                    .build();
+
+            RetryOperationHelper.retryOperation(
+                    context,
+                    "wrapped",
+                    (ctx, attempt) -> {
+                        assertSame(childContext, ctx);
+                        return "verified";
+                    },
+                    config);
+        }
+
+        @Test
+        void rethrowsLastExceptionWhenAllRetriesExhausted() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
+                    .build();
+
+            var thrown = assertThrows(
+                    RuntimeException.class,
+                    () -> RetryOperationHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                throw new RuntimeException("attempt-" + attempt);
+                            },
+                            config));
+
+            // The last attempt's exception is rethrown
+            assertEquals("attempt-3", thrown.getMessage());
+        }
+    }
+}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java
new file mode 100644
index 000000000..f5893a37c
--- /dev/null
+++ b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java
@@ -0,0 +1,52 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.util;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.lambda.durable.DurableContext;
+
+class RetryableOperationTest {
+
+    @Test
+    void canBeImplementedAsLambda() {
+        RetryableOperation operation = (ctx, attempt) -> "result-" + attempt;
+        var context = mock(DurableContext.class);
+
+        assertEquals("result-1", operation.execute(context, 1));
+        assertEquals("result-3", operation.execute(context, 3));
+    }
+
+    @Test
+    void receivesContextAndAttempt() {
+        var context = mock(DurableContext.class);
+        RetryableOperation operation = (ctx, attempt) -> {
+            assertSame(context, ctx);
+            return "attempt-" + attempt;
+        };
+
+        assertEquals("attempt-1", operation.execute(context, 1));
+        assertEquals("attempt-2", operation.execute(context, 2));
+    }
+
+    @Test
+    void canThrowExceptions() {
+        RetryableOperation operation = (ctx, attempt) -> {
+            throw new RuntimeException("failed on attempt " + attempt);
+        };
+        var context = mock(DurableContext.class);
+
+        var exception = assertThrows(RuntimeException.class, () -> operation.execute(context, 1));
+        assertEquals("failed on attempt 1", exception.getMessage());
+    }
+
+    @Test
+    void canReturnNull() {
+        RetryableOperation operation = (ctx, attempt) -> null;
+        var context = mock(DurableContext.class);
+
+        assertNull(operation.execute(context, 1));
+    }
+}

From 254dab1792047f1218d3e8b5a40c60970fa1fd40 Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Wed, 22 Apr 2026 15:48:47 -0700
Subject: [PATCH 13/19] feat: add example tests and integration tests for new
 RetryableOperation util

---
 .../callback/RetryWaitForCallbackExample.java |  50 ++++
 .../examples/invoke/RetryInvokeExample.java   |  41 +++
 .../RetryWaitForCallbackExampleTest.java      | 148 ++++++++++
 .../invoke/RetryInvokeExampleTest.java        | 130 +++++++++
 .../durable/RetryInvokeIntegrationTest.java   | 198 ++++++++++++++
 .../RetryWaitForCallbackIntegrationTest.java  | 253 ++++++++++++++++++
 .../durable/util/RetryOperationHelper.java    |   8 +
 .../util/RetryOperationHelperTest.java        |  42 +++
 8 files changed, 870 insertions(+)
 create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
 create mode 100644 examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
 create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExampleTest.java
 create mode 100644 examples/src/test/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExampleTest.java
 create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
 create mode 100644 sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java

diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
new file mode 100644
index 000000000..e16ca8f43
--- /dev/null
+++ b/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.RetryOperationConfig;
+import software.amazon.lambda.durable.examples.types.ApprovalRequest;
+import software.amazon.lambda.durable.retry.RetryDecision;
+import software.amazon.lambda.durable.util.RetryOperationHelper;
+
+/**
+ * Example demonstrating {@link RetryOperationHelper} with {@code context.waitForCallback}.
+ *
+ * 
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.
+ *
+ * 
Each attempt uses a unique callback name ({@code "approval-1"}, {@code "approval-2"}, etc.) so the execution
+ * history stays clean and replay-safe. The anonymous form is used, so attempts run directly in the caller's context.
+ */
+public class RetryWaitForCallbackExample extends DurableHandler {
+
+    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 = RetryOperationHelper.retryOperation(
+                context,
+                (ctx, attempt) -> ctx.waitForCallback(
+                        "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
+                                .info("Attempt {}: sending callback {} to approval system", attempt, callbackId)),
+                RetryOperationConfig.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);
+    }
+}
diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
new file mode 100644
index 000000000..774c7c681
--- /dev/null
+++ b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
@@ -0,0 +1,41 @@
+// 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.RetryOperationConfig;
+import software.amazon.lambda.durable.examples.types.GreetingRequest;
+import software.amazon.lambda.durable.retry.RetryDecision;
+import software.amazon.lambda.durable.util.RetryOperationHelper;
+
+/**
+ * Example demonstrating {@link RetryOperationHelper} with {@code context.invoke}.
+ *
+ * 
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.
+ *
+ * 
The anonymous form is used, so attempts run directly in the caller's context without child-context wrapping.
+ */
+public class RetryInvokeExample extends DurableHandler {
+
+    private static final int MAX_ATTEMPTS = 3;
+
+    @Override
+    public String handleRequest(GreetingRequest input, DurableContext context) {
+        return RetryOperationHelper.retryOperation(
+                context,
+                (ctx, attempt) -> ctx.invoke(
+                        "call-greeting-" + attempt,
+                        "simple-step-example" + input.getName() + ":$LATEST",
+                        input,
+                        String.class),
+                RetryOperationConfig.builder()
+                        .retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS
+                                ? RetryDecision.retry(Duration.ofSeconds(2))
+                                : RetryDecision.fail())
+                        .build());
+    }
+}
diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExampleTest.java
new file mode 100644
index 000000000..ea5e8a4a8
--- /dev/null
+++ b/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());
+    }
+}
diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExampleTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExampleTest.java
new file mode 100644
index 000000000..7d8e66104
--- /dev/null
+++ b/examples/src/test/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExampleTest.java
@@ -0,0 +1,130 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples.invoke;
+
+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.GreetingRequest;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+
+class RetryInvokeExampleTest {
+
+    @Test
+    void succeedsOnFirstAttempt() {
+        var handler = new RetryInvokeExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+        var input = new GreetingRequest("world");
+
+        // First run — starts the invoke, suspends waiting for result
+        var result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete the first invoke attempt
+        runner.completeChainedInvoke("call-greeting-1", "\"hello world\"");
+        result = runner.run(input);
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("hello world", result.getResult(String.class));
+    }
+
+    @Test
+    void retriesAfterFirstAttemptFails() {
+        var handler = new RetryInvokeExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+        var input = new GreetingRequest("world");
+
+        // First run — starts invoke attempt 1
+        var result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail the first invoke attempt
+        runner.failChainedInvoke(
+                "call-greeting-1",
+                ErrorObject.builder()
+                        .errorType("TransientError")
+                        .errorMessage("Service unavailable")
+                        .build());
+
+        // Second run — processes the failure, does backoff wait, starts invoke attempt 2
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past the backoff wait
+        runner.advanceTime();
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete the second invoke attempt
+        runner.completeChainedInvoke("call-greeting-2", "\"hello on retry\"");
+        result = runner.run(input);
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("hello on retry", result.getResult(String.class));
+    }
+
+    @Test
+    void failsAfterAllRetriesExhausted() {
+        var handler = new RetryInvokeExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+        var input = new GreetingRequest("world");
+
+        // First run — starts invoke attempt 1
+        var result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail attempt 1
+        runner.failChainedInvoke(
+                "call-greeting-1",
+                ErrorObject.builder()
+                        .errorType("TransientError")
+                        .errorMessage("fail 1")
+                        .build());
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff wait 1
+        runner.advanceTime();
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail attempt 2
+        runner.failChainedInvoke(
+                "call-greeting-2",
+                ErrorObject.builder()
+                        .errorType("TransientError")
+                        .errorMessage("fail 2")
+                        .build());
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff wait 2
+        runner.advanceTime();
+        result = runner.run(input);
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail attempt 3 — this is the last attempt, retryStrategy returns fail()
+        runner.failChainedInvoke(
+                "call-greeting-3",
+                ErrorObject.builder()
+                        .errorType("TransientError")
+                        .errorMessage("fail 3")
+                        .build());
+        result = runner.run(input);
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+    }
+
+    @Test
+    void suspendsOnFirstRun() {
+        var handler = new RetryInvokeExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+        var input = new GreetingRequest("test");
+
+        var result = runner.run(input);
+
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+    }
+}
diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
new file mode 100644
index 000000000..1d7185b95
--- /dev/null
+++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
@@ -0,0 +1,198 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import java.time.Duration;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.services.lambda.model.ErrorObject;
+import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.exception.InvokeFailedException;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.retry.RetryDecision;
+import software.amazon.lambda.durable.retry.RetryStrategies;
+import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+import software.amazon.lambda.durable.util.RetryOperationHelper;
+
+class RetryInvokeIntegrationTest {
+
+    @Test
+    void invokeSucceedsOnFirstAttempt() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
+                                .build()));
+
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.completeChainedInvoke("invoke-1", "\"success\"");
+        result = runner.run("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("success", result.getResult(String.class));
+    }
+
+    @Test
+    void invokeRetriesAfterFailure() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
+                                .build()));
+
+        // First run — invoke attempt 1 starts
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail attempt 1
+        runner.failChainedInvoke(
+                "invoke-1",
+                ErrorObject.builder()
+                        .errorType("TransientError")
+                        .errorMessage("service unavailable")
+                        .build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff wait
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete attempt 2
+        runner.completeChainedInvoke("invoke-2", "\"recovered\"");
+        result = runner.run("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("recovered", result.getResult(String.class));
+    }
+
+    @Test
+    void invokeFailsAfterAllRetriesExhausted() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        RetryOperationConfig.builder()
+                                .retryStrategy((error, attempt) ->
+                                        attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
+                                .build()));
+
+        // Attempt 1
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.failChainedInvoke(
+                "invoke-1", ErrorObject.builder().errorMessage("fail 1").build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Attempt 2 — last attempt
+        runner.failChainedInvoke(
+                "invoke-2", ErrorObject.builder().errorMessage("fail 2").build());
+        result = runner.run("test");
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+    }
+
+    @Test
+    void invokeRetryWithCustomBackoffDelay() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        RetryOperationConfig.builder()
+                                .retryStrategy((error, attempt) -> attempt < 3
+                                        ? RetryDecision.retry(Duration.ofSeconds(attempt * 5L))
+                                        : RetryDecision.fail())
+                                .build()));
+
+        // Attempt 1 fails
+        var result = runner.run("test");
+        runner.failChainedInvoke(
+                "invoke-1", ErrorObject.builder().errorMessage("fail").build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past first backoff (5s)
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Attempt 2 succeeds
+        runner.completeChainedInvoke("invoke-2", "\"ok\"");
+        result = runner.run("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("ok", result.getResult(String.class));
+    }
+
+    @Test
+    void invokeRetryWithStepsBeforeAndAfter() {
+        var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
+            var prefix = context.step("prepare", String.class, stepCtx -> "prepared");
+
+            var invokeResult = RetryOperationHelper.retryOperation(
+                    context,
+                    (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                    RetryOperationConfig.builder()
+                            .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
+                            .build());
+
+            return context.step("finalize", String.class, stepCtx -> prefix + " -> " + invokeResult + " -> done");
+        });
+
+        // First run — prepare step completes, invoke starts
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete invoke
+        runner.completeChainedInvoke("invoke-1", "\"invoked\"");
+        result = runner.runUntilComplete("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("prepared -> invoked -> done", result.getResult(String.class));
+    }
+
+    @Test
+    void invokeRetryPreservesOriginalExceptionType() {
+        var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
+            try {
+                return RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                                .build());
+            } catch (InvokeFailedException e) {
+                assertEquals("invoke failed", e.getMessage());
+                throw e;
+            }
+        });
+
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.failChainedInvoke(
+                "invoke-1", ErrorObject.builder().errorMessage("invoke failed").build());
+        result = runner.run("test");
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+    }
+}
diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
new file mode 100644
index 000000000..f7806d577
--- /dev/null
+++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
@@ -0,0 +1,253 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import java.time.Duration;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.services.lambda.model.ErrorObject;
+import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.retry.RetryDecision;
+import software.amazon.lambda.durable.retry.RetryStrategies;
+import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+import software.amazon.lambda.durable.util.RetryOperationHelper;
+
+class RetryWaitForCallbackIntegrationTest {
+
+    @Test
+    void waitForCallbackSucceedsOnFirstAttempt() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.waitForCallback(
+                                "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
+                                        .info("Submitting callback {}", callbackId)),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
+                                .build()));
+
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // waitForCallback("approval-1", ...) creates "approval-1-callback" internally
+        var callbackId = runner.getCallbackId("approval-1-callback");
+        assertNotNull(callbackId);
+        runner.completeCallback(callbackId, "\"approved\"");
+
+        result = runner.runUntilComplete("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("approved", result.getResult(String.class));
+    }
+
+    @Test
+    void waitForCallbackRetriesAfterFailure() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) -> ctx.waitForCallback(
+                                "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
+                                        .info("Attempt {} callback {}", attempt, callbackId)),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
+                                .build()));
+
+        // First run — starts waitForCallback attempt 1
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Fail attempt 1
+        var callbackId1 = runner.getCallbackId("approval-1-callback");
+        assertNotNull(callbackId1);
+        runner.failCallback(
+                callbackId1,
+                ErrorObject.builder()
+                        .errorType("Rejected")
+                        .errorMessage("denied by reviewer")
+                        .build());
+
+        // Run — processes failure, hits backoff wait
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete attempt 2
+        var callbackId2 = runner.getCallbackId("approval-2-callback");
+        assertNotNull(callbackId2);
+        runner.completeCallback(callbackId2, "\"approved on retry\"");
+
+        result = runner.runUntilComplete("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("approved on retry", result.getResult(String.class));
+    }
+
+    @Test
+    void waitForCallbackFailsAfterAllRetriesExhausted() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) ->
+                                ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
+                        RetryOperationConfig.builder()
+                                .retryStrategy((error, attempt) ->
+                                        attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
+                                .build()));
+
+        // Attempt 1
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.failCallback(
+                runner.getCallbackId("approval-1-callback"),
+                ErrorObject.builder().errorMessage("fail 1").build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Attempt 2 — last attempt
+        runner.failCallback(
+                runner.getCallbackId("approval-2-callback"),
+                ErrorObject.builder().errorMessage("fail 2").build());
+        result = runner.run("test");
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+    }
+
+    @Test
+    void waitForCallbackRetryWithStepsBeforeAndAfter() {
+        var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
+            var prefix = context.step("prepare", String.class, stepCtx -> "prepared");
+
+            var callbackResult = RetryOperationHelper.retryOperation(
+                    context,
+                    (ctx, attempt) ->
+                            ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
+                    RetryOperationConfig.builder()
+                            .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
+                            .build());
+
+            return context.step("finalize", String.class, stepCtx -> prefix + " -> " + callbackResult + " -> done");
+        });
+
+        // First run — prepare completes, waitForCallback starts
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Complete callback
+        var callbackId = runner.getCallbackId("approval-1-callback");
+        runner.completeCallback(callbackId, "\"approved\"");
+
+        result = runner.runUntilComplete("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("prepared -> approved -> done", result.getResult(String.class));
+    }
+
+    @Test
+    void waitForCallbackRetryMultipleFailuresThenSuccess() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) ->
+                                ctx.waitForCallback("cb-" + attempt, String.class, (callbackId, stepCtx) -> {}),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.fixedDelay(4, Duration.ofSeconds(1)))
+                                .build()));
+
+        // Attempt 1 — fail
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.failCallback(
+                runner.getCallbackId("cb-1-callback"),
+                ErrorObject.builder().errorMessage("fail 1").build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Attempt 2 — fail
+        runner.failCallback(
+                runner.getCallbackId("cb-2-callback"),
+                ErrorObject.builder().errorMessage("fail 2").build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Attempt 3 — succeed
+        runner.completeCallback(runner.getCallbackId("cb-3-callback"), "\"third time's the charm\"");
+        result = runner.runUntilComplete("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("third time's the charm", result.getResult(String.class));
+    }
+
+    @Test
+    void waitForCallbackRetryWithSubmitterLogic() {
+        // Verify the submitter runs on each retry attempt
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> RetryOperationHelper.retryOperation(
+                        context,
+                        (ctx, attempt) ->
+                                ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {
+                                    // Submitter runs each attempt — in a real scenario this would
+                                    // send the callbackId to an external system
+                                    stepCtx.getLogger().info("Attempt {} submitting {}", attempt, callbackId);
+                                }),
+                        RetryOperationConfig.builder()
+                                .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
+                                .build()));
+
+        // Attempt 1
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Verify submitter step was created for attempt 1
+        var submitterOp = runner.getOperation("approval-1-submitter");
+        assertNotNull(submitterOp, "Submitter step should exist for attempt 1");
+
+        // Fail attempt 1
+        runner.failCallback(
+                runner.getCallbackId("approval-1-callback"),
+                ErrorObject.builder().errorMessage("rejected").build());
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance past backoff, start attempt 2
+        runner.advanceTime();
+        result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Verify submitter step was created for attempt 2
+        var submitterOp2 = runner.getOperation("approval-2-submitter");
+        assertNotNull(submitterOp2, "Submitter step should exist for attempt 2");
+
+        // Complete attempt 2
+        runner.completeCallback(runner.getCallbackId("approval-2-callback"), "\"approved\"");
+        result = runner.runUntilComplete("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("approved", result.getResult(String.class));
+    }
+}
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
index c63f0aad0..f4eaf8c28 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
@@ -7,6 +7,8 @@
 import software.amazon.lambda.durable.DurableContext;
 import software.amazon.lambda.durable.TypeToken;
 import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
+import software.amazon.lambda.durable.execution.SuspendExecutionException;
 import software.amazon.lambda.durable.retry.RetryDecision;
 
 /**
@@ -112,6 +114,9 @@ public static  T retryOperation(
     /**
      * Core retry loop. Replay-safe because every side-effect is a durable operation: the user's operation calls durable
      * primitives, and backoff uses {@code context.wait()}.
+     *
+     * 
{@link SuspendExecutionException} and {@link UnrecoverableDurableExecutionException} are never retried — they
+     * are internal SDK control flow signals that must propagate immediately.
      */
     private static  T executeRetryLoop(
             DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) {
@@ -119,6 +124,9 @@ private static  T executeRetryLoop(
         while (true) {
             try {
                 return operation.execute(context, attempt);
+            } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) {
+                // Internal SDK control flow — never retry, always propagate
+                throw e;
             } catch (Exception e) {
                 RetryDecision decision = config.retryStrategy().makeRetryDecision(e, attempt);
                 if (!decision.shouldRetry()) {
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
index 9558f3c83..265fa968c 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
@@ -15,6 +15,8 @@
 import software.amazon.lambda.durable.DurableContext;
 import software.amazon.lambda.durable.TypeToken;
 import software.amazon.lambda.durable.config.RetryOperationConfig;
+import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
+import software.amazon.lambda.durable.execution.SuspendExecutionException;
 import software.amazon.lambda.durable.retry.RetryDecision;
 import software.amazon.lambda.durable.retry.RetryStrategies;
 
@@ -416,5 +418,45 @@ void rethrowsLastExceptionWhenAllRetriesExhausted() {
             // The last attempt's exception is rethrown
             assertEquals("attempt-3", thrown.getMessage());
         }
+
+        @Test
+        void propagatesSuspendExecutionExceptionWithoutRetrying() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1)))
+                    .build();
+
+            assertThrows(
+                    SuspendExecutionException.class,
+                    () -> RetryOperationHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                throw new SuspendExecutionException();
+                            },
+                            config));
+
+            // Should never reach the wait — SuspendExecutionException propagates immediately
+            verify(context, never()).wait(anyString(), any(Duration.class));
+        }
+
+        @Test
+        void propagatesUnrecoverableDurableExecutionExceptionWithoutRetrying() {
+            var config = RetryOperationConfig.builder()
+                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1)))
+                    .build();
+
+            assertThrows(
+                    UnrecoverableDurableExecutionException.class,
+                    () -> RetryOperationHelper.retryOperation(
+                            context,
+                            (ctx, attempt) -> {
+                                throw new UnrecoverableDurableExecutionException(
+                                        software.amazon.awssdk.services.lambda.model.ErrorObject.builder()
+                                                .errorMessage("unrecoverable")
+                                                .build());
+                            },
+                            config));
+
+            verify(context, never()).wait(anyString(), any(Duration.class));
+        }
     }
 }

From fadce47fcb4a67096f10de869c04c299e309ba62 Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Thu, 30 Apr 2026 10:54:44 -0700
Subject: [PATCH 14/19] chore: use virtual context when wrapInChildContext is
 false

---
 .../amazon/lambda/durable/DurableContext.java |  26 +-
 .../durable/config/RetryOperationConfig.java  | 103 ----
 .../durable/context/DurableContextImpl.java   |  89 +++-
 .../durable/model/OperationSubType.java       |   3 +-
 .../operation/ChildContextOperation.java      |   2 +-
 .../durable/util/RetryOperationHelper.java    | 143 ------
 .../durable/util/RetryableOperation.java      |  26 -
 .../config/RetryOperationConfigTest.java      |  80 ---
 .../context/DurableContextWithRetryTest.java  | 167 ++++---
 .../util/RetryOperationHelperTest.java        | 462 ------------------
 .../durable/util/RetryableOperationTest.java  |  52 --
 11 files changed, 181 insertions(+), 972 deletions(-)
 delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java
 delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
 delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java
 delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
 delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
 delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java

diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
index 9518e5060..53a17a2c2 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
@@ -740,9 +740,10 @@  DurableFuture waitForConditionAsync(
      * 
Every side-effect in the loop is a durable operation, so the loop is replay-safe by construction. On replay,
      * completed operations return cached results instantly and the loop fast-forwards to the current attempt.
      *
-     * 
By default, the retry loop runs directly on the caller's context. If
-     * {@link WithRetryConfig#wrapInChildContext()} is enabled, the loop is wrapped in a child context so all attempts
-     * are grouped under a single named operation in execution history.
+     * 
The retry loop always runs in a child context to provide an isolated operation ID namespace. If
+     * {@link WithRetryConfig#wrapInChildContext()} is enabled, the child context is checkpointed (persisted) so all
+     * attempts are grouped under a single named operation in execution history. Otherwise, a virtual child context is
+     * used — no checkpointing overhead, but the child re-executes on replay.
      *
      * @param  the result type
      * @param name operation name (used for backoff wait names, and as the child context name when wrapping)
@@ -755,9 +756,10 @@  DurableFuture waitForConditionAsync(
     /**
      * Replay-safe retry loop for any durable operation (anonymous form, sync).
      *
-     * 
By default, the retry loop runs directly on the caller's context. If
-     * {@link WithRetryConfig#wrapInChildContext()} is enabled, the loop is wrapped in a child context with a default
-     * name so all attempts are grouped under a single operation in execution history.
+     * 
The retry loop always runs in a child context to provide an isolated operation ID namespace. If
+     * {@link WithRetryConfig#wrapInChildContext()} is enabled, the child context is checkpointed (persisted) so all
+     * attempts are grouped under a single named operation in execution history. Otherwise, a virtual child context is
+     * used — no checkpointing overhead, but the child re-executes on replay.
      *
      * @param  the result type
      * @param operation the retryable operation — receives the context and 1-based attempt number
@@ -769,8 +771,10 @@  DurableFuture waitForConditionAsync(
     /**
      * Replay-safe retry loop for any durable operation (named form, async).
      *
-     * 
Wraps the retry loop in {@code runInChildContextAsync} so all attempts are grouped under a single named
-     * operation in execution history, and returns a {@link DurableFuture} that can be composed or blocked on.
+     * 
The retry loop always runs in a child context to provide an isolated operation ID namespace. If
+     * {@link WithRetryConfig#wrapInChildContext()} is enabled, the child context is checkpointed (persisted) so all
+     * attempts are grouped under a single named operation in execution history. Otherwise, a virtual child context is
+     * used — no checkpointing overhead, but the child re-executes on replay.
      *
      * @param  the result type
      * @param name operation name (used for child context and backoff wait names)
@@ -783,8 +787,10 @@  DurableFuture waitForConditionAsync(
     /**
      * Replay-safe retry loop for any durable operation (anonymous form, async).
      *
-     * 
Wraps the retry loop in {@code runInChildContextAsync} with a default name so all attempts are grouped under a
-     * single operation in execution history, and returns a {@link DurableFuture} that can be composed or blocked on.
+     * 
The retry loop always runs in a child context to provide an isolated operation ID namespace. If
+     * {@link WithRetryConfig#wrapInChildContext()} is enabled, the child context is checkpointed (persisted) so all
+     * attempts are grouped under a single named operation in execution history. Otherwise, a virtual child context is
+     * used — no checkpointing overhead, but the child re-executes on replay.
      *
      * @param  the result type
      * @param operation the retryable operation — receives the context and 1-based attempt number
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java
deleted file mode 100644
index 696d1d253..000000000
--- a/sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java
+++ /dev/null
@@ -1,103 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.config;
-
-import software.amazon.lambda.durable.retry.RetryStrategy;
-
-/**
- * Configuration for {@link software.amazon.lambda.durable.util.RetryOperationHelper#retryOperation}.
- *
- * 
Uses the same {@link RetryStrategy} shape that developers already know from {@link StepConfig}, so there are zero
- * new retry concepts to learn.
- */
-public class RetryOperationConfig {
-    private final RetryStrategy retryStrategy;
-    private final boolean wrapInChildContext;
-
-    private RetryOperationConfig(Builder builder) {
-        this.retryStrategy = builder.retryStrategy;
-        this.wrapInChildContext = builder.wrapInChildContext;
-    }
-
-    /**
-     * Returns the retry strategy. Same type as {@link StepConfig#retryStrategy()}.
-     *
-     * @return the retry strategy, never null
-     */
-    public RetryStrategy retryStrategy() {
-        return retryStrategy;
-    }
-
-    /**
-     * Whether to wrap the retry loop in {@code runInChildContext} so all attempts are grouped under a single named
-     * operation in execution history. Only applies when a name is provided to the named form of {@code retryOperation}.
-     * Defaults to {@code true}.
-     *
-     * @return true if child-context wrapping is enabled
-     */
-    public boolean wrapInChildContext() {
-        return wrapInChildContext;
-    }
-
-    /**
-     * Creates a new builder for {@code RetryOperationConfig}.
-     *
-     * @return a new builder instance
-     */
-    public static Builder builder() {
-        return new Builder();
-    }
-
-    /** Builder for creating {@link RetryOperationConfig} instances. */
-    public static class Builder {
-        private RetryStrategy retryStrategy;
-        private boolean wrapInChildContext = true;
-
-        private Builder() {}
-
-        /**
-         * Sets the retry strategy. Required.
-         *
-         * 
Reuses the exact same {@link RetryStrategy} interface from {@link StepConfig}. All existing factory
-         * methods ({@link software.amazon.lambda.durable.retry.RetryStrategies#exponentialBackoff},
-         * {@link software.amazon.lambda.durable.retry.RetryStrategies#fixedDelay}, presets, and custom lambdas) work
-         * without modification.
-         *
-         * @param retryStrategy the retry strategy to use
-         * @return this builder for method chaining
-         */
-        public Builder retryStrategy(RetryStrategy retryStrategy) {
-            this.retryStrategy = retryStrategy;
-            return this;
-        }
-
-        /**
-         * Controls whether the retry loop is wrapped in a child context. Only meaningful for the named form of
-         * {@code retryOperation}. Defaults to {@code true}.
-         *
-         * 
When {@code true}, all attempts and backoff waits are grouped under a single named operation in execution
-         * history, providing a cleaner view and isolated operation ID space. Set to {@code false} to flatten attempts
-         * into the parent context.
-         *
-         * @param wrapInChildContext whether to wrap in a child context
-         * @return this builder for method chaining
-         */
-        public Builder wrapInChildContext(boolean wrapInChildContext) {
-            this.wrapInChildContext = wrapInChildContext;
-            return this;
-        }
-
-        /**
-         * Builds the {@link RetryOperationConfig} instance.
-         *
-         * @return a new config with the configured options
-         * @throws IllegalArgumentException if retryStrategy is not set
-         */
-        public RetryOperationConfig build() {
-            if (retryStrategy == null) {
-                throw new IllegalArgumentException("retryStrategy is required");
-            }
-            return new RetryOperationConfig(this);
-        }
-    }
-}
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
index 82fa29bce..1274e1c20 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
@@ -238,6 +238,25 @@ private  DurableFuture runInChildContextAsync(
             Function func,
             RunInChildContextConfig config,
             OperationSubType subType) {
+        return runInChildContextAsync(name, resultType, func, config, subType, false);
+    }
+
+    private  DurableFuture runInVirtualChildContextAsync(
+            String name,
+            TypeToken resultType,
+            Function func,
+            RunInChildContextConfig config,
+            OperationSubType subType) {
+        return runInChildContextAsync(name, resultType, func, config, subType, true);
+    }
+
+    private  DurableFuture runInChildContextAsync(
+            String name,
+            TypeToken resultType,
+            Function func,
+            RunInChildContextConfig config,
+            OperationSubType subType,
+            boolean isVirtual) {
         Objects.requireNonNull(resultType, "resultType cannot be null");
         Objects.requireNonNull(config, "RunInChildContextConfig cannot be null");
         ParameterValidator.validateOperationName(name);
@@ -253,7 +272,9 @@ private  DurableFuture runInChildContextAsync(
                 func,
                 resultType,
                 config,
-                this);
+                this,
+                isVirtual,
+                null);
 
         operation.execute();
         return operation;
@@ -386,10 +407,21 @@ public  T withRetry(String name, WithRetry operation, WithRetryConfig conf
         Objects.requireNonNull(config, "config cannot be null");
 
         if (config.wrapInChildContext()) {
-            return (T) runInChildContext(
-                    name, new TypeToken() {}, childCtx -> executeRetryLoop(childCtx, name, operation, config));
+            return (T) runInChildContextAsync(
+                            name,
+                            new TypeToken() {},
+                            childCtx -> executeRetryLoop(childCtx, name, operation, config),
+                            RunInChildContextConfig.builder().build(),
+                            OperationSubType.WITH_RETRY)
+                    .get();
         }
-        return executeRetryLoop(this, name, operation, config);
+        return (T) runInVirtualChildContextAsync(
+                        name,
+                        new TypeToken() {},
+                        childCtx -> executeRetryLoop(childCtx, name, operation, config),
+                        RunInChildContextConfig.builder().build(),
+                        OperationSubType.WITH_RETRY)
+                .get();
     }
 
     @Override
@@ -399,12 +431,21 @@ public  T withRetry(WithRetry operation, WithRetryConfig config) {
         Objects.requireNonNull(config, "config cannot be null");
 
         if (config.wrapInChildContext()) {
-            return (T) runInChildContext(
-                    ANONYMOUS_CHILD_CONTEXT_NAME,
-                    new TypeToken() {},
-                    childCtx -> executeRetryLoop(childCtx, null, operation, config));
+            return (T) runInChildContextAsync(
+                            ANONYMOUS_CHILD_CONTEXT_NAME,
+                            new TypeToken() {},
+                            childCtx -> executeRetryLoop(childCtx, null, operation, config),
+                            RunInChildContextConfig.builder().build(),
+                            OperationSubType.WITH_RETRY)
+                    .get();
         }
-        return executeRetryLoop(this, null, operation, config);
+        return (T) runInVirtualChildContextAsync(
+                        ANONYMOUS_CHILD_CONTEXT_NAME,
+                        new TypeToken() {},
+                        childCtx -> executeRetryLoop(childCtx, null, operation, config),
+                        RunInChildContextConfig.builder().build(),
+                        OperationSubType.WITH_RETRY)
+                .get();
     }
 
     @Override
@@ -414,8 +455,20 @@ public  DurableFuture withRetryAsync(String name, WithRetry operation,
         Objects.requireNonNull(operation, "operation cannot be null");
         Objects.requireNonNull(config, "config cannot be null");
 
-        return (DurableFuture) runInChildContextAsync(
-                name, new TypeToken() {}, childCtx -> executeRetryLoop(childCtx, name, operation, config));
+        if (config.wrapInChildContext()) {
+            return (DurableFuture) runInChildContextAsync(
+                    name,
+                    new TypeToken() {},
+                    childCtx -> executeRetryLoop(childCtx, name, operation, config),
+                    RunInChildContextConfig.builder().build(),
+                    OperationSubType.WITH_RETRY);
+        }
+        return (DurableFuture) runInVirtualChildContextAsync(
+                name,
+                new TypeToken() {},
+                childCtx -> executeRetryLoop(childCtx, name, operation, config),
+                RunInChildContextConfig.builder().build(),
+                OperationSubType.WITH_RETRY);
     }
 
     @Override
@@ -424,10 +477,20 @@ public  DurableFuture withRetryAsync(WithRetry operation, WithRetryConf
         Objects.requireNonNull(operation, "operation cannot be null");
         Objects.requireNonNull(config, "config cannot be null");
 
-        return (DurableFuture) runInChildContextAsync(
+        if (config.wrapInChildContext()) {
+            return (DurableFuture) runInChildContextAsync(
+                    ANONYMOUS_CHILD_CONTEXT_NAME,
+                    new TypeToken() {},
+                    childCtx -> executeRetryLoop(childCtx, null, operation, config),
+                    RunInChildContextConfig.builder().build(),
+                    OperationSubType.WITH_RETRY);
+        }
+        return (DurableFuture) runInVirtualChildContextAsync(
                 ANONYMOUS_CHILD_CONTEXT_NAME,
                 new TypeToken() {},
-                childCtx -> executeRetryLoop(childCtx, null, operation, config));
+                childCtx -> executeRetryLoop(childCtx, null, operation, config),
+                RunInChildContextConfig.builder().build(),
+                OperationSubType.WITH_RETRY);
     }
 
     /**
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/model/OperationSubType.java b/sdk/src/main/java/software/amazon/lambda/durable/model/OperationSubType.java
index 90c351487..416798f9f 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/model/OperationSubType.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/model/OperationSubType.java
@@ -15,7 +15,8 @@ public enum OperationSubType {
     PARALLEL("Parallel"),
     PARALLEL_BRANCH("ParallelBranch"),
     WAIT_FOR_CALLBACK("WaitForCallback"),
-    WAIT_FOR_CONDITION("WaitForCondition");
+    WAIT_FOR_CONDITION("WaitForCondition"),
+    WITH_RETRY("WithRetry");
 
     private final String value;
 
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java
index 39fd6e4eb..524de5958 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java
@@ -240,7 +240,7 @@ private Throwable translateException(Operation op, ErrorObject errorObject) {
             case WAIT_FOR_CALLBACK -> handleWaitForCallbackFailure();
             case MAP_ITERATION -> new MapIterationFailedException(op);
             case PARALLEL_BRANCH -> new ParallelBranchFailedException(op);
-            case RUN_IN_CHILD_CONTEXT -> new ChildContextFailedException(op);
+            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());
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
deleted file mode 100644
index f4eaf8c28..000000000
--- a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryOperationHelper.java
+++ /dev/null
@@ -1,143 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.util;
-
-import java.time.Duration;
-import java.util.Objects;
-import software.amazon.lambda.durable.DurableContext;
-import software.amazon.lambda.durable.TypeToken;
-import software.amazon.lambda.durable.config.RetryOperationConfig;
-import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
-import software.amazon.lambda.durable.execution.SuspendExecutionException;
-import software.amazon.lambda.durable.retry.RetryDecision;
-
-/**
- * Replay-safe retry loop for any durable operation.
- *
- * 
Provides the same retry-with-backoff pattern that {@code context.step()} has built in, but for operations that
- * cannot live inside a step ({@code waitForCallback}, {@code invoke}, {@code waitForCondition}, etc.).
- *
- * 
Every side-effect in the loop is a durable operation, so the loop is replay-safe by construction. On replay,
- * completed operations return cached results instantly and the loop fast-forwards to the current attempt.
- *
- * 
Usage — callback retry

- *
- * 
{@code
- * var result = RetryOperationHelper.retryOperation(
- *     context,
- *     "approval",
- *     (ctx, attempt) -> ctx.waitForCallback(
- *         "approval-" + attempt,
- *         String.class,
- *         (callbackId, stepCtx) -> sendApprovalEmail(approverEmail, callbackId)
- *     ),
- *     RetryOperationConfig.builder()
- *         .retryStrategy(RetryStrategies.exponentialBackoff(
- *             3, Duration.ofSeconds(2), Duration.ofSeconds(30), 2.0, JitterStrategy.FULL))
- *         .build()
- * );
- * }

- *
- * 
Usage — invoke retry (anonymous form)

- *
- * 
{@code
- * var result = RetryOperationHelper.retryOperation(
- *     context,
- *     (ctx, attempt) -> ctx.invoke(
- *         "charge-" + attempt, paymentFnArn, new ChargeRequest(orderId), String.class),
- *     RetryOperationConfig.builder()
- *         .retryStrategy((err, att) -> att < 3
- *             ? RetryDecision.retry(Duration.ofSeconds(1))
- *             : RetryDecision.fail())
- *         .build()
- * );
- * }

- */
-public final class RetryOperationHelper {
-
-    private static final Duration DEFAULT_BACKOFF_DELAY = Duration.ofSeconds(1);
-    private static final String BACKOFF_SUFFIX = "-backoff-";
-    private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-";
-
-    private RetryOperationHelper() {
-        // utility class
-    }
-
-    /**
-     * Named form — wraps the retry loop in {@code runInChildContext} by default so all attempts are grouped under a
-     * single named operation in execution history.
-     *
-     * 
The child-context wrapping can be disabled via
-     * {@link RetryOperationConfig.Builder#wrapInChildContext(boolean)}.
-     *
-     * @param  the result type
-     * @param context the durable context
-     * @param name operation name (used for child context and backoff wait names)
-     * @param operation the retryable operation — receives the context and 1-based attempt number
-     * @param config retry configuration including the retry strategy
-     * @return the operation result
-     */
-    @SuppressWarnings("unchecked")
-    public static  T retryOperation(
-            DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) {
-        Objects.requireNonNull(context, "context cannot be null");
-        Objects.requireNonNull(name, "name cannot be null");
-        Objects.requireNonNull(operation, "operation cannot be null");
-        Objects.requireNonNull(config, "config cannot be null");
-
-        if (config.wrapInChildContext()) {
-            return (T) context.runInChildContext(
-                    name, new TypeToken() {}, childCtx -> executeRetryLoop(childCtx, name, operation, config));
-        }
-        return executeRetryLoop(context, name, operation, config);
-    }
-
-    /**
-     * Anonymous form — runs the retry loop directly in the caller's context. No child-context wrapping is applied
-     * regardless of the {@code wrapInChildContext} config setting.
-     *
-     * @param  the result type
-     * @param context the durable context
-     * @param operation the retryable operation — receives the context and 1-based attempt number
-     * @param config retry configuration including the retry strategy
-     * @return the operation result
-     */
-    public static  T retryOperation(
-            DurableContext context, RetryableOperation operation, RetryOperationConfig config) {
-        Objects.requireNonNull(context, "context cannot be null");
-        Objects.requireNonNull(operation, "operation cannot be null");
-        Objects.requireNonNull(config, "config cannot be null");
-
-        return executeRetryLoop(context, null, operation, config);
-    }
-
-    /**
-     * Core retry loop. Replay-safe because every side-effect is a durable operation: the user's operation calls durable
-     * primitives, and backoff uses {@code context.wait()}.
-     *
-     * 
{@link SuspendExecutionException} and {@link UnrecoverableDurableExecutionException} are never retried — they
-     * are internal SDK control flow signals that must propagate immediately.
-     */
-    private static  T executeRetryLoop(
-            DurableContext context, String name, RetryableOperation operation, RetryOperationConfig config) {
-        var attempt = 1;
-        while (true) {
-            try {
-                return operation.execute(context, attempt);
-            } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) {
-                // Internal SDK control flow — never retry, always propagate
-                throw e;
-            } catch (Exception e) {
-                RetryDecision decision = config.retryStrategy().makeRetryDecision(e, attempt);
-                if (!decision.shouldRetry()) {
-                    throw e;
-                }
-
-                var delay = decision.delay().isZero() ? DEFAULT_BACKOFF_DELAY : decision.delay();
-                var waitName = name != null ? name + BACKOFF_SUFFIX + attempt : ANONYMOUS_BACKOFF_PREFIX + attempt;
-                context.wait(waitName, delay);
-                attempt++;
-            }
-        }
-    }
-}
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java
deleted file mode 100644
index d6b2496c1..000000000
--- a/sdk/src/main/java/software/amazon/lambda/durable/util/RetryableOperation.java
+++ /dev/null
@@ -1,26 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.util;
-
-import software.amazon.lambda.durable.DurableContext;
-
-/**
- * A durable operation that can be retried end-to-end by {@link RetryOperationHelper}.
- *
- * 
Receives the durable context and the 1-based attempt number so callers can generate unique operation names per
- * attempt (e.g., {@code "approval-" + attempt}).
- *
- * @param  the result type
- */
-@FunctionalInterface
-public interface RetryableOperation {
-
-    /**
-     * Executes the durable operation.
-     *
-     * @param context the durable context to use for durable operations
-     * @param attempt the current attempt number (1-based: first attempt is 1)
-     * @return the operation result
-     */
-    T execute(DurableContext context, int attempt);
-}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
deleted file mode 100644
index 726da5b61..000000000
--- a/sdk/src/test/java/software/amazon/lambda/durable/config/RetryOperationConfigTest.java
+++ /dev/null
@@ -1,80 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.config;
-
-import static org.junit.jupiter.api.Assertions.*;
-
-import org.junit.jupiter.api.Test;
-import software.amazon.lambda.durable.retry.RetryDecision;
-import software.amazon.lambda.durable.retry.RetryStrategies;
-
-class RetryOperationConfigTest {
-
-    @Test
-    void builderWithRetryStrategy() {
-        var strategy = RetryStrategies.Presets.DEFAULT;
-
-        var config = RetryOperationConfig.builder().retryStrategy(strategy).build();
-
-        assertEquals(strategy, config.retryStrategy());
-    }
-
-    @Test
-    void builderWithoutRetryStrategy_shouldThrow() {
-        var exception = assertThrows(IllegalArgumentException.class, () -> RetryOperationConfig.builder()
-                .build());
-
-        assertEquals("retryStrategy is required", exception.getMessage());
-    }
-
-    @Test
-    void wrapInChildContext_defaultsToTrue() {
-        var config = RetryOperationConfig.builder()
-                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                .build();
-
-        assertTrue(config.wrapInChildContext());
-    }
-
-    @Test
-    void wrapInChildContext_canBeSetToFalse() {
-        var config = RetryOperationConfig.builder()
-                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                .wrapInChildContext(false)
-                .build();
-
-        assertFalse(config.wrapInChildContext());
-    }
-
-    @Test
-    void wrapInChildContext_canBeSetToTrue() {
-        var config = RetryOperationConfig.builder()
-                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                .wrapInChildContext(true)
-                .build();
-
-        assertTrue(config.wrapInChildContext());
-    }
-
-    @Test
-    void builderChaining() {
-        var strategy = RetryStrategies.Presets.DEFAULT;
-
-        var config = RetryOperationConfig.builder()
-                .retryStrategy(strategy)
-                .wrapInChildContext(false)
-                .build();
-
-        assertEquals(strategy, config.retryStrategy());
-        assertFalse(config.wrapInChildContext());
-    }
-
-    @Test
-    void builderWithCustomLambdaRetryStrategy() {
-        var config = RetryOperationConfig.builder()
-                .retryStrategy((error, attempt) -> RetryDecision.fail())
-                .build();
-
-        assertNotNull(config.retryStrategy());
-    }
-}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
index 8a3e1bade..89b251cb5 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
@@ -23,7 +23,6 @@
 import software.amazon.lambda.durable.model.WithRetry;
 import software.amazon.lambda.durable.retry.RetryDecision;
 import software.amazon.lambda.durable.retry.RetryStrategies;
-import software.amazon.lambda.durable.util.CompletedDurableFuture;
 
 @SuppressWarnings("unchecked")
 class DurableContextWithRetryTest {
@@ -42,9 +41,12 @@ void setUp() {
     /**
      * Stubs the withRetry/withRetryAsync methods on a mock DurableContext so they execute the real retry loop logic.
      * This is needed because withRetry/withRetryAsync are abstract interface methods, and mocks return null by default.
+     *
+     * 
All forms always run in a child context (matching the real implementation which uses a virtual child context
+     * when {@code wrapInChildContext} is false, and a checkpointed child context when true).
      */
     private void stubWithRetryMethods(DurableContext mock) {
-        // Named sync form
+        // Named sync form — always runs in a child context
         when(mock.withRetry(any(), nullable(WithRetry.class), nullable(WithRetryConfig.class)))
                 .thenAnswer(invocation -> {
                     String name = invocation.getArgument(0);
@@ -53,29 +55,25 @@ private void stubWithRetryMethods(DurableContext mock) {
                     Objects.requireNonNull(name, "name cannot be null");
                     Objects.requireNonNull(operation, "operation cannot be null");
                     Objects.requireNonNull(config, "config cannot be null");
-                    if (config.wrapInChildContext()) {
-                        return mock.runInChildContext(
-                                name,
-                                new TypeToken() {},
-                                childCtx -> executeRetryLoop(childCtx, name, operation, config));
-                    }
-                    return executeRetryLoop(mock, name, operation, config);
+                    return mock.runInChildContextAsync(
+                                    name,
+                                    new TypeToken() {},
+                                    childCtx -> executeRetryLoop(childCtx, name, operation, config))
+                            .get();
                 });
 
-        // Anonymous sync form
+        // Anonymous sync form — always runs in a child context
         when(mock.withRetry(nullable(WithRetry.class), nullable(WithRetryConfig.class)))
                 .thenAnswer(invocation -> {
                     WithRetry operation = invocation.getArgument(0);
                     WithRetryConfig config = invocation.getArgument(1);
                     Objects.requireNonNull(operation, "operation cannot be null");
                     Objects.requireNonNull(config, "config cannot be null");
-                    if (config.wrapInChildContext()) {
-                        return mock.runInChildContext(
-                                "retry",
-                                new TypeToken() {},
-                                childCtx -> executeRetryLoop(childCtx, null, operation, config));
-                    }
-                    return executeRetryLoop(mock, null, operation, config);
+                    return mock.runInChildContextAsync(
+                                    "retry",
+                                    new TypeToken() {},
+                                    childCtx -> executeRetryLoop(childCtx, null, operation, config))
+                            .get();
                 });
 
         // Named async form
@@ -134,7 +132,8 @@ private void stubChildContext(String name) {
         when(context.runInChildContextAsync(eq(name), any(TypeToken.class), any()))
                 .thenAnswer(invocation -> {
                     Function func = invocation.getArgument(2);
-                    return new CompletedDurableFuture<>(func.apply(childContext));
+                    var result = func.apply(childContext);
+                    return (DurableFuture) () -> result;
                 });
     }
 
@@ -143,25 +142,24 @@ private void stubChildContextAnyName() {
         when(context.runInChildContextAsync(anyString(), any(TypeToken.class), any()))
                 .thenAnswer(invocation -> {
                     Function func = invocation.getArgument(2);
-                    return new CompletedDurableFuture<>(func.apply(childContext));
+                    var result = func.apply(childContext);
+                    return (DurableFuture) () -> result;
                 });
     }
 
-    /** Stubs runInChildContext (sync) to immediately execute the function with childContext. */
-    private void stubChildContextSync(String name) {
-        when(context.runInChildContext(eq(name), any(TypeToken.class), any())).thenAnswer(invocation -> {
-            Function func = invocation.getArgument(2);
-            return func.apply(childContext);
-        });
-    }
-
     // --- Named form tests ---
 
     @Nested
     class NamedForm {
 
+        @BeforeEach
+        void setUpChildContext() {
+            stubChildContext("my-op");
+            stubChildContext("name");
+        }
+
         @Test
-        void successOnFirstAttempt_runsDirectlyOnContext() {
+        void successOnFirstAttempt() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
@@ -169,18 +167,18 @@ void successOnFirstAttempt_runsDirectlyOnContext() {
             var result = context.withRetry("my-op", (ctx, attempt) -> "success", config);
 
             assertEquals("success", result);
-            verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
+            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
         }
 
         @Test
-        void doesNotWrapInChildContext() {
+        void alwaysUsesChildContext() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
             context.withRetry("my-op", (ctx, attempt) -> "result", config);
 
-            verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
+            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
         }
 
         @Test
@@ -204,10 +202,10 @@ void retriesWithBackoffWaits_namedForm() {
 
             assertEquals("success-on-3", result);
             assertEquals(3, callCount[0]);
-            // Backoff waits happen on the caller's context directly
-            verify(context).wait("my-op-backoff-1", Duration.ofSeconds(5));
-            verify(context).wait("my-op-backoff-2", Duration.ofSeconds(5));
-            verify(context, times(2)).wait(anyString(), any(Duration.class));
+            // Backoff waits happen on the child context
+            verify(childContext).wait("my-op-backoff-1", Duration.ofSeconds(5));
+            verify(childContext).wait("my-op-backoff-2", Duration.ofSeconds(5));
+            verify(childContext, times(2)).wait(anyString(), any(Duration.class));
         }
 
         @Test
@@ -226,7 +224,7 @@ void rethrowsWhenRetryStrategyReturnsFail() {
                             config));
 
             assertEquals("terminal", exception.getMessage());
-            verify(context, never()).wait(anyString(), any(Duration.class));
+            verify(childContext, never()).wait(anyString(), any(Duration.class));
         }
 
         @Test
@@ -250,7 +248,7 @@ void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
 
             assertEquals("ok", result);
             // Zero delay should be replaced with 1-second default
-            verify(context).wait("my-op-backoff-1", Duration.ofSeconds(1));
+            verify(childContext).wait("my-op-backoff-1", Duration.ofSeconds(1));
         }
 
         @Test
@@ -288,7 +286,7 @@ void operationReturnsNull() {
         }
 
         @Test
-        void passesCallerContextToOperation() {
+        void passesChildContextToOperation() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
@@ -296,7 +294,7 @@ void passesCallerContextToOperation() {
             context.withRetry(
                     "my-op",
                     (ctx, attempt) -> {
-                        assertSame(context, ctx);
+                        assertSame(childContext, ctx);
                         return "verified";
                     },
                     config);
@@ -308,6 +306,11 @@ void passesCallerContextToOperation() {
     @Nested
     class AnonymousForm {
 
+        @BeforeEach
+        void setUpChildContext() {
+            stubChildContext("retry");
+        }
+
         @Test
         void successOnFirstAttempt() {
             var config = WithRetryConfig.builder()
@@ -320,14 +323,14 @@ void successOnFirstAttempt() {
         }
 
         @Test
-        void doesNotWrapInChildContext() {
+        void alwaysUsesChildContext() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
             context.withRetry((ctx, attempt) -> "result", config);
 
-            verify(context, never()).runInChildContextAsync(anyString(), any(TypeToken.class), any());
+            verify(context).runInChildContextAsync(eq("retry"), any(TypeToken.class), any());
         }
 
         @Test
@@ -347,8 +350,8 @@ void retriesWithAnonymousBackoffNames() {
                     config);
 
             assertEquals("done", result);
-            verify(context).wait("retry-backoff-1", Duration.ofSeconds(2));
-            verify(context).wait("retry-backoff-2", Duration.ofSeconds(2));
+            verify(childContext).wait("retry-backoff-1", Duration.ofSeconds(2));
+            verify(childContext).wait("retry-backoff-2", Duration.ofSeconds(2));
         }
 
         @Test
@@ -411,18 +414,18 @@ void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
                     config);
 
             assertEquals("ok", result);
-            verify(context).wait("retry-backoff-1", Duration.ofSeconds(1));
+            verify(childContext).wait("retry-backoff-1", Duration.ofSeconds(1));
         }
 
         @Test
-        void passesCallerContextToOperation() {
+        void passesChildContextToOperation() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
             context.withRetry(
                     (ctx, attempt) -> {
-                        assertSame(context, ctx);
+                        assertSame(childContext, ctx);
                         return "verified";
                     },
                     config);
@@ -610,6 +613,11 @@ void syncAndAsyncProduceSameResult() {
     @Nested
     class RetryBehavior {
 
+        @BeforeEach
+        void setUpChildContext() {
+            stubChildContextAnyName();
+        }
+
         @Test
         void passesCorrectAttemptNumberToOperation() {
             var attempts = new ArrayList();
@@ -700,26 +708,26 @@ void respectsCustomDelayFromRetryDecision() {
                     },
                     config);
 
-            verify(context).wait("retry-backoff-1", Duration.ofSeconds(10));
-            verify(context).wait("retry-backoff-2", Duration.ofSeconds(20));
+            verify(childContext).wait("retry-backoff-1", Duration.ofSeconds(10));
+            verify(childContext).wait("retry-backoff-2", Duration.ofSeconds(20));
         }
 
         @Test
-        void anonymousFormPassesCallerContextToOperation() {
+        void anonymousFormPassesChildContextToOperation() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
             context.withRetry(
                     (ctx, attempt) -> {
-                        assertSame(context, ctx);
+                        assertSame(childContext, ctx);
                         return "verified";
                     },
                     config);
         }
 
         @Test
-        void namedFormPassesCallerContextToOperation() {
+        void namedFormPassesChildContextToOperation() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
@@ -727,7 +735,7 @@ void namedFormPassesCallerContextToOperation() {
             context.withRetry(
                     "wrapped",
                     (ctx, attempt) -> {
-                        assertSame(context, ctx);
+                        assertSame(childContext, ctx);
                         return "verified";
                     },
                     config);
@@ -766,7 +774,7 @@ void propagatesSuspendExecutionExceptionWithoutRetrying() {
                             config));
 
             // Should never reach the wait — SuspendExecutionException propagates immediately
-            verify(context, never()).wait(anyString(), any(Duration.class));
+            verify(childContext, never()).wait(anyString(), any(Duration.class));
         }
 
         @Test
@@ -786,7 +794,7 @@ void propagatesUnrecoverableDurableExecutionExceptionWithoutRetrying() {
                             },
                             config));
 
-            verify(context, never()).wait(anyString(), any(Duration.class));
+            verify(childContext, never()).wait(anyString(), any(Duration.class));
         }
 
         @Test
@@ -808,7 +816,7 @@ void propagatesSuspendExecutionExceptionOnLaterAttempt() {
                             config));
 
             // First attempt retried (one backoff wait), second attempt suspended immediately
-            verify(context, times(1)).wait(anyString(), any(Duration.class));
+            verify(childContext, times(1)).wait(anyString(), any(Duration.class));
         }
 
         @Test
@@ -831,7 +839,7 @@ void propagatesUnrecoverableDurableExecutionExceptionOnLaterAttempt() {
                             },
                             config));
 
-            verify(context, times(1)).wait(anyString(), any(Duration.class));
+            verify(childContext, times(1)).wait(anyString(), any(Duration.class));
         }
 
         @Test
@@ -860,10 +868,13 @@ void preservesCheckedExceptionSubclassType() {
     @Nested
     class WrapInChildContext {
 
-        @Test
-        void namedForm_wrapsInChildContextWhenEnabled() {
-            stubChildContextSync("my-op");
+        @BeforeEach
+        void setUpChildContext() {
+            stubChildContextAnyName();
+        }
 
+        @Test
+        void namedForm_alwaysUsesChildContextWhenEnabled() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .wrapInChildContext(true)
@@ -872,37 +883,36 @@ void namedForm_wrapsInChildContextWhenEnabled() {
             var result = context.withRetry("my-op", (ctx, attempt) -> "wrapped", config);
 
             assertEquals("wrapped", result);
-            verify(context).runInChildContext(eq("my-op"), any(TypeToken.class), any());
+            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
         }
 
         @Test
-        void namedForm_doesNotWrapWhenDisabled() {
+        void namedForm_usesChildContextWhenDisabled() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .wrapInChildContext(false)
                     .build();
 
-            var result = context.withRetry("my-op", (ctx, attempt) -> "direct", config);
+            var result = context.withRetry("my-op", (ctx, attempt) -> "virtual", config);
 
-            assertEquals("direct", result);
-            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+            assertEquals("virtual", result);
+            // Still uses a child context (virtual in real impl)
+            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
         }
 
         @Test
-        void namedForm_doesNotWrapByDefault() {
+        void namedForm_usesChildContextByDefault() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            context.withRetry("my-op", (ctx, attempt) -> "direct", config);
+            context.withRetry("my-op", (ctx, attempt) -> "virtual", config);
 
-            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
         }
 
         @Test
-        void anonymousForm_wrapsInChildContextWhenEnabled() {
-            stubChildContextSync("retry");
-
+        void anonymousForm_alwaysUsesChildContextWhenEnabled() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .wrapInChildContext(true)
@@ -911,26 +921,25 @@ void anonymousForm_wrapsInChildContextWhenEnabled() {
             var result = context.withRetry((ctx, attempt) -> "wrapped", config);
 
             assertEquals("wrapped", result);
-            verify(context).runInChildContext(eq("retry"), any(TypeToken.class), any());
+            verify(context).runInChildContextAsync(eq("retry"), any(TypeToken.class), any());
         }
 
         @Test
-        void anonymousForm_doesNotWrapWhenDisabled() {
+        void anonymousForm_usesChildContextWhenDisabled() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .wrapInChildContext(false)
                     .build();
 
-            var result = context.withRetry((ctx, attempt) -> "direct", config);
+            var result = context.withRetry((ctx, attempt) -> "virtual", config);
 
-            assertEquals("direct", result);
-            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
+            assertEquals("virtual", result);
+            // Still uses a child context (virtual in real impl)
+            verify(context).runInChildContextAsync(eq("retry"), any(TypeToken.class), any());
         }
 
         @Test
         void namedForm_passesChildContextToOperationWhenWrapped() {
-            stubChildContextSync("my-op");
-
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .wrapInChildContext(true)
@@ -947,8 +956,6 @@ void namedForm_passesChildContextToOperationWhenWrapped() {
 
         @Test
         void namedForm_retriesWithBackoffOnChildContextWhenWrapped() {
-            stubChildContextSync("my-op");
-
             var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) ->
                             attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
@@ -972,8 +979,6 @@ void namedForm_retriesWithBackoffOnChildContextWhenWrapped() {
 
         @Test
         void anonymousForm_retriesWithBackoffOnChildContextWhenWrapped() {
-            stubChildContextSync("retry");
-
             var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) ->
                             attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
deleted file mode 100644
index 265fa968c..000000000
--- a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryOperationHelperTest.java
+++ /dev/null
@@ -1,462 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.util;
-
-import static org.junit.jupiter.api.Assertions.*;
-import static org.mockito.ArgumentMatchers.*;
-import static org.mockito.Mockito.*;
-
-import java.time.Duration;
-import java.util.ArrayList;
-import java.util.function.Function;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Nested;
-import org.junit.jupiter.api.Test;
-import software.amazon.lambda.durable.DurableContext;
-import software.amazon.lambda.durable.TypeToken;
-import software.amazon.lambda.durable.config.RetryOperationConfig;
-import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
-import software.amazon.lambda.durable.execution.SuspendExecutionException;
-import software.amazon.lambda.durable.retry.RetryDecision;
-import software.amazon.lambda.durable.retry.RetryStrategies;
-
-class RetryOperationHelperTest {
-
-    private DurableContext context;
-
-    @BeforeEach
-    void setUp() {
-        context = mock(DurableContext.class);
-    }
-
-    // --- Named form tests ---
-
-    @Nested
-    class NamedForm {
-
-        @Test
-        void successOnFirstAttempt_wrapsInChildContext() {
-            // runInChildContext should be called; delegate to the function immediately
-            when(context.runInChildContext(eq("my-op"), any(TypeToken.class), any()))
-                    .thenAnswer(invocation -> {
-                        Function func = invocation.getArgument(2);
-                        return func.apply(context);
-                    });
-
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var result = RetryOperationHelper.retryOperation(context, "my-op", (ctx, attempt) -> "success", config);
-
-            assertEquals("success", result);
-            verify(context).runInChildContext(eq("my-op"), any(TypeToken.class), any());
-        }
-
-        @Test
-        void successOnFirstAttempt_noChildContext_whenDisabled() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(false)
-                    .build();
-
-            var result = RetryOperationHelper.retryOperation(context, "my-op", (ctx, attempt) -> "direct", config);
-
-            assertEquals("direct", result);
-            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
-        }
-
-        @Test
-        void retriesWithBackoffWaits_namedForm() {
-            // Disable child context wrapping so we can directly verify wait calls
-            var callCount = new int[] {0};
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) ->
-                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
-                    .wrapInChildContext(false)
-                    .build();
-
-            var result = RetryOperationHelper.retryOperation(
-                    context,
-                    "my-op",
-                    (ctx, attempt) -> {
-                        callCount[0]++;
-                        if (attempt < 3) {
-                            throw new RuntimeException("fail-" + attempt);
-                        }
-                        return "success-on-3";
-                    },
-                    config);
-
-            assertEquals("success-on-3", result);
-            assertEquals(3, callCount[0]);
-            verify(context).wait("my-op-backoff-1", Duration.ofSeconds(5));
-            verify(context).wait("my-op-backoff-2", Duration.ofSeconds(5));
-            verify(context, times(2)).wait(anyString(), any(Duration.class));
-        }
-
-        @Test
-        void rethrowsWhenRetryStrategyReturnsFail() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(false)
-                    .build();
-
-            var exception = assertThrows(
-                    RuntimeException.class,
-                    () -> RetryOperationHelper.retryOperation(
-                            context,
-                            "my-op",
-                            (ctx, attempt) -> {
-                                throw new RuntimeException("terminal");
-                            },
-                            config));
-
-            assertEquals("terminal", exception.getMessage());
-            verify(context, never()).wait(anyString(), any(Duration.class));
-        }
-
-        @Test
-        void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(
-                            (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail())
-                    .wrapInChildContext(false)
-                    .build();
-
-            var callCount = new int[] {0};
-            var result = RetryOperationHelper.retryOperation(
-                    context,
-                    "my-op",
-                    (ctx, attempt) -> {
-                        callCount[0]++;
-                        if (attempt == 1) {
-                            throw new RuntimeException("fail");
-                        }
-                        return "ok";
-                    },
-                    config);
-
-            assertEquals("ok", result);
-            // Zero delay should be replaced with 1-second default
-            verify(context).wait("my-op-backoff-1", Duration.ofSeconds(1));
-        }
-
-        @Test
-        void nullContext_shouldThrow() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(null, "name", (ctx, a) -> "x", config));
-        }
-
-        @Test
-        void nullName_shouldThrow() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, null, (ctx, a) -> "x", config));
-        }
-
-        @Test
-        void nullOperation_shouldThrow() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, "name", null, config));
-        }
-
-        @Test
-        void nullConfig_shouldThrow() {
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, "name", (ctx, a) -> "x", null));
-        }
-    }
-
-    // --- Anonymous form tests ---
-
-    @Nested
-    class AnonymousForm {
-
-        @Test
-        void successOnFirstAttempt() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var result = RetryOperationHelper.retryOperation(context, (ctx, attempt) -> "anonymous-success", config);
-
-            assertEquals("anonymous-success", result);
-            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
-        }
-
-        @Test
-        void retriesWithAnonymousBackoffNames() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) ->
-                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
-                    .build();
-
-            var result = RetryOperationHelper.retryOperation(
-                    context,
-                    (ctx, attempt) -> {
-                        if (attempt < 3) {
-                            throw new RuntimeException("fail");
-                        }
-                        return "done";
-                    },
-                    config);
-
-            assertEquals("done", result);
-            verify(context).wait("retry-backoff-1", Duration.ofSeconds(2));
-            verify(context).wait("retry-backoff-2", Duration.ofSeconds(2));
-        }
-
-        @Test
-        void neverWrapsInChildContext_evenWhenConfigSaysTrue() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(true) // should be ignored for anonymous form
-                    .build();
-
-            RetryOperationHelper.retryOperation(context, (ctx, attempt) -> "result", config);
-
-            verify(context, never()).runInChildContext(anyString(), any(TypeToken.class), any());
-        }
-
-        @Test
-        void rethrowsOriginalException() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var original = new IllegalStateException("original error");
-            var thrown = assertThrows(
-                    IllegalStateException.class,
-                    () -> RetryOperationHelper.retryOperation(
-                            context,
-                            (ctx, attempt) -> {
-                                throw original;
-                            },
-                            config));
-
-            assertSame(original, thrown);
-        }
-
-        @Test
-        void nullContext_shouldThrow() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(null, (ctx, a) -> "x", config));
-        }
-
-        @Test
-        void nullOperation_shouldThrow() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, (RetryableOperation) null, config));
-        }
-
-        @Test
-        void nullConfig_shouldThrow() {
-            assertThrows(
-                    NullPointerException.class,
-                    () -> RetryOperationHelper.retryOperation(context, (ctx, a) -> "x", null));
-        }
-    }
-
-    // --- Retry behavior tests ---
-
-    @Nested
-    class RetryBehavior {
-
-        @Test
-        void passesCorrectAttemptNumberToOperation() {
-            var attempts = new ArrayList();
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) ->
-                            attempt < 4 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
-                    .wrapInChildContext(false)
-                    .build();
-
-            RetryOperationHelper.retryOperation(
-                    context,
-                    "track",
-                    (ctx, attempt) -> {
-                        attempts.add(attempt);
-                        if (attempt < 4) {
-                            throw new RuntimeException("not yet");
-                        }
-                        return "done";
-                    },
-                    config);
-
-            assertEquals(4, attempts.size());
-            assertEquals(1, attempts.get(0));
-            assertEquals(2, attempts.get(1));
-            assertEquals(3, attempts.get(2));
-            assertEquals(4, attempts.get(3));
-        }
-
-        @Test
-        void passesErrorToRetryStrategy() {
-            var errors = new ArrayList();
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) -> {
-                        errors.add(error);
-                        return attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail();
-                    })
-                    .build();
-
-            assertThrows(
-                    RuntimeException.class,
-                    () -> RetryOperationHelper.retryOperation(
-                            context,
-                            (ctx, attempt) -> {
-                                throw new RuntimeException("error-" + attempt);
-                            },
-                            config));
-
-            assertEquals(3, errors.size());
-            assertEquals("error-1", errors.get(0).getMessage());
-            assertEquals("error-2", errors.get(1).getMessage());
-            assertEquals("error-3", errors.get(2).getMessage());
-        }
-
-        @Test
-        void respectsCustomDelayFromRetryDecision() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(attempt * 10L)))
-                    .build();
-
-            RetryOperationHelper.retryOperation(
-                    context,
-                    (ctx, attempt) -> {
-                        if (attempt <= 2) {
-                            throw new RuntimeException("fail");
-                        }
-                        return "ok";
-                    },
-                    config);
-
-            verify(context).wait("retry-backoff-1", Duration.ofSeconds(10));
-            verify(context).wait("retry-backoff-2", Duration.ofSeconds(20));
-        }
-
-        @Test
-        void passesContextToOperation() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            RetryOperationHelper.retryOperation(
-                    context,
-                    (ctx, attempt) -> {
-                        assertSame(context, ctx);
-                        return "verified";
-                    },
-                    config);
-        }
-
-        @Test
-        void namedFormPassesChildContextToOperation_whenWrapped() {
-            var childContext = mock(DurableContext.class);
-            when(context.runInChildContext(eq("wrapped"), any(TypeToken.class), any()))
-                    .thenAnswer(invocation -> {
-                        Function func = invocation.getArgument(2);
-                        return func.apply(childContext);
-                    });
-
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(true)
-                    .build();
-
-            RetryOperationHelper.retryOperation(
-                    context,
-                    "wrapped",
-                    (ctx, attempt) -> {
-                        assertSame(childContext, ctx);
-                        return "verified";
-                    },
-                    config);
-        }
-
-        @Test
-        void rethrowsLastExceptionWhenAllRetriesExhausted() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
-                    .build();
-
-            var thrown = assertThrows(
-                    RuntimeException.class,
-                    () -> RetryOperationHelper.retryOperation(
-                            context,
-                            (ctx, attempt) -> {
-                                throw new RuntimeException("attempt-" + attempt);
-                            },
-                            config));
-
-            // The last attempt's exception is rethrown
-            assertEquals("attempt-3", thrown.getMessage());
-        }
-
-        @Test
-        void propagatesSuspendExecutionExceptionWithoutRetrying() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1)))
-                    .build();
-
-            assertThrows(
-                    SuspendExecutionException.class,
-                    () -> RetryOperationHelper.retryOperation(
-                            context,
-                            (ctx, attempt) -> {
-                                throw new SuspendExecutionException();
-                            },
-                            config));
-
-            // Should never reach the wait — SuspendExecutionException propagates immediately
-            verify(context, never()).wait(anyString(), any(Duration.class));
-        }
-
-        @Test
-        void propagatesUnrecoverableDurableExecutionExceptionWithoutRetrying() {
-            var config = RetryOperationConfig.builder()
-                    .retryStrategy((error, attempt) -> RetryDecision.retry(Duration.ofSeconds(1)))
-                    .build();
-
-            assertThrows(
-                    UnrecoverableDurableExecutionException.class,
-                    () -> RetryOperationHelper.retryOperation(
-                            context,
-                            (ctx, attempt) -> {
-                                throw new UnrecoverableDurableExecutionException(
-                                        software.amazon.awssdk.services.lambda.model.ErrorObject.builder()
-                                                .errorMessage("unrecoverable")
-                                                .build());
-                            },
-                            config));
-
-            verify(context, never()).wait(anyString(), any(Duration.class));
-        }
-    }
-}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java
deleted file mode 100644
index f5893a37c..000000000
--- a/sdk/src/test/java/software/amazon/lambda/durable/util/RetryableOperationTest.java
+++ /dev/null
@@ -1,52 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.util;
-
-import static org.junit.jupiter.api.Assertions.*;
-import static org.mockito.Mockito.*;
-
-import org.junit.jupiter.api.Test;
-import software.amazon.lambda.durable.DurableContext;
-
-class RetryableOperationTest {
-
-    @Test
-    void canBeImplementedAsLambda() {
-        RetryableOperation operation = (ctx, attempt) -> "result-" + attempt;
-        var context = mock(DurableContext.class);
-
-        assertEquals("result-1", operation.execute(context, 1));
-        assertEquals("result-3", operation.execute(context, 3));
-    }
-
-    @Test
-    void receivesContextAndAttempt() {
-        var context = mock(DurableContext.class);
-        RetryableOperation operation = (ctx, attempt) -> {
-            assertSame(context, ctx);
-            return "attempt-" + attempt;
-        };
-
-        assertEquals("attempt-1", operation.execute(context, 1));
-        assertEquals("attempt-2", operation.execute(context, 2));
-    }
-
-    @Test
-    void canThrowExceptions() {
-        RetryableOperation operation = (ctx, attempt) -> {
-            throw new RuntimeException("failed on attempt " + attempt);
-        };
-        var context = mock(DurableContext.class);
-
-        var exception = assertThrows(RuntimeException.class, () -> operation.execute(context, 1));
-        assertEquals("failed on attempt 1", exception.getMessage());
-    }
-
-    @Test
-    void canReturnNull() {
-        RetryableOperation operation = (ctx, attempt) -> null;
-        var context = mock(DurableContext.class);
-
-        assertNull(operation.execute(context, 1));
-    }
-}

From 9af4a42a8c1ff39b7669b133ecd0a443d36145f9 Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Thu, 30 Apr 2026 15:35:23 -0700
Subject: [PATCH 15/19] chore: refactor withRetry to use withRetryAsync

---
 .../amazon/lambda/durable/DurableContext.java |  8 ++-
 .../durable/context/DurableContextImpl.java   | 49 -------------------
 2 files changed, 6 insertions(+), 51 deletions(-)

diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
index 53a17a2c2..f2ec53a04 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
@@ -751,7 +751,9 @@  DurableFuture waitForConditionAsync(
      * @param config retry configuration including the retry strategy and child context wrapping
      * @return the operation result
      */
-     T withRetry(String name, WithRetry operation, WithRetryConfig config);
+    default  T withRetry(String name, WithRetry operation, WithRetryConfig config) {
+        return withRetryAsync(name, operation, config).get();
+    }
 
     /**
      * Replay-safe retry loop for any durable operation (anonymous form, sync).
@@ -766,7 +768,9 @@  DurableFuture waitForConditionAsync(
      * @param config retry configuration including the retry strategy and child context wrapping
      * @return the operation result
      */
-     T withRetry(WithRetry operation, WithRetryConfig config);
+    default  T withRetry(WithRetry operation, WithRetryConfig config) {
+        return withRetryAsync(operation, config).get();
+    }
 
     /**
      * Replay-safe retry loop for any durable operation (named form, async).
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
index 1274e1c20..4a8b26aff 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
@@ -399,55 +399,6 @@ public  DurableFuture waitForConditionAsync(
     private static final String ANONYMOUS_CHILD_CONTEXT_NAME = "retry";
     private static final String ANONYMOUS_BACKOFF_PREFIX = "retry-backoff-";
 
-    @Override
-    @SuppressWarnings("unchecked")
-    public  T withRetry(String name, WithRetry operation, WithRetryConfig config) {
-        Objects.requireNonNull(name, "name cannot be null");
-        Objects.requireNonNull(operation, "operation cannot be null");
-        Objects.requireNonNull(config, "config cannot be null");
-
-        if (config.wrapInChildContext()) {
-            return (T) runInChildContextAsync(
-                            name,
-                            new TypeToken() {},
-                            childCtx -> executeRetryLoop(childCtx, name, operation, config),
-                            RunInChildContextConfig.builder().build(),
-                            OperationSubType.WITH_RETRY)
-                    .get();
-        }
-        return (T) runInVirtualChildContextAsync(
-                        name,
-                        new TypeToken() {},
-                        childCtx -> executeRetryLoop(childCtx, name, operation, config),
-                        RunInChildContextConfig.builder().build(),
-                        OperationSubType.WITH_RETRY)
-                .get();
-    }
-
-    @Override
-    @SuppressWarnings("unchecked")
-    public  T withRetry(WithRetry operation, WithRetryConfig config) {
-        Objects.requireNonNull(operation, "operation cannot be null");
-        Objects.requireNonNull(config, "config cannot be null");
-
-        if (config.wrapInChildContext()) {
-            return (T) runInChildContextAsync(
-                            ANONYMOUS_CHILD_CONTEXT_NAME,
-                            new TypeToken() {},
-                            childCtx -> executeRetryLoop(childCtx, null, operation, config),
-                            RunInChildContextConfig.builder().build(),
-                            OperationSubType.WITH_RETRY)
-                    .get();
-        }
-        return (T) runInVirtualChildContextAsync(
-                        ANONYMOUS_CHILD_CONTEXT_NAME,
-                        new TypeToken() {},
-                        childCtx -> executeRetryLoop(childCtx, null, operation, config),
-                        RunInChildContextConfig.builder().build(),
-                        OperationSubType.WITH_RETRY)
-                .get();
-    }
-
     @Override
     @SuppressWarnings("unchecked")
     public  DurableFuture withRetryAsync(String name, WithRetry operation, WithRetryConfig config) {

From f54d8a4a7b1c006ab7cd863e16d65bb52c0150e4 Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Thu, 30 Apr 2026 16:12:22 -0700
Subject: [PATCH 16/19] chore: remove anonymous form function override

---
 .../callback/RetryWaitForCallbackExample.java |   3 +-
 .../examples/invoke/RetryInvokeExample.java   |   3 +-
 .../durable/RetryInvokeIntegrationTest.java   |   6 +
 .../RetryWaitForCallbackIntegrationTest.java  |   6 +
 .../amazon/lambda/durable/DurableContext.java |  42 +-
 .../durable/context/DurableContextImpl.java   |  64 +-
 .../context/DurableContextWithRetryTest.java  | 774 +++++-------------
 7 files changed, 214 insertions(+), 684 deletions(-)

diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
index 953366050..c08d44937 100644
--- a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
+++ b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
@@ -17,7 +17,7 @@
  * each time.
  *
  * 
Each attempt uses a unique callback name ({@code "approval-1"}, {@code "approval-2"}, etc.) so the execution
- * history stays clean and replay-safe. The anonymous form is used, so attempts are grouped under a default-named child
+ * 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 {
@@ -34,6 +34,7 @@ public String handleRequest(ApprovalRequest input, DurableContext context) {
 
         // Step 2: waitForCallback with retry — if the external system fails, try again with a fresh callback
         var approvalResult = context.withRetry(
+                null,
                 (ctx, attempt) -> ctx.waitForCallback(
                         "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
                                 .info("Attempt {}: sending callback {} to approval system", attempt, callbackId)),
diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
index 7850b2d7b..8f577036f 100644
--- a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
+++ b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
@@ -16,7 +16,7 @@
  * uses a unique operation name ({@code "call-greeting-1"}, {@code "call-greeting-2"}, etc.) so the execution history
  * stays clean and replay-safe.
  *
- * 
The anonymous form is used, so attempts are grouped under a default-named child context.
+ * 
A {@code null} name is used, so attempts are grouped under a default-named child context.
  */
 public class RetryInvokeExample extends DurableHandler {
 
@@ -25,6 +25,7 @@ public class RetryInvokeExample extends DurableHandler
     @Override
     public String handleRequest(GreetingRequest input, DurableContext context) {
         return context.withRetry(
+                null,
                 (ctx, attempt) -> ctx.invoke(
                         "call-greeting-" + attempt,
                         "simple-step-example" + input.getName() + ":$LATEST",
diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
index fd206831d..5aa873692 100644
--- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
+++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
@@ -21,6 +21,7 @@ void invokeSucceedsOnFirstAttempt() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
                 (input, context) -> context.withRetry(
+                        null,
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
@@ -41,6 +42,7 @@ void invokeRetriesAfterFailure() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
                 (input, context) -> context.withRetry(
+                        null,
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
@@ -78,6 +80,7 @@ void invokeFailsAfterAllRetriesExhausted() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
                 (input, context) -> context.withRetry(
+                        null,
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy((error, attempt) ->
@@ -111,6 +114,7 @@ void invokeRetryWithCustomBackoffDelay() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
                 (input, context) -> context.withRetry(
+                        null,
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy((error, attempt) -> attempt < 3
@@ -144,6 +148,7 @@ void invokeRetryWithStepsBeforeAndAfter() {
             var prefix = context.step("prepare", String.class, stepCtx -> "prepared");
 
             var invokeResult = context.withRetry(
+                    null,
                     (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                     WithRetryConfig.builder()
                             .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
@@ -169,6 +174,7 @@ void invokeRetryPreservesOriginalExceptionType() {
         var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
             try {
                 return context.withRetry(
+                        null,
                         (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy(RetryStrategies.Presets.NO_RETRY)
diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
index 1cd7b7e50..d07d944c4 100644
--- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
+++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
@@ -20,6 +20,7 @@ void waitForCallbackSucceedsOnFirstAttempt() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
                 (input, context) -> context.withRetry(
+                        null,
                         (ctx, attempt) -> ctx.waitForCallback(
                                 "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
                                         .info("Submitting callback {}", callbackId)),
@@ -46,6 +47,7 @@ void waitForCallbackRetriesAfterFailure() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
                 (input, context) -> context.withRetry(
+                        null,
                         (ctx, attempt) -> ctx.waitForCallback(
                                 "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
                                         .info("Attempt {} callback {}", attempt, callbackId)),
@@ -92,6 +94,7 @@ void waitForCallbackFailsAfterAllRetriesExhausted() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
                 (input, context) -> context.withRetry(
+                        null,
                         (ctx, attempt) ->
                                 ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
                         WithRetryConfig.builder()
@@ -129,6 +132,7 @@ void waitForCallbackRetryWithStepsBeforeAndAfter() {
             var prefix = context.step("prepare", String.class, stepCtx -> "prepared");
 
             var callbackResult = context.withRetry(
+                    null,
                     (ctx, attempt) ->
                             ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
                     WithRetryConfig.builder()
@@ -157,6 +161,7 @@ void waitForCallbackRetryMultipleFailuresThenSuccess() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
                 (input, context) -> context.withRetry(
+                        null,
                         (ctx, attempt) ->
                                 ctx.waitForCallback("cb-" + attempt, String.class, (callbackId, stepCtx) -> {}),
                         WithRetryConfig.builder()
@@ -202,6 +207,7 @@ void waitForCallbackRetryWithSubmitterLogic() {
         var runner = LocalDurableTestRunner.create(
                 String.class,
                 (input, context) -> context.withRetry(
+                        null,
                         (ctx, attempt) ->
                                 ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {
                                     // Submitter runs each attempt — in a real scenario this would
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
index f2ec53a04..ca17109ce 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
@@ -732,7 +732,7 @@  DurableFuture waitForConditionAsync(
     // =============== withRetry ================
 
     /**
-     * Replay-safe retry loop for any durable operation (named form, sync).
+     * Replay-safe retry loop for any durable operation (sync).
      *
      * 
Provides the same retry-with-backoff pattern that {@code step()} has built in, but for operations that cannot
      * live inside a step ({@code waitForCallback}, {@code invoke}, {@code waitForCondition}, etc.).
@@ -746,7 +746,8 @@  DurableFuture waitForConditionAsync(
      * used — no checkpointing overhead, but the child re-executes on replay.
      *
      * @param  the result type
-     * @param name operation name (used for backoff wait names, and as the child context name when wrapping)
+     * @param name operation name (used for backoff wait names, and as the child context name when wrapping); pass
+     *     {@code null} for an anonymous retry whose backoff waits use default names
      * @param operation the retryable operation — receives the context and 1-based attempt number
      * @param config retry configuration including the retry strategy and child context wrapping
      * @return the operation result
@@ -756,7 +757,7 @@ default  T withRetry(String name, WithRetry operation, WithRetryConfig con
     }
 
     /**
-     * Replay-safe retry loop for any durable operation (anonymous form, sync).
+     * Replay-safe retry loop for any durable operation (async).
      *
      * 
The retry loop always runs in a child context to provide an isolated operation ID namespace. If
      * {@link WithRetryConfig#wrapInChildContext()} is enabled, the child context is checkpointed (persisted) so all
@@ -764,45 +765,14 @@ default  T withRetry(String name, WithRetry operation, WithRetryConfig con
      * used — no checkpointing overhead, but the child re-executes on replay.
      *
      * @param  the result type
-     * @param operation the retryable operation — receives the context and 1-based attempt number
-     * @param config retry configuration including the retry strategy and child context wrapping
-     * @return the operation result
-     */
-    default  T withRetry(WithRetry operation, WithRetryConfig config) {
-        return withRetryAsync(operation, config).get();
-    }
-
-    /**
-     * Replay-safe retry loop for any durable operation (named form, async).
-     *
-     * 
The retry loop always runs in a child context to provide an isolated operation ID namespace. If
-     * {@link WithRetryConfig#wrapInChildContext()} is enabled, the child context is checkpointed (persisted) so all
-     * attempts are grouped under a single named operation in execution history. Otherwise, a virtual child context is
-     * used — no checkpointing overhead, but the child re-executes on replay.
-     *
-     * @param  the result type
-     * @param name operation name (used for child context and backoff wait names)
+     * @param name operation name (used for child context and backoff wait names); pass {@code null} for an anonymous
+     *     retry whose backoff waits use default names
      * @param operation the retryable operation — receives the context and 1-based attempt number
      * @param config retry configuration including the retry strategy
      * @return a future representing the operation result
      */
      DurableFuture withRetryAsync(String name, WithRetry operation, WithRetryConfig config);
 
-    /**
-     * Replay-safe retry loop for any durable operation (anonymous form, async).
-     *
-     * 
The retry loop always runs in a child context to provide an isolated operation ID namespace. If
-     * {@link WithRetryConfig#wrapInChildContext()} is enabled, the child context is checkpointed (persisted) so all
-     * attempts are grouped under a single named operation in execution history. Otherwise, a virtual child context is
-     * used — no checkpointing overhead, but the child re-executes on replay.
-     *
-     * @param  the result type
-     * @param operation the retryable operation — receives the context and 1-based attempt number
-     * @param config retry configuration including the retry strategy
-     * @return a future representing the operation result
-     */
-     DurableFuture withRetryAsync(WithRetry operation, WithRetryConfig config);
-
     /**
      * Function applied to each item in a map operation.
      *
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
index 8487969a4..56e2ab36b 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
@@ -238,25 +238,6 @@ private  DurableFuture runInChildContextAsync(
             Function func,
             RunInChildContextConfig config,
             OperationSubType subType) {
-        return runInChildContextAsync(name, resultType, func, config, subType, false);
-    }
-
-    private  DurableFuture runInVirtualChildContextAsync(
-            String name,
-            TypeToken resultType,
-            Function func,
-            RunInChildContextConfig config,
-            OperationSubType subType) {
-        return runInChildContextAsync(name, resultType, func, config, subType, true);
-    }
-
-    private  DurableFuture runInChildContextAsync(
-            String name,
-            TypeToken resultType,
-            Function func,
-            RunInChildContextConfig config,
-            OperationSubType subType,
-            boolean isVirtual) {
         Objects.requireNonNull(resultType, "resultType cannot be null");
         Objects.requireNonNull(config, "RunInChildContextConfig cannot be null");
         ParameterValidator.validateOperationName(name);
@@ -272,9 +253,7 @@ private  DurableFuture runInChildContextAsync(
                 func,
                 resultType,
                 config,
-                this,
-                isVirtual,
-                null);
+                this);
 
         operation.execute();
         return operation;
@@ -402,45 +381,18 @@ public  DurableFuture waitForConditionAsync(
     @Override
     @SuppressWarnings("unchecked")
     public  DurableFuture withRetryAsync(String name, WithRetry operation, WithRetryConfig config) {
-        Objects.requireNonNull(name, "name cannot be null");
         Objects.requireNonNull(operation, "operation cannot be null");
         Objects.requireNonNull(config, "config cannot be null");
 
-        if (config.wrapInChildContext()) {
-            return (DurableFuture) runInChildContextAsync(
-                    name,
-                    new TypeToken() {},
-                    childCtx -> executeRetryLoop(childCtx, name, operation, config),
-                    RunInChildContextConfig.builder().build(),
-                    OperationSubType.WITH_RETRY);
-        }
-        return (DurableFuture) runInVirtualChildContextAsync(
-                name,
-                new TypeToken() {},
-                childCtx -> executeRetryLoop(childCtx, name, operation, config),
-                RunInChildContextConfig.builder().build(),
-                OperationSubType.WITH_RETRY);
-    }
-
-    @Override
-    @SuppressWarnings("unchecked")
-    public  DurableFuture withRetryAsync(WithRetry operation, WithRetryConfig config) {
-        Objects.requireNonNull(operation, "operation cannot be null");
-        Objects.requireNonNull(config, "config cannot be null");
+        var childContextName = name != null ? name : ANONYMOUS_CHILD_CONTEXT_NAME;
 
-        if (config.wrapInChildContext()) {
-            return (DurableFuture) runInChildContextAsync(
-                    ANONYMOUS_CHILD_CONTEXT_NAME,
-                    new TypeToken() {},
-                    childCtx -> executeRetryLoop(childCtx, null, operation, config),
-                    RunInChildContextConfig.builder().build(),
-                    OperationSubType.WITH_RETRY);
-        }
-        return (DurableFuture) runInVirtualChildContextAsync(
-                ANONYMOUS_CHILD_CONTEXT_NAME,
+        return (DurableFuture) runInChildContextAsync(
+                childContextName,
                 new TypeToken() {},
-                childCtx -> executeRetryLoop(childCtx, null, operation, config),
-                RunInChildContextConfig.builder().build(),
+                childCtx -> executeRetryLoop(childCtx, name, operation, config),
+                RunInChildContextConfig.builder()
+                        .isVirtual(!config.wrapInChildContext())
+                        .build(),
                 OperationSubType.WITH_RETRY);
     }
 
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
index 89b251cb5..5dfb24eee 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
@@ -46,63 +46,36 @@ void setUp() {
      * when {@code wrapInChildContext} is false, and a checkpointed child context when true).
      */
     private void stubWithRetryMethods(DurableContext mock) {
-        // Named sync form — always runs in a child context
+        // Sync form — always runs in a child context
         when(mock.withRetry(any(), nullable(WithRetry.class), nullable(WithRetryConfig.class)))
                 .thenAnswer(invocation -> {
                     String name = invocation.getArgument(0);
                     WithRetry operation = invocation.getArgument(1);
                     WithRetryConfig config = invocation.getArgument(2);
-                    Objects.requireNonNull(name, "name cannot be null");
                     Objects.requireNonNull(operation, "operation cannot be null");
                     Objects.requireNonNull(config, "config cannot be null");
+                    var childContextName = name != null ? name : "retry";
                     return mock.runInChildContextAsync(
-                                    name,
+                                    childContextName,
                                     new TypeToken() {},
                                     childCtx -> executeRetryLoop(childCtx, name, operation, config))
                             .get();
                 });
 
-        // Anonymous sync form — always runs in a child context
-        when(mock.withRetry(nullable(WithRetry.class), nullable(WithRetryConfig.class)))
-                .thenAnswer(invocation -> {
-                    WithRetry operation = invocation.getArgument(0);
-                    WithRetryConfig config = invocation.getArgument(1);
-                    Objects.requireNonNull(operation, "operation cannot be null");
-                    Objects.requireNonNull(config, "config cannot be null");
-                    return mock.runInChildContextAsync(
-                                    "retry",
-                                    new TypeToken() {},
-                                    childCtx -> executeRetryLoop(childCtx, null, operation, config))
-                            .get();
-                });
-
-        // Named async form
+        // Async form
         when(mock.withRetryAsync(any(), nullable(WithRetry.class), nullable(WithRetryConfig.class)))
                 .thenAnswer(invocation -> {
                     String name = invocation.getArgument(0);
                     WithRetry operation = invocation.getArgument(1);
                     WithRetryConfig config = invocation.getArgument(2);
-                    Objects.requireNonNull(name, "name cannot be null");
                     Objects.requireNonNull(operation, "operation cannot be null");
                     Objects.requireNonNull(config, "config cannot be null");
+                    var childContextName = name != null ? name : "retry";
                     return mock.runInChildContextAsync(
-                            name,
+                            childContextName,
                             new TypeToken() {},
                             childCtx -> executeRetryLoop(childCtx, name, operation, config));
                 });
-
-        // Anonymous async form
-        when(mock.withRetryAsync(nullable(WithRetry.class), nullable(WithRetryConfig.class)))
-                .thenAnswer(invocation -> {
-                    WithRetry operation = invocation.getArgument(0);
-                    WithRetryConfig config = invocation.getArgument(1);
-                    Objects.requireNonNull(operation, "operation cannot be null");
-                    Objects.requireNonNull(config, "config cannot be null");
-                    return mock.runInChildContextAsync(
-                            "retry",
-                            new TypeToken() {},
-                            childCtx -> executeRetryLoop(childCtx, null, operation, config));
-                });
     }
 
     /** Replicates the retry loop logic from DurableContextImpl for test stubbing. */
@@ -147,504 +120,115 @@ private void stubChildContextAnyName() {
                 });
     }
 
-    // --- Named form tests ---
-
-    @Nested
-    class NamedForm {
-
-        @BeforeEach
-        void setUpChildContext() {
-            stubChildContext("my-op");
-            stubChildContext("name");
-        }
-
-        @Test
-        void successOnFirstAttempt() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var result = context.withRetry("my-op", (ctx, attempt) -> "success", config);
-
-            assertEquals("success", result);
-            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
-        }
-
-        @Test
-        void alwaysUsesChildContext() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            context.withRetry("my-op", (ctx, attempt) -> "result", config);
-
-            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
-        }
-
-        @Test
-        void retriesWithBackoffWaits_namedForm() {
-            var callCount = new int[] {0};
-            var config = WithRetryConfig.builder()
-                    .retryStrategy((error, attempt) ->
-                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
-                    .build();
-
-            var result = context.withRetry(
-                    "my-op",
-                    (ctx, attempt) -> {
-                        callCount[0]++;
-                        if (attempt < 3) {
-                            throw new RuntimeException("fail-" + attempt);
-                        }
-                        return "success-on-3";
-                    },
-                    config);
-
-            assertEquals("success-on-3", result);
-            assertEquals(3, callCount[0]);
-            // Backoff waits happen on the child context
-            verify(childContext).wait("my-op-backoff-1", Duration.ofSeconds(5));
-            verify(childContext).wait("my-op-backoff-2", Duration.ofSeconds(5));
-            verify(childContext, times(2)).wait(anyString(), any(Duration.class));
-        }
-
-        @Test
-        void rethrowsWhenRetryStrategyReturnsFail() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var exception = assertThrows(
-                    RuntimeException.class,
-                    () -> context.withRetry(
-                            "my-op",
-                            (ctx, attempt) -> {
-                                throw new RuntimeException("terminal");
-                            },
-                            config));
-
-            assertEquals("terminal", exception.getMessage());
-            verify(childContext, never()).wait(anyString(), any(Duration.class));
-        }
-
-        @Test
-        void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(
-                            (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail())
-                    .build();
-
-            var callCount = new int[] {0};
-            var result = context.withRetry(
-                    "my-op",
-                    (ctx, attempt) -> {
-                        callCount[0]++;
-                        if (attempt == 1) {
-                            throw new RuntimeException("fail");
-                        }
-                        return "ok";
-                    },
-                    config);
-
-            assertEquals("ok", result);
-            // Zero delay should be replaced with 1-second default
-            verify(childContext).wait("my-op-backoff-1", Duration.ofSeconds(1));
-        }
-
-        @Test
-        void nullName_shouldThrow() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(NullPointerException.class, () -> context.withRetry(null, (ctx, a) -> "x", config));
-        }
-
-        @Test
-        void nullOperation_shouldThrow() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(NullPointerException.class, () -> context.withRetry("name", null, config));
-        }
-
-        @Test
-        void nullConfig_shouldThrow() {
-            assertThrows(NullPointerException.class, () -> context.withRetry("name", (ctx, a) -> "x", null));
-        }
-
-        @Test
-        void operationReturnsNull() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var result = context.withRetry("my-op", (WithRetry) (ctx, attempt) -> null, config);
-
-            assertNull(result);
-        }
-
-        @Test
-        void passesChildContextToOperation() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            context.withRetry(
-                    "my-op",
-                    (ctx, attempt) -> {
-                        assertSame(childContext, ctx);
-                        return "verified";
-                    },
-                    config);
-        }
-    }
-
-    // --- Anonymous form tests ---
-
-    @Nested
-    class AnonymousForm {
-
-        @BeforeEach
-        void setUpChildContext() {
-            stubChildContext("retry");
-        }
-
-        @Test
-        void successOnFirstAttempt() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var result = context.withRetry((ctx, attempt) -> "anonymous-success", config);
-
-            assertEquals("anonymous-success", result);
-        }
-
-        @Test
-        void alwaysUsesChildContext() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            context.withRetry((ctx, attempt) -> "result", config);
-
-            verify(context).runInChildContextAsync(eq("retry"), any(TypeToken.class), any());
-        }
-
-        @Test
-        void retriesWithAnonymousBackoffNames() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy((error, attempt) ->
-                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
-                    .build();
-
-            var result = context.withRetry(
-                    (ctx, attempt) -> {
-                        if (attempt < 3) {
-                            throw new RuntimeException("fail");
-                        }
-                        return "done";
-                    },
-                    config);
-
-            assertEquals("done", result);
-            verify(childContext).wait("retry-backoff-1", Duration.ofSeconds(2));
-            verify(childContext).wait("retry-backoff-2", Duration.ofSeconds(2));
-        }
-
-        @Test
-        void rethrowsOriginalException() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var original = new IllegalStateException("original error");
-            var thrown = assertThrows(
-                    IllegalStateException.class,
-                    () -> context.withRetry(
-                            (ctx, attempt) -> {
-                                throw original;
-                            },
-                            config));
-
-            assertSame(original, thrown);
-        }
-
-        @Test
-        void nullOperation_shouldThrow() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            assertThrows(NullPointerException.class, () -> context.withRetry((WithRetry) null, config));
-        }
-
-        @Test
-        void nullConfig_shouldThrow() {
-            assertThrows(NullPointerException.class, () -> context.withRetry((ctx, a) -> "x", null));
-        }
-
-        @Test
-        void operationReturnsNull() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var result = context.withRetry((WithRetry) (ctx, attempt) -> null, config);
-
-            assertNull(result);
-        }
-
-        @Test
-        void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(
-                            (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail())
-                    .build();
-
-            var result = context.withRetry(
-                    (ctx, attempt) -> {
-                        if (attempt == 1) {
-                            throw new RuntimeException("fail");
-                        }
-                        return "ok";
-                    },
-                    config);
-
-            assertEquals("ok", result);
-            verify(childContext).wait("retry-backoff-1", Duration.ofSeconds(1));
-        }
-
-        @Test
-        void passesChildContextToOperation() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            context.withRetry(
-                    (ctx, attempt) -> {
-                        assertSame(childContext, ctx);
-                        return "verified";
-                    },
-                    config);
-        }
-    }
-
-    // --- Async form tests ---
-
-    @Nested
-    class AsyncForm {
-
-        @Test
-        void namedAsyncReturnsDurableFuture() {
-            stubChildContext("my-op");
-
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            DurableFuture future = context.withRetryAsync("my-op", (ctx, attempt) -> "async-result", config);
-
-            assertNotNull(future);
-            assertEquals("async-result", future.get());
-        }
-
-        @Test
-        void namedAsyncAlwaysWrapsInChildContext() {
-            stubChildContext("my-op");
-
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            context.withRetryAsync("my-op", (ctx, attempt) -> "result", config);
-
-            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
-        }
-
-        @Test
-        void anonymousAsyncReturnsDurableFuture() {
-            stubChildContext("retry");
-
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            DurableFuture future = context.withRetryAsync((ctx, attempt) -> "anon-async", config);
-
-            assertNotNull(future);
-            assertEquals("anon-async", future.get());
-        }
-
-        @Test
-        void anonymousAsyncWrapsInChildContext() {
-            stubChildContext("retry");
-
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            context.withRetryAsync((ctx, attempt) -> "result", config);
-
-            verify(context).runInChildContextAsync(eq("retry"), any(TypeToken.class), any());
-        }
-
-        @Test
-        void namedAsyncRetriesWithBackoff() {
-            stubChildContext("my-op");
-
-            var config = WithRetryConfig.builder()
-                    .retryStrategy((error, attempt) ->
-                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
-                    .build();
-
-            var callCount = new int[] {0};
-            DurableFuture future = context.withRetryAsync(
-                    "my-op",
-                    (ctx, attempt) -> {
-                        callCount[0]++;
-                        if (attempt < 3) {
-                            throw new RuntimeException("fail-" + attempt);
-                        }
-                        return "success-on-3";
-                    },
-                    config);
-
-            assertEquals("success-on-3", future.get());
-            assertEquals(3, callCount[0]);
-            verify(childContext).wait("my-op-backoff-1", Duration.ofSeconds(5));
-            verify(childContext).wait("my-op-backoff-2", Duration.ofSeconds(5));
-        }
-
-        @Test
-        void anonymousAsyncRetriesWithBackoff() {
-            stubChildContext("retry");
-
-            var config = WithRetryConfig.builder()
-                    .retryStrategy((error, attempt) ->
-                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
-                    .build();
-
-            DurableFuture future = context.withRetryAsync(
-                    (ctx, attempt) -> {
-                        if (attempt < 3) {
-                            throw new RuntimeException("fail");
-                        }
-                        return "done";
-                    },
-                    config);
-
-            assertEquals("done", future.get());
-            verify(childContext).wait("retry-backoff-1", Duration.ofSeconds(2));
-            verify(childContext).wait("retry-backoff-2", Duration.ofSeconds(2));
-        }
+    // --- Core retry logic (uses named form; retry behavior is identical for all forms) ---
 
-        @Test
-        void namedAsyncNullName_shouldThrow() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
+    @Nested
+    class CoreRetryLogic {
 
-            assertThrows(NullPointerException.class, () -> context.withRetryAsync(null, (ctx, a) -> "x", config));
+        @BeforeEach
+        void setUpChildContext() {
+            stubChildContextAnyName();
         }
 
         @Test
-        void namedAsyncNullOperation_shouldThrow() {
+        void successOnFirstAttempt() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            assertThrows(NullPointerException.class, () -> context.withRetryAsync("name", null, config));
-        }
+            var result = context.withRetry("my-op", (ctx, attempt) -> "success", config);
 
-        @Test
-        void namedAsyncNullConfig_shouldThrow() {
-            assertThrows(NullPointerException.class, () -> context.withRetryAsync("name", (ctx, a) -> "x", null));
+            assertEquals("success", result);
         }
 
         @Test
-        void anonymousAsyncNullOperation_shouldThrow() {
+        void retriesWithBackoffWaits() {
+            var callCount = new int[] {0};
             var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .retryStrategy((error, attempt) ->
+                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
                     .build();
 
-            assertThrows(NullPointerException.class, () -> context.withRetryAsync((WithRetry) null, config));
-        }
+            var result = context.withRetry(
+                    "my-op",
+                    (ctx, attempt) -> {
+                        callCount[0]++;
+                        if (attempt < 3) {
+                            throw new RuntimeException("fail-" + attempt);
+                        }
+                        return "success-on-3";
+                    },
+                    config);
 
-        @Test
-        void anonymousAsyncNullConfig_shouldThrow() {
-            assertThrows(NullPointerException.class, () -> context.withRetryAsync((ctx, a) -> "x", null));
+            assertEquals("success-on-3", result);
+            assertEquals(3, callCount[0]);
+            verify(childContext).wait("my-op-backoff-1", Duration.ofSeconds(5));
+            verify(childContext).wait("my-op-backoff-2", Duration.ofSeconds(5));
+            verify(childContext, times(2)).wait(anyString(), any(Duration.class));
         }
 
         @Test
-        void asyncOperationReturnsNull() {
-            stubChildContext("my-op");
-
+        void rethrowsWhenRetryStrategyReturnsFail() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            DurableFuture future =
-                    context.withRetryAsync("my-op", (WithRetry) (ctx, attempt) -> null, config);
+            var exception = assertThrows(
+                    RuntimeException.class,
+                    () -> context.withRetry(
+                            "my-op",
+                            (ctx, attempt) -> {
+                                throw new RuntimeException("terminal");
+                            },
+                            config));
 
-            assertNull(future.get());
+            assertEquals("terminal", exception.getMessage());
+            verify(childContext, never()).wait(anyString(), any(Duration.class));
         }
 
         @Test
-        void syncAndAsyncProduceSameResult() {
-            stubChildContextAnyName();
-
+        void rethrowsLastExceptionWhenAllRetriesExhausted() {
             var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
                     .build();
 
-            var syncResult = context.withRetry("op", (ctx, attempt) -> "value", config);
-            var asyncResult = context.withRetryAsync("op", (ctx, attempt) -> "value", config)
-                    .get();
-
-            assertEquals(syncResult, asyncResult);
-        }
-    }
-
-    // --- Retry behavior tests ---
-
-    @Nested
-    class RetryBehavior {
+            var thrown = assertThrows(
+                    RuntimeException.class,
+                    () -> context.withRetry(
+                            "my-op",
+                            (ctx, attempt) -> {
+                                throw new RuntimeException("attempt-" + attempt);
+                            },
+                            config));
 
-        @BeforeEach
-        void setUpChildContext() {
-            stubChildContextAnyName();
+            assertEquals("attempt-3", thrown.getMessage());
         }
 
         @Test
-        void passesCorrectAttemptNumberToOperation() {
-            var attempts = new ArrayList();
+        void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
             var config = WithRetryConfig.builder()
-                    .retryStrategy((error, attempt) ->
-                            attempt < 4 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
+                    .retryStrategy(
+                            (error, attempt) -> attempt < 2 ? RetryDecision.retry(Duration.ZERO) : RetryDecision.fail())
                     .build();
 
-            context.withRetry(
+            var callCount = new int[] {0};
+            var result = context.withRetry(
+                    "my-op",
                     (ctx, attempt) -> {
-                        attempts.add(attempt);
-                        if (attempt < 4) {
-                            throw new RuntimeException("not yet");
+                        callCount[0]++;
+                        if (attempt == 1) {
+                            throw new RuntimeException("fail");
                         }
-                        return "done";
+                        return "ok";
                     },
                     config);
 
-            assertEquals(4, attempts.size());
-            assertEquals(1, attempts.get(0));
-            assertEquals(2, attempts.get(1));
-            assertEquals(3, attempts.get(2));
-            assertEquals(4, attempts.get(3));
+            assertEquals("ok", result);
+            verify(childContext).wait("my-op-backoff-1", Duration.ofSeconds(1));
         }
 
         @Test
-        void passesCorrectAttemptNumberToOperation_namedForm() {
+        void passesCorrectAttemptNumberToOperation() {
             var attempts = new ArrayList();
             var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) ->
@@ -682,6 +266,7 @@ void passesErrorToRetryStrategy() {
             assertThrows(
                     RuntimeException.class,
                     () -> context.withRetry(
+                            "my-op",
                             (ctx, attempt) -> {
                                 throw new RuntimeException("error-" + attempt);
                             },
@@ -700,6 +285,7 @@ void respectsCustomDelayFromRetryDecision() {
                     .build();
 
             context.withRetry(
+                    "my-op",
                     (ctx, attempt) -> {
                         if (attempt <= 2) {
                             throw new RuntimeException("fail");
@@ -708,17 +294,18 @@ void respectsCustomDelayFromRetryDecision() {
                     },
                     config);
 
-            verify(childContext).wait("retry-backoff-1", Duration.ofSeconds(10));
-            verify(childContext).wait("retry-backoff-2", Duration.ofSeconds(20));
+            verify(childContext).wait("my-op-backoff-1", Duration.ofSeconds(10));
+            verify(childContext).wait("my-op-backoff-2", Duration.ofSeconds(20));
         }
 
         @Test
-        void anonymousFormPassesChildContextToOperation() {
+        void passesChildContextToOperation() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
             context.withRetry(
+                    "my-op",
                     (ctx, attempt) -> {
                         assertSame(childContext, ctx);
                         return "verified";
@@ -727,36 +314,46 @@ void anonymousFormPassesChildContextToOperation() {
         }
 
         @Test
-        void namedFormPassesChildContextToOperation() {
+        void operationReturnsNull() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            context.withRetry(
-                    "wrapped",
-                    (ctx, attempt) -> {
-                        assertSame(childContext, ctx);
-                        return "verified";
-                    },
-                    config);
+            var result = context.withRetry("my-op", (WithRetry) (ctx, attempt) -> null, config);
+
+            assertNull(result);
         }
 
         @Test
-        void rethrowsLastExceptionWhenAllRetriesExhausted() {
+        void preservesCheckedExceptionSubclassType() {
             var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
+            var original = new SerDesException("deserialization failed", new RuntimeException("bad json"));
+
             var thrown = assertThrows(
-                    RuntimeException.class,
+                    SerDesException.class,
                     () -> context.withRetry(
+                            "my-op",
                             (ctx, attempt) -> {
-                                throw new RuntimeException("attempt-" + attempt);
+                                throw original;
                             },
                             config));
 
-            // The last attempt's exception is rethrown
-            assertEquals("attempt-3", thrown.getMessage());
+            assertSame(original, thrown);
+            assertEquals("deserialization failed", thrown.getMessage());
+        }
+    }
+
+    // --- Exception propagation (SuspendExecution, Unrecoverable) ---
+
+    @Nested
+    class ExceptionPropagation {
+
+        @BeforeEach
+        void setUpChildContext() {
+            stubChildContextAnyName();
         }
 
         @Test
@@ -768,12 +365,12 @@ void propagatesSuspendExecutionExceptionWithoutRetrying() {
             assertThrows(
                     SuspendExecutionException.class,
                     () -> context.withRetry(
+                            "my-op",
                             (ctx, attempt) -> {
                                 throw new SuspendExecutionException();
                             },
                             config));
 
-            // Should never reach the wait — SuspendExecutionException propagates immediately
             verify(childContext, never()).wait(anyString(), any(Duration.class));
         }
 
@@ -786,6 +383,7 @@ void propagatesUnrecoverableDurableExecutionExceptionWithoutRetrying() {
             assertThrows(
                     UnrecoverableDurableExecutionException.class,
                     () -> context.withRetry(
+                            "my-op",
                             (ctx, attempt) -> {
                                 throw new UnrecoverableDurableExecutionException(
                                         software.amazon.awssdk.services.lambda.model.ErrorObject.builder()
@@ -806,16 +404,15 @@ void propagatesSuspendExecutionExceptionOnLaterAttempt() {
             assertThrows(
                     SuspendExecutionException.class,
                     () -> context.withRetry(
+                            "my-op",
                             (ctx, attempt) -> {
                                 if (attempt == 1) {
                                     throw new RuntimeException("transient");
                                 }
-                                // Second attempt triggers suspend — must propagate, not retry
                                 throw new SuspendExecutionException();
                             },
                             config));
 
-            // First attempt retried (one backoff wait), second attempt suspended immediately
             verify(childContext, times(1)).wait(anyString(), any(Duration.class));
         }
 
@@ -828,6 +425,7 @@ void propagatesUnrecoverableDurableExecutionExceptionOnLaterAttempt() {
             assertThrows(
                     UnrecoverableDurableExecutionException.class,
                     () -> context.withRetry(
+                            "my-op",
                             (ctx, attempt) -> {
                                 if (attempt == 1) {
                                     throw new RuntimeException("transient");
@@ -841,128 +439,112 @@ void propagatesUnrecoverableDurableExecutionExceptionOnLaterAttempt() {
 
             verify(childContext, times(1)).wait(anyString(), any(Duration.class));
         }
-
-        @Test
-        void preservesCheckedExceptionSubclassType() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
-
-            var original = new SerDesException("deserialization failed", new RuntimeException("bad json"));
-
-            var thrown = assertThrows(
-                    SerDesException.class,
-                    () -> context.withRetry(
-                            (ctx, attempt) -> {
-                                throw original;
-                            },
-                            config));
-
-            assertSame(original, thrown);
-            assertEquals("deserialization failed", thrown.getMessage());
-        }
     }
 
-    // --- WrapInChildContext tests ---
+    // --- Naming: named form uses the provided name, null-name form defaults to "retry" ---
 
     @Nested
-    class WrapInChildContext {
-
-        @BeforeEach
-        void setUpChildContext() {
-            stubChildContextAnyName();
-        }
+    class Naming {
 
         @Test
-        void namedForm_alwaysUsesChildContextWhenEnabled() {
+        void namedFormUsesProvidedNameForChildContextAndBackoff() {
+            stubChildContext("my-op");
             var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(true)
+                    .retryStrategy((error, attempt) ->
+                            attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
                     .build();
 
-            var result = context.withRetry("my-op", (ctx, attempt) -> "wrapped", config);
+            context.withRetry(
+                    "my-op",
+                    (ctx, attempt) -> {
+                        if (attempt < 2) {
+                            throw new RuntimeException("fail");
+                        }
+                        return "ok";
+                    },
+                    config);
 
-            assertEquals("wrapped", result);
             verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
+            verify(childContext).wait("my-op-backoff-1", Duration.ofSeconds(5));
         }
 
         @Test
-        void namedForm_usesChildContextWhenDisabled() {
+        void nullNameFormDefaultsToRetryForChildContextAndBackoff() {
+            stubChildContext("retry");
             var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(false)
+                    .retryStrategy((error, attempt) ->
+                            attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
                     .build();
 
-            var result = context.withRetry("my-op", (ctx, attempt) -> "virtual", config);
+            context.withRetry(
+                    null,
+                    (ctx, attempt) -> {
+                        if (attempt < 2) {
+                            throw new RuntimeException("fail");
+                        }
+                        return "ok";
+                    },
+                    config);
 
-            assertEquals("virtual", result);
-            // Still uses a child context (virtual in real impl)
-            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
+            verify(context).runInChildContextAsync(eq("retry"), any(TypeToken.class), any());
+            verify(childContext).wait("retry-backoff-1", Duration.ofSeconds(2));
         }
+    }
 
-        @Test
-        void namedForm_usesChildContextByDefault() {
-            var config = WithRetryConfig.builder()
-                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .build();
+    // --- Sync vs async: sync returns value, async returns DurableFuture ---
 
-            context.withRetry("my-op", (ctx, attempt) -> "virtual", config);
+    @Nested
+    class SyncVsAsync {
 
-            verify(context).runInChildContextAsync(eq("my-op"), any(TypeToken.class), any());
+        @BeforeEach
+        void setUpChildContext() {
+            stubChildContextAnyName();
         }
 
         @Test
-        void anonymousForm_alwaysUsesChildContextWhenEnabled() {
+        void syncReturnsValueDirectly() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(true)
                     .build();
 
-            var result = context.withRetry((ctx, attempt) -> "wrapped", config);
+            var result = context.withRetry("my-op", (ctx, attempt) -> "sync-value", config);
 
-            assertEquals("wrapped", result);
-            verify(context).runInChildContextAsync(eq("retry"), any(TypeToken.class), any());
+            assertEquals("sync-value", result);
         }
 
         @Test
-        void anonymousForm_usesChildContextWhenDisabled() {
+        void asyncReturnsDurableFuture() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(false)
                     .build();
 
-            var result = context.withRetry((ctx, attempt) -> "virtual", config);
+            DurableFuture future = context.withRetryAsync("my-op", (ctx, attempt) -> "async-value", config);
 
-            assertEquals("virtual", result);
-            // Still uses a child context (virtual in real impl)
-            verify(context).runInChildContextAsync(eq("retry"), any(TypeToken.class), any());
+            assertNotNull(future);
+            assertEquals("async-value", future.get());
         }
 
         @Test
-        void namedForm_passesChildContextToOperationWhenWrapped() {
+        void syncAndAsyncProduceSameResult() {
             var config = WithRetryConfig.builder()
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
-                    .wrapInChildContext(true)
                     .build();
 
-            context.withRetry(
-                    "my-op",
-                    (ctx, attempt) -> {
-                        assertSame(childContext, ctx);
-                        return "verified";
-                    },
-                    config);
+            var syncResult = context.withRetry("op", (ctx, attempt) -> "value", config);
+            var asyncResult = context.withRetryAsync("op", (ctx, attempt) -> "value", config)
+                    .get();
+
+            assertEquals(syncResult, asyncResult);
         }
 
         @Test
-        void namedForm_retriesWithBackoffOnChildContextWhenWrapped() {
+        void asyncRetriesWithBackoff() {
             var config = WithRetryConfig.builder()
                     .retryStrategy((error, attempt) ->
                             attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(5)) : RetryDecision.fail())
-                    .wrapInChildContext(true)
                     .build();
 
-            var result = context.withRetry(
+            DurableFuture future = context.withRetryAsync(
                     "my-op",
                     (ctx, attempt) -> {
                         if (attempt < 3) {
@@ -972,31 +554,43 @@ void namedForm_retriesWithBackoffOnChildContextWhenWrapped() {
                     },
                     config);
 
-            assertEquals("success-on-3", result);
+            assertEquals("success-on-3", future.get());
             verify(childContext).wait("my-op-backoff-1", Duration.ofSeconds(5));
             verify(childContext).wait("my-op-backoff-2", Duration.ofSeconds(5));
         }
+    }
+
+    // --- Null guards ---
+
+    @Nested
+    class NullGuards {
 
         @Test
-        void anonymousForm_retriesWithBackoffOnChildContextWhenWrapped() {
+        void syncNullOperationThrows() {
             var config = WithRetryConfig.builder()
-                    .retryStrategy((error, attempt) ->
-                            attempt < 3 ? RetryDecision.retry(Duration.ofSeconds(2)) : RetryDecision.fail())
-                    .wrapInChildContext(true)
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var result = context.withRetry(
-                    (ctx, attempt) -> {
-                        if (attempt < 3) {
-                            throw new RuntimeException("fail");
-                        }
-                        return "done";
-                    },
-                    config);
+            assertThrows(NullPointerException.class, () -> context.withRetry("name", null, config));
+        }
 
-            assertEquals("done", result);
-            verify(childContext).wait("retry-backoff-1", Duration.ofSeconds(2));
-            verify(childContext).wait("retry-backoff-2", Duration.ofSeconds(2));
+        @Test
+        void syncNullConfigThrows() {
+            assertThrows(NullPointerException.class, () -> context.withRetry("name", (ctx, a) -> "x", null));
+        }
+
+        @Test
+        void asyncNullOperationThrows() {
+            var config = WithRetryConfig.builder()
+                    .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                    .build();
+
+            assertThrows(NullPointerException.class, () -> context.withRetryAsync("name", null, config));
+        }
+
+        @Test
+        void asyncNullConfigThrows() {
+            assertThrows(NullPointerException.class, () -> context.withRetryAsync("name", (ctx, a) -> "x", null));
         }
     }
 }

From b9438b787121e2c498a87a8e439b06b7e79afbbb Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Thu, 30 Apr 2026 16:26:20 -0700
Subject: [PATCH 17/19] chore: add cloud based integration tests for withRetry,
 add missing resources to template.yml

---
 .../examples/CloudBasedIntegrationTest.java   |  83 +++++++
 examples/template.yaml                        | 224 ++++++++++++++++++
 2 files changed, 307 insertions(+)

diff --git a/examples/src/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java b/examples/src/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java
index a83109e89..d9df8bc67 100644
--- a/examples/src/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java
+++ b/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(
diff --git a/examples/template.yaml b/examples/template.yaml
index c5a5049b5..1a256c5c3 100644
--- a/examples/template.yaml
+++ b/examples/template.yaml
@@ -432,6 +432,186 @@ Resources:
                 - lambda:GetDurableExecutionState
               Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:concurrent-wait-for-condition-example-${JavaVersion}-runtime"
 
+  RetryWaitForCallbackExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      FunctionName: !Join
+        - '-'
+        - - 'retry-wait-for-callback-example'
+          - !Ref JavaVersion
+          - runtime
+      Handler: "software.amazon.lambda.durable.examples.callback.RetryWaitForCallbackExample"
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:retry-wait-for-callback-example-${JavaVersion}-runtime"
+
+  WaitForCallbackFailedExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      FunctionName: !Join
+        - '-'
+        - - 'wait-for-callback-failed-example'
+          - !Ref JavaVersion
+          - runtime
+      Handler: "software.amazon.lambda.durable.examples.callback.WaitForCallbackFailedExample"
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:wait-for-callback-failed-example-${JavaVersion}-runtime"
+
+  CustomPollingExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      FunctionName: !Join
+        - '-'
+        - - 'custom-polling-example'
+          - !Ref JavaVersion
+          - runtime
+      Handler: "software.amazon.lambda.durable.examples.general.CustomPollingExample"
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+                - lambda:InvokeFunction
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:custom-polling-example-${JavaVersion}-runtime"
+            - Effect: Allow
+              Action:
+                - lambda:InvokeFunction
+              Resource: '*'
+
+  RetryInvokeExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      FunctionName: !Join
+        - '-'
+        - - 'retry-invoke-example'
+          - !Ref JavaVersion
+          - runtime
+      Handler: "software.amazon.lambda.durable.examples.invoke.RetryInvokeExample"
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+                - lambda:InvokeFunction
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:retry-invoke-example-${JavaVersion}-runtime"
+            - Effect: Allow
+              Action:
+                - lambda:InvokeFunction
+              Resource: '*'
+
+  DeserializationFailedMapExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      FunctionName: !Join
+        - '-'
+        - - 'deserialization-failed-map-example'
+          - !Ref JavaVersion
+          - runtime
+      Handler: "software.amazon.lambda.durable.examples.map.DeserializationFailedMapExample"
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:deserialization-failed-map-example-${JavaVersion}-runtime"
+
+  ParallelExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      FunctionName: !Join
+        - '-'
+        - - 'parallel-example'
+          - !Ref JavaVersion
+          - runtime
+      Handler: "software.amazon.lambda.durable.examples.parallel.ParallelExample"
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:parallel-example-${JavaVersion}-runtime"
+
+  ParallelFailureToleranceExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      FunctionName: !Join
+        - '-'
+        - - 'parallel-failure-tolerance-example'
+          - !Ref JavaVersion
+          - runtime
+      Handler: "software.amazon.lambda.durable.examples.parallel.ParallelFailureToleranceExample"
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:parallel-failure-tolerance-example-${JavaVersion}-runtime"
+
+  ParallelWithWaitExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      FunctionName: !Join
+        - '-'
+        - - 'parallel-with-wait-example'
+          - !Ref JavaVersion
+          - runtime
+      Handler: "software.amazon.lambda.durable.examples.parallel.ParallelWithWaitExample"
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:parallel-with-wait-example-${JavaVersion}-runtime"
+
+  DeserializationFailedParallelExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      FunctionName: !Join
+        - '-'
+        - - 'deserialization-failed-parallel-example'
+          - !Ref JavaVersion
+          - runtime
+      Handler: "software.amazon.lambda.durable.examples.parallel.DeserializationFailedParallelExample"
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:deserialization-failed-parallel-example-${JavaVersion}-runtime"
+
+  DeserializationFailureExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      FunctionName: !Join
+        - '-'
+        - - 'deserialization-failure-example'
+          - !Ref JavaVersion
+          - runtime
+      Handler: "software.amazon.lambda.durable.examples.step.DeserializationFailureExample"
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:deserialization-failure-example-${JavaVersion}-runtime"
+
   ManyAsyncStepsVirtualThreadPoolExampleFunction:
     Type: AWS::Serverless::Function
     Condition: IsJava21OrLater
@@ -515,6 +695,10 @@ Outputs:
     Description: Child Context Example Function ARN
     Value: !GetAtt ChildContextExampleFunction.Arn
 
+  VirtualChildContextExampleFunction:
+    Description: Virtual Child Context Example Function ARN
+    Value: !GetAtt VirtualChildContextExampleFunction.Arn
+
   WaitAsyncExampleFunction:
     Description: Wait Async Example Function ARN
     Value: !GetAtt WaitAsyncExampleFunction.Arn
@@ -543,6 +727,46 @@ Outputs:
     Description: Concurrent Wait For Condition Example Function ARN
     Value: !GetAtt ConcurrentWaitForConditionExampleFunction.Arn
 
+  RetryWaitForCallbackExampleFunction:
+    Description: Retry Wait For Callback Example Function ARN
+    Value: !GetAtt RetryWaitForCallbackExampleFunction.Arn
+
+  WaitForCallbackFailedExampleFunction:
+    Description: Wait For Callback Failed Example Function ARN
+    Value: !GetAtt WaitForCallbackFailedExampleFunction.Arn
+
+  CustomPollingExampleFunction:
+    Description: Custom Polling Example Function ARN
+    Value: !GetAtt CustomPollingExampleFunction.Arn
+
+  RetryInvokeExampleFunction:
+    Description: Retry Invoke Example Function ARN
+    Value: !GetAtt RetryInvokeExampleFunction.Arn
+
+  DeserializationFailedMapExampleFunction:
+    Description: Deserialization Failed Map Example Function ARN
+    Value: !GetAtt DeserializationFailedMapExampleFunction.Arn
+
+  ParallelExampleFunction:
+    Description: Parallel Example Function ARN
+    Value: !GetAtt ParallelExampleFunction.Arn
+
+  ParallelFailureToleranceExampleFunction:
+    Description: Parallel Failure Tolerance Example Function ARN
+    Value: !GetAtt ParallelFailureToleranceExampleFunction.Arn
+
+  ParallelWithWaitExampleFunction:
+    Description: Parallel With Wait Example Function ARN
+    Value: !GetAtt ParallelWithWaitExampleFunction.Arn
+
+  DeserializationFailedParallelExampleFunction:
+    Description: Deserialization Failed Parallel Example Function ARN
+    Value: !GetAtt DeserializationFailedParallelExampleFunction.Arn
+
+  DeserializationFailureExampleFunction:
+    Description: Deserialization Failure Example Function ARN
+    Value: !GetAtt DeserializationFailureExampleFunction.Arn
+
   ManyAsyncStepsVirtualThreadPoolExampleFunction:
     Condition: IsJava21OrLater
     Description: Many Async Steps Virtual Thread Pool Example Function ARN

From fbcd329ed7d9e19331763c0277b1882c7442290b Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Thu, 30 Apr 2026 16:37:44 -0700
Subject: [PATCH 18/19] chore: make config optional by adding override with
 config parameter for withRetry and withRetryAsync

---
 .../amazon/lambda/durable/DurableContext.java | 36 +++++++++
 .../context/DurableContextWithRetryTest.java  | 79 ++++++++++++++++++-
 2 files changed, 113 insertions(+), 2 deletions(-)

diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
index ca17109ce..7528ffba2 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
@@ -731,6 +731,24 @@  DurableFuture waitForConditionAsync(
 
     // =============== withRetry ================
 
+    /**
+     * Replay-safe retry loop for any durable operation (sync) with default configuration.
+     *
+     * 
Uses {@link WithRetryConfig} defaults:
+     * {@link software.amazon.lambda.durable.retry.RetryStrategies.Presets#DEFAULT} retry strategy and no child context
+     * wrapping.
+     *
+     * @param  the result type
+     * @param name operation name (used for backoff wait names, and as the child context name when wrapping); pass
+     *     {@code null} for an anonymous retry whose backoff waits use default names
+     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @return the operation result
+     * @see #withRetry(String, WithRetry, WithRetryConfig)
+     */
+    default  T withRetry(String name, WithRetry operation) {
+        return withRetry(name, operation, WithRetryConfig.builder().build());
+    }
+
     /**
      * Replay-safe retry loop for any durable operation (sync).
      *
@@ -756,6 +774,24 @@ default  T withRetry(String name, WithRetry operation, WithRetryConfig con
         return withRetryAsync(name, operation, config).get();
     }
 
+    /**
+     * Replay-safe retry loop for any durable operation (async) with default configuration.
+     *
+     * 
Uses {@link WithRetryConfig} defaults:
+     * {@link software.amazon.lambda.durable.retry.RetryStrategies.Presets#DEFAULT} retry strategy and no child context
+     * wrapping.
+     *
+     * @param  the result type
+     * @param name operation name (used for child context and backoff wait names); pass {@code null} for an anonymous
+     *     retry whose backoff waits use default names
+     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @return a future representing the operation result
+     * @see #withRetryAsync(String, WithRetry, WithRetryConfig)
+     */
+    default  DurableFuture withRetryAsync(String name, WithRetry operation) {
+        return withRetryAsync(name, operation, WithRetryConfig.builder().build());
+    }
+
     /**
      * Replay-safe retry loop for any durable operation (async).
      *
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
index 5dfb24eee..15465631b 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
@@ -46,7 +46,7 @@ void setUp() {
      * when {@code wrapInChildContext} is false, and a checkpointed child context when true).
      */
     private void stubWithRetryMethods(DurableContext mock) {
-        // Sync form — always runs in a child context
+        // Sync form with config — always runs in a child context
         when(mock.withRetry(any(), nullable(WithRetry.class), nullable(WithRetryConfig.class)))
                 .thenAnswer(invocation -> {
                     String name = invocation.getArgument(0);
@@ -62,7 +62,14 @@ private void stubWithRetryMethods(DurableContext mock) {
                             .get();
                 });
 
-        // Async form
+        // Sync form without config — delegates to the 3-arg form with default config
+        when(mock.withRetry(any(), nullable(WithRetry.class))).thenAnswer(invocation -> {
+            String name = invocation.getArgument(0);
+            WithRetry operation = invocation.getArgument(1);
+            return mock.withRetry(name, operation, WithRetryConfig.builder().build());
+        });
+
+        // Async form with config
         when(mock.withRetryAsync(any(), nullable(WithRetry.class), nullable(WithRetryConfig.class)))
                 .thenAnswer(invocation -> {
                     String name = invocation.getArgument(0);
@@ -76,6 +83,14 @@ private void stubWithRetryMethods(DurableContext mock) {
                             new TypeToken() {},
                             childCtx -> executeRetryLoop(childCtx, name, operation, config));
                 });
+
+        // Async form without config — delegates to the 3-arg form with default config
+        when(mock.withRetryAsync(any(), nullable(WithRetry.class))).thenAnswer(invocation -> {
+            String name = invocation.getArgument(0);
+            WithRetry operation = invocation.getArgument(1);
+            return mock.withRetryAsync(
+                    name, operation, WithRetryConfig.builder().build());
+        });
     }
 
     /** Replicates the retry loop logic from DurableContextImpl for test stubbing. */
@@ -593,4 +608,64 @@ void asyncNullConfigThrows() {
             assertThrows(NullPointerException.class, () -> context.withRetryAsync("name", (ctx, a) -> "x", null));
         }
     }
+
+    // --- Default config overloads (no WithRetryConfig parameter) ---
+
+    @Nested
+    class DefaultConfigOverloads {
+
+        @BeforeEach
+        void setUpChildContext() {
+            stubChildContextAnyName();
+        }
+
+        @Test
+        void syncWithRetryWithoutConfigSucceedsOnFirstAttempt() {
+            var result = context.withRetry("my-op", (ctx, attempt) -> "default-config-result");
+
+            assertEquals("default-config-result", result);
+        }
+
+        @Test
+        void syncWithRetryWithoutConfigRetriesOnFailure() {
+            var callCount = new int[] {0};
+
+            var result = context.withRetry("my-op", (ctx, attempt) -> {
+                callCount[0]++;
+                if (attempt == 1) {
+                    throw new RuntimeException("transient");
+                }
+                return "recovered";
+            });
+
+            assertEquals("recovered", result);
+            assertEquals(2, callCount[0]);
+            verify(childContext).wait(eq("my-op-backoff-1"), any(Duration.class));
+        }
+
+        @Test
+        void asyncWithRetryWithoutConfigSucceedsOnFirstAttempt() {
+            DurableFuture future = context.withRetryAsync("my-op", (ctx, attempt) -> "async-default");
+
+            assertNotNull(future);
+            assertEquals("async-default", future.get());
+        }
+
+        @Test
+        void asyncWithRetryWithoutConfigRetriesOnFailure() {
+            var callCount = new int[] {0};
+
+            DurableFuture future = context.withRetryAsync("my-op", (ctx, attempt) -> {
+                callCount[0]++;
+                if (attempt == 1) {
+                    throw new RuntimeException("transient");
+                }
+                return "async-recovered";
+            });
+
+            assertEquals("async-recovered", future.get());
+            assertEquals(2, callCount[0]);
+            verify(childContext).wait(eq("my-op-backoff-1"), any(Duration.class));
+        }
+    }
 }

From 3b98eec8b6d486aad28ae0ae68f89d2fb26b4cb3 Mon Sep 17 00:00:00 2001
From: hsilan 
Date: Fri, 1 May 2026 11:21:46 -0700
Subject: [PATCH 19/19] chore: use BiFunction
 instead of WithRetry

---
 .../callback/RetryWaitForCallbackExample.java |  2 +-
 .../examples/invoke/RetryInvokeExample.java   |  2 +-
 .../durable/RetryInvokeIntegrationTest.java   | 12 +--
 .../RetryWaitForCallbackIntegrationTest.java  | 12 +--
 .../amazon/lambda/durable/DurableContext.java | 22 ++---
 .../durable/context/DurableContextImpl.java   | 11 ++-
 .../lambda/durable/model/WithRetry.java       | 27 ------
 .../context/DurableContextWithRetryTest.java  | 82 ++++++++++---------
 .../lambda/durable/model/WithRetryTest.java   | 52 ------------
 9 files changed, 75 insertions(+), 147 deletions(-)
 delete mode 100644 sdk/src/main/java/software/amazon/lambda/durable/model/WithRetry.java
 delete mode 100644 sdk/src/test/java/software/amazon/lambda/durable/model/WithRetryTest.java

diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
index c08d44937..1698e9a0c 100644
--- a/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
+++ b/examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java
@@ -35,7 +35,7 @@ public String handleRequest(ApprovalRequest input, DurableContext context) {
         // Step 2: waitForCallback with retry — if the external system fails, try again with a fresh callback
         var approvalResult = context.withRetry(
                 null,
-                (ctx, attempt) -> ctx.waitForCallback(
+                (attempt, ctx) -> ctx.waitForCallback(
                         "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
                                 .info("Attempt {}: sending callback {} to approval system", attempt, callbackId)),
                 WithRetryConfig.builder()
diff --git a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
index 8f577036f..28cec8dc0 100644
--- a/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
+++ b/examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java
@@ -26,7 +26,7 @@ public class RetryInvokeExample extends DurableHandler
     public String handleRequest(GreetingRequest input, DurableContext context) {
         return context.withRetry(
                 null,
-                (ctx, attempt) -> ctx.invoke(
+                (attempt, ctx) -> ctx.invoke(
                         "call-greeting-" + attempt,
                         "simple-step-example" + input.getName() + ":$LATEST",
                         input,
diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
index 5aa873692..7acacc988 100644
--- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
+++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java
@@ -22,7 +22,7 @@ void invokeSucceedsOnFirstAttempt() {
                 String.class,
                 (input, context) -> context.withRetry(
                         null,
-                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        (attempt, ctx) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
                                 .build()));
@@ -43,7 +43,7 @@ void invokeRetriesAfterFailure() {
                 String.class,
                 (input, context) -> context.withRetry(
                         null,
-                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        (attempt, ctx) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
                                 .build()));
@@ -81,7 +81,7 @@ void invokeFailsAfterAllRetriesExhausted() {
                 String.class,
                 (input, context) -> context.withRetry(
                         null,
-                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        (attempt, ctx) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy((error, attempt) ->
                                         attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
@@ -115,7 +115,7 @@ void invokeRetryWithCustomBackoffDelay() {
                 String.class,
                 (input, context) -> context.withRetry(
                         null,
-                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        (attempt, ctx) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy((error, attempt) -> attempt < 3
                                         ? RetryDecision.retry(Duration.ofSeconds(attempt * 5L))
@@ -149,7 +149,7 @@ void invokeRetryWithStepsBeforeAndAfter() {
 
             var invokeResult = context.withRetry(
                     null,
-                    (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                    (attempt, ctx) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                     WithRetryConfig.builder()
                             .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
                             .build());
@@ -175,7 +175,7 @@ void invokeRetryPreservesOriginalExceptionType() {
             try {
                 return context.withRetry(
                         null,
-                        (ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
+                        (attempt, ctx) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
                         WithRetryConfig.builder()
                                 .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                                 .build());
diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
index d07d944c4..d6ad39c0c 100644
--- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
+++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java
@@ -21,7 +21,7 @@ void waitForCallbackSucceedsOnFirstAttempt() {
                 String.class,
                 (input, context) -> context.withRetry(
                         null,
-                        (ctx, attempt) -> ctx.waitForCallback(
+                        (attempt, ctx) -> ctx.waitForCallback(
                                 "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
                                         .info("Submitting callback {}", callbackId)),
                         WithRetryConfig.builder()
@@ -48,7 +48,7 @@ void waitForCallbackRetriesAfterFailure() {
                 String.class,
                 (input, context) -> context.withRetry(
                         null,
-                        (ctx, attempt) -> ctx.waitForCallback(
+                        (attempt, ctx) -> ctx.waitForCallback(
                                 "approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
                                         .info("Attempt {} callback {}", attempt, callbackId)),
                         WithRetryConfig.builder()
@@ -95,7 +95,7 @@ void waitForCallbackFailsAfterAllRetriesExhausted() {
                 String.class,
                 (input, context) -> context.withRetry(
                         null,
-                        (ctx, attempt) ->
+                        (attempt, ctx) ->
                                 ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
                         WithRetryConfig.builder()
                                 .retryStrategy((error, attempt) ->
@@ -133,7 +133,7 @@ void waitForCallbackRetryWithStepsBeforeAndAfter() {
 
             var callbackResult = context.withRetry(
                     null,
-                    (ctx, attempt) ->
+                    (attempt, ctx) ->
                             ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
                     WithRetryConfig.builder()
                             .retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
@@ -162,7 +162,7 @@ void waitForCallbackRetryMultipleFailuresThenSuccess() {
                 String.class,
                 (input, context) -> context.withRetry(
                         null,
-                        (ctx, attempt) ->
+                        (attempt, ctx) ->
                                 ctx.waitForCallback("cb-" + attempt, String.class, (callbackId, stepCtx) -> {}),
                         WithRetryConfig.builder()
                                 .retryStrategy(RetryStrategies.fixedDelay(4, Duration.ofSeconds(1)))
@@ -208,7 +208,7 @@ void waitForCallbackRetryWithSubmitterLogic() {
                 String.class,
                 (input, context) -> context.withRetry(
                         null,
-                        (ctx, attempt) ->
+                        (attempt, ctx) ->
                                 ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {
                                     // Submitter runs each attempt — in a real scenario this would
                                     // send the callbackId to an external system
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
index 7528ffba2..38065b749 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java
@@ -20,7 +20,6 @@
 import software.amazon.lambda.durable.context.BaseContext;
 import software.amazon.lambda.durable.model.MapResult;
 import software.amazon.lambda.durable.model.WaitForConditionResult;
-import software.amazon.lambda.durable.model.WithRetry;
 
 public interface DurableContext extends BaseContext {
     /**
@@ -741,11 +740,11 @@  DurableFuture waitForConditionAsync(
      * @param  the result type
      * @param name operation name (used for backoff wait names, and as the child context name when wrapping); pass
      *     {@code null} for an anonymous retry whose backoff waits use default names
-     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param operation the retryable operation — receives the 1-based attempt number and the durable context
      * @return the operation result
-     * @see #withRetry(String, WithRetry, WithRetryConfig)
+     * @see #withRetry(String, BiFunction, WithRetryConfig)
      */
-    default  T withRetry(String name, WithRetry operation) {
+    default  T withRetry(String name, BiFunction operation) {
         return withRetry(name, operation, WithRetryConfig.builder().build());
     }
 
@@ -766,11 +765,11 @@ default  T withRetry(String name, WithRetry operation) {
      * @param  the result type
      * @param name operation name (used for backoff wait names, and as the child context name when wrapping); pass
      *     {@code null} for an anonymous retry whose backoff waits use default names
-     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param operation the retryable operation — receives the 1-based attempt number and the durable context
      * @param config retry configuration including the retry strategy and child context wrapping
      * @return the operation result
      */
-    default  T withRetry(String name, WithRetry operation, WithRetryConfig config) {
+    default  T withRetry(String name, BiFunction operation, WithRetryConfig config) {
         return withRetryAsync(name, operation, config).get();
     }
 
@@ -784,11 +783,11 @@ default  T withRetry(String name, WithRetry operation, WithRetryConfig con
      * @param  the result type
      * @param name operation name (used for child context and backoff wait names); pass {@code null} for an anonymous
      *     retry whose backoff waits use default names
-     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param operation the retryable operation — receives the 1-based attempt number and the durable context
      * @return a future representing the operation result
-     * @see #withRetryAsync(String, WithRetry, WithRetryConfig)
+     * @see #withRetryAsync(String, BiFunction, WithRetryConfig)
      */
-    default  DurableFuture withRetryAsync(String name, WithRetry operation) {
+    default  DurableFuture withRetryAsync(String name, BiFunction operation) {
         return withRetryAsync(name, operation, WithRetryConfig.builder().build());
     }
 
@@ -803,11 +802,12 @@ default  DurableFuture withRetryAsync(String name, WithRetry operation)
      * @param  the result type
      * @param name operation name (used for child context and backoff wait names); pass {@code null} for an anonymous
      *     retry whose backoff waits use default names
-     * @param operation the retryable operation — receives the context and 1-based attempt number
+     * @param operation the retryable operation — receives the 1-based attempt number and the durable context
      * @param config retry configuration including the retry strategy
      * @return a future representing the operation result
      */
-     DurableFuture withRetryAsync(String name, WithRetry operation, WithRetryConfig config);
+     DurableFuture withRetryAsync(
+            String name, BiFunction operation, WithRetryConfig config);
 
     /**
      * Function applied to each item in a map operation.
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
index 56e2ab36b..62414aba9 100644
--- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
+++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java
@@ -38,7 +38,6 @@
 import software.amazon.lambda.durable.model.OperationIdentifier;
 import software.amazon.lambda.durable.model.OperationSubType;
 import software.amazon.lambda.durable.model.WaitForConditionResult;
-import software.amazon.lambda.durable.model.WithRetry;
 import software.amazon.lambda.durable.operation.CallbackOperation;
 import software.amazon.lambda.durable.operation.ChildContextOperation;
 import software.amazon.lambda.durable.operation.InvokeOperation;
@@ -380,7 +379,8 @@ public  DurableFuture waitForConditionAsync(
 
     @Override
     @SuppressWarnings("unchecked")
-    public  DurableFuture withRetryAsync(String name, WithRetry operation, WithRetryConfig config) {
+    public  DurableFuture withRetryAsync(
+            String name, BiFunction operation, WithRetryConfig config) {
         Objects.requireNonNull(operation, "operation cannot be null");
         Objects.requireNonNull(config, "config cannot be null");
 
@@ -404,11 +404,14 @@ public  DurableFuture withRetryAsync(String name, WithRetry operation,
      * are internal SDK control flow signals that must propagate immediately.
      */
     private static  T executeRetryLoop(
-            DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
+            DurableContext context,
+            String name,
+            BiFunction operation,
+            WithRetryConfig config) {
         var attempt = 1;
         while (true) {
             try {
-                return operation.execute(context, attempt);
+                return operation.apply(attempt, context);
             } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) {
                 // Internal SDK control flow — never retry, always propagate
                 throw e;
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/model/WithRetry.java b/sdk/src/main/java/software/amazon/lambda/durable/model/WithRetry.java
deleted file mode 100644
index 3a5675eaa..000000000
--- a/sdk/src/main/java/software/amazon/lambda/durable/model/WithRetry.java
+++ /dev/null
@@ -1,27 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.model;
-
-import software.amazon.lambda.durable.DurableContext;
-
-/**
- * A durable operation that can be retried end-to-end by
- * {@link software.amazon.lambda.durable.DurableContext#withRetry}.
- *
- * 
Receives the durable context and the 1-based attempt number so callers can generate unique operation names per
- * attempt (e.g., {@code "approval-" + attempt}).
- *
- * @param  the result type
- */
-@FunctionalInterface
-public interface WithRetry {
-
-    /**
-     * Executes the durable operation.
-     *
-     * @param context the durable context to use for durable operations
-     * @param attempt the current attempt number (1-based: first attempt is 1)
-     * @return the operation result
-     */
-    T execute(DurableContext context, int attempt);
-}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java b/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
index 15465631b..5f0c2c9aa 100644
--- a/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
+++ b/sdk/src/test/java/software/amazon/lambda/durable/context/DurableContextWithRetryTest.java
@@ -9,6 +9,7 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Objects;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Nested;
@@ -20,7 +21,6 @@
 import software.amazon.lambda.durable.exception.SerDesException;
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.execution.SuspendExecutionException;
-import software.amazon.lambda.durable.model.WithRetry;
 import software.amazon.lambda.durable.retry.RetryDecision;
 import software.amazon.lambda.durable.retry.RetryStrategies;
 
@@ -47,10 +47,10 @@ void setUp() {
      */
     private void stubWithRetryMethods(DurableContext mock) {
         // Sync form with config — always runs in a child context
-        when(mock.withRetry(any(), nullable(WithRetry.class), nullable(WithRetryConfig.class)))
+        when(mock.withRetry(any(), nullable(BiFunction.class), nullable(WithRetryConfig.class)))
                 .thenAnswer(invocation -> {
                     String name = invocation.getArgument(0);
-                    WithRetry operation = invocation.getArgument(1);
+                    BiFunction operation = invocation.getArgument(1);
                     WithRetryConfig config = invocation.getArgument(2);
                     Objects.requireNonNull(operation, "operation cannot be null");
                     Objects.requireNonNull(config, "config cannot be null");
@@ -63,17 +63,17 @@ private void stubWithRetryMethods(DurableContext mock) {
                 });
 
         // Sync form without config — delegates to the 3-arg form with default config
-        when(mock.withRetry(any(), nullable(WithRetry.class))).thenAnswer(invocation -> {
+        when(mock.withRetry(any(), nullable(BiFunction.class))).thenAnswer(invocation -> {
             String name = invocation.getArgument(0);
-            WithRetry operation = invocation.getArgument(1);
+            BiFunction operation = invocation.getArgument(1);
             return mock.withRetry(name, operation, WithRetryConfig.builder().build());
         });
 
         // Async form with config
-        when(mock.withRetryAsync(any(), nullable(WithRetry.class), nullable(WithRetryConfig.class)))
+        when(mock.withRetryAsync(any(), nullable(BiFunction.class), nullable(WithRetryConfig.class)))
                 .thenAnswer(invocation -> {
                     String name = invocation.getArgument(0);
-                    WithRetry operation = invocation.getArgument(1);
+                    BiFunction operation = invocation.getArgument(1);
                     WithRetryConfig config = invocation.getArgument(2);
                     Objects.requireNonNull(operation, "operation cannot be null");
                     Objects.requireNonNull(config, "config cannot be null");
@@ -85,9 +85,9 @@ private void stubWithRetryMethods(DurableContext mock) {
                 });
 
         // Async form without config — delegates to the 3-arg form with default config
-        when(mock.withRetryAsync(any(), nullable(WithRetry.class))).thenAnswer(invocation -> {
+        when(mock.withRetryAsync(any(), nullable(BiFunction.class))).thenAnswer(invocation -> {
             String name = invocation.getArgument(0);
-            WithRetry operation = invocation.getArgument(1);
+            BiFunction operation = invocation.getArgument(1);
             return mock.withRetryAsync(
                     name, operation, WithRetryConfig.builder().build());
         });
@@ -95,11 +95,14 @@ private void stubWithRetryMethods(DurableContext mock) {
 
     /** Replicates the retry loop logic from DurableContextImpl for test stubbing. */
     private static  T executeRetryLoop(
-            DurableContext context, String name, WithRetry operation, WithRetryConfig config) {
+            DurableContext context,
+            String name,
+            BiFunction operation,
+            WithRetryConfig config) {
         var attempt = 1;
         while (true) {
             try {
-                return operation.execute(context, attempt);
+                return operation.apply(attempt, context);
             } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) {
                 throw e;
             } catch (Exception e) {
@@ -151,7 +154,7 @@ void successOnFirstAttempt() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var result = context.withRetry("my-op", (ctx, attempt) -> "success", config);
+            var result = context.withRetry("my-op", (attempt, ctx) -> "success", config);
 
             assertEquals("success", result);
         }
@@ -166,7 +169,7 @@ void retriesWithBackoffWaits() {
 
             var result = context.withRetry(
                     "my-op",
-                    (ctx, attempt) -> {
+                    (attempt, ctx) -> {
                         callCount[0]++;
                         if (attempt < 3) {
                             throw new RuntimeException("fail-" + attempt);
@@ -192,7 +195,7 @@ void rethrowsWhenRetryStrategyReturnsFail() {
                     RuntimeException.class,
                     () -> context.withRetry(
                             "my-op",
-                            (ctx, attempt) -> {
+                            (attempt, ctx) -> {
                                 throw new RuntimeException("terminal");
                             },
                             config));
@@ -211,7 +214,7 @@ void rethrowsLastExceptionWhenAllRetriesExhausted() {
                     RuntimeException.class,
                     () -> context.withRetry(
                             "my-op",
-                            (ctx, attempt) -> {
+                            (attempt, ctx) -> {
                                 throw new RuntimeException("attempt-" + attempt);
                             },
                             config));
@@ -229,7 +232,7 @@ void usesDefaultDelayWhenRetryDecisionDelayIsZero() {
             var callCount = new int[] {0};
             var result = context.withRetry(
                     "my-op",
-                    (ctx, attempt) -> {
+                    (attempt, ctx) -> {
                         callCount[0]++;
                         if (attempt == 1) {
                             throw new RuntimeException("fail");
@@ -252,7 +255,7 @@ void passesCorrectAttemptNumberToOperation() {
 
             context.withRetry(
                     "track",
-                    (ctx, attempt) -> {
+                    (attempt, ctx) -> {
                         attempts.add(attempt);
                         if (attempt < 4) {
                             throw new RuntimeException("not yet");
@@ -282,7 +285,7 @@ void passesErrorToRetryStrategy() {
                     RuntimeException.class,
                     () -> context.withRetry(
                             "my-op",
-                            (ctx, attempt) -> {
+                            (attempt, ctx) -> {
                                 throw new RuntimeException("error-" + attempt);
                             },
                             config));
@@ -301,7 +304,7 @@ void respectsCustomDelayFromRetryDecision() {
 
             context.withRetry(
                     "my-op",
-                    (ctx, attempt) -> {
+                    (attempt, ctx) -> {
                         if (attempt <= 2) {
                             throw new RuntimeException("fail");
                         }
@@ -321,7 +324,7 @@ void passesChildContextToOperation() {
 
             context.withRetry(
                     "my-op",
-                    (ctx, attempt) -> {
+                    (attempt, ctx) -> {
                         assertSame(childContext, ctx);
                         return "verified";
                     },
@@ -334,7 +337,8 @@ void operationReturnsNull() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var result = context.withRetry("my-op", (WithRetry) (ctx, attempt) -> null, config);
+            var result = context.withRetry(
+                    "my-op", (BiFunction) (attempt, ctx) -> null, config);
 
             assertNull(result);
         }
@@ -351,7 +355,7 @@ void preservesCheckedExceptionSubclassType() {
                     SerDesException.class,
                     () -> context.withRetry(
                             "my-op",
-                            (ctx, attempt) -> {
+                            (attempt, ctx) -> {
                                 throw original;
                             },
                             config));
@@ -381,7 +385,7 @@ void propagatesSuspendExecutionExceptionWithoutRetrying() {
                     SuspendExecutionException.class,
                     () -> context.withRetry(
                             "my-op",
-                            (ctx, attempt) -> {
+                            (attempt, ctx) -> {
                                 throw new SuspendExecutionException();
                             },
                             config));
@@ -399,7 +403,7 @@ void propagatesUnrecoverableDurableExecutionExceptionWithoutRetrying() {
                     UnrecoverableDurableExecutionException.class,
                     () -> context.withRetry(
                             "my-op",
-                            (ctx, attempt) -> {
+                            (attempt, ctx) -> {
                                 throw new UnrecoverableDurableExecutionException(
                                         software.amazon.awssdk.services.lambda.model.ErrorObject.builder()
                                                 .errorMessage("unrecoverable")
@@ -420,7 +424,7 @@ void propagatesSuspendExecutionExceptionOnLaterAttempt() {
                     SuspendExecutionException.class,
                     () -> context.withRetry(
                             "my-op",
-                            (ctx, attempt) -> {
+                            (attempt, ctx) -> {
                                 if (attempt == 1) {
                                     throw new RuntimeException("transient");
                                 }
@@ -441,7 +445,7 @@ void propagatesUnrecoverableDurableExecutionExceptionOnLaterAttempt() {
                     UnrecoverableDurableExecutionException.class,
                     () -> context.withRetry(
                             "my-op",
-                            (ctx, attempt) -> {
+                            (attempt, ctx) -> {
                                 if (attempt == 1) {
                                     throw new RuntimeException("transient");
                                 }
@@ -471,7 +475,7 @@ void namedFormUsesProvidedNameForChildContextAndBackoff() {
 
             context.withRetry(
                     "my-op",
-                    (ctx, attempt) -> {
+                    (attempt, ctx) -> {
                         if (attempt < 2) {
                             throw new RuntimeException("fail");
                         }
@@ -493,7 +497,7 @@ void nullNameFormDefaultsToRetryForChildContextAndBackoff() {
 
             context.withRetry(
                     null,
-                    (ctx, attempt) -> {
+                    (attempt, ctx) -> {
                         if (attempt < 2) {
                             throw new RuntimeException("fail");
                         }
@@ -522,7 +526,7 @@ void syncReturnsValueDirectly() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var result = context.withRetry("my-op", (ctx, attempt) -> "sync-value", config);
+            var result = context.withRetry("my-op", (attempt, ctx) -> "sync-value", config);
 
             assertEquals("sync-value", result);
         }
@@ -533,7 +537,7 @@ void asyncReturnsDurableFuture() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            DurableFuture future = context.withRetryAsync("my-op", (ctx, attempt) -> "async-value", config);
+            DurableFuture future = context.withRetryAsync("my-op", (attempt, ctx) -> "async-value", config);
 
             assertNotNull(future);
             assertEquals("async-value", future.get());
@@ -545,8 +549,8 @@ void syncAndAsyncProduceSameResult() {
                     .retryStrategy(RetryStrategies.Presets.NO_RETRY)
                     .build();
 
-            var syncResult = context.withRetry("op", (ctx, attempt) -> "value", config);
-            var asyncResult = context.withRetryAsync("op", (ctx, attempt) -> "value", config)
+            var syncResult = context.withRetry("op", (attempt, ctx) -> "value", config);
+            var asyncResult = context.withRetryAsync("op", (attempt, ctx) -> "value", config)
                     .get();
 
             assertEquals(syncResult, asyncResult);
@@ -561,7 +565,7 @@ void asyncRetriesWithBackoff() {
 
             DurableFuture future = context.withRetryAsync(
                     "my-op",
-                    (ctx, attempt) -> {
+                    (attempt, ctx) -> {
                         if (attempt < 3) {
                             throw new RuntimeException("fail-" + attempt);
                         }
@@ -591,7 +595,7 @@ void syncNullOperationThrows() {
 
         @Test
         void syncNullConfigThrows() {
-            assertThrows(NullPointerException.class, () -> context.withRetry("name", (ctx, a) -> "x", null));
+            assertThrows(NullPointerException.class, () -> context.withRetry("name", (a, ctx) -> "x", null));
         }
 
         @Test
@@ -605,7 +609,7 @@ void asyncNullOperationThrows() {
 
         @Test
         void asyncNullConfigThrows() {
-            assertThrows(NullPointerException.class, () -> context.withRetryAsync("name", (ctx, a) -> "x", null));
+            assertThrows(NullPointerException.class, () -> context.withRetryAsync("name", (a, ctx) -> "x", null));
         }
     }
 
@@ -621,7 +625,7 @@ void setUpChildContext() {
 
         @Test
         void syncWithRetryWithoutConfigSucceedsOnFirstAttempt() {
-            var result = context.withRetry("my-op", (ctx, attempt) -> "default-config-result");
+            var result = context.withRetry("my-op", (attempt, ctx) -> "default-config-result");
 
             assertEquals("default-config-result", result);
         }
@@ -630,7 +634,7 @@ void syncWithRetryWithoutConfigSucceedsOnFirstAttempt() {
         void syncWithRetryWithoutConfigRetriesOnFailure() {
             var callCount = new int[] {0};
 
-            var result = context.withRetry("my-op", (ctx, attempt) -> {
+            var result = context.withRetry("my-op", (attempt, ctx) -> {
                 callCount[0]++;
                 if (attempt == 1) {
                     throw new RuntimeException("transient");
@@ -645,7 +649,7 @@ void syncWithRetryWithoutConfigRetriesOnFailure() {
 
         @Test
         void asyncWithRetryWithoutConfigSucceedsOnFirstAttempt() {
-            DurableFuture future = context.withRetryAsync("my-op", (ctx, attempt) -> "async-default");
+            DurableFuture future = context.withRetryAsync("my-op", (attempt, ctx) -> "async-default");
 
             assertNotNull(future);
             assertEquals("async-default", future.get());
@@ -655,7 +659,7 @@ void asyncWithRetryWithoutConfigSucceedsOnFirstAttempt() {
         void asyncWithRetryWithoutConfigRetriesOnFailure() {
             var callCount = new int[] {0};
 
-            DurableFuture future = context.withRetryAsync("my-op", (ctx, attempt) -> {
+            DurableFuture future = context.withRetryAsync("my-op", (attempt, ctx) -> {
                 callCount[0]++;
                 if (attempt == 1) {
                     throw new RuntimeException("transient");
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/model/WithRetryTest.java b/sdk/src/test/java/software/amazon/lambda/durable/model/WithRetryTest.java
deleted file mode 100644
index db6c83953..000000000
--- a/sdk/src/test/java/software/amazon/lambda/durable/model/WithRetryTest.java
+++ /dev/null
@@ -1,52 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-package software.amazon.lambda.durable.model;
-
-import static org.junit.jupiter.api.Assertions.*;
-import static org.mockito.Mockito.*;
-
-import org.junit.jupiter.api.Test;
-import software.amazon.lambda.durable.DurableContext;
-
-class WithRetryTest {
-
-    @Test
-    void canBeImplementedAsLambda() {
-        WithRetry operation = (ctx, attempt) -> "result-" + attempt;
-        var context = mock(DurableContext.class);
-
-        assertEquals("result-1", operation.execute(context, 1));
-        assertEquals("result-3", operation.execute(context, 3));
-    }
-
-    @Test
-    void receivesContextAndAttempt() {
-        var context = mock(DurableContext.class);
-        WithRetry operation = (ctx, attempt) -> {
-            assertSame(context, ctx);
-            return "attempt-" + attempt;
-        };
-
-        assertEquals("attempt-1", operation.execute(context, 1));
-        assertEquals("attempt-2", operation.execute(context, 2));
-    }
-
-    @Test
-    void canThrowExceptions() {
-        WithRetry operation = (ctx, attempt) -> {
-            throw new RuntimeException("failed on attempt " + attempt);
-        };
-        var context = mock(DurableContext.class);
-
-        var exception = assertThrows(RuntimeException.class, () -> operation.execute(context, 1));
-        assertEquals("failed on attempt 1", exception.getMessage());
-    }
-
-    @Test
-    void canReturnNull() {
-        WithRetry operation = (ctx, attempt) -> null;
-        var context = mock(DurableContext.class);
-
-        assertNull(operation.execute(context, 1));
-    }
-}