feat: add replay aware state and logging

- check if it is replaying in ExecutionState
- add should_replay logging in default logger

feat: update logging extra field to CamelCase

Author: Alex Wang

src/aws_durable_execution_sdk_python/concurrency/executor.py

@@ -381,6 +381,7 @@ def _execute_item_in_child_context(
             executor_context._parent_id,  # noqa: SLF001
             name,
         )
+        child_context.state.track_replay(operation_id=operation_id)
 
         def run_in_child_handler():
             return self.execute_item(child_context, executable)

src/aws_durable_execution_sdk_python/context.py

@@ -176,7 +176,8 @@ def __init__(
         self._step_counter: OrderedCounter = OrderedCounter()
 
         log_info = LogInfo(
-            execution_arn=state.durable_execution_arn, parent_id=parent_id
+            execution_state=state,
+            parent_id=parent_id,
         )
         self._log_info = log_info
         self.logger: Logger = logger or Logger.from_log_info(
@@ -205,7 +206,8 @@ def create_child_context(self, parent_id: str) -> DurableContext:
             parent_id=parent_id,
             logger=self.logger.with_log_info(
                 LogInfo(
-                    execution_arn=self.state.durable_execution_arn, parent_id=parent_id
+                    execution_state=self.state,
+                    parent_id=parent_id,
                 )
             ),
         )
@@ -269,6 +271,7 @@ def create_callback(
         if not config:
             config = CallbackConfig()
         operation_id: str = self._create_step_id()
+        self.state.track_replay(operation_id=operation_id)
         callback_id: str = create_callback_handler(
             state=self.state,
             operation_identifier=OperationIdentifier(
@@ -302,12 +305,14 @@ def invoke(
         Returns:
             The result of the invoked function
         """
+        operation_id = self._create_step_id()
+        self.state.track_replay(operation_id=operation_id)
         return invoke_handler(
             function_name=function_name,
             payload=payload,
             state=self.state,
             operation_identifier=OperationIdentifier(
-                operation_id=self._create_step_id(),
+                operation_id=operation_id,
                 parent_id=self._parent_id,
                 name=name,
             ),
@@ -325,6 +330,7 @@ def map(
         map_name: str | None = self._resolve_step_name(name, func)
 
         operation_id = self._create_step_id()
+        self.state.track_replay(operation_id=operation_id)
         operation_identifier = OperationIdentifier(
             operation_id=operation_id, parent_id=self._parent_id, name=map_name
         )
@@ -367,6 +373,7 @@ def parallel(
         """Execute multiple callables in parallel."""
         # _create_step_id() is thread-safe. rest of method is safe, since using local copy of parent id
         operation_id = self._create_step_id()
+        self.state.track_replay(operation_id=operation_id)
         parallel_context = self.create_child_context(parent_id=operation_id)
         operation_identifier = OperationIdentifier(
             operation_id=operation_id, parent_id=self._parent_id, name=name
@@ -420,6 +427,7 @@ def run_in_child_context(
         step_name: str | None = self._resolve_step_name(name, func)
         # _create_step_id() is thread-safe. rest of method is safe, since using local copy of parent id
         operation_id = self._create_step_id()
+        self.state.track_replay(operation_id=operation_id)
 
         def callable_with_child_context():
             return func(self.create_child_context(parent_id=operation_id))
@@ -441,13 +449,15 @@ def step(
     ) -> T:
         step_name = self._resolve_step_name(name, func)
         logger.debug("Step name: %s", step_name)
+        operation_id = self._create_step_id()
+        self.state.track_replay(operation_id=operation_id)
 
         return step_handler(
             func=func,
             config=config,
             state=self.state,
             operation_identifier=OperationIdentifier(
-                operation_id=self._create_step_id(),
+                operation_id=operation_id,
                 parent_id=self._parent_id,
                 name=step_name,
             ),
@@ -465,11 +475,13 @@ def wait(self, duration: Duration, name: str | None = None) -> None:
         if seconds < 1:
             msg = "duration must be at least 1 second"
             raise ValidationError(msg)
+        operation_id = self._create_step_id()
+        self.state.track_replay(operation_id=operation_id)
         wait_handler(
             seconds=seconds,
             state=self.state,
             operation_identifier=OperationIdentifier(
-                operation_id=self._create_step_id(),
+                operation_id=operation_id,
                 parent_id=self._parent_id,
                 name=name,
             ),
@@ -515,12 +527,14 @@ def wait_for_condition(
             msg = "`config` is required for wait_for_condition"
             raise ValidationError(msg)
 
+        operation_id = self._create_step_id()
+        self.state.track_replay(operation_id=operation_id)
         return wait_for_condition_handler(
             check=check,
             config=config,
             state=self.state,
             operation_identifier=OperationIdentifier(
-                operation_id=self._create_step_id(),
+                operation_id=operation_id,
                 parent_id=self._parent_id,
                 name=name,
             ),

src/aws_durable_execution_sdk_python/execution.py

@@ -9,7 +9,7 @@
 from enum import Enum
 from typing import TYPE_CHECKING, Any
 
-from aws_durable_execution_sdk_python.context import DurableContext, ExecutionState
+from aws_durable_execution_sdk_python.context import DurableContext
 from aws_durable_execution_sdk_python.exceptions import (
     BackgroundThreadError,
     BotoClientError,
@@ -27,6 +27,7 @@
     OperationType,
     OperationUpdate,
 )
+from aws_durable_execution_sdk_python.state import ExecutionState, ReplayStatus
 
 if TYPE_CHECKING:
     from collections.abc import Callable, MutableMapping
@@ -268,6 +269,9 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
             initial_checkpoint_token=invocation_input.checkpoint_token,
             operations={},
             service_client=service_client,
+            replay_status=ReplayStatus.REPLAY
+            if len(invocation_input.initial_execution_state.operations) > 0
+            else ReplayStatus.NEW,
         )
 
         execution_state.fetch_paginated_operations(

src/aws_durable_execution_sdk_python/logger.py

@@ -8,26 +8,30 @@
 from aws_durable_execution_sdk_python.types import LoggerInterface
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping, MutableMapping
+    from collections.abc import Callable, Mapping, MutableMapping
 
+    from aws_durable_execution_sdk_python.context import ExecutionState
     from aws_durable_execution_sdk_python.identifier import OperationIdentifier
 
 
 @dataclass(frozen=True)
 class LogInfo:
-    execution_arn: str
+    execution_state: ExecutionState
     parent_id: str | None = None
     operation_id: str | None = None
     name: str | None = None
     attempt: int | None = None
 
     @classmethod
     def from_operation_identifier(
-        cls, execution_arn: str, op_id: OperationIdentifier, attempt: int | None = None
+        cls,
+        execution_state: ExecutionState,
+        op_id: OperationIdentifier,
+        attempt: int | None = None,
     ) -> LogInfo:
         """Create new log info from an execution arn, OperationIdentifier and attempt."""
         return cls(
-            execution_arn=execution_arn,
+            execution_state=execution_state,
             parent_id=op_id.parent_id,
             operation_id=op_id.operation_id,
             name=op_id.name,
@@ -37,7 +41,7 @@ def from_operation_identifier(
     def with_parent_id(self, parent_id: str) -> LogInfo:
         """Clone the log info with a new parent id."""
         return LogInfo(
-            execution_arn=self.execution_arn,
+            execution_state=self.execution_state,
             parent_id=parent_id,
             operation_id=self.operation_id,
             name=self.name,
@@ -47,25 +51,33 @@ def with_parent_id(self, parent_id: str) -> LogInfo:
 
 class Logger(LoggerInterface):
     def __init__(
-        self, logger: LoggerInterface, default_extra: Mapping[str, object]
+        self,
+        logger: LoggerInterface,
+        default_extra: Mapping[str, object],
+        execution_state: ExecutionState,
     ) -> None:
         self._logger = logger
         self._default_extra = default_extra
+        self._execution_state = execution_state
 
     @classmethod
     def from_log_info(cls, logger: LoggerInterface, info: LogInfo) -> Logger:
         """Create a new logger with the given LogInfo."""
-        extra: MutableMapping[str, object] = {"execution_arn": info.execution_arn}
+        extra: MutableMapping[str, object] = {
+            "executionArn": info.execution_state.durable_execution_arn
+        }
         if info.parent_id:
-            extra["parent_id"] = info.parent_id
+            extra["parentId"] = info.parent_id
         if info.name:
             # Use 'operation_name' instead of 'name' as key because the stdlib LogRecord internally reserved 'name' parameter
-            extra["operation_name"] = info.name
+            extra["operationName"] = info.name
         if info.attempt is not None:
             extra["attempt"] = info.attempt + 1
         if info.operation_id:
-            extra["operation_id"] = info.operation_id
-        return cls(logger, extra)
+            extra["operationId"] = info.operation_id
+        return cls(
+            logger=logger, default_extra=extra, execution_state=info.execution_state
+        )
 
     def with_log_info(self, info: LogInfo) -> Logger:
         """Clone the existing logger with new LogInfo."""
@@ -81,29 +93,39 @@ def get_logger(self) -> LoggerInterface:
     def debug(
         self, msg: object, *args: object, extra: Mapping[str, object] | None = None
     ) -> None:
-        merged_extra = {**self._default_extra, **(extra or {})}
-        self._logger.debug(msg, *args, extra=merged_extra)
+        self._log(self._logger.debug, msg, *args, extra=extra)
 
     def info(
         self, msg: object, *args: object, extra: Mapping[str, object] | None = None
     ) -> None:
-        merged_extra = {**self._default_extra, **(extra or {})}
-        self._logger.info(msg, *args, extra=merged_extra)
+        self._log(self._logger.info, msg, *args, extra=extra)
 
     def warning(
         self, msg: object, *args: object, extra: Mapping[str, object] | None = None
     ) -> None:
-        merged_extra = {**self._default_extra, **(extra or {})}
-        self._logger.warning(msg, *args, extra=merged_extra)
+        self._log(self._logger.warning, msg, *args, extra=extra)
 
     def error(
         self, msg: object, *args: object, extra: Mapping[str, object] | None = None
     ) -> None:
-        merged_extra = {**self._default_extra, **(extra or {})}
-        self._logger.error(msg, *args, extra=merged_extra)
+        self._log(self._logger.error, msg, *args, extra=extra)
 
     def exception(
         self, msg: object, *args: object, extra: Mapping[str, object] | None = None
     ) -> None:
+        self._log(self._logger.exception, msg, *args, extra=extra)
+
+    def _log(
+        self,
+        log_func: Callable,
+        msg: object,
+        *args: object,
+        extra: Mapping[str, object] | None = None,
+    ):
+        if not self._should_log():
+            return
         merged_extra = {**self._default_extra, **(extra or {})}
-        self._logger.exception(msg, *args, extra=merged_extra)
+        log_func(msg, *args, extra=merged_extra)
+
+    def _should_log(self) -> bool:
+        return not self._execution_state.is_replaying()

src/aws_durable_execution_sdk_python/operation/step.py

@@ -128,7 +128,7 @@ def step_handler(
     step_context = StepContext(
         logger=context_logger.with_log_info(
             LogInfo.from_operation_identifier(
-                execution_arn=state.durable_execution_arn,
+                execution_state=state,
                 op_id=operation_identifier,
                 attempt=attempt,
             )

src/aws_durable_execution_sdk_python/operation/wait_for_condition.py

@@ -133,7 +133,7 @@ def wait_for_condition_handler(
         check_context = WaitForConditionCheckContext(
             logger=context_logger.with_log_info(
                 LogInfo.from_operation_identifier(
-                    execution_arn=state.durable_execution_arn,
+                    execution_state=state,
                     op_id=operation_identifier,
                     attempt=attempt,
                 )

src/aws_durable_execution_sdk_python/state.py

@@ -8,6 +8,7 @@
 import threading
 import time
 from dataclasses import dataclass
+from enum import Enum
 from threading import Lock
 from typing import TYPE_CHECKING
 
@@ -210,6 +211,13 @@ def get_next_attempt_timestamp(self) -> datetime.datetime | None:
 CHECKPOINT_NOT_FOUND = CheckpointedResult.create_not_found()
 
 
+class ReplayStatus(Enum):
+    """Status indicating whether execution is replaying or executing new operations."""
+
+    REPLAY = "replay"
+    NEW = "new"
+
+
 class ExecutionState:
     """Get, set and maintain execution state. This is mutable. Create and check checkpoints."""
 
@@ -220,6 +228,7 @@ def __init__(
         operations: MutableMapping[str, Operation],
         service_client: DurableServiceClient,
         batcher_config: CheckpointBatcherConfig | None = None,
+        replay_status: ReplayStatus = ReplayStatus.NEW,
     ):
         self.durable_execution_arn: str = durable_execution_arn
         self._current_checkpoint_token: str = initial_checkpoint_token
@@ -247,6 +256,8 @@ def __init__(
 
         # Protects parent_to_children and parent_done
         self._parent_done_lock: Lock = Lock()
+        self._replay_status: ReplayStatus = replay_status
+        self._replay_status_lock: Lock = Lock()
 
     def fetch_paginated_operations(
         self,
@@ -277,6 +288,42 @@ def fetch_paginated_operations(
         with self._operations_lock:
             self.operations.update({op.operation_id: op for op in all_operations})
 
+    def track_replay(self, operation_id: str) -> None:
+        """Check if operation exists with completed status; if not, transition to NEW status.
+
+        This method is called before each operation (step, wait, invoke, etc.) to determine
+        if we've reached the replay boundary. Once we encounter an operation that doesn't
+        exist or isn't completed, we transition from REPLAY to NEW status, which enables
+        logging for all subsequent code.
+
+        Args:
+            operation_id: The operation ID to check
+        """
+        with self._replay_status_lock:
+            if self._replay_status == ReplayStatus.REPLAY:
+                operation = self.operations.get(operation_id)
+                # Transition if operation doesn't exist OR isn't in a completed state
+                if not operation or operation.status not in {
+                    OperationStatus.SUCCEEDED,
+                    OperationStatus.FAILED,
+                    OperationStatus.CANCELLED,
+                    OperationStatus.STOPPED,
+                }:
+                    logger.debug(
+                        "Transitioning from REPLAY to NEW status at operation %s",
+                        operation_id,
+                    )
+                    self._replay_status = ReplayStatus.NEW
+
+    def is_replaying(self) -> bool:
+        """Check if execution is currently in replay mode.
+
+        Returns:
+            True if in REPLAY status, False if in NEW status
+        """
+        with self._replay_status_lock:
+            return self._replay_status is ReplayStatus.REPLAY
+
     def get_checkpoint_result(self, checkpoint_id: str) -> CheckpointedResult:
         """Get checkpoint result.
 

tests/e2e/execution_int_test.py

@@ -217,10 +217,10 @@ def mock_checkpoint(
             123,
             "str",
             extra={
-                "execution_arn": "test-arn",
-                "operation_name": "mystep",
+                "executionArn": "test-arn",
+                "operationName": "mystep",
                 "attempt": 1,
-                "operation_id": operation_id,
+                "operationId": operation_id,
             },
         )