[feature] add an optional stepContext parameter (#130)

* add an optional stepContext parameter

* rename parent context id to context id

* rename parameter

* update loggin example

* add some comments

* update DurableLogger to use StepContext class

* update docs and examples to use child context logger

* add more javadoc and minor code cleanups

* add more comments and remove registration from context

* fix of test cases

Author: zhongkechen

README.md

@@ -354,8 +354,8 @@ The SDK provides a `DurableLogger` via `ctx.getLogger()` that automatically incl
 protected OrderResult handleRequest(Order order, DurableContext ctx) {
     ctx.getLogger().info("Processing order: {}", order.getId());
 
-    var result = ctx.step("validate", String.class, () -> {
-        ctx.getLogger().debug("Validating order details");
+    var result = ctx.step("validate", String.class, stepCtx -> {
+        stepCtx.getLogger().debug("Validating order details");
         return validate(order);
     });
 

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

@@ -47,13 +47,13 @@ public String handleRequest(ApprovalRequest input, DurableContext context) {
         var callback = context.createCallback("approval", String.class, config);
 
         // Step 2.5: Log AWS CLI command to complete the callback
-        context.step("log-callback-command", Void.class, () -> {
+        context.step("log-callback-command", Void.class, ctx -> {
             var callbackId = callback.callbackId();
             // The result must be base64-encoded JSON
             var command = String.format(
                     "aws lambda send-durable-execution-callback-success --callback-id %s --result $(echo -n '\"approved\"' | base64)",
                     callbackId);
-            context.getLogger().info("To complete this callback, run: {}", command);
+            ctx.getLogger().info("To complete this callback, run: {}", command);
             return null;
         });
 

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

@@ -31,7 +31,7 @@ public String handleRequest(GreetingRequest input, DurableContext context) {
         // Child context 1: Order validation — step + wait + step
         var orderFuture = context.runInChildContextAsync("order-validation", String.class, child -> {
             var prepared = child.step("prepare-order", String.class, () -> "Order for " + name);
-            context.getLogger().info("Order prepared, waiting for validation");
+            child.getLogger().info("Order prepared, waiting for validation");
 
             child.wait("validation-delay", Duration.ofSeconds(5));
 
@@ -41,7 +41,7 @@ public String handleRequest(GreetingRequest input, DurableContext context) {
         // Child context 2: Inventory check — step + wait + step
         var inventoryFuture = context.runInChildContextAsync("inventory-check", String.class, child -> {
             var stock = child.step("check-stock", String.class, () -> "Stock available for " + name);
-            context.getLogger().info("Stock checked, waiting for confirmation");
+            child.getLogger().info("Stock checked, waiting for confirmation");
 
             child.wait("confirmation-delay", Duration.ofSeconds(3));
 

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

@@ -15,24 +15,22 @@ public class LoggingExample extends DurableHandler<GreetingRequest, String> {
 
     @Override
     public String handleRequest(GreetingRequest input, DurableContext context) {
-        var logger = context.getLogger();
-
         // Log at execution level (outside any step)
-        logger.info("Processing greeting for: {}", input.getName());
+        context.getLogger().info("Processing greeting for: {}", input.getName());
 
         // Step 1: Create greeting - logs inside step include operation context
-        var greeting = context.step("create-greeting", String.class, () -> {
-            logger.info("Creating greeting message");
+        var greeting = context.step("create-greeting", String.class, ctx -> {
+            ctx.getLogger().info("Creating greeting message");
             return "Hello, " + input.getName();
         });
 
         // Step 2: Transform
-        var result = context.step("transform", String.class, () -> {
-            logger.info("Transforming greeting to uppercase");
+        var result = context.step("transform", String.class, ctx -> {
+            ctx.getLogger().info("Transforming greeting to uppercase");
             return greeting.toUpperCase() + "!";
         });
 
-        logger.info("Completed processing, result: {}", result);
+        context.getLogger().info("Completed processing, result: {}", result);
         return result;
     }
 }

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

@@ -31,7 +31,7 @@ public class RetryExample extends DurableHandler<Object, String> {
     @Override
     public String handleRequest(Object input, DurableContext context) {
         // Step 1: Record start time
-        startTime = context.step("record-start-time", Instant.class, Instant::now);
+        startTime = context.step("record-start-time", Instant.class, () -> Instant.now());
         logger.info("Recorded start time: {}", startTime);
 
         // Step 2: Call that never retries (fails immediately)

sdk/src/main/java/software/amazon/lambda/durable/BaseContext.java

@@ -0,0 +1,89 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import software.amazon.lambda.durable.execution.ExecutionManager;
+import software.amazon.lambda.durable.logging.DurableLogger;
+
+public abstract class BaseContext {
+    protected final ExecutionManager executionManager;
+    private final DurableConfig durableConfig;
+    private final Context lambdaContext;
+    private final ExecutionContext executionContext;
+    private final String contextId;
+    private boolean isReplaying;
+
+    /** Creates a new BaseContext instance. */
+    protected BaseContext(
+            ExecutionManager executionManager, DurableConfig durableConfig, Context lambdaContext, String contextId) {
+        this.executionManager = executionManager;
+        this.durableConfig = durableConfig;
+        this.lambdaContext = lambdaContext;
+        this.contextId = contextId;
+        this.executionContext = new ExecutionContext(executionManager.getDurableExecutionArn());
+        this.isReplaying = executionManager.hasOperationsForContext(contextId);
+    }
+
+    // =============== accessors ================
+    /**
+     * Gets a logger with additional information of the current execution context.
+     *
+     * @return a DurableLogger instance
+     */
+    public abstract DurableLogger getLogger();
+
+    /**
+     * Returns the AWS Lambda runtime context.
+     *
+     * @return the Lambda context
+     */
+    public Context getLambdaContext() {
+        return lambdaContext;
+    }
+
+    /**
+     * Returns metadata about the current durable execution.
+     *
+     * <p>The execution context provides information that remains constant throughout the execution lifecycle, such as
+     * the durable execution ARN. This is useful for tracking execution progress, correlating logs, and referencing this
+     * execution in external systems.
+     *
+     * @return the execution context
+     */
+    public ExecutionContext getExecutionContext() {
+        return executionContext;
+    }
+
+    /**
+     * Returns the configuration for durable execution behavior.
+     *
+     * @return the durable configuration
+     */
+    public DurableConfig getDurableConfig() {
+        return durableConfig;
+    }
+
+    // ============= internal utilities ===============
+
+    /** Gets the context ID for this context. Null for root context, set for child contexts. */
+    public String getContextId() {
+        return contextId;
+    }
+
+    public ExecutionManager getExecutionManager() {
+        return executionManager;
+    }
+
+    /** Returns whether this context is currently in replay mode. */
+    boolean isReplaying() {
+        return isReplaying;
+    }
+
+    /**
+     * Transitions this context from replay to execution mode. Called when the first un-cached operation is encountered.
+     */
+    void setExecutionMode() {
+        this.isReplaying = false;
+    }
+}