fix(core): reject checkpoints at virtual child ids

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

@@ -51,7 +51,10 @@
     CallbackOperationExecutor,
     wait_for_callback_handler,
 )
-from aws_durable_execution_sdk_python.operation.child import child_handler
+from aws_durable_execution_sdk_python.operation.child import (
+    ChildOperationExecutor,
+    child_handler,
+)
 from aws_durable_execution_sdk_python.operation.invoke import InvokeOperationExecutor
 from aws_durable_execution_sdk_python.operation.map import map_handler
 from aws_durable_execution_sdk_python.operation.parallel import parallel_handler
@@ -515,12 +518,18 @@ def _next_operation_exists(self) -> bool:
 
     @contextmanager
     def _replay_aware(
-        self, operation_identifier: OperationIdentifier | None = None
+        self,
+        operation_identifier: OperationIdentifier | None = None,
+        *,
+        is_virtual: bool = False,
     ) -> Iterator[None]:
         """Wrap a single operation with replay-boundary detection.
 
         The operation kind is inferred from its own checkpoint (when one
-        exists), so no per-call-site flag is needed.
+        exists). The one per-call-site input is `is_virtual`: a virtual child
+        context never checkpoints itself, so with `is_virtual=True` any
+        checkpoint at its id is rejected right after the identity check,
+        before replay and update hooks run.
 
         The boundary has these parts:
 
@@ -555,6 +564,10 @@ def _replay_aware(
         if was_replaying and next_operation is not None:
             if operation_identifier is not None:
                 operation_identifier.validate_checkpoint(next_operation)
+                if is_virtual:
+                    ChildOperationExecutor.validate_virtual_checkpoint(
+                        operation_identifier, next_operation
+                    )
 
             # While replaying, an operation that already has a checkpoint was
             # observed in a prior invocation. If the backend says the operation
@@ -601,6 +614,7 @@ def _operation_replay_aware(
         name: str | None = None,
         *,
         operation_type: OperationType,
+        is_virtual: bool = False,
     ) -> Iterator[OperationIdentifier]:
         """Allocate an operation and validate its replay identity before hooks."""
         operation_identifier = OperationIdentifier(
@@ -610,7 +624,7 @@ def _operation_replay_aware(
             name=name,
             operation_type=operation_type,
         )
-        with self._replay_aware(operation_identifier):
+        with self._replay_aware(operation_identifier, is_virtual=is_virtual):
             yield operation_identifier
 
     # endregion replay status
@@ -814,14 +828,15 @@ def run_in_child_context(
             if config and config.sub_type
             else OperationSubType.RUN_IN_CHILD_CONTEXT
         )
+        is_virtual: bool = config.is_virtual if config else False
 
         with self._operation_replay_aware(
             sub_type,
             step_name,
             operation_type=OperationType.CONTEXT,
+            is_virtual=is_virtual,
         ) as operation_identifier:
             operation_id = operation_identifier.operation_id
-            is_virtual: bool = config.is_virtual if config else False
 
             def callable_with_child_context():
                 return func(

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

@@ -10,11 +10,13 @@
     ChildContextError,
     ExecutionError,
     InvocationError,
+    NonDeterministicExecutionError,
     SuspendExecution,
 )
 from aws_durable_execution_sdk_python.lambda_service import (
     ContextOptions,
     ErrorObject,
+    Operation,
     OperationStatus,
     OperationSubType,
     OperationUpdate,
@@ -71,6 +73,28 @@ def __init__(
         self.is_virtual: bool = config.is_virtual
         self.sub_type = config.sub_type or OperationSubType.RUN_IN_CHILD_CONTEXT
 
+    @staticmethod
+    def validate_virtual_checkpoint(
+        operation_identifier: OperationIdentifier,
+        checkpoint: Operation | None,
+    ) -> None:
+        """Reject a checkpoint at the operation id of a virtual child context.
+
+        A virtual child never checkpoints itself, so a checkpoint at its id was
+        written by a non-virtual child in an earlier invocation. Consuming it
+        would skip the virtual body and leave its inner operations unvalidated.
+        """
+        if not isinstance(checkpoint, Operation):
+            return
+        msg = (
+            "Non-deterministic child context at "
+            f"id={operation_identifier.operation_id!r}: checkpoint contains a "
+            "non-virtual child context but current child context is virtual"
+        )
+        raise NonDeterministicExecutionError(
+            msg, step_id=operation_identifier.operation_id
+        )
+
     def check_result_status(self) -> CheckResult[T]:
         """Check operation status and create START checkpoint if needed.
 
@@ -82,8 +106,14 @@ def check_result_status(self) -> CheckResult[T]:
 
         Raises:
             ChildContextError: For FAILED operations
+            NonDeterministicExecutionError: For a virtual child whose operation
+                id already has a checkpoint
         """
         checkpointed_result = self._get_checkpoint_result()
+        if self.is_virtual:
+            self.validate_virtual_checkpoint(
+                self.operation_identifier, checkpointed_result.operation
+            )
 
         # Terminal success without replay_children - deserialize and return
         if (

packages/aws-durable-execution-sdk-python/tests/context_test.py

@@ -3145,6 +3145,67 @@ def on_operation_end(self, info):
     assert step_body_calls == []
 
 
+@pytest.mark.parametrize(
+    ("checkpoint_status", "updated"),
+    [
+        (OperationStatus.STARTED, False),
+        (OperationStatus.SUCCEEDED, True),
+    ],
+)
+def test_virtual_child_checkpoint_is_rejected_before_replay_hooks(
+    checkpoint_status: OperationStatus,
+    updated: bool,
+):
+    """A virtual child fails on a container checkpoint before replay/update hooks."""
+    captured: list[str] = []
+
+    class _CapturingPlugin(DurableInstrumentationPlugin):
+        def on_operation_start(self, info):
+            captured.append(f"start:{info.operation_id}")
+
+        def on_operation_end(self, info):
+            captured.append(f"end:{info.operation_id}")
+
+    plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()])
+    child_body_calls: list[bool] = []
+    with plugin_executor.run():
+        state = ExecutionState(
+            durable_execution_arn="arn",
+            initial_checkpoint_token="token",  # noqa: S106
+            operations={},
+            service_client=Mock(),
+            plugin_executor=plugin_executor,
+            updated_operation_ids=[],
+        )
+        ctx = DurableContext(
+            state=state,
+            execution_context=ExecutionContext(durable_execution_arn="arn"),
+            replay_status=ReplayStatus.REPLAY,
+        )
+        next_id = ctx._peek_next_operation_id()  # noqa: SLF001
+        state._operations[next_id] = Operation(  # noqa: SLF001
+            operation_id=next_id,
+            operation_type=OperationType.CONTEXT,
+            status=checkpoint_status,
+            sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT,
+            name="child",
+        )
+        if updated:
+            state._updated_operation_ids.add(next_id)  # noqa: SLF001
+
+        with pytest.raises(
+            NonDeterministicExecutionError, match="current child context is virtual"
+        ):
+            ctx.run_in_child_context(
+                lambda _child_context: child_body_calls.append(True),
+                name="child",
+                config=ChildConfig(is_virtual=True),
+            )
+
+    assert captured == []
+    assert child_body_calls == []
+
+
 def test_replay_aware_emits_update_hook_for_operation_updated_since_last_invocation():
     """Updated terminal operations emit operation_end, not replay start+end."""
     captured: list[tuple[str, str, bool, OperationStatus]] = []

packages/aws-durable-execution-sdk-python/tests/execution_test.py

@@ -2957,6 +2957,86 @@ def test_handler(event: Any, context: DurableContext) -> str:
     assert child_body_calls == []
 
 
+@pytest.mark.parametrize(
+    "child_status",
+    [OperationStatus.SUCCEEDED, OperationStatus.FAILED, OperationStatus.STARTED],
+)
+def test_durable_execution_rejects_non_virtual_child_checkpoint_in_virtual_replay(
+    child_status: OperationStatus,
+):
+    """Replay rejects a child changed from non-virtual to virtual.
+
+    A virtual child never checkpoints itself, so a CONTEXT checkpoint at its id
+    comes from a non-virtual child. Replay fails before the cached result or
+    error is used and before the child body runs.
+    """
+    mock_client = Mock(spec=DurableServiceClient)
+    child_id = OperationIdNamespace().create_id_for_step(1)
+    execution_operation = Operation(
+        operation_id="exec1",
+        operation_type=OperationType.EXECUTION,
+        status=OperationStatus.STARTED,
+        execution_details=ExecutionDetails(input_payload="{}"),
+    )
+    child_operation = Operation(
+        operation_id=child_id,
+        operation_type=OperationType.CONTEXT,
+        status=child_status,
+        sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT,
+        name="child",
+        context_details=ContextDetails(
+            result=(
+                json.dumps("cached")
+                if child_status is OperationStatus.SUCCEEDED
+                else None
+            ),
+            error=(
+                ErrorObject(
+                    message="child failed",
+                    type="ValueError",
+                    data=None,
+                    stack_trace=None,
+                )
+                if child_status is OperationStatus.FAILED
+                else None
+            ),
+        ),
+    )
+    invocation_input = DurableExecutionInvocationInputWithClient(
+        durable_execution_arn="arn:test:execution/exec1",
+        checkpoint_token="token123",  # noqa: S106
+        initial_execution_state=InitialExecutionState(
+            operations=[execution_operation, child_operation],
+            next_marker="",
+        ),
+        service_client=mock_client,
+    )
+    child_body_calls: list[bool] = []
+
+    def child_body(_child: DurableContext) -> str:
+        child_body_calls.append(True)
+        return "executed"
+
+    @durable_execution
+    def test_handler(event: Any, context: DurableContext) -> str:
+        return context.run_in_child_context(
+            child_body,
+            name="child",
+            config=ChildConfig(is_virtual=True),
+        )
+
+    result = test_handler(invocation_input, _make_lambda_context())
+
+    assert result["Status"] == InvocationStatus.FAILED.value
+    assert (
+        result["Error"]["ErrorType"]
+        == "aws_durable_execution_sdk_python.exceptions.NonDeterministicExecutionError"
+    )
+    assert "current child context is virtual" in result["Error"]["ErrorMessage"]
+    assert child_body_calls == []
+    mock_client.checkpoint.assert_not_called()
+
+
 @pytest.mark.parametrize("operation_kind", ["map", "parallel"])
 @pytest.mark.parametrize("parent_replay_children", [False, True])
 @pytest.mark.parametrize(

packages/aws-durable-execution-sdk-python/tests/operation/child_test.py

@@ -15,21 +15,24 @@
     DurableApiErrorCategory,
     ExecutionError,
     InvocationError,
+    NonDeterministicExecutionError,
     RetryableSerDesError,
     SerDesError,
     StepError,
 )
 from aws_durable_execution_sdk_python.identifier import OperationIdentifier
 from aws_durable_execution_sdk_python.lambda_service import (
+    ContextDetails,
     ErrorObject,
+    Operation,
     OperationAction,
     OperationStatus,
     OperationSubType,
     OperationType,
 )
 from aws_durable_execution_sdk_python.operation.child import child_handler
 from aws_durable_execution_sdk_python.serdes import SerDes, SerDesContext
-from aws_durable_execution_sdk_python.state import ExecutionState
+from aws_durable_execution_sdk_python.state import CheckpointedResult, ExecutionState
 from aws_durable_execution_sdk_python.types import SummaryGenerator
 from tests.serdes_test import (
     CustomDictSerDes,
@@ -1315,6 +1318,81 @@ def setup_mocks():
     assert mock_state2.create_checkpoint.call_count == 0  # No checkpoints
 
 
+@pytest.mark.parametrize(
+    ("status", "replay_children"),
+    [
+        (OperationStatus.SUCCEEDED, False),
+        (OperationStatus.SUCCEEDED, True),
+        (OperationStatus.FAILED, False),
+        (OperationStatus.STARTED, False),
+    ],
+)
+def test_child_handler_is_virtual_rejects_container_checkpoint(
+    status: OperationStatus, replay_children: bool
+):
+    """Reject a CONTEXT checkpoint at the id of a virtual child.
+
+    A virtual child never checkpoints itself, so the checkpoint was written by
+    a non-virtual child in an earlier invocation. Returning its result or
+    re-raising its error would skip the virtual body, so replay must fail
+    before the checkpoint is used.
+    """
+    mock_state = Mock(spec=ExecutionState)
+    mock_state.durable_execution_arn = "test_arn"
+    mock_state.get_checkpoint_result.return_value = (
+        CheckpointedResult.create_from_operation(
+            Operation(
+                operation_id="op1",
+                operation_type=OperationType.CONTEXT,
+                status=status,
+                sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT,
+                name="test_name",
+                context_details=ContextDetails(
+                    replay_children=replay_children,
+                    result=(
+                        json.dumps("cached_result")
+                        if status is OperationStatus.SUCCEEDED
+                        else None
+                    ),
+                    error=(
+                        ErrorObject(
+                            message="Previous failure",
+                            type="ValueError",
+                            data=None,
+                            stack_trace=None,
+                        )
+                        if status is OperationStatus.FAILED
+                        else None
+                    ),
+                ),
+            )
+        )
+    )
+    mock_callable = Mock(return_value="fresh_result")
+
+    with pytest.raises(
+        NonDeterministicExecutionError, match="current child context is virtual"
+    ) as exc_info:
+        child_handler(
+            mock_callable,
+            mock_state,
+            OperationIdentifier(
+                "op1",
+                OperationSubType.RUN_IN_CHILD_CONTEXT,
+                None,
+                "test_name",
+                operation_type=OperationType.CONTEXT,
+            ),
+            ChildConfig(is_virtual=True),
+        )
+
+    assert exc_info.value.step_id == "op1"
+    mock_callable.assert_not_called()
+    mock_state.wrap_user_function.assert_not_called()
+    mock_state.create_checkpoint.assert_not_called()
+    mock_state.emit_child_context_end_hook.assert_not_called()
+
+
 # region first-run round-trip
 def test_child_handler_first_run_returns_round_tripped_result():
     """First-run result must equal the replay (deserialized-from-checkpoint) result.