fix(plugin): surface missing state fields on hook info records

otel-plugin/src/test/java/software/amazon/lambda/durable/otel/MdcSpanEnricherTest.java

@@ -8,6 +8,7 @@
 import io.opentelemetry.sdk.trace.SdkTracerProvider;
 import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
 import java.time.Instant;
+import java.util.Map;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 import org.slf4j.MDC;
@@ -63,26 +64,27 @@ void plugin_withMdcEnabled_setsFieldsInMdc() {
                         .enableMdc(true)
                         .build());
 
-        plugin.onInvocationStart(new InvocationInfo("req-1", "arn:exec-mdc-test", true, Instant.now()));
+        plugin.onInvocationStart(
+                new InvocationInfo("req-1", "arn:exec-mdc-test", true, Instant.now(), Map.of(), Map.of()));
 
         plugin.onUserFunctionStart(
-                new UserFunctionStartInfo("op-1", "step", "STEP", "Step", null, Instant.now(), false, 1));
+                new UserFunctionStartInfo("op-1", "step", "STEP", "Step", null, Instant.now(), false, false, 1));
 
         // MDC should have trace fields after onUserFunctionStart
         assertNotNull(MDC.get(MdcSpanEnricher.MDC_TRACE_ID));
         assertNotNull(MDC.get(MdcSpanEnricher.MDC_SPAN_ID));
         assertNotNull(MDC.get(MdcSpanEnricher.MDC_TRACE_SAMPLED));
 
         plugin.onUserFunctionEnd(new UserFunctionEndInfo(
-                "op-1", "step", "STEP", "Step", null, Instant.now(), Instant.now(), false, 1, true, null));
+                "op-1", "step", "STEP", "Step", null, Instant.now(), Instant.now(), false, false, 1, true, null));
 
         // After onUserFunctionEnd: span_id is cleared, but trace_id remains for handler-level logs between steps
         assertNotNull(MDC.get(MdcSpanEnricher.MDC_TRACE_ID), "trace_id should persist between steps");
         assertNull(MDC.get(MdcSpanEnricher.MDC_SPAN_ID), "span_id should be cleared after step");
         assertNotNull(MDC.get(MdcSpanEnricher.MDC_TRACE_SAMPLED), "trace_flags should persist between steps");
 
-        plugin.onInvocationEnd(
-                new InvocationEndInfo("req-1", "arn:exec-mdc-test", true, InvocationStatus.SUCCEEDED, null));
+        plugin.onInvocationEnd(new InvocationEndInfo(
+                "req-1", "arn:exec-mdc-test", true, Instant.now(), Map.of(), InvocationStatus.SUCCEEDED, null));
 
         // After onInvocationEnd: all MDC fields are cleared
         assertNull(MDC.get(MdcSpanEnricher.MDC_TRACE_ID));

sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java

@@ -28,6 +28,7 @@
 import software.amazon.lambda.durable.plugin.InvocationEndInfo;
 import software.amazon.lambda.durable.plugin.InvocationInfo;
 import software.amazon.lambda.durable.plugin.InvocationStatus;
+import software.amazon.lambda.durable.plugin.PluginInfoConverter;
 import software.amazon.lambda.durable.plugin.PluginRunner;
 import software.amazon.lambda.durable.serde.SerDes;
 import software.amazon.lambda.durable.util.ExceptionHelper;
@@ -67,11 +68,19 @@ public static <I, O> DurableExecutionOutput execute(
                         // onInvocationStart runs on the user thread so plugins can
                         // inject ThreadLocal objects, update MDC, etc.
                         // executionStartTime comes from the initial EXECUTION operation in the first backend event.
+                        // The operation maps are snapshots of the state delivered for this invocation; on a replay
+                        // invocation updatedOperations names the operations the backend completed while suspended.
                         pluginRunner.onInvocationStart(new InvocationInfo(
                                 requestId,
                                 executionArn,
                                 isFirstInvocation,
-                                executionManager.getExecutionOperation().startTimestamp()));
+                                executionManager.getExecutionOperation().startTimestamp(),
+                                PluginInfoConverter.toOperationItemMap(
+                                        executionManager.getOperationsSnapshot(),
+                                        executionManager.getInitialOperationIds()),
+                                PluginInfoConverter.toOperationItemMap(
+                                        executionManager.getUpdatedOperationsSnapshot(),
+                                        executionManager.getInitialOperationIds())));
 
                         var userInput = extractUserInput(
                                 executionManager.getExecutionOperation(), config.getSerDes(), inputType);
@@ -99,6 +108,7 @@ public static <I, O> DurableExecutionOutput execute(
                                 if (cause instanceof SuspendExecutionException) {
                                     fireOnInvocationEnd(
                                             pluginRunner,
+                                            executionManager,
                                             requestId,
                                             executionArn,
                                             isFirstInvocation,
@@ -115,6 +125,7 @@ public static <I, O> DurableExecutionOutput execute(
                                         && unrecoverableDurableExecutionException.isRetryable()) {
                                     fireOnInvocationEnd(
                                             pluginRunner,
+                                            executionManager,
                                             requestId,
                                             executionArn,
                                             isFirstInvocation,
@@ -127,6 +138,7 @@ public static <I, O> DurableExecutionOutput execute(
                                 logger.debug("Execution failed: {}", cause.getMessage());
                                 fireOnInvocationEnd(
                                         pluginRunner,
+                                        executionManager,
                                         requestId,
                                         executionArn,
                                         isFirstInvocation,
@@ -141,6 +153,7 @@ public static <I, O> DurableExecutionOutput execute(
                                     DurableExecutionOutput.success(handleLargePayload(executionManager, outputPayload));
                             fireOnInvocationEnd(
                                     pluginRunner,
+                                    executionManager,
                                     requestId,
                                     executionArn,
                                     isFirstInvocation,
@@ -159,12 +172,24 @@ public static <I, O> DurableExecutionOutput execute(
 
     private static void fireOnInvocationEnd(
             PluginRunner pluginRunner,
+            ExecutionManager executionManager,
             String requestId,
             String executionArn,
             boolean isFirstInvocation,
             InvocationStatus status,
             Throwable error) {
-        pluginRunner.onInvocationEnd(new InvocationEndInfo(requestId, executionArn, isFirstInvocation, status, error));
+        // The end info repeats the start info's identity surface (execution start time, operation snapshot) so an
+        // invocation-end hook never has to correlate back to the start hook. The snapshot is taken at end time, so
+        // unlike the start info it also contains operations created during this invocation.
+        pluginRunner.onInvocationEnd(new InvocationEndInfo(
+                requestId,
+                executionArn,
+                isFirstInvocation,
+                executionManager.getExecutionOperation().startTimestamp(),
+                PluginInfoConverter.toOperationItemMap(
+                        executionManager.getOperationsSnapshot(), executionManager.getInitialOperationIds()),
+                status,
+                error));
     }
 
     private static String handleLargePayload(ExecutionManager executionManager, String outputPayload) {

sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java

@@ -5,6 +5,7 @@
 import com.amazonaws.services.lambda.runtime.Context;
 import java.time.Instant;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
@@ -63,6 +64,7 @@ public class ExecutionManager implements SafeCloseable {
     private final AtomicReference<ExecutionMode> executionMode;
     private final DurableConfig durableConfig;
     private final Set<String> updatedOperationIdsSinceLastInvocation;
+    private final Set<String> initialOperationIds;
 
     // ===== Thread Coordination =====
     private final Map<String, BaseDurableOperation> registeredOperations = new ConcurrentHashMap<>();
@@ -89,6 +91,11 @@ public ExecutionManager(DurableExecutionInput input, DurableConfig config, Conte
         this.operationStorage = checkpointManager.fetchAllPages(input.initialExecutionState()).stream()
                 .collect(Collectors.toConcurrentMap(Operation::id, op -> op));
 
+        // The ids delivered in this invocation's initial state. Everything else in operationStorage is created during
+        // this invocation, so this set is what distinguishes replayed operations from freshly-started ones for the
+        // plugin hooks' isReplay indicators.
+        this.initialOperationIds = Set.copyOf(operationStorage.keySet());
+
         // Start in REPLAY mode if we have more than just the initial EXECUTION operation
         this.executionMode =
                 new AtomicReference<>(operationStorage.size() > 1 ? ExecutionMode.REPLAY : ExecutionMode.EXECUTION);
@@ -132,6 +139,47 @@ public boolean isOperationUpdatedSinceLastInvocation(String operationId) {
         return updatedOperationIdsSinceLastInvocation.contains(operationId);
     }
 
+    /**
+     * Returns {@code true} if the given operation was present in the checkpointed state delivered at the start of this
+     * invocation, i.e. it predates this invocation and is being replayed rather than started fresh. Unlike
+     * {@link #getOperationAndUpdateReplayState(String)} this does not mutate the execution's replay mode, so it is safe
+     * to call from plugin-hook firing sites.
+     *
+     * @param operationId the operation ID to check
+     * @return true if the operation was delivered in this invocation's initial state
+     */
+    public boolean wasObservedAtInvocationStart(String operationId) {
+        return initialOperationIds.contains(operationId);
+    }
+
+    /** Returns the ids of the operations delivered in this invocation's initial state. */
+    public Set<String> getInitialOperationIds() {
+        return initialOperationIds;
+    }
+
+    /**
+     * Returns an immutable snapshot of the operations currently tracked for this execution, including the initial
+     * EXECUTION operation. Non-mutating; intended for the invocation-level plugin hooks.
+     *
+     * @return a snapshot of the tracked operations
+     */
+    public Collection<Operation> getOperationsSnapshot() {
+        return List.copyOf(operationStorage.values());
+    }
+
+    /**
+     * Returns the subset of {@link #getOperationsSnapshot()} whose ids the backend reported as updated since the last
+     * successful invocation. Empty on the first invocation. Ids without a corresponding tracked operation are skipped.
+     *
+     * @return a snapshot of the externally-updated operations
+     */
+    public Collection<Operation> getUpdatedOperationsSnapshot() {
+        return updatedOperationIdsSinceLastInvocation.stream()
+                .map(operationStorage::get)
+                .filter(Objects::nonNull)
+                .toList();
+    }
+
     /** Registers an operation so it can receive checkpoint completion notifications. */
     public void registerOperation(BaseDurableOperation operation) {
         registeredOperations.put(operation.getOperationId(), operation);
@@ -162,7 +210,11 @@ private void onCheckpointComplete(List<Operation> newOperations) {
             durableConfig
                     .getPluginRunner()
                     .onOperationChange(PluginInfoConverter.toOperationChangeInfo(
-                            requestId, durableExecutionArn, updatedOperations, operationStorage.values()));
+                            requestId,
+                            durableExecutionArn,
+                            updatedOperations,
+                            operationStorage.values(),
+                            initialOperationIds));
         }
     }
 

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

@@ -350,7 +350,11 @@ protected void runUserHandler(Runnable runnable, ThreadType threadType) {
     protected <T> T runUserFunction(Integer attempt, Supplier<T> userFunction) {
         var pluginRunner = getPluginRunner();
         var startInfo = PluginInfoConverter.toUserFunctionStartInfo(
-                operationIdentifier, durableContext.getParentId(), durableContext.isReplaying(), attempt);
+                operationIdentifier,
+                durableContext.getParentId(),
+                executionManager.wasObservedAtInvocationStart(getOperationId()),
+                durableContext.isReplaying(),
+                attempt);
         pluginRunner.onUserFunctionStart(startInfo);
         try {
             T result = userFunction.get();

sdk/src/main/java/software/amazon/lambda/durable/plugin/InvocationEndInfo.java

@@ -2,12 +2,23 @@
 // SPDX-License-Identifier: Apache-2.0
 package software.amazon.lambda.durable.plugin;
 
+import java.time.Instant;
+import java.util.Map;
+
 /**
  * Information provided at the end of a Lambda invocation.
  *
+ * <p>Carries the same invocation-identity surface as {@link InvocationInfo} so an invocation-end hook never has to
+ * correlate back to the start hook to learn the execution start time or the operation state.
+ *
  * @param requestId the Lambda request ID for this invocation
  * @param durableExecutionArn the durable execution ARN
  * @param isFirstInvocation true if this is the first invocation of the execution
+ * @param executionStartTime the start timestamp of the durable execution, taken from the initial EXECUTION operation in
+ *     the first event delivered by the backend. Stable across all invocations of the same execution.
+ * @param operations a snapshot of the checkpointed operations known when the invocation ended, keyed by operation ID.
+ *     Unlike {@link InvocationInfo#operations()} this includes operations created during this invocation. Includes the
+ *     initial EXECUTION operation. Empty-but-never-null.
  * @param invocationStatus the invocation outcome (SUCCEEDED, FAILED, or PENDING)
  * @param executionError non-null if the execution failed
  * @deprecated This is a preview API that is experimental and may be changed or removed in future releases.
@@ -17,5 +28,7 @@ public record InvocationEndInfo(
         String requestId,
         String durableExecutionArn,
         boolean isFirstInvocation,
+        Instant executionStartTime,
+        Map<String, OperationChangeItemInfo> operations,
         InvocationStatus invocationStatus,
         Throwable executionError) {}

sdk/src/main/java/software/amazon/lambda/durable/plugin/InvocationInfo.java

@@ -3,6 +3,7 @@
 package software.amazon.lambda.durable.plugin;
 
 import java.time.Instant;
+import java.util.Map;
 
 /**
  * Invocation-level information available to plugin hooks.
@@ -12,8 +13,19 @@
  * @param isFirstInvocation true if this is the first invocation of the execution (not a replay invocation)
  * @param executionStartTime the start timestamp of the durable execution, taken from the initial EXECUTION operation in
  *     the first event delivered by the backend. Stable across all invocations of the same execution.
+ * @param operations a snapshot of the checkpointed operations delivered at the start of this invocation, keyed by
+ *     operation ID. Includes the initial EXECUTION operation. Empty-but-never-null.
+ * @param updatedOperations the subset of {@code operations} that changed externally between the previous invocation and
+ *     this one (a wait timer expired, a callback was received, a chained invoke completed), keyed by operation ID.
+ *     Sourced from the {@code UpdatedOperationIds} field of the durable invocation input, so it is empty on the first
+ *     invocation. Empty-but-never-null.
  * @deprecated This is a preview API that is experimental and may be changed or removed in future releases.
  */
 @Deprecated
 public record InvocationInfo(
-        String requestId, String durableExecutionArn, boolean isFirstInvocation, Instant executionStartTime) {}
+        String requestId,
+        String durableExecutionArn,
+        boolean isFirstInvocation,
+        Instant executionStartTime,
+        Map<String, OperationChangeItemInfo> operations,
+        Map<String, OperationChangeItemInfo> updatedOperations) {}

sdk/src/main/java/software/amazon/lambda/durable/plugin/OperationChangeItemInfo.java

@@ -6,7 +6,11 @@
 import software.amazon.awssdk.services.lambda.model.OperationStatus;
 
 /**
- * Operation-level information for a single operation within an {@link OperationChangeInfo}.
+ * Operation-level information for a single operation within an {@link OperationChangeInfo}, and the snapshot record
+ * used for the operation maps carried on {@link InvocationInfo} / {@link InvocationEndInfo}.
+ *
+ * <p>Carries the full operation field surface, mirroring {@link OperationEndInfo}, so a plugin observing an operation
+ * through a change delta or an invocation-level map sees the same fields it would see through the per-operation hooks.
  *
  * @param id operation ID
  * @param name human-readable operation name (may be null)
@@ -15,8 +19,11 @@
  * @param parentId parent operation ID (null for root-level operations)
  * @param startTimestamp when the operation started
  * @param endTimestamp when the operation ended
- * @param error non-null if the operation failed
  * @param status operation status
+ * @param attempt the attempt number for retriable operations (STEP, WAIT_FOR_CONDITION) — null for others
+ * @param isReplay true if this operation was already present in the checkpointed state delivered at the start of the
+ *     current invocation (i.e. it predates this invocation) rather than being created during it
+ * @param error non-null if the operation failed
  * @deprecated This is a preview API that is experimental and may be changed or removed in future releases.
  */
 @Deprecated
@@ -28,5 +35,7 @@ public record OperationChangeItemInfo(
         String parentId,
         Instant startTimestamp,
         Instant endTimestamp,
-        Throwable error,
-        OperationStatus status) {}
+        OperationStatus status,
+        Integer attempt,
+        boolean isReplay,
+        Throwable error) {}