fix(plugin): set OperationStartInfo start_time from operation (#559)

Author: nvasiu

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

@@ -4,10 +4,10 @@
 import datetime
 import functools
 import logging
+from collections.abc import Mapping, Sequence
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from enum import Enum
-from collections.abc import Mapping, Sequence
 from typing import Any, Callable, MutableMapping, cast
 
 from aws_durable_execution_sdk_python.identifier import OperationIdentifier
@@ -381,7 +381,9 @@ def on_user_function_end(self, start_info: UserFunctionStartInfo, error) -> None
             UserFunctionEndInfo.from_start_info(start_info, error), sync=True
         )
 
-    def on_operation_action(self, update: OperationUpdate):
+    def on_operation_action(
+        self, update: OperationUpdate, operation: Operation | None = None
+    ):
         """Execute any registered plugins for a given operation when an update is checkpointed
 
         Args:
@@ -398,7 +400,7 @@ def on_operation_action(self, update: OperationUpdate):
                     sub_type=update.sub_type,
                     name=update.name,
                     parent_id=update.parent_id,
-                    start_time=datetime.datetime.now(datetime.UTC),
+                    start_time=operation.start_timestamp if operation else None,
                     is_replayed=False,
                     status=OperationStatus.STARTED,
                 ),

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

@@ -751,7 +751,9 @@ def checkpoint_batches_forever(self) -> None:
                         output.new_execution_state.next_marker,
                     )
                     for update in updates:
-                        self._plugin_executor.on_operation_action(update)
+                        self._plugin_executor.on_operation_action(
+                            update, self.operations.get(update.operation_id)
+                        )
 
                     self._plugin_executor.on_operation_update(
                         updated_operations,

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

@@ -554,6 +554,36 @@ def on_operation_start(self, info: OperationStartInfo) -> None:
         self.assertIn("operation_start:op-1", self.plugin.calls)
         self.assertEqual(captured[0].status, OperationStatus.STARTED)
 
+    def test_start_action_uses_server_start_timestamp(self):
+        captured: list[OperationStartInfo] = []
+
+        class _CapturingPlugin(_TrackingPlugin):
+            def on_operation_start(self, info: OperationStartInfo) -> None:
+                super().on_operation_start(info)
+                captured.append(info)
+
+        self.plugin = _CapturingPlugin()
+        self.executor = PluginExecutor(plugins=[self.plugin])
+        update = MagicMock()
+        update.action = OperationAction.START
+        update.operation_id = "op-1"
+        update.operation_type = OperationType.STEP
+        update.sub_type = OperationSubType.STEP
+        update.name = "my-step"
+        update.parent_id = "parent-1"
+
+        operation = Operation(
+            operation_id="op-1",
+            operation_type=OperationType.STEP,
+            status=OperationStatus.STARTED,
+            start_timestamp=START_TS,
+        )
+
+        with self.executor.run():
+            self.executor.on_operation_action(update, operation)
+
+        self.assertEqual(captured[0].start_time, START_TS)
+
     def test_non_start_action_does_not_fire(self):
         update = MagicMock()
         update.action = OperationAction.SUCCEED