fix(core): validate replay operation identity

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,21 @@ 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 _build_items_snapshot(self) -> tuple[CompletionItemStatus, ...]:
         """Build the per-branch status snapshot for the custom predicate.
 
@@ -239,7 +256,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 +500,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 +530,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 +557,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 +592,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 +604,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 +614,31 @@ 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:
+                if branch_checkpoint.is_existent():
+                    operation_identifier.validate_checkpoint(
+                        branch_checkpoint.operation
+                    )
+                    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
+                    )
+            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
                 )
@@ -650,12 +711,13 @@ 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
         )
+        operation_identifier.validate_checkpoint(checkpoint.operation)
         if checkpoint.is_succeeded():
             result: ResultType = self._execute_item_in_child_context(
                 executor_context, executable
@@ -670,6 +732,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 +760,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)
+            operation_identifier.validate_checkpoint(checkpoint.operation)
 
             result: ResultType | None = None
             error = None

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

@@ -5,7 +5,11 @@
 import hashlib
 from dataclasses import dataclass
 
+from aws_durable_execution_sdk_python.exceptions import (
+    NonDeterministicExecutionError,
+)
 from aws_durable_execution_sdk_python.lambda_service import (
+    Operation,
     OperationType,
     OperationSubType,
 )
@@ -40,3 +44,42 @@ class OperationIdentifier:
     @property
     def type(self) -> OperationType:
         return OperationType.from_sub_type(self.sub_type)
+
+    def validate_checkpoint(self, checkpoint: Operation | None) -> None:
+        """Ensure replay history belongs to this operation before it is consumed."""
+        if not isinstance(checkpoint, Operation):
+            return
+
+        expected_name = self.name or None
+        checkpoint_name = checkpoint.name or None
+        expected_parent_id = self.parent_id or None
+        checkpoint_parent_id = checkpoint.parent_id or None
+        mismatches: list[str] = []
+
+        if checkpoint.operation_type is not self.type:
+            mismatches.append(
+                f"type checkpoint={checkpoint.operation_type.value!r} current={self.type.value!r}"
+            )
+        if checkpoint.sub_type is not self.sub_type:
+            checkpoint_sub_type = (
+                checkpoint.sub_type.value if checkpoint.sub_type is not None else None
+            )
+            mismatches.append(
+                f"subtype checkpoint={checkpoint_sub_type!r} current={self.sub_type.value!r}"
+            )
+        if checkpoint_name != expected_name:
+            mismatches.append(
+                f"name checkpoint={checkpoint_name!r} current={expected_name!r}"
+            )
+        if checkpoint_parent_id != expected_parent_id:
+            mismatches.append(
+                f"parent_id checkpoint={checkpoint_parent_id!r} current={expected_parent_id!r}"
+            )
+
+        if mismatches:
+            mismatch_details = ", ".join(mismatches)
+            msg = (
+                "Non-deterministic operation identity at "
+                f"id={self.operation_id!r}: {mismatch_details}"
+            )
+            raise NonDeterministicExecutionError(msg, step_id=self.operation_id)

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

@@ -9,7 +9,11 @@
 from aws_durable_execution_sdk_python.exceptions import InvalidStateError
 
 if TYPE_CHECKING:
-    from aws_durable_execution_sdk_python.state import CheckpointedResult
+    from aws_durable_execution_sdk_python.identifier import OperationIdentifier
+    from aws_durable_execution_sdk_python.state import (
+        CheckpointedResult,
+        ExecutionState,
+    )
 
 T = TypeVar("T")
 
@@ -103,6 +107,17 @@ class OperationExecutor(ABC, Generic[T]):
     - execute(): Execute the operation logic with checkpoint data
     """
 
+    state: ExecutionState
+    operation_identifier: OperationIdentifier
+
+    def _get_checkpoint_result(self) -> CheckpointedResult:
+        """Return this operation's checkpoint after validating replay identity."""
+        checkpointed_result = self.state.get_checkpoint_result(
+            self.operation_identifier.operation_id
+        )
+        self.operation_identifier.validate_checkpoint(checkpointed_result.operation)
+        return checkpointed_result
+
     @abstractmethod
     def check_result_status(self) -> CheckResult[T]:
         """Check operation status and create START checkpoint if needed.

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

@@ -82,9 +82,7 @@ def check_result_status(self) -> CheckResult[str]:
         Raises:
             CallbackError: If callback_details are missing from checkpoint
         """
-        checkpointed_result: CheckpointedResult = self.state.get_checkpoint_result(
-            self.operation_identifier.operation_id
-        )
+        checkpointed_result = self._get_checkpoint_result()
 
         # CRITICAL: Do NOT raise on FAILED - defer error to Callback.result()
         # If checkpoint exists (any status including FAILED), return ready to execute

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

@@ -8,6 +8,7 @@
 from aws_durable_execution_sdk_python.config import ChildConfig
 from aws_durable_execution_sdk_python.exceptions import (
     ChildContextError,
+    ExecutionError,
     InvocationError,
     SuspendExecution,
 )
@@ -82,9 +83,7 @@ def check_result_status(self) -> CheckResult[T]:
         Raises:
             ChildContextError: For FAILED operations
         """
-        checkpointed_result: CheckpointedResult = self.state.get_checkpoint_result(
-            self.operation_identifier.operation_id
-        )
+        checkpointed_result = self._get_checkpoint_result()
 
         # Terminal success without replay_children - deserialize and return
         if (
@@ -271,6 +270,11 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T:
         except SuspendExecution:
             # Don't checkpoint SuspendExecution - let it bubble up
             raise
+        except ExecutionError:
+            # Execution-terminal SDK errors (including nondeterminism) must
+            # escape unchanged without mutating history or being wrapped as a
+            # child failure.
+            raise
         except Exception as e:
             # Retryable InvocationError: re-raise with no FAIL checkpoint so the
             # backend retry re-runs. Non-retryable falls through to FAIL + wrap.

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

@@ -84,9 +84,7 @@ def check_result_status(self) -> CheckResult[R]:
             InvokeError: For FAILED, TIMED_OUT, or STOPPED operations
             SuspendExecution: For STARTED operations waiting for completion
         """
-        checkpointed_result: CheckpointedResult = self.state.get_checkpoint_result(
-            self.operation_identifier.operation_id
-        )
+        checkpointed_result = self._get_checkpoint_result()
 
         # Terminal success - deserialize and return
         if checkpointed_result.is_succeeded():

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

@@ -124,6 +124,7 @@ def map_handler(
     checkpoint: CheckpointedResult = execution_state.get_checkpoint_result(
         operation_identifier.operation_id
     )
+    operation_identifier.validate_checkpoint(checkpoint.operation)
     if checkpoint.is_succeeded():
         # if we've reached this point, then not only is the step succeeded, but it is also `replay_children`.
         return executor.replay(execution_state, map_context, checkpoint)

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

@@ -114,6 +114,7 @@ def parallel_handler(
     checkpoint = execution_state.get_checkpoint_result(
         operation_identifier.operation_id
     )
+    operation_identifier.validate_checkpoint(checkpoint.operation)
     if checkpoint.is_succeeded():
         return executor.replay(execution_state, parallel_context, checkpoint)
     return executor.execute(execution_state, executor_context=parallel_context)

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

@@ -93,9 +93,7 @@ def check_result_status(self) -> CheckResult[T]:
             StepInterruptedError: For interrupted AT_MOST_ONCE operations
             SuspendExecution: For PENDING operations waiting for retry
         """
-        checkpointed_result: CheckpointedResult = self.state.get_checkpoint_result(
-            self.operation_identifier.operation_id
-        )
+        checkpointed_result = self._get_checkpoint_result()
 
         # Terminal success - deserialize and return
         if checkpointed_result.is_succeeded():
@@ -175,9 +173,7 @@ def check_result_status(self) -> CheckResult[T]:
             # After creating sync checkpoint, check the status
             if is_sync:
                 # Refresh checkpoint result to check for immediate response
-                refreshed_result: CheckpointedResult = self.state.get_checkpoint_result(
-                    self.operation_identifier.operation_id
-                )
+                refreshed_result = self._get_checkpoint_result()
 
                 # START checkpoint only returns STARTED status
                 # Any errors would be thrown as runtime exceptions during checkpoint creation