feat: context aware replay status (#488)

- Add per-context replay status
- Add replay status for map and parallel branches
- Use per-context replay for hooks
- Add logging assertion for examples

---------

Co-authored-by: Alex Wang <wangyb@amazon.com>

Author: wangyb-A

packages/aws-durable-execution-sdk-python-examples/examples-catalog.json

@@ -280,6 +280,36 @@
       },
       "path": "./src/logger_example/logger_example.py"
     },
+    {
+      "name": "Replay Logging",
+      "description": "Demonstrating replay-aware logger de-duplication across a wait/replay boundary",
+      "handler": "replay_logging.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "loggingConfig": {
+        "ApplicationLogLevel": "INFO",
+        "LogFormat": "JSON"
+      },
+      "path": "./src/logger_example/replay_logging.py"
+    },
+    {
+      "name": "Replay Logging Concurrent",
+      "description": "Demonstrating per-branch replay-aware logger de-duplication across concurrent parallel branch waits",
+      "handler": "replay_logging_concurrent.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "loggingConfig": {
+        "ApplicationLogLevel": "INFO",
+        "LogFormat": "JSON"
+      },
+      "path": "./src/logger_example/replay_logging_concurrent.py"
+    },
     {
       "name": "Steps with Retry",
       "description": "Multiple steps with retry logic in a polling pattern",

packages/aws-durable-execution-sdk-python-examples/src/logger_example/replay_logging.py

@@ -0,0 +1,93 @@
+"""Example demonstrating replay-aware logging across a wait boundary."""
+
+from typing import Any
+
+from aws_durable_execution_sdk_python.config import Duration
+from aws_durable_execution_sdk_python.context import (
+    DurableContext,
+    StepContext,
+    durable_step,
+    durable_with_child_context,
+)
+from aws_durable_execution_sdk_python.execution import durable_execution
+
+
+@durable_step
+def prepare(step_context: StepContext, item: str) -> str:
+    """A step that runs before the wait.
+
+    Its log is emitted on the first invocation. On replay this step is not
+    re-executed (it returns its checkpointed result), so this log does not
+    repeat.
+    """
+    step_context.logger.info("Preparing item", extra={"item": item})
+    return f"prepared:{item}"
+
+
+@durable_step
+def finalize(step_context: StepContext, prepared: str) -> str:
+    """A step that runs after the wait (new work on the replay invocation)."""
+    step_context.logger.info("Finalizing item", extra={"prepared": prepared})
+    return f"done:{prepared}"
+
+
+@durable_with_child_context
+def audit(child_ctx: DurableContext, prepared: str) -> str:
+    """Child context with its own logger and its own replay status."""
+    child_ctx.logger.info(
+        "Auditing in child context (before child wait)",
+        extra={"prepared": prepared, "child_is_replaying": child_ctx.is_replaying()},
+    )
+
+    # The child's own replay boundary.
+    child_ctx.wait(duration=Duration.from_seconds(5), name="audit_cooldown")
+
+    # After the child's wait: emitted as new work on the child's replay.
+    child_ctx.logger.info(
+        "Resumed in child context (after child wait)",
+        extra={"child_is_replaying": child_ctx.is_replaying()},
+    )
+
+    return child_ctx.step(lambda _: f"audited:{prepared}", name="record_audit")
+
+
+@durable_execution
+def handler(event: Any, context: DurableContext) -> dict[str, Any]:
+    """Handler demonstrating replay-aware logging across a wait."""
+    item: str = event.get("item", "widget") if isinstance(event, dict) else "widget"
+
+    # --- Before the wait ---
+    # On the replay invocation these lines are de-duplicated by the replay-aware
+    # logger because the context is still replaying when it reaches them.
+    context.logger.info(
+        "Workflow started (before wait)",
+        extra={"item": item, "is_replaying": context.is_replaying()},
+    )
+
+    prepared: str = context.step(prepare(item), name="prepare")
+
+    context.logger.info(
+        "Prepared, about to wait",
+        extra={"prepared": prepared, "is_replaying": context.is_replaying()},
+    )
+
+    # --- The replay boundary ---
+    # The wait suspends the execution. When it resumes, the handler replays from
+    # the top; everything above is de-duplicated, and everything below is new.
+    context.wait(duration=Duration.from_seconds(5), name="cooldown")
+
+    # --- After the wait ---
+    # These logs are emitted on the replay invocation because the context has
+    # crossed its replay boundary and is no longer replaying.
+    context.logger.info(
+        "Resumed after wait",
+        extra={"is_replaying": context.is_replaying()},
+    )
+
+    audited: str = context.run_in_child_context(audit(prepared), name="audit")
+
+    result: str = context.step(finalize(audited), name="finalize")
+
+    context.logger.info("Workflow completed", extra={"result": result})
+
+    return {"result": result, "item": item}

packages/aws-durable-execution-sdk-python-examples/src/logger_example/replay_logging_concurrent.py

@@ -0,0 +1,83 @@
+"""Replay-aware logging across concurrent (parallel) branch replays.
+
+Each parallel branch runs in its own child context with its own replay status.
+A branch that contains a `wait` suspends and is replayed independently of the
+other branches and of the parent. This example demonstrates that the
+replay-aware logger de-duplicates per branch:
+
+- A branch's "before wait" log is emitted once (on the invocation where the
+  branch first reaches its wait) and de-duplicated on the branch's replay.
+- A branch's "after wait" log is emitted once, as new work, when that branch
+  resumes.
+- The branches stagger their waits so they resume on different invocations,
+  making the per-branch independence visible in CloudWatch.
+
+Deploy and invoke asynchronously, then inspect the logs. Each `branch` field
+identifies the emitting branch so you can confirm each line appears exactly
+once across all invocations.
+"""
+
+from typing import Any
+
+from aws_durable_execution_sdk_python.config import Duration, ParallelConfig
+from aws_durable_execution_sdk_python.context import (
+    DurableContext,
+    durable_parallel_branch,
+)
+from aws_durable_execution_sdk_python.execution import durable_execution
+
+
+def _make_branch(name: str, wait_seconds: int):
+    """Build a parallel branch that logs around its own wait boundary."""
+
+    @durable_parallel_branch(name=name)
+    def branch(ctx: DurableContext) -> str:
+        # Before the branch's wait. On this branch's replay invocation this is
+        # de-duplicated because the branch is still replaying when it reaches
+        # here. This line is the load-bearing case: there is no inner operation
+        # before it, so the branch's child context relies on starting in the
+        # correct replay status.
+        ctx.logger.info(
+            "branch start (before wait)",
+            extra={"branch": name, "is_replaying": ctx.is_replaying()},
+        )
+
+        # The branch's own replay boundary. Different per branch so the branches
+        # resume on different invocations.
+        ctx.wait(duration=Duration.from_seconds(wait_seconds), name=f"{name}_wait")
+
+        # After the branch's wait: emitted as new work on the branch's replay.
+        ctx.logger.info(
+            "branch resumed (after wait)",
+            extra={"branch": name, "is_replaying": ctx.is_replaying()},
+        )
+
+        return ctx.step(lambda _: f"{name}-done", name=f"{name}_finalize")
+
+    return branch
+
+
+@durable_execution
+def handler(event: Any, context: DurableContext) -> dict[str, Any]:
+    """Run concurrent branches that each log around their own wait."""
+    context.logger.info(
+        "workflow started",
+        extra={"is_replaying": context.is_replaying()},
+    )
+
+    results: list[str] = context.parallel(
+        functions=[
+            _make_branch("alpha", 3)(),
+            _make_branch("bravo", 6)(),
+            _make_branch("charlie", 9)(),
+        ],
+        name="concurrent_branches",
+        config=ParallelConfig(max_concurrency=3),
+    ).get_results()
+
+    context.logger.info(
+        "workflow completed",
+        extra={"results": results, "is_replaying": context.is_replaying()},
+    )
+
+    return {"results": results}

packages/aws-durable-execution-sdk-python-examples/template.yaml

@@ -492,6 +492,42 @@
         }
       }
     },
+    "ReplayLogging": {
+      "Type": "AWS::Serverless::Function",
+      "Properties": {
+        "CodeUri": "build/",
+        "Handler": "replay_logging.handler",
+        "Description": "Demonstrating replay-aware logger de-duplication across a wait/replay boundary",
+        "Role": {
+          "Fn::GetAtt": [
+            "DurableFunctionRole",
+            "Arn"
+          ]
+        },
+        "DurableConfig": {
+          "RetentionPeriodInDays": 7,
+          "ExecutionTimeout": 300
+        }
+      }
+    },
+    "ReplayLoggingConcurrent": {
+      "Type": "AWS::Serverless::Function",
+      "Properties": {
+        "CodeUri": "build/",
+        "Handler": "replay_logging_concurrent.handler",
+        "Description": "Demonstrating per-branch replay-aware logger de-duplication across concurrent parallel branch waits",
+        "Role": {
+          "Fn::GetAtt": [
+            "DurableFunctionRole",
+            "Arn"
+          ]
+        },
+        "DurableConfig": {
+          "RetentionPeriodInDays": 7,
+          "ExecutionTimeout": 300
+        }
+      }
+    },
     "StepsWithRetry": {
       "Type": "AWS::Serverless::Function",
       "Properties": {

packages/aws-durable-execution-sdk-python-examples/test/logger_example/test_logger_example.py

@@ -1,9 +1,11 @@
 """Tests for logger_example."""
 
+import logging
+
 import pytest
+
 from aws_durable_execution_sdk_python.execution import InvocationStatus
 from aws_durable_execution_sdk_python.lambda_service import OperationType
-
 from src.logger_example import logger_example
 from test.conftest import deserialize_operation_payload
 
@@ -33,3 +35,46 @@ def test_logger_example(durable_runner):
         op for op in result.operations if op.operation_type.value == "CONTEXT"
     ]
     assert len(context_ops) >= 1
+
+
+@pytest.mark.example
+@pytest.mark.durable_execution(
+    handler=logger_example.handler,
+    lambda_function_name="logger example",
+)
+def test_logger_example_emits_expected_logs(durable_runner, caplog):
+    """Verify the durable logger emits the expected messages and enriched extras.
+
+    The SDK logger wraps the stdlib root logger, so ``caplog`` captures every
+    record emitted from both the top-level context and from inside steps. Step
+    logs are enriched with the operation name and attempt number, which we
+    assert on via the LogRecord attributes.
+
+    Log capture only works in local mode (the handler runs in-process); in cloud
+    mode the handler runs in a deployed Lambda, so this test is skipped there.
+    """
+    if durable_runner.mode != "local":
+        pytest.skip("Log capture is only available in local (in-process) mode")
+
+    with caplog.at_level(logging.INFO):
+        with durable_runner:
+            result = durable_runner.run(input={"id": "test-123"}, timeout=10)
+
+    assert result.status is InvocationStatus.SUCCEEDED
+
+    messages = [record.getMessage() for record in caplog.records]
+    # Top-level context logs (no step_id) and in-step logs both appear.
+    assert "Starting workflow" in messages
+    assert "Hello from my_step" in messages
+    assert "Workflow completed" in messages
+
+    # The warning emitted inside my_step is enriched with the step's extra
+    # (my_arg) and with the operation name the SDK attaches automatically.
+    warning = next(
+        record
+        for record in caplog.records
+        if record.getMessage() == "Warning from my_step"
+    )
+    assert warning.levelno == logging.WARNING
+    assert warning.my_arg == 123
+    assert warning.operationName == "my_step"