test: add error preservation regression tests

packages/aws-durable-execution-sdk-python/tests/e2e/checkpoint_response_int_test.py

@@ -11,9 +11,13 @@
 
 import pytest
 
+from aws_durable_execution_sdk_python.concurrency.models import BatchResult
 from aws_durable_execution_sdk_python.config import ChildConfig, Duration
 from aws_durable_execution_sdk_python.context import DurableContext, durable_step
-from aws_durable_execution_sdk_python.exceptions import InvocationError
+from aws_durable_execution_sdk_python.exceptions import (
+    ChildContextError,
+    InvocationError,
+)
 from aws_durable_execution_sdk_python.execution import (
     InvocationStatus,
     durable_execution,
@@ -762,3 +766,75 @@ def my_handler(event, context: DurableContext) -> str:
             if hasattr(op, "action") and op.action.value == "FAIL"
         ]
         assert len(fail_updates) == 0
+
+
+def test_parallel_and_child_context_agree_on_custom_error_type():
+    """ctx.parallel() and ctx.run_in_child_context() surface the same
+    error_type for an identical custom exception: the caller's original type,
+    not the ChildContextError wrapper class name.
+    """
+
+    class PermanentFailure(Exception):
+        pass
+
+    def branch(child_context: DurableContext) -> str:
+        msg: str = "Invalid input data"
+        raise PermanentFailure(msg)
+
+    captured: dict[str, str | None] = {}
+
+    @durable_execution
+    def my_handler(event, context: DurableContext) -> str:
+        try:
+            context.run_in_child_context(branch)
+        except ChildContextError as e:
+            captured["child"] = e.error_type
+
+        result: BatchResult[str] = context.parallel([branch])
+        try:
+            result.throw_if_error()
+        except ChildContextError as e:
+            captured["parallel"] = e.error_type
+
+        return "handled"
+
+    with patch(
+        "aws_durable_execution_sdk_python.execution.LambdaClient"
+    ) as mock_client_class:
+        mock_client = Mock()
+        mock_client_class.initialize_client.return_value = mock_client
+
+        mock_checkpoint, _ = create_mock_checkpoint_with_operations()
+        mock_client.checkpoint = mock_checkpoint
+
+        event = {
+            "DurableExecutionArn": "test-arn/execution-1",
+            "CheckpointToken": "test-token",
+            "InitialExecutionState": {
+                "Operations": [
+                    {
+                        "Id": "execution-1",
+                        "Type": "EXECUTION",
+                        "Status": "STARTED",
+                        "ExecutionDetails": {"InputPayload": "{}"},
+                    }
+                ],
+                "NextMarker": "",
+            },
+            "LocalRunner": True,
+        }
+
+        lambda_context = Mock()
+        lambda_context.aws_request_id = "test-request-id"
+        lambda_context.client_context = None
+        lambda_context.identity = None
+        lambda_context._epoch_deadline_time_in_ms = 0  # noqa: SLF001
+        lambda_context.invoked_function_arn = "test-arn"
+        lambda_context.tenant_id = None
+
+        my_handler(event, lambda_context)
+
+    # Both operations surface the caller's original type, and they agree.
+    assert captured["child"] == "PermanentFailure"
+    assert captured["parallel"] == "PermanentFailure"
+    assert captured["child"] == captured["parallel"]

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

@@ -15,6 +15,7 @@
     DurableApiErrorCategory,
     ExecutionError,
     InvocationError,
+    StepError,
 )
 from aws_durable_execution_sdk_python.identifier import OperationIdentifier
 from aws_durable_execution_sdk_python.lambda_service import (
@@ -472,6 +473,75 @@ def test_child_handler_non_retryable_invocation_error_wrapped():
     assert fail_operation.action is OperationAction.FAIL
 
 
+def test_child_handler_preserves_data_across_nested_boundaries():
+    """Error data and stack_trace survive >=2 nested run_in_child_context
+    boundaries.
+
+    Each boundary carries the escaping error's data and stack_trace forward,
+    both on the raised ChildContextError and on the FAIL checkpoint that a
+    replay reconstructs from.
+    """
+    mock_state = Mock(spec=ExecutionState)
+    mock_state.durable_execution_arn = "test_arn"
+    mock_result = Mock()
+    mock_result.is_succeeded.return_value = False
+    mock_result.is_failed.return_value = False
+    mock_result.is_started.return_value = False
+    mock_result.is_replay_children.return_value = False
+    mock_result.is_existent.return_value = False
+    mock_state.get_checkpoint_result.return_value = mock_result
+    # Return each function unchanged so both nested bodies actually run.
+    mock_state.wrap_user_function.side_effect = lambda func, *args, **kwargs: func
+
+    # Innermost failure carries a serialized payload and stack trace, as an
+    # operation error reconstructed from a checkpoint would.
+    inner_error: StepError = StepError(
+        "inner step failed",
+        error_type="ValueError",
+        data="serialized-payload",
+        stack_trace=["frame-a", "frame-b"],
+    )
+
+    def inner_body():
+        raise inner_error
+
+    def outer_body():
+        return child_handler(
+            inner_body,
+            mock_state,
+            OperationIdentifier(
+                "inner", OperationSubType.RUN_IN_CHILD_CONTEXT, None, "inner_name"
+            ),
+            None,
+        )
+
+    with pytest.raises(ChildContextError) as exc_info:
+        child_handler(
+            outer_body,
+            mock_state,
+            OperationIdentifier(
+                "outer", OperationSubType.RUN_IN_CHILD_CONTEXT, None, "outer_name"
+            ),
+            None,
+        )
+
+    # Data and stack_trace survive both boundaries on the surfaced error.
+    assert exc_info.value.data == "serialized-payload"
+    assert exc_info.value.stack_trace == ["frame-a", "frame-b"]
+
+    # Both FAIL checkpoints (inner and outer) recorded the same payload, so the
+    # value also survives a replay that rebuilds from either checkpoint.
+    fail_operations = [
+        call.kwargs["operation_update"]
+        for call in mock_state.create_checkpoint.call_args_list
+        if call.kwargs["operation_update"].action is OperationAction.FAIL
+    ]
+    assert len(fail_operations) == 2
+    for fail_operation in fail_operations:
+        assert fail_operation.error.data == "serialized-payload"
+        assert fail_operation.error.stack_trace == ["frame-a", "frame-b"]
+
+
 def test_child_handler_with_config():
     """Test child_handler with config parameter."""
     mock_state = Mock(spec=ExecutionState)