fix(plugin): Propagate error to onOperationEnd for failed operations

Author: ayushiahjolia

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

@@ -202,6 +202,52 @@ void waitCompletedDuringSuspension_producesOperationSpan() {
                                 .toList());
     }
 
+    @Test
+    void invokeFailedDuringSuspension_producesErrorSpan() {
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, ctx) -> {
+                    try {
+                        ctx.invoke("call-target", "target-fn", "{}", String.class);
+                    } catch (Exception e) {
+                        // expected failure
+                    }
+                    return "handled";
+                },
+                otelConfig);
+
+        // First invocation: invoke starts, suspends waiting for target
+        var result1 = runner.run("input");
+        assertEquals(ExecutionStatus.PENDING, result1.getStatus());
+
+        // Target fails while Lambda is frozen
+        runner.failChainedInvoke(
+                "call-target",
+                software.amazon.awssdk.services.lambda.model.ErrorObject.builder()
+                        .errorType("TargetError")
+                        .errorMessage("target function failed")
+                        .build());
+
+        // Resume
+        var result2 = runner.run("input");
+        assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus());
+
+        // The invoke operation span from the second invocation should have ERROR status
+        var allSpans = spanExporter.getFinishedSpanItems();
+        var invokeSpans = allSpans.stream()
+                .filter(s -> s.getName().contains("call-target"))
+                .toList();
+        assertTrue(invokeSpans.size() >= 2, "Should have at least 2 invoke spans (one PENDING, one ended with error)");
+
+        // The span from the second invocation should be marked as error
+        var errorSpan = invokeSpans.stream()
+                .filter(s -> s.getStatus().getStatusCode() == io.opentelemetry.api.trace.StatusCode.ERROR)
+                .findFirst();
+        assertTrue(
+                errorSpan.isPresent(),
+                "Should have an invoke span with ERROR status when invoke fails during suspension");
+    }
+
     @Test
     void childContext_producesNestedSpans() {
         var runner = LocalDurableTestRunner.create(

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

@@ -10,6 +10,7 @@
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.services.lambda.model.ErrorObject;
 import software.amazon.lambda.durable.config.StepConfig;
 import software.amazon.lambda.durable.model.ExecutionStatus;
 import software.amazon.lambda.durable.model.WaitForConditionResult;
@@ -235,6 +236,100 @@ void plugin_operationEnd_firedOnceForStepCompletingInCurrentInvocation() {
         assertEquals(1, step1EndCount, "step1 onOperationEnd should fire exactly once");
     }
 
+    @Test
+    void plugin_operationEnd_includesError_whenInvokeFailsDuringSuspension() {
+        var plugin = new RecordingPlugin();
+        var config = DurableConfig.builder().withPlugins(plugin).build();
+
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> {
+                    try {
+                        context.invoke("call-target", "target-fn", "{}", String.class);
+                    } catch (Exception e) {
+                        // expected
+                    }
+                    return "done";
+                },
+                config);
+
+        // First invocation: invoke starts, suspends waiting for target
+        var result1 = runner.run("input");
+        assertEquals(ExecutionStatus.PENDING, result1.getStatus());
+
+        // Target fails while Lambda is frozen
+        runner.failChainedInvoke(
+                "call-target",
+                ErrorObject.builder()
+                        .errorType("TargetError")
+                        .errorMessage("target function failed")
+                        .build());
+
+        // Clear plugin state to only track second invocation
+        plugin.operationEnds.clear();
+
+        var result2 = runner.run("input");
+        assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus());
+
+        // onOperationEnd for "call-target" should fire with error info
+        var invokeEnd = plugin.operationEnds.stream()
+                .filter(info -> "call-target".equals(info.name()))
+                .findFirst()
+                .orElse(null);
+        assertNotNull(invokeEnd, "onOperationEnd should fire for failed invoke during suspension");
+        assertNotNull(invokeEnd.error(), "error should be propagated for failed invoke");
+        assertEquals("target function failed", invokeEnd.error().getMessage());
+    }
+
+    @Test
+    void plugin_operationEnd_includesError_whenStepFailsViaCheckpoint() {
+        var plugin = new RecordingPlugin();
+        var config = DurableConfig.builder().withPlugins(plugin).build();
+
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> context.step(
+                        "failing-step",
+                        String.class,
+                        stepCtx -> {
+                            throw new RuntimeException("step exploded");
+                        },
+                        StepConfig.builder()
+                                .retryStrategy(RetryStrategies.Presets.NO_RETRY)
+                                .build()),
+                config);
+
+        var result = runner.run("input");
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+
+        // onOperationEnd for "failing-step" should fire with error info
+        var stepEnd = plugin.operationEnds.stream()
+                .filter(info -> "failing-step".equals(info.name()))
+                .findFirst()
+                .orElse(null);
+        assertNotNull(stepEnd, "onOperationEnd should fire for failed step");
+        assertNotNull(stepEnd.error(), "error should be propagated for failed step");
+        assertTrue(stepEnd.error().getMessage().contains("step exploded"));
+    }
+
+    @Test
+    void plugin_operationEnd_noError_whenOperationSucceeds() {
+        var plugin = new RecordingPlugin();
+        var config = DurableConfig.builder().withPlugins(plugin).build();
+
+        var runner = LocalDurableTestRunner.create(
+                String.class, (input, context) -> context.step("ok-step", String.class, stepCtx -> "success"), config);
+
+        runner.runUntilComplete("input");
+
+        var stepEnd = plugin.operationEnds.stream()
+                .filter(info -> "ok-step".equals(info.name()))
+                .findFirst()
+                .orElse(null);
+        assertNotNull(stepEnd, "onOperationEnd should fire for successful step");
+        assertNull(stepEnd.error(), "error should be null for successful step");
+    }
+
     // ─── User function hooks ─────────────────────────────────────────────
 
     @Test

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

@@ -10,10 +10,13 @@
 import java.util.concurrent.atomic.AtomicReference;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.services.lambda.model.ErrorObject;
 import software.amazon.awssdk.services.lambda.model.Operation;
+import software.amazon.awssdk.services.lambda.model.OperationStatus;
 import software.amazon.awssdk.services.lambda.model.OperationType;
 import software.amazon.awssdk.services.lambda.model.OperationUpdate;
 import software.amazon.lambda.durable.context.DurableContextImpl;
+import software.amazon.lambda.durable.exception.DurableOperationException;
 import software.amazon.lambda.durable.exception.IllegalDurableOperationException;
 import software.amazon.lambda.durable.exception.NonDeterministicExecutionException;
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
@@ -141,7 +144,7 @@ public void execute() {
                 // This enables the OTel plugin to emit spans for operations that transitioned while Lambda was frozen.
                 if (replayCompletedOperation.get()
                         && executionManager.isOperationUpdatedSinceLastInvocation(getOperationId())) {
-                    fireOnOperationEnd(existing, null);
+                    fireOnOperationEnd(existing, extractErrorFromOperation(existing));
                 }
 
                 replay(existing);
@@ -354,7 +357,7 @@ public void onCheckpointComplete(Operation operation) {
 
             // Fire onOperationEnd plugin hook — operation reached terminal status for the first time (not replay)
             if (!replayCompletedOperation.get()) {
-                fireOnOperationEnd(operation, null);
+                fireOnOperationEnd(operation, extractErrorFromOperation(operation));
             }
 
             markCompletionFutureCompleted();
@@ -522,4 +525,43 @@ private void fireOnOperationEnd(Operation operation, Throwable error) {
                 operation, operationIdentifier, durableContext.getParentId(), error);
         getPluginRunner().onOperationEnd(info);
     }
+
+    /**
+     * Extracts the error from a terminal operation as a Throwable. Returns null if the operation succeeded or has no
+     * error details.
+     */
+    private Throwable extractErrorFromOperation(Operation operation) {
+        if (operation.status() != OperationStatus.FAILED
+                && operation.status() != OperationStatus.TIMED_OUT
+                && operation.status() != OperationStatus.STOPPED) {
+            return null;
+        }
+        var errorObject = getErrorObject(operation);
+        if (errorObject == null) {
+            return null;
+        }
+        return new DurableOperationException(operation, errorObject);
+    }
+
+    /** Extracts the ErrorObject from an operation based on its type. */
+    private static ErrorObject getErrorObject(Operation operation) {
+        if (operation.type() == null) {
+            return null;
+        }
+        return switch (operation.type()) {
+            case STEP ->
+                operation.stepDetails() != null ? operation.stepDetails().error() : null;
+            case CHAINED_INVOKE ->
+                operation.chainedInvokeDetails() != null
+                        ? operation.chainedInvokeDetails().error()
+                        : null;
+            case CALLBACK ->
+                operation.callbackDetails() != null
+                        ? operation.callbackDetails().error()
+                        : null;
+            case CONTEXT ->
+                operation.contextDetails() != null ? operation.contextDetails().error() : null;
+            default -> null;
+        };
+    }
 }