test: add replay identity conformance case

packages/aws-durable-execution-sdk-python-conformance-tests/handlers/general/__init__.py

@@ -0,0 +1 @@
+"""General cross-cutting conformance handlers."""

packages/aws-durable-execution-sdk-python-conformance-tests/handlers/general/replay_operation_identity.py

@@ -0,0 +1,29 @@
+"""11-1: Replay rejects an operation-type mismatch."""
+
+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,
+)
+from aws_durable_execution_sdk_python.execution import durable_execution
+
+
+@durable_step
+def unexpected_step(step_context: StepContext) -> str:
+    step_context.logger.info("DETERMINISM_STEP_BODY_EXECUTED")
+    return "unexpected"
+
+
+@durable_execution
+def handler(_event: Any, context: DurableContext) -> str | None:
+    replay_logger = context.logger.with_is_replaying(lambda: False)
+
+    if context.is_replaying():
+        replay_logger.info("DETERMINISM_REPLAY_CANARY")
+        return context.step(unexpected_step(), name="identity-slot")
+
+    context.wait(Duration.from_seconds(1), name="identity-slot")
+    return None

packages/aws-durable-execution-sdk-python-conformance-tests/template_general.yaml

@@ -0,0 +1,49 @@
+AWSTemplateFormatVersion: '2010-09-09'
+Transform: AWS::Serverless-2016-10-31
+Globals:
+  Function:
+    Runtime: python3.13
+    Timeout: 60
+    MemorySize: 128
+    LoggingConfig:
+      LogFormat: JSON
+      ApplicationLogLevel: INFO
+      SystemLogLevel: INFO
+Resources:
+  DurableFunctionRole:
+    Type: AWS::IAM::Role
+    Properties:
+      AssumeRolePolicyDocument:
+        Version: '2012-10-17'
+        Statement:
+        - Effect: Allow
+          Principal:
+            Service: lambda.amazonaws.com
+          Action: sts:AssumeRole
+      ManagedPolicyArns:
+      - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
+      Policies:
+      - PolicyName: DurableExecutionPolicy
+        PolicyDocument:
+          Version: '2012-10-17'
+          Statement:
+          - Effect: Allow
+            Action:
+            - lambda:CheckpointDurableExecution
+            - lambda:GetDurableExecutionState
+            Resource: '*'
+  ReplayOperationIdentity:
+    Type: AWS::Serverless::Function
+    TestingMetadata:
+      TestDescription: ["11-1"]
+    Properties:
+      CodeUri: lambda-build/
+      Handler: general.replay_operation_identity.handler
+      Description: Replay rejects an operation-type mismatch
+      Role:
+        Fn::GetAtt:
+        - DurableFunctionRole
+        - Arn
+      DurableConfig:
+        RetentionPeriodInDays: 7
+        ExecutionTimeout: 300

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

@@ -31,8 +31,10 @@
 )
 from aws_durable_execution_sdk_python.exceptions import (
     DurableOperationError,
+    ExecutionError,
     InvalidStateError,
     InvocationError,
+    NonDeterministicExecutionError,
     OrphanedChildException,
     SuspendExecution,
     TimedSuspendExecution,
@@ -154,6 +156,59 @@ def get_iteration_name(self, index: int) -> str:
         name: str | None = self.executables[index].name
         return name if name is not None else f"{self.name_prefix}{index}"
 
+    def _get_iteration_operation_identifier(
+        self,
+        executor_context: DurableContext,
+        executable: Executable[CallableType],
+    ) -> OperationIdentifier:
+        """Build the stable operation identity for one branch or iteration."""
+        return OperationIdentifier(
+            operation_id=self.operation_id_namespace.create_id_for_step(
+                executable.index
+            ),
+            sub_type=self.sub_type_iteration,
+            parent_id=executor_context._parent_id,  # noqa: SLF001
+            name=self.get_iteration_name(executable.index),
+        )
+
+    def _validate_branch_checkpoint_nesting(
+        self,
+        operation_identifier: OperationIdentifier,
+        checkpoint: CheckpointedResult,
+    ) -> None:
+        """Validate branch identity and reject NESTED history in FLAT mode."""
+        operation_identifier.validate_checkpoint(checkpoint.operation)
+        if self.nesting_type is NestingType.FLAT and checkpoint.is_existent():
+            msg = (
+                "Non-deterministic branch nesting at "
+                f"id={operation_identifier.operation_id!r}: "
+                "checkpoint contains a NESTED branch context but current "
+                "nesting is FLAT"
+            )
+            raise NonDeterministicExecutionError(
+                msg, step_id=operation_identifier.operation_id
+            )
+
+    @staticmethod
+    def _find_flat_branch_child_operation_id(
+        execution_state: ExecutionState,
+        executor_context: DurableContext,
+        operation_identifier: OperationIdentifier,
+    ) -> str | None:
+        """Return a branch-namespaced operation parented directly to the batch."""
+        operations = execution_state.operations
+        branch_namespace = OperationIdNamespace(operation_identifier.operation_id)
+        expected_flat_parent_id = executor_context._parent_id  # noqa: SLF001
+        for logical_step in range(1, len(operations) + 1):
+            child_operation_id = branch_namespace.create_id_for_step(logical_step)
+            child_operation = operations.get(child_operation_id)
+            if (
+                child_operation is not None
+                and child_operation.parent_id == expected_flat_parent_id
+            ):
+                return child_operation_id
+        return None
+
     def _build_items_snapshot(self) -> tuple[CompletionItemStatus, ...]:
         """Build the per-branch status snapshot for the custom predicate.
 
@@ -239,7 +294,11 @@ def execute(
         def submit(branch: Branch[CallableType, ResultType]) -> None:
             branch.start()
             pool.submit(
-                self._branch_worker, executor_context, events, branch.executable
+                self._branch_worker,
+                execution_state,
+                executor_context,
+                events,
+                branch.executable,
             )
 
         try:
@@ -479,15 +538,17 @@ def _create_result(
 
     def _branch_worker(
         self,
+        execution_state: ExecutionState,
         executor_context: DurableContext,
         events: queue.Queue[BranchEvent[ResultType]],
         executable: Executable[CallableType],
     ) -> None:
         """Worker-thread body: run one branch and report its outcome.
 
-        Converts every outcome into a :class:`BranchEvent` on the queue and
-        never raises into the pool. The coordinator loop is the sole
-        consumer of the events.
+        Converts every outcome into a :class:`BranchEvent` on the queue. Fatal
+        errors are also re-raised into the pool after posting their event; the
+        coordinator loop consumes the event and propagates the error on the
+        calling thread.
         """
         try:
             result: ResultType = self._execute_item_in_child_context(
@@ -507,6 +568,22 @@ def _branch_worker(
                 executable.index,
             )
             events.put(BranchEvent.orphaned(executable.index))
+        except ExecutionError as e:
+            # Execution-terminal SDK errors (including nondeterminism) must
+            # bypass branch failure tolerance and custom completion policies.
+            parent_operation_id: str | None = executor_context._parent_id  # noqa: SLF001
+            if (
+                parent_operation_id is not None
+                and execution_state.record_branch_fatal_error(parent_operation_id, e)
+                is False
+            ):
+                logger.debug(
+                    "Ignoring fatal error from orphaned branch %s",
+                    executable.index,
+                )
+                return
+            events.put(BranchEvent.fatal(executable.index, e))
+            raise
         except Exception as e:  # noqa: BLE001
             # A retryable error (e.g. RetryableSerDesError) escapes the batch:
             # the coordinator re-raises it so the invocation fails and the
@@ -518,6 +595,17 @@ def _branch_worker(
             # Post a fatal event so the coordinator re-raises it on the
             # calling thread instead of blocking forever on the queue, then
             # let the exception propagate to the worker thread.
+            parent_operation_id = executor_context._parent_id  # noqa: SLF001
+            if (
+                parent_operation_id is not None
+                and execution_state.record_branch_fatal_error(parent_operation_id, e)
+                is False
+            ):
+                logger.debug(
+                    "Ignoring fatal error from orphaned branch %s",
+                    executable.index,
+                )
+                return
             events.put(BranchEvent.fatal(executable.index, e))
             raise
         else:
@@ -542,10 +630,10 @@ def _execute_item_in_child_context(
         and execution-order invariant.
         """
 
-        operation_id: str = self.operation_id_namespace.create_id_for_step(
-            executable.index
+        operation_identifier = self._get_iteration_operation_identifier(
+            executor_context, executable
         )
-        name: str = self.get_iteration_name(executable.index)
+        operation_id = operation_identifier.operation_id
         is_virtual: bool = self.nesting_type is NestingType.FLAT
 
         child_context: DurableContext = executor_context.create_child_context(
@@ -554,13 +642,6 @@ def _execute_item_in_child_context(
         # For NESTED this is for branch's START/SUCCEED/FAIL checkpoints (not the children of the branch).
         # For FLAT `child_handler` skips checkpoints, so not used.
         # Construct it unconditionally to keep the call simple.
-        operation_identifier = OperationIdentifier(
-            operation_id=operation_id,
-            sub_type=self.sub_type_iteration,
-            parent_id=executor_context._parent_id,  # noqa: SLF001
-            name=name,
-        )
-
         # The branch/iteration container op is resolved here via child_handler,
         # bypassing context.run_in_child_context and therefore the parent's
         # `_replay_aware`. Replicate the two things `_replay_aware` would have
@@ -571,13 +652,21 @@ def _execute_item_in_child_context(
         #      de-duplicated during a map/parallel replay.
         #   2. Replay hook: a branch that already has a checkpoint was observed
         #      in a prior invocation, so emit the plugin replay hook (once).
-        # Virtual (FLAT) branches do not checkpoint themselves, so neither
-        # applies; their inner operations still self-correct via `_replay_aware`.
-        if not is_virtual and child_context.is_replaying():
-            branch_checkpoint = child_context.state.get_checkpoint_result(operation_id)
-            if not branch_checkpoint.is_existent():
+        # Virtual (FLAT) branches do not checkpoint themselves. Therefore an
+        # existing branch-container checkpoint proves that replay changed from
+        # NESTED and must be rejected before child_handler can consume it.
+        if child_context.is_replaying():
+            branch_checkpoint = child_context.state.get_checkpoint_result(
+                operation_identifier.operation_id
+            )
+            if is_virtual:
+                self._validate_branch_checkpoint_nesting(
+                    operation_identifier, branch_checkpoint
+                )
+            elif not branch_checkpoint.is_existent():
                 child_context._set_replay_status_new()  # noqa: SLF001
             elif branch_checkpoint.operation is not None:
+                operation_identifier.validate_checkpoint(branch_checkpoint.operation)
                 child_context.state.emit_operation_replay_hook(
                     branch_checkpoint.operation
                 )
@@ -627,6 +716,35 @@ def replay(
             if executable.index >= record.started_total:
                 continue
             if executable.index in record.started_indexes:
+                operation_identifier = self._get_iteration_operation_identifier(
+                    executor_context, executable
+                )
+                checkpoint = execution_state.get_checkpoint_result(
+                    operation_identifier.operation_id
+                )
+                self._validate_branch_checkpoint_nesting(
+                    operation_identifier, checkpoint
+                )
+                if (
+                    self.nesting_type is NestingType.NESTED
+                    and not checkpoint.is_existent()
+                ):
+                    flat_child_operation_id = self._find_flat_branch_child_operation_id(
+                        execution_state,
+                        executor_context,
+                        operation_identifier,
+                    )
+                    if flat_child_operation_id is not None:
+                        msg = (
+                            "Non-deterministic branch nesting at "
+                            f"id={operation_identifier.operation_id!r}: "
+                            "recorded STARTED NESTED branch has no container "
+                            "checkpoint but contains a child operation from FLAT "
+                            f"history at id={flat_child_operation_id!r}"
+                        )
+                        raise NonDeterministicExecutionError(
+                            msg, step_id=operation_identifier.operation_id
+                        )
                 items.append(BatchItem(executable.index, BatchItemStatus.STARTED))
                 continue
             items.append(
@@ -650,12 +768,26 @@ def _replay_terminal_item(
         themselves, so re-executing the branch body over its inner
         operations' checkpoints discriminates success from failure.
         """
-        operation_id: str = self.operation_id_namespace.create_id_for_step(
-            executable.index
+        operation_identifier = self._get_iteration_operation_identifier(
+            executor_context, executable
         )
         checkpoint: CheckpointedResult = execution_state.get_checkpoint_result(
-            operation_id
+            operation_identifier.operation_id
         )
+        self._validate_branch_checkpoint_nesting(operation_identifier, checkpoint)
+        if self.nesting_type is NestingType.NESTED and not checkpoint.is_terminal():
+            checkpoint_status = (
+                checkpoint.status.value if checkpoint.status is not None else None
+            )
+            msg = (
+                "Non-deterministic branch nesting at "
+                f"id={operation_identifier.operation_id!r}: "
+                "recorded terminal branch requires a terminal NESTED branch "
+                f"context checkpoint, got status={checkpoint_status!r}"
+            )
+            raise NonDeterministicExecutionError(
+                msg, step_id=operation_identifier.operation_id
+            )
         if checkpoint.is_succeeded():
             result: ResultType = self._execute_item_in_child_context(
                 executor_context, executable
@@ -670,6 +802,10 @@ def _replay_terminal_item(
                 flat_result: ResultType = self._execute_item_in_child_context(
                     executor_context, executable
                 )
+            except ExecutionError:
+                # Nondeterminism and other execution-terminal SDK errors must
+                # not be downgraded to a failed FLAT item.
+                raise
             except Exception as e:  # noqa: BLE001
                 if isinstance(e, InvocationError) and e.is_retryable():
                     # Escape the batch so the invocation fails and the backend
@@ -694,10 +830,13 @@ def _replay_from_checkpoints(
         """
         items: list[BatchItem[ResultType]] = []
         for executable in self.executables:
-            operation_id = self.operation_id_namespace.create_id_for_step(
-                executable.index
+            operation_identifier = self._get_iteration_operation_identifier(
+                executor_context, executable
+            )
+            checkpoint = execution_state.get_checkpoint_result(
+                operation_identifier.operation_id
             )
-            checkpoint = execution_state.get_checkpoint_result(operation_id)
+            self._validate_branch_checkpoint_nesting(operation_identifier, checkpoint)
 
             result: ResultType | None = None
             error = None