feat(plugin): surface execution input and result on invocation hooks

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

@@ -4,11 +4,15 @@
 
 import static org.junit.jupiter.api.Assertions.*;
 
+import java.io.IOException;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.jupiter.api.Test;
 import software.amazon.awssdk.services.lambda.model.ErrorObject;
 import software.amazon.awssdk.services.lambda.model.OperationStatus;
@@ -19,7 +23,10 @@
 import software.amazon.lambda.durable.model.WaitForConditionResult;
 import software.amazon.lambda.durable.plugin.*;
 import software.amazon.lambda.durable.retry.RetryStrategies;
+import software.amazon.lambda.durable.serde.JacksonSerDes;
+import software.amazon.lambda.durable.serde.SerDes;
 import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+import software.amazon.lambda.durable.util.ExceptionHelper;
 
 /** Integration tests verifying plugin hooks fire correctly during durable execution lifecycle. */
 class PluginIntegrationTest {
@@ -122,6 +129,167 @@ void plugin_receivesInvocationEnd_withFailedStatus_onError() {
         assertNotNull(plugin.invocationEnds.get(0).executionError());
     }
 
+    @Test
+    void plugin_invocationHooks_carryExecutionInputAndResult_onSuccess() {
+        var plugin = new RecordingPlugin();
+        var config = DurableConfig.builder().withPlugins(plugin).build();
+
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> context.step("greet", String.class, stepCtx -> "Hello " + input),
+                config);
+
+        var result = runner.runUntilComplete("World");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+
+        // The deserialized handler input reaches the start hook ...
+        assertEquals("World", plugin.invocationStarts.get(0).executionInput());
+        // ... and the end hook, alongside the value the handler returned.
+        var end = plugin.invocationEnds.get(0);
+        assertEquals("World", end.executionInput());
+        assertEquals("Hello World", end.executionResult());
+    }
+
+    @Test
+    void plugin_invocationEnd_omitsExecutionResult_onFailure() {
+        var plugin = new RecordingPlugin();
+        var config = DurableConfig.builder().withPlugins(plugin).build();
+
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> context.step(
+                        "failing",
+                        String.class,
+                        stepCtx -> {
+                            throw new RuntimeException("boom");
+                        },
+                        StepConfig.builder()
+                                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                                .build()),
+                config);
+
+        var result = runner.run("input");
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+
+        var end = plugin.invocationEnds.get(0);
+        assertEquals("input", end.executionInput());
+        assertNull(end.executionResult(), "a failed execution produced no result");
+    }
+
+    @Test
+    void plugin_invocationEnd_omitsExecutionResult_onSuspension() {
+        var plugin = new RecordingPlugin();
+        var config = DurableConfig.builder().withPlugins(plugin).build();
+
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> {
+                    context.wait("pause", Duration.ofMinutes(5));
+                    return "complete";
+                },
+                config);
+
+        var result = runner.run("input");
+
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        var end = plugin.invocationEnds.get(0);
+        assertEquals(InvocationStatus.PENDING, end.invocationStatus());
+        assertEquals("input", end.executionInput());
+        assertNull(end.executionResult(), "a suspended invocation has not produced a result yet");
+    }
+
+    @Test
+    void plugin_executionInput_isDeserializedOnce_andSharedWithHandler() {
+        var serDes = new CountingSerDes();
+        var plugin = new RecordingPlugin();
+        var config =
+                DurableConfig.builder().withPlugins(plugin).withSerDes(serDes).build();
+        var handlerInput = new AtomicReference<Object>();
+
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> {
+                    handlerInput.set(input);
+                    return context.step("greet", String.class, stepCtx -> "Hello " + input);
+                },
+                config);
+
+        var result = runner.runUntilComplete("World");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        // The handler input is deserialized once and that same instance reaches the hooks — a second
+        // deserialization would double the cost and re-run side effects in a stateful SerDes.
+        assertEquals(1, serDes.inputDeserializations("World"));
+        assertSame(handlerInput.get(), plugin.invocationStarts.get(0).executionInput());
+        assertSame(handlerInput.get(), plugin.invocationEnds.get(0).executionInput());
+    }
+
+    /** SerDes that counts how many times each payload is deserialized. */
+    static class CountingSerDes implements SerDes {
+        private final JacksonSerDes delegate = new JacksonSerDes();
+        private final Map<String, AtomicInteger> deserializations = new ConcurrentHashMap<>();
+
+        @Override
+        public String serialize(Object value) {
+            return delegate.serialize(value);
+        }
+
+        @Override
+        public <T> T deserialize(String data, TypeToken<T> typeToken) {
+            deserializations
+                    .computeIfAbsent(data, ignored -> new AtomicInteger())
+                    .incrementAndGet();
+            return delegate.deserialize(data, typeToken);
+        }
+
+        int inputDeserializations(String value) {
+            var counter = deserializations.get(delegate.serialize(value));
+            return counter == null ? 0 : counter.get();
+        }
+    }
+
+    @Test
+    void plugin_hooksStayPaired_whenSerDesSneakyThrowsCheckedException() {
+        var plugin = new RecordingPlugin();
+        var config = DurableConfig.builder()
+                .withPlugins(plugin)
+                .withSerDes(new SneakyThrowingSerDes())
+                .build();
+
+        var runner = LocalDurableTestRunner.create(String.class, (input, context) -> "unreachable", config);
+
+        var result = runner.run("input");
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+        // SerDes.deserialize declares no checked exceptions, so an implementation can sneaky-throw one. The start
+        // hook must still fire, otherwise a plugin keying state on it sees an end with no start.
+        assertEquals(1, plugin.invocationStarts.size(), "onInvocationStart must fire before the failure surfaces");
+        assertNull(plugin.invocationStarts.get(0).executionInput(), "input could not be deserialized");
+        assertEquals(1, plugin.invocationEnds.size());
+        assertEquals(InvocationStatus.FAILED, plugin.invocationEnds.get(0).invocationStatus());
+    }
+
+    /**
+     * SerDes that sneaky-throws a checked exception from deserialize, as DurableInputOutputSerDes does on IO errors.
+     */
+    static class SneakyThrowingSerDes implements SerDes {
+        private final JacksonSerDes delegate = new JacksonSerDes();
+
+        @Override
+        public String serialize(Object value) {
+            return delegate.serialize(value);
+        }
+
+        @Override
+        public <T> T deserialize(String data, TypeToken<T> typeToken) {
+            ExceptionHelper.sneakyThrow(new IOException("sneaky checked failure"));
+            return null;
+        }
+    }
+
     // ─── Operation-level hooks ───────────────────────────────────────────
 
     @Test

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

@@ -7,6 +7,7 @@
 import java.nio.charset.StandardCharsets;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -60,21 +61,41 @@ public static <I, O> DurableExecutionOutput execute(
             var executionArn = input.durableExecutionArn();
 
             executionManager.registerActiveThread(null);
+            // Captured for onInvocationEnd, which runs outside the handler thread below.
+            var pluginExecutionInput = new AtomicReference<>();
             var handlerFuture = CompletableFuture.supplyAsync(
                     () -> {
                         executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT));
 
+                        // Deserialize once and share the value with the plugin hooks and the handler below. A second
+                        // deserialization would double the cost, hand plugins a different object than the handler, and
+                        // re-run any side effects in a stateful custom SerDes. A failure is captured rather than thrown
+                        // so onInvocationStart still fires before it surfaces, keeping the start/end hooks paired.
+                        // SerDes is a public extension point whose deserialize declares no checked exceptions, so an
+                        // implementation may sneaky-throw one; capture every Throwable and rethrow it unchanged.
+                        I userInput = null;
+                        Throwable inputFailure = null;
+                        try {
+                            userInput = extractUserInput(
+                                    executionManager.getExecutionOperation(), config.getSerDes(), inputType);
+                        } catch (Throwable t) {
+                            inputFailure = t;
+                        }
+                        pluginExecutionInput.set(userInput);
+
                         // 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.
                         pluginRunner.onInvocationStart(new InvocationInfo(
                                 requestId,
                                 executionArn,
                                 isFirstInvocation,
-                                executionManager.getExecutionOperation().startTimestamp()));
+                                executionManager.getExecutionOperation().startTimestamp(),
+                                userInput));
+                        if (inputFailure != null) {
+                            ExceptionHelper.sneakyThrow(inputFailure);
+                        }
 
-                        var userInput = extractUserInput(
-                                executionManager.getExecutionOperation(), config.getSerDes(), inputType);
                         var context = DurableContextImpl.createRootContext(executionManager, config, lambdaContext);
                         DurableContextImpl.setCurrentContext(context);
                         // use a try-with-resources to clear logger properties
@@ -103,6 +124,8 @@ public static <I, O> DurableExecutionOutput execute(
                                             executionArn,
                                             isFirstInvocation,
                                             InvocationStatus.PENDING,
+                                            null,
+                                            pluginExecutionInput.get(),
                                             null);
                                     return DurableExecutionOutput.pending();
                                 }
@@ -119,7 +142,9 @@ public static <I, O> DurableExecutionOutput execute(
                                             executionArn,
                                             isFirstInvocation,
                                             InvocationStatus.RETRYING,
-                                            cause);
+                                            cause,
+                                            pluginExecutionInput.get(),
+                                            null);
                                     throw unrecoverableDurableExecutionException;
                                 }
 
@@ -131,7 +156,9 @@ public static <I, O> DurableExecutionOutput execute(
                                         executionArn,
                                         isFirstInvocation,
                                         InvocationStatus.FAILED,
-                                        cause);
+                                        cause,
+                                        pluginExecutionInput.get(),
+                                        null);
                                 return DurableExecutionOutput.failure(buildErrorObject(cause, config.getSerDes()));
                             }
                             // user handler complete successfully
@@ -145,7 +172,9 @@ public static <I, O> DurableExecutionOutput execute(
                                     executionArn,
                                     isFirstInvocation,
                                     InvocationStatus.SUCCEEDED,
-                                    null);
+                                    null,
+                                    pluginExecutionInput.get(),
+                                    result);
                             return output;
                         })
                         .join();
@@ -163,8 +192,11 @@ private static void fireOnInvocationEnd(
             String executionArn,
             boolean isFirstInvocation,
             InvocationStatus status,
-            Throwable error) {
-        pluginRunner.onInvocationEnd(new InvocationEndInfo(requestId, executionArn, isFirstInvocation, status, error));
+            Throwable error,
+            Object executionInput,
+            Object executionResult) {
+        pluginRunner.onInvocationEnd(new InvocationEndInfo(
+                requestId, executionArn, isFirstInvocation, status, error, executionInput, executionResult));
     }
 
     private static String handleLargePayload(ExecutionManager executionManager, String outputPayload) {

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

@@ -12,10 +12,52 @@
  * @param isFirstInvocation true if this is the first invocation of the execution
  * @param invocationStatus the invocation outcome (SUCCEEDED, FAILED, or PENDING)
  * @param executionError non-null if the execution failed; this component is experimental
+ * @param executionInput the deserialized execution input passed to the user handler, or null when no plugins are
+ *     registered or the input could not be deserialized; this component is experimental
+ * @param executionResult the value the user handler returned, or null unless the invocation completed the execution
+ *     successfully; this component is experimental
  */
 public record InvocationEndInfo(
         String requestId,
         String durableExecutionArn,
         boolean isFirstInvocation,
         InvocationStatus invocationStatus,
-        @Experimental Throwable executionError) {}
+        @Experimental Throwable executionError,
+        @Experimental Object executionInput,
+        @Experimental Object executionResult) {
+
+    /**
+     * Creates invocation-end information without the execution input or result.
+     *
+     * <p>Retained so callers written before {@code executionInput} and {@code executionResult} were added keep
+     * compiling; both resolve to null.
+     *
+     * @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 invocationStatus the invocation outcome
+     * @param executionError non-null if the execution failed
+     */
+    public InvocationEndInfo(
+            String requestId,
+            String durableExecutionArn,
+            boolean isFirstInvocation,
+            InvocationStatus invocationStatus,
+            Throwable executionError) {
+        this(requestId, durableExecutionArn, isFirstInvocation, invocationStatus, executionError, null, null);
+    }
+
+    /**
+     * Returns a representation that omits {@code executionInput} and {@code executionResult}.
+     *
+     * <p>The generated representation would render both payloads, so plugins that log this object whole would start
+     * emitting customer inputs and results, potentially including secrets or personal data. Read the components
+     * explicitly to record them.
+     */
+    @Override
+    public String toString() {
+        return "InvocationEndInfo[requestId=" + requestId + ", durableExecutionArn=" + durableExecutionArn
+                + ", isFirstInvocation=" + isFirstInvocation + ", invocationStatus=" + invocationStatus
+                + ", executionError=" + executionError + "]";
+    }
+}