fix(plugin): Fire onOperationStart for all non-terminal operations during replay

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

@@ -139,13 +139,12 @@ public void execute() {
                 validateReplay(existing);
                 if (ExecutionManager.isTerminalStatus(existing.status())) {
                     replayCompletedOperation.set(true);
-                } else if (getType() == OperationType.STEP || getType() == OperationType.CONTEXT) {
-                    // Non-terminal STEP/CONTEXT operations are being re-executed (user code runs again).
-                    // Fire onOperationStart so the OTel plugin can create a parent span for attempt spans.
+                } else {
+                    // Non-terminal operations are being replayed. Fire onOperationStart so plugins
+                    // can observe all in-progress operations during replay, including WAIT/INVOKE/CALLBACK
+                    // that are still pending.
                     fireOnOperationStart(existing);
                 }
-                // WAIT/INVOKE/CALLBACK in non-terminal status just poll — no onOperationStart needed.
-                // They'll get a continuation span via onOperationEnd when they complete.
 
                 // Fire onOperationEnd for operations that completed during suspension (between invocations).
                 // The OTel plugin handles the missing onOperationStart by creating a continuation span linked

sdk/src/test/java/software/amazon/lambda/durable/operation/BaseDurableOperationPluginTest.java

@@ -0,0 +1,166 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.operation;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState;
+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.lambda.durable.DurableConfig;
+import software.amazon.lambda.durable.TestUtils;
+import software.amazon.lambda.durable.context.DurableContextImpl;
+import software.amazon.lambda.durable.execution.ExecutionManager;
+import software.amazon.lambda.durable.execution.ThreadContext;
+import software.amazon.lambda.durable.execution.ThreadType;
+import software.amazon.lambda.durable.model.DurableExecutionInput;
+import software.amazon.lambda.durable.model.OperationIdentifier;
+import software.amazon.lambda.durable.model.OperationSubType;
+import software.amazon.lambda.durable.plugin.DurableExecutionPlugin;
+import software.amazon.lambda.durable.plugin.OperationInfo;
+
+/**
+ * Unit tests verifying that BaseDurableOperation.execute() fires onOperationStart with isReplay=true for all
+ * non-terminal operations during replay, regardless of operation type.
+ *
+ * <p>This mirrors the Python SDK's TestPluginExecutorOnOperationReplay tests.
+ */
+class BaseDurableOperationPluginTest {
+
+    private static final String EXECUTION_OP_ID = "exec-123";
+    private static final String EXECUTION_ARN =
+            "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/exec-name/" + EXECUTION_OP_ID;
+    private static final String OPERATION_ID = TestUtils.hashOperationId("1");
+    private static final String OPERATION_NAME = "test-op";
+
+    @Test
+    void execute_firesOnOperationStart_withIsReplayTrue_forNonTerminalWait() {
+        var plugin = new RecordingPlugin();
+        var waitOp = Operation.builder()
+                .id(OPERATION_ID)
+                .name(OPERATION_NAME)
+                .type(OperationType.WAIT)
+                .subType("Wait")
+                .status(OperationStatus.STARTED)
+                .build();
+
+        var executionManager = createExecutionManager(List.of(waitOp), plugin);
+        var durableContext = mockDurableContext(executionManager, plugin);
+
+        var operation = new WaitOperation(
+                OperationIdentifier.of(OPERATION_ID, OPERATION_NAME, OperationSubType.WAIT),
+                Duration.ofMinutes(5),
+                durableContext);
+
+        operation.execute();
+
+        // onOperationStart should fire with isReplay=true for the non-terminal WAIT
+        assertEquals(1, plugin.operationStarts.size(), "Should fire exactly one onOperationStart");
+        var info = plugin.operationStarts.get(0);
+        assertEquals(OPERATION_NAME, info.name());
+        assertEquals("WAIT", info.type());
+        assertTrue(info.isReplay(), "isReplay should be true for a replayed non-terminal wait");
+    }
+
+    @ParameterizedTest
+    @EnumSource(
+            value = OperationStatus.class,
+            names = {"SUCCEEDED", "FAILED", "TIMED_OUT", "STOPPED"})
+    void execute_doesNotFireOnOperationStart_forTerminalOperation(OperationStatus terminalStatus) {
+        var plugin = new RecordingPlugin();
+        var waitOp = Operation.builder()
+                .id(OPERATION_ID)
+                .name(OPERATION_NAME)
+                .type(OperationType.WAIT)
+                .subType("Wait")
+                .status(terminalStatus)
+                .build();
+
+        var executionManager = createExecutionManager(List.of(waitOp), plugin);
+        var durableContext = mockDurableContext(executionManager, plugin);
+
+        var operation = new WaitOperation(
+                OperationIdentifier.of(OPERATION_ID, OPERATION_NAME, OperationSubType.WAIT),
+                Duration.ofMinutes(5),
+                durableContext);
+
+        operation.execute();
+
+        // onOperationStart should NOT fire for terminal operations during replay
+        assertTrue(
+                plugin.operationStarts.isEmpty(),
+                "Should not fire onOperationStart for terminal " + terminalStatus + " operation");
+    }
+
+    @Test
+    void execute_firesOnOperationStart_withIsReplayFalse_forFirstExecution() {
+        var plugin = new RecordingPlugin();
+        // No existing operations — first execution
+        var executionManager = createExecutionManager(List.of(), plugin);
+        var durableContext = mockDurableContext(executionManager, plugin);
+
+        var operation = new WaitOperation(
+                OperationIdentifier.of(OPERATION_ID, OPERATION_NAME, OperationSubType.WAIT),
+                Duration.ofMinutes(5),
+                durableContext);
+
+        operation.execute();
+
+        // onOperationStart should fire with isReplay=false on first execution
+        assertEquals(1, plugin.operationStarts.size(), "Should fire exactly one onOperationStart");
+        var info = plugin.operationStarts.get(0);
+        assertEquals(OPERATION_NAME, info.name());
+        assertFalse(info.isReplay(), "isReplay should be false on first execution");
+    }
+
+    // ─── Helpers ─────────────────────────────────────────────────────────
+
+    private ExecutionManager createExecutionManager(List<Operation> additionalOps, RecordingPlugin plugin) {
+        var client = TestUtils.createMockClient();
+        var operations = new ArrayList<Operation>();
+        operations.add(Operation.builder()
+                .id(EXECUTION_OP_ID)
+                .type(OperationType.EXECUTION)
+                .status(OperationStatus.STARTED)
+                .build());
+        operations.addAll(additionalOps);
+        var initialState =
+                CheckpointUpdatedExecutionState.builder().operations(operations).build();
+        var config = DurableConfig.builder()
+                .withDurableExecutionClient(client)
+                .withPlugins(plugin)
+                .build();
+        var executionManager = new ExecutionManager(
+                new DurableExecutionInput(EXECUTION_ARN, "test-token", initialState), config, null);
+        executionManager.setCurrentThreadContext(new ThreadContext("Root", ThreadType.CONTEXT));
+        return executionManager;
+    }
+
+    private DurableContextImpl mockDurableContext(ExecutionManager executionManager, RecordingPlugin plugin) {
+        var durableContext = mock(DurableContextImpl.class);
+        when(durableContext.getExecutionManager()).thenReturn(executionManager);
+        when(durableContext.getDurableConfig())
+                .thenReturn(DurableConfig.builder().withPlugins(plugin).build());
+        return durableContext;
+    }
+
+    /** Plugin that records onOperationStart calls. */
+    private static class RecordingPlugin implements DurableExecutionPlugin {
+        final List<OperationInfo> operationStarts = Collections.synchronizedList(new ArrayList<>());
+
+        @Override
+        public void onOperationStart(OperationInfo info) {
+            operationStarts.add(info);
+        }
+    }
+}