feat: Add replay mode detection for default logger

- Add visited_operation to logger
- Replace execution_arn with execution_state
- Add is_replay method so that we will not log when it is in replay mode

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.logger.visit_operation(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.logger.visit_operation(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.logger.visit_operation(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.logger.visit_operation(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.logger.visit_operation(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.logger.visit_operation(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.logger.visit_operation(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,14 @@ 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.logger.visit_operation(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 +528,15 @@ def wait_for_condition(
             msg = "`config` is required for wait_for_condition"
             raise ValidationError(msg)
 
+        operation_id = self._create_step_id()
+        self.logger.visit_operation(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/logger.py

@@ -5,29 +5,34 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
+from aws_durable_execution_sdk_python.lambda_service import OperationType
 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 +42,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,15 +52,23 @@ 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,
+        visited_operations: set[str] | None = None,
     ) -> None:
         self._logger = logger
         self._default_extra = default_extra
+        self._execution_state = execution_state
+        self._visited_operations = visited_operations or set()
 
     @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] = {
+            "execution_arn": info.execution_state.durable_execution_arn
+        }
         if info.parent_id:
             extra["parent_id"] = info.parent_id
         if info.name:
@@ -65,7 +78,9 @@ def from_log_info(cls, logger: LoggerInterface, info: LogInfo) -> Logger:
             extra["attempt"] = info.attempt + 1
         if info.operation_id:
             extra["operation_id"] = info.operation_id
-        return cls(logger, extra)
+        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 +96,52 @@ 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 visit_operation(self, operation_id: str):
+        self._visited_operations.add(operation_id)
+
+    def _is_replay(self) -> bool:
+        if not self._execution_state or not self._execution_state.operations:
+            return False
+
+        return any(
+            operation_id in self._visited_operations
+            for operation_id, operation in self._execution_state.operations.items()
+            if operation.operation_type != OperationType.EXECUTION
+        )
+
+    def _should_log(self) -> bool:
+        return not self._is_replay()

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,
                 )