fix(plugin): correct replay callbacks and operation span timing

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py

@@ -68,7 +68,7 @@ class InvocationOtelPlugin(DurableInstrumentationPlugin):
     Operation IDs are converted into deterministic span IDs. The first observed
     span for an operation uses that deterministic ID; later continuation spans
     use newly generated span IDs and link back to the deterministic span ID so
-    trace viewers can relate retries and replay-created terminal spans to the
+    trace viewers can relate retries and cross-invocation completions to the
     original logical operation.
 
     Args:
@@ -388,7 +388,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None:
             operation_id=info.operation_id,
             name=info.name or info.operation_id,
             attributes=attributes,
-            start_time=info.start_time,
+            start_time=datetime.datetime.now(datetime.UTC),
             parent_span=parent_span,
             existed=info.is_replayed,
         )
@@ -407,7 +407,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None:
             # Context operations are tracked using on_user_function_end.
             return
         span = self._get_span(info.operation_id)
-        if not span:
+        if span is None:
             # the span was not started in the current invocation, so we need to
             # create a new one that links to the previous one
             parent_span = self._resolve_parent_span(info.parent_id)
@@ -416,7 +416,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None:
                 operation_id=info.operation_id,
                 name=info.name or info.operation_id,
                 attributes=attributes,
-                start_time=info.start_time,
+                start_time=datetime.datetime.now(datetime.UTC),
                 parent_span=parent_span,
                 existed=True,
             )
@@ -431,10 +431,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None:
         else:
             span.set_status(StatusCode.OK)
 
-        end_timestamp = info.end_time
-        if end_timestamp is not None and end_timestamp == info.start_time:
-            end_timestamp += datetime.timedelta(microseconds=1)
-        self._end_span(info.operation_id, end_timestamp)
+        self._end_span(info.operation_id)
 
     def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
         """Called when a context or step operation starts user code.

packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py

@@ -209,7 +209,7 @@ def test_cross_invocation_operation_end_uses_deterministic_span_id():
     plugin, exporter = _create_plugin()
     plugin.on_invocation_start(_invocation_start_info())
 
-    # Operation end with no matching start (started in a prior invocation).
+    # Backend-updated completion for an operation started in a prior invocation.
     plugin.on_operation_end(
         OperationEndInfo(
             operation_id="step-earlier",
@@ -218,7 +218,7 @@ def test_cross_invocation_operation_end_uses_deterministic_span_id():
             name="earlier-step",
             parent_id=None,
             start_time=START_TIME,
-            is_replayed=True,
+            is_replayed=False,
             status=OperationStatus.SUCCEEDED,
             end_time=END_TIME,
             error=None,

packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import time
 from concurrent.futures import ThreadPoolExecutor
 from datetime import UTC, datetime
 
@@ -218,6 +219,9 @@ def test_operation_callbacks_emit_child_span_with_deterministic_span_id():
     )
     active_wait_span = plugin._get_span(operation_id)
     assert active_wait_span is not None
+    invocation_span = plugin._get_span(None)
+    assert invocation_span is not None
+    assert invocation_span.start_time <= active_wait_span.start_time
     assert (
         active_wait_span.attributes["durable.operation.status"]
         == OperationStatus.STARTED.value
@@ -250,6 +254,12 @@ def test_operation_callbacks_emit_child_span_with_deterministic_span_id():
         EXECUTION_ARN, operation_id
     )
     assert wait_span.parent.span_id == invocation_span.context.span_id
+    assert (
+        invocation_span.start_time
+        <= wait_span.start_time
+        <= wait_span.end_time
+        <= invocation_span.end_time
+    )
     assert wait_span.attributes["durable.operation.id"] == operation_id
     assert wait_span.attributes["durable.operation.type"] == OperationType.WAIT.value
     assert (
@@ -297,10 +307,13 @@ def test_operation_end_without_start_emits_continuation_span_with_link():
     )
 
 
-def test_continuation_span_uses_recorded_start_and_end_times():
-    """Continuation spans use the recorded operation start/end times."""
+def test_continuation_span_uses_current_start_and_end_times():
+    """Continuation spans use current times within the invocation."""
     plugin, exporter = _create_plugin()
     plugin.on_invocation_start(_invocation_start_info())
+    invocation_span = plugin._get_span(None)
+    assert invocation_span is not None
+    before_callback = time.time_ns()
 
     plugin.on_operation_end(
         OperationEndInfo(
@@ -316,59 +329,11 @@ def test_continuation_span_uses_recorded_start_and_end_times():
             error=None,
         )
     )
+    after_callback = time.time_ns()
 
     span = exporter.get_finished_spans()[0]
-    expected_start = int(START_TIME.timestamp() * 1_000_000_000)
-    expected_end = int(END_TIME.timestamp() * 1_000_000_000)
-    assert span.start_time == expected_start
-    assert span.end_time == expected_end
-    # Duration must be non-negative.
-    assert span.end_time >= span.start_time
-
-
-def test_replayed_operation_start_emits_continuation_span_with_link():
-    """Replayed operation spans should not reuse the original deterministic span ID."""
-    plugin, exporter = _create_plugin()
-    plugin.on_invocation_start(_invocation_start_info())
-    operation_id = "wait-replayed"
-    random_span_id = int("abcdef1234567890", 16)
-    plugin._id_generator._fallback_id_generator.generate_span_id = lambda: (
-        random_span_id
-    )
-
-    plugin.on_operation_start(
-        OperationStartInfo(
-            operation_id=operation_id,
-            operation_type=OperationType.WAIT,
-            sub_type=OperationSubType.WAIT,
-            name="replayed-wait",
-            parent_id=None,
-            start_time=START_TIME,
-            is_replayed=True,
-            status=OperationStatus.SUCCEEDED,
-        )
-    )
-    plugin.on_operation_end(
-        OperationEndInfo(
-            operation_id=operation_id,
-            operation_type=OperationType.WAIT,
-            sub_type=OperationSubType.WAIT,
-            name="replayed-wait",
-            parent_id=None,
-            start_time=START_TIME,
-            is_replayed=True,
-            status=OperationStatus.SUCCEEDED,
-            end_time=END_TIME,
-            error=None,
-        )
-    )
-
-    span = exporter.get_finished_spans()[0]
-    assert span.name == "replayed-wait"
-    assert span.context.span_id == random_span_id
-    assert span.links[0].context.span_id == operation_id_to_span_id(
-        EXECUTION_ARN, operation_id
-    )
+    assert invocation_span.start_time <= span.start_time
+    assert before_callback <= span.start_time <= span.end_time <= after_callback
 
 
 def test_retried_operation_start_emits_continuation_span_with_link():

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py

@@ -213,9 +213,9 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
     def on_operation_start(self, info: OperationStartInfo) -> None:
         """
         Called before an operation's START checkpoint is queued, or when a
-        prior operation is replayed. This guarantees that it strictly precedes
-        ``on_user_function_start``. This is called NOT within the thread that
-        runs operation.
+        prior non-terminal operation is replayed. This guarantees that it
+        strictly precedes ``on_user_function_start``. This is called NOT within
+        the thread that runs operation.
 
         Args:
             info: Information about the operation.
@@ -225,8 +225,9 @@ def on_operation_start(self, info: OperationStartInfo) -> None:
 
     def on_operation_end(self, info: OperationEndInfo) -> None:
         """
-        Called when an operation checkpoints a terminal status, or when a prior
-        terminal operation is replayed. This is called NOT within the thread that runs operation.
+        Called when an operation checkpoints a terminal status. Terminal
+        operations are not emitted again during replay. This is called NOT
+        within the thread that runs operation.
 
         Args:
             info: Information about the operation.
@@ -407,7 +408,10 @@ def on_operation_action(
             )
 
     def on_operation_replay(self, operation: Operation) -> None:
-        """Execute plugins for a checkpointed operation observed during replay."""
+        """Execute plugins for a non-terminal operation observed during replay."""
+        if self._is_terminal_status(operation.status):
+            return
+
         start_info = OperationStartInfo(
             operation_id=operation.operation_id,
             operation_type=operation.operation_type,
@@ -420,29 +424,6 @@ def on_operation_replay(self, operation: Operation) -> None:
         )
         self.execute_plugins(start_info, sync=True)
 
-        if self._is_terminal_status(operation.status):
-            self.execute_plugins(
-                OperationEndInfo(
-                    operation_id=operation.operation_id,
-                    operation_type=operation.operation_type,
-                    sub_type=operation.sub_type,
-                    name=operation.name,
-                    parent_id=operation.parent_id,
-                    start_time=operation.start_timestamp,
-                    end_time=operation.end_timestamp,
-                    result=_extract_result(operation),
-                    status=operation.status,
-                    error=self._extract_error(operation),
-                    attempt=(
-                        operation.step_details.attempt
-                        if operation.step_details
-                        else None
-                    ),
-                    is_replayed=True,
-                ),
-                sync=True,
-            )
-
     def on_operation_update(
         self,
         operation_or_operations: Operation | Sequence[Operation] | None,

packages/aws-durable-execution-sdk-python/tests/plugin_test.py

@@ -642,6 +642,48 @@ def test_fail_action_does_not_fire(self):
         self.assertEqual(self.plugin.calls, [])
 
 
+class TestPluginExecutorOnOperationReplay(unittest.TestCase):
+    """Tests for PluginExecutor.on_operation_replay."""
+
+    def test_terminal_operation_does_not_fire_callbacks(self):
+        terminal_statuses = [
+            OperationStatus.SUCCEEDED,
+            OperationStatus.FAILED,
+            OperationStatus.TIMED_OUT,
+            OperationStatus.CANCELLED,
+            OperationStatus.STOPPED,
+        ]
+
+        for status in terminal_statuses:
+            with self.subTest(status=status):
+                plugin = _TrackingPlugin()
+                executor = PluginExecutor(plugins=[plugin])
+                operation = Operation(
+                    operation_id="op-1",
+                    operation_type=OperationType.STEP,
+                    status=status,
+                )
+
+                with executor.run():
+                    executor.on_operation_replay(operation)
+
+                self.assertEqual(plugin.calls, [])
+
+    def test_non_terminal_operation_fires_operation_start(self):
+        plugin = _TrackingPlugin()
+        executor = PluginExecutor(plugins=[plugin])
+        operation = Operation(
+            operation_id="op-1",
+            operation_type=OperationType.WAIT,
+            status=OperationStatus.STARTED,
+        )
+
+        with executor.run():
+            executor.on_operation_replay(operation)
+
+        self.assertEqual(plugin.calls, ["operation_start:op-1"])
+
+
 class TestPluginExecutorOnOperationUpdate(unittest.TestCase):
     """Tests for PluginExecutor.on_operation_update."""
 

packages/aws-durable-execution-sdk-python/tests/state_test.py

@@ -4921,14 +4921,26 @@ def test_plugin_executor_not_called_for_pending_operations():
     assert len(operation_end_calls) == 0
 
 
-def test_emit_operation_replay_hook_fires_start_and_end_for_terminal_operation():
-    """emit_operation_replay_hook emits start+end (is_replayed=True) for terminal ops."""
+@pytest.mark.parametrize(
+    "terminal_status",
+    [
+        OperationStatus.SUCCEEDED,
+        OperationStatus.FAILED,
+        OperationStatus.TIMED_OUT,
+        OperationStatus.CANCELLED,
+        OperationStatus.STOPPED,
+    ],
+)
+def test_emit_operation_replay_hook_skips_terminal_operation(
+    terminal_status: OperationStatus,
+):
+    """Terminal operations do not emit plugin callbacks during replay."""
     start_time = datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC)
     end_time = datetime.datetime(2025, 1, 2, tzinfo=datetime.UTC)
     operation = Operation(
         operation_id="step-1",
         operation_type=OperationType.STEP,
-        status=OperationStatus.SUCCEEDED,
+        status=terminal_status,
         parent_id="parent-1",
         name="my-step",
         start_timestamp=start_time,
@@ -4959,10 +4971,7 @@ def on_operation_end(self, info):
         state.emit_operation_replay_hook(operation)
         state.emit_operation_replay_hook(operation)
 
-    assert captured == [
-        ("start", "step-1", True, OperationStatus.SUCCEEDED),
-        ("end", "step-1", True, OperationStatus.SUCCEEDED),
-    ]
+    assert captured == []
 
 
 def test_emit_operation_replay_hook_fires_only_start_for_non_terminal_operation():