fix: reject orphan checkpoint race

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

@@ -308,7 +308,8 @@ def __init__(
         # Operations whose parent has completed
         self._parent_done: set[str] = set()
 
-        # Protects parent_to_children and parent_done
+        # Protects parent_to_children and parent_done. When both state locks are
+        # required, acquire _completion_lock before _parent_done_lock.
         self._parent_done_lock: Lock = Lock()
 
         # Branch thread pools created by concurrency coordinators. A pool
@@ -538,6 +539,23 @@ def _reject_if_execution_completed(
             operation_id=operation_id,
         )
 
+    def _reject_if_parent_done(self, operation_update: OperationUpdate) -> None:
+        """Raise OrphanedChildException when the operation is orphaned.
+
+        Must be called while holding _parent_done_lock so the check can be
+        linearized with parent completion and, at the final call site, enqueue.
+        """
+        if operation_update.operation_id not in self._parent_done:
+            return
+        logger.debug(
+            "Rejecting checkpoint for operation %s - parent is done",
+            operation_update.operation_id,
+        )
+        raise OrphanedChildException(
+            "Parent context completed, child operation cannot checkpoint",
+            operation_id=operation_update.operation_id,
+        )
+
     def create_checkpoint(
         self,
         operation_update: OperationUpdate | None = None,
@@ -631,18 +649,7 @@ def create_checkpoint(
                     self._mark_orphans(operation_update.operation_id)
 
                 # Check if this operation's parent is done
-                if operation_update.operation_id in self._parent_done:
-                    logger.debug(
-                        "Rejecting checkpoint for operation %s - parent is done",
-                        operation_update.operation_id,
-                    )
-                    error_msg = (
-                        "Parent context completed, child operation cannot checkpoint"
-                    )
-                    raise OrphanedChildException(
-                        error_msg,
-                        operation_id=operation_update.operation_id,
-                    )
+                self._reject_if_parent_done(operation_update)
 
         # Check if background checkpointing has failed
         if self._checkpointing_failed.is_set():
@@ -675,14 +682,20 @@ def create_checkpoint(
         # drains the queue - on completion via _settle_after_execution_completed, or
         # on failure in the exception handler. Re-check both terminal conditions
         # inside the lock so a checkpoint is never enqueued after a drain and left
-        # with a waiter that blocks forever.
+        # with a waiter that blocks forever. Acquire _parent_done_lock inside
+        # _completion_lock and hold both through queue insertion so parent completion
+        # cannot mark the operation orphaned between the final check and enqueue.
         with self._completion_lock:
             if self._checkpointing_failed.is_set():
                 # Raises the stored BackgroundThreadError.
                 self._checkpointing_failed.wait()
             self._reject_if_execution_completed(operation_update)
-            # Enqueue the wrapper object (operation_update can be None for empty checkpoints)
-            self._checkpoint_queue.put(queued_op)
+            if operation_update is None:
+                self._checkpoint_queue.put(queued_op)
+            else:
+                with self._parent_done_lock:
+                    self._reject_if_parent_done(operation_update)
+                    self._checkpoint_queue.put(queued_op)
 
         # Conditionally wait for completion based on is_sync parameter
         if is_sync:

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

@@ -1426,6 +1426,87 @@ def test_rejection_of_operations_from_completed_parents():
     assert exc_info.value.operation_id == "child_1"
 
 
+def test_parent_done_recheck_rejects_checkpoint_racing_parent_completion():
+    """A checkpoint that passed validation is rejected if its parent completes.
+
+    The child pauses in the operation hook after releasing _parent_done_lock and
+    before acquiring _completion_lock. Completing the parent during that pause
+    deterministically reproduces the validation-to-enqueue race without sleeps.
+    """
+    mock_lambda_client = Mock(spec=LambdaClient)
+    mock_plugin_executor = Mock(spec=PluginExecutor)
+    state = ExecutionState(
+        durable_execution_arn="test_arn",
+        initial_checkpoint_token="token123",  # noqa: S106
+        operations={},
+        service_client=mock_lambda_client,
+        plugin_executor=mock_plugin_executor,
+    )
+
+    child_reached_hook = threading.Event()
+    release_child_hook = threading.Event()
+
+    def block_child_completion(operation_update: OperationUpdate, **_: object) -> None:
+        if (
+            operation_update.operation_id == "child_1"
+            and operation_update.action == OperationAction.SUCCEED
+        ):
+            child_reached_hook.set()
+            assert release_child_hook.wait(timeout=2.0), (
+                "child checkpoint was not released"
+            )
+
+    mock_plugin_executor.on_operation_action.side_effect = block_child_completion
+
+    state.create_checkpoint(
+        OperationUpdate(
+            operation_id="parent_1",
+            operation_type=OperationType.CONTEXT,
+            action=OperationAction.START,
+        ),
+        is_sync=False,
+    )
+    state.create_checkpoint(
+        OperationUpdate(
+            operation_id="child_1",
+            operation_type=OperationType.CONTEXT,
+            action=OperationAction.START,
+            parent_id="parent_1",
+        ),
+        is_sync=False,
+    )
+    child_complete = OperationUpdate(
+        operation_id="child_1",
+        operation_type=OperationType.CONTEXT,
+        action=OperationAction.SUCCEED,
+        parent_id="parent_1",
+    )
+
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        child_future = executor.submit(state.create_checkpoint, child_complete, False)
+        assert child_reached_hook.wait(timeout=2.0), (
+            "child checkpoint did not reach the hook"
+        )
+
+        try:
+            state.create_checkpoint(
+                OperationUpdate(
+                    operation_id="parent_1",
+                    operation_type=OperationType.CONTEXT,
+                    action=OperationAction.SUCCEED,
+                ),
+                is_sync=False,
+            )
+            assert "child_1" in state._parent_done
+        finally:
+            release_child_hook.set()
+
+        with pytest.raises(OrphanedChildException) as exc_info:
+            child_future.result(timeout=2.0)
+
+    assert exc_info.value.operation_id == "child_1"
+
+
 def test_nested_parallel_operations_deep_hierarchy():
     """Test that nested parallel operations handle deep hierarchies correctly."""
     mock_lambda_client = Mock(spec=LambdaClient)