fix: checkpoint the wrapped error for failed contexts

packages/aws-durable-execution-sdk-python-examples/test/wait_for_callback/test_wait_for_callback_submitter_failure.py

@@ -24,13 +24,10 @@ def test_fail_after_exhausting_retries_when_submitter_always_fails(durable_runne
     # Execution should fail after retries are exhausted
     assert result.status is InvocationStatus.FAILED
 
-    # A failed submitter step surfaces as CallbackSubmitterError (the submitter's
-    # StepError is translated at the wait_for_callback call site). The type is
-    # identical on first run and replay, so the execution-level record pins it.
+    # User code sees CallbackSubmitterError, and the execution records the
+    # error it wraps: the submitter's own exception. The type is identical on
+    # first run and replay, so the execution-level record pins it.
     error = result.error
     assert error is not None
-    assert (
-        error.type
-        == "aws_durable_execution_sdk_python.exceptions.CallbackSubmitterError"
-    )
+    assert error.type == "Exception"
     assert "Simulated submitter failure" in error.message

packages/aws-durable-execution-sdk-python-extras/README.md

@@ -225,8 +225,8 @@ pipeline                       child context
 
 A failed `vm.invoke` raises its MicroVM error inside the handler. The handler can catch it and send another job to the same MicroVM. An error that leaves the handler ends the session:
 
-1. A MicroVM error that the handler raises itself reaches the caller with its own class, such as `MicrovmJobFailedError`.
-2. Any other error reaches the caller as the SDK's `ChildContextError`, as from `run_in_child_context`. This includes a MicroVM error that leaves a child context that the handler created: that child context records `ChildContextError` as its error type. The caller then gets the MicroVM error's message, but not its class. Inside the handler, the `ChildContextError` of that child context has the MicroVM error as its `__cause__`. An SDK error that leaves the handler, such as the `StepError` of a failed step, also reaches the caller as `ChildContextError`, with the SDK error as its `__cause__`. This rule matches `run_in_child_context`. The JavaScript operation differs here: it passes an SDK error to the caller with its own class.
+1. A MicroVM error that leaves the handler reaches the caller with its own class, such as `MicrovmJobFailedError`. This includes a MicroVM error that leaves a child context that the handler created: the SDK records the error that a `ChildContextError` wraps, not the `ChildContextError`. Inside the handler, the `ChildContextError` of that child context has the MicroVM error as its `__cause__`.
+2. Any other error reaches the caller as the SDK's `ChildContextError`, as from `run_in_child_context`. An SDK error that leaves the handler, such as the `StepError` of a failed step, also reaches the caller as `ChildContextError`. Its `error_type` is the type of the error that failed the step, for example `LookupError`. This rule matches `run_in_child_context`. The JavaScript operation differs here: it passes an SDK error to the caller with its own class.
 
 An invalid job raises `ValidationError` or `TypeError` from `vm.invoke`, before the job's durable operations start. The handler can catch it too.
 

packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/session.py

@@ -113,13 +113,13 @@ def microvm_session(
     error that leaves the handler ends the session, and the MicroVM is
     terminated:
 
-    1. A MicroVM error, such as :class:`MicrovmJobFailedError`, that the
-       handler raises itself reaches the caller with its own class.
+    1. A MicroVM error, such as :class:`MicrovmJobFailedError`, that leaves
+       the handler reaches the caller with its own class. This includes a
+       MicroVM error that leaves a child context that the handler created:
+       the SDK records the error that the child context's
+       ``ChildContextError`` wraps.
     2. Any other error reaches the caller as ``ChildContextError``, as from
-       ``run_in_child_context``. This includes a MicroVM error that leaves
-       a child context that the handler created: that child context records
-       ``ChildContextError`` as its error type. The caller then gets the
-       MicroVM error's message, but not its class.
+       ``run_in_child_context``.
 
     The session's value is the handler's return value, checkpointed as the
     child context result. So it must be serializable by the SDK's default

packages/aws-durable-execution-sdk-python-extras/tests/microvm/session_test.py

@@ -1016,14 +1016,13 @@ def body(vm: MicrovmSession, ctx: DurableContext) -> Any:
     ]
 
 
-def test_wraps_a_job_error_that_leaves_a_child_context_of_the_handler(
+def test_keeps_a_job_error_that_leaves_a_child_context_of_the_handler(
     run: Any, fake_client: Any, fake_endpoint: Any
 ) -> None:
-    # Only an error that the handler raises itself keeps its MicroVM class.
-    # The SDK records one error type per child context. So an error from the
-    # handler's own child context reaches the caller as ChildContextError,
-    # with the MicroVM error's message. Inside the handler, the
-    # ChildContextError's __cause__ is still the MicroVM class.
+    # The SDK records the error a ChildContextError wraps. So a MicroVM error
+    # that leaves the handler's own child context reaches the caller with its
+    # MicroVM class. Inside the handler, the ChildContextError's __cause__ is
+    # the MicroVM class.
     client = fake_client(endpoint=fake_endpoint().url)
     config = session_config(client)
     inside: list[str] = []
@@ -1042,14 +1041,14 @@ def body(vm: MicrovmSession, ctx: DurableContext) -> Any:
 
         try:
             microvm_session(context, "pipeline", config=config, handler=body)
-        except ChildContextError as error:
-            return [type(error.__cause__).__name__, error.message]
+        except MicrovmJobFailedError as error:
+            return [type(error).__name__, error.message]
         return None
 
     result = run(handler, None, [("build", fail_with("BuildError", "failed"))])
 
     assert outcome(result) == [
-        "ChildContextError",
+        "MicrovmJobFailedError",
         'MicroVM "build": the job failed (BuildError): failed',
     ]
     assert "MicrovmJobFailedError" in inside
@@ -1111,12 +1110,17 @@ def body(vm: MicrovmSession, ctx: DurableContext) -> Any:
         try:
             microvm_session(context, "pipeline", config=config, handler=body)
         except ChildContextError as error:
-            return [type(error.__cause__).__name__, error.message]
+            return [
+                type(error.__cause__).__name__,
+                error.error_type,
+                error.message,
+            ]
         return None
 
     result = run(handler)
 
-    assert outcome(result) == ["StepError", "no artifact"]
+    # The session context records the step's LookupError, not StepError.
+    assert outcome(result) == ["DurableOperationError", "LookupError", "no artifact"]
     assert len(client.requests("terminate_microvm")) == 1
 
 

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

@@ -305,25 +305,45 @@ def result(self) -> T | None:
     def _raise_terminal_failure(
         self, checkpointed_result: CheckpointedResult
     ) -> NoReturn:
-        """Raise the graded error for a terminal callback failure.
+        _raise_callback_terminal_failure(checkpointed_result)
 
-        TIMED_OUT -> CallbackTimeoutError; any other non-succeeded terminal
-        (external FAILED, cancelled, stopped) -> CallbackExternalError.
-        Carries the message/data/stack_trace; the external error's own type
-        stays on the callback operation's details.
-        """
-        error: ErrorObject | None = checkpointed_result.error
-        is_timeout: bool = checkpointed_result.is_timed_out()
-        error_cls: type[CallbackError] = (
-            CallbackTimeoutError if is_timeout else CallbackExternalError
-        )
-        default_message: str = "Callback timed out" if is_timeout else "Callback failed"
-        message: str = error.message if error and error.message else default_message
-        raise error_cls(
-            message=message,
-            data=error.data if error else None,
-            stack_trace=error.stack_trace if error else None,
-        )
+
+def _is_blank_type(error_type: str | None) -> bool:
+    """A missing type, or one made only of ASCII whitespace."""
+    return error_type is None or all(ch in " \t\n\r\f\v" for ch in error_type)
+
+
+def _raise_callback_terminal_failure(
+    checkpointed_result: CheckpointedResult,
+) -> NoReturn:
+    """Raise the graded error for a terminal callback failure.
+
+    TIMED_OUT -> CallbackTimeoutError; any other non-succeeded terminal
+    (external FAILED, cancelled, stopped) -> CallbackExternalError.
+    Carries the message/data/stack_trace. An external failure also carries
+    the sender's ErrorType as ``error_type``.
+    """
+    error: ErrorObject | None = checkpointed_result.error
+    is_timeout: bool = checkpointed_result.is_timed_out()
+    error_cls: type[CallbackError] = (
+        CallbackTimeoutError if is_timeout else CallbackExternalError
+    )
+    default_message: str = "Callback timed out" if is_timeout else "Callback failed"
+    message: str = error.message if error and error.message else default_message
+    # An external failure keeps the sender's ErrorType, the explicit type
+    # "Error" included, so it is recorded and surfaced as that type. Only a
+    # missing or blank type falls back to the SDK class.
+    error_type: str | None = (
+        error.type
+        if not is_timeout and error is not None and not _is_blank_type(error.type)
+        else None
+    )
+    raise error_cls(
+        message=message,
+        error_type=error_type,
+        data=error.data if error else None,
+        stack_trace=error.stack_trace if error else None,
+    )
 
 
 class DurableContext(DurableContextProtocol):
@@ -908,29 +928,65 @@ def wait_for_callback(
         def wait_in_child_context(context: DurableContext):
             return wait_for_callback_handler(context, submitter, step_name, config)
 
-        # run_in_child_context is generic and raises ChildContextError with the
-        # inner failure on __cause__; translate to callback-typed errors here, at
-        # the call site, so it re-fires identically on replay.
+        # The context records the wrapped error's type, so the callback-typed
+        # error is recovered from which child failed: the submitter step or the
+        # callback. Both are read from checkpoints, so first run and replay
+        # raise the same error.
+        context_operation_id: str = self._peek_next_operation_id()
         try:
             return self.run_in_child_context(
                 wait_in_child_context,
                 step_name,
                 ChildConfig(sub_type=OperationSubType.WAIT_FOR_CALLBACK),
             )
         except ChildContextError as e:
-            inner: BaseException | None = e.__cause__
-            # Callback failures (external/timeout/base) pass through.
-            if isinstance(inner, CallbackError):
-                raise inner
-            # A failed submitter step surfaces as CallbackSubmitterError, keeping
-            # the step's data/stack_trace (the StepError stays as its cause).
-            if isinstance(inner, StepError):
-                raise CallbackSubmitterError(
-                    inner.message or "Callback submitter failed",
-                    data=inner.data,
-                    stack_trace=inner.stack_trace,
-                ) from inner
-            raise
+            self._raise_wait_for_callback_error(e, context_operation_id)
+
+    def _raise_wait_for_callback_error(
+        self, error: ChildContextError, context_operation_id: str
+    ) -> NoReturn:
+        """Translate a failed wait_for_callback context to a callback error.
+
+        A failed submitter step surfaces as CallbackSubmitterError and a
+        callback failure as CallbackExternalError or CallbackTimeoutError.
+        Checkpoints written before the context recorded the wrapped error hold
+        the callback error or StepError itself; those surface as before.
+        """
+        inner: BaseException | None = error.__cause__
+        # The context holds a callback error type: a timeout, a callback
+        # failure without a sender type, or the base CallbackError.
+        if isinstance(inner, CallbackError):
+            raise inner
+        if isinstance(inner, StepError):
+            # The context holds StepError itself.
+            raise CallbackSubmitterError(
+                inner.message or "Callback submitter failed",
+                data=inner.data,
+                stack_trace=inner.stack_trace,
+            ) from inner
+        namespace = OperationIdNamespace(prefix=context_operation_id)
+        # wait_for_callback_handler creates the callback first, then the step.
+        callback_result: CheckpointedResult = self.state.get_checkpoint_result(
+            namespace.create_id_for_step(1)
+        )
+        submitter_result: CheckpointedResult = self.state.get_checkpoint_result(
+            namespace.create_id_for_step(2)
+        )
+        if submitter_result.is_failed():
+            raise CallbackSubmitterError(
+                error.message or "Callback submitter failed",
+                error_type=error.error_type,
+                data=error.data,
+                stack_trace=error.stack_trace,
+            ) from inner
+        if (
+            callback_result.is_failed()
+            or callback_result.is_cancelled()
+            or callback_result.is_timed_out()
+            or callback_result.is_stopped()
+        ):
+            _raise_callback_terminal_failure(callback_result)
+        raise error
 
     def wait_for_condition(
         self,

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

@@ -414,6 +414,28 @@ class CallbackSubmitterError(CallbackError):
 }
 
 
+# SDK wrapper types: each re-raises the failure that escaped an operation and
+# carries that failure's fields (error_type, message, data, stack_trace). A
+# checkpoint records the wrapped error, never the wrapper. Matched by exact
+# class, so a user subclass is never unwrapped. CallbackTimeoutError, the base
+# CallbackError, SerDesError, BatchCompletionError and ValidationError are SDK
+# errors in their own right and are not wrappers.
+SDK_WRAPPER_ERROR_TYPES: frozenset[type[DurableOperationError]] = frozenset(
+    {
+        DurableOperationError,
+        StepError,
+        InvokeError,
+        ChildContextError,
+        WaitForConditionError,
+        CallbackExternalError,
+        CallbackSubmitterError,
+    }
+)
+SDK_WRAPPER_ERROR_TYPE_NAMES: frozenset[str] = frozenset(
+    f"{cls.__module__}.{cls.__qualname__}" for cls in SDK_WRAPPER_ERROR_TYPES
+)
+
+
 def register_operation_error(error_cls: type[DurableOperationError]) -> None:
     """Register an SDK-owned operation-error subclass for replay reconstruction.
 

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

@@ -17,6 +17,8 @@
     CheckpointError,
     DurableOperationError,
     GetExecutionStateError,
+    SDK_WRAPPER_ERROR_TYPE_NAMES,
+    SDK_WRAPPER_ERROR_TYPES,
     SerDesError,
 )
 
@@ -200,6 +202,15 @@ def _qualified_error_type(exception: BaseException) -> str:
     return f"{module}{cls_.__qualname__}"
 
 
+# Bound on the SDK wrappers unwrapped when recording an error.
+_MAX_WRAPPER_UNWRAPS: int = 16
+
+
+def _is_sdk_wrapper(exception: BaseException) -> bool:
+    # Exact class match: a user subclass of an SDK wrapper is the user's error.
+    return type(exception) in SDK_WRAPPER_ERROR_TYPES
+
+
 @dataclass(frozen=True)
 class ErrorObject:
     message: str | None
@@ -218,6 +229,46 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> ErrorObject:
 
     @classmethod
     def from_exception(cls, exception: Exception) -> ErrorObject:
+        """Build the ErrorObject that records ``exception``.
+
+        An SDK wrapper (StepError, ChildContextError, InvokeError,
+        CallbackExternalError, ...) is unwrapped first, so the record carries
+        the type, message, data and stack trace of the error it wraps rather
+        than the wrapper's class. A wrapper keeps those fields itself: its
+        ``error_type`` names the wrapped error. When ``error_type`` names
+        another wrapper, the walk continues with ``__cause__``. The walk takes
+        at most ``_MAX_WRAPPER_UNWRAPS`` steps and stops at a cause it has
+        already visited; a wrapper it stops at is recorded by its
+        ``error_type``.
+        """
+        current: Exception = exception
+        visited: set[int] = {id(current)}
+        for _ in range(_MAX_WRAPPER_UNWRAPS):
+            if not _is_sdk_wrapper(current):
+                break
+            wrapper = cast(DurableOperationError, current)
+            if wrapper.error_type not in SDK_WRAPPER_ERROR_TYPE_NAMES:
+                return cls._from_wrapper(wrapper)
+            cause = wrapper.__cause__
+            if not isinstance(cause, Exception) or id(cause) in visited:
+                break
+            visited.add(id(cause))
+            current = cause
+        if _is_sdk_wrapper(current):
+            return cls._from_wrapper(cast(DurableOperationError, current))
+        return cls._from_unwrapped_exception(current)
+
+    @classmethod
+    def _from_wrapper(cls, wrapper: DurableOperationError) -> ErrorObject:
+        return cls(
+            message=wrapper.message,
+            type=wrapper.error_type,
+            data=wrapper.data,
+            stack_trace=wrapper.stack_trace,
+        )
+
+    @classmethod
+    def _from_unwrapped_exception(cls, exception: Exception) -> ErrorObject:
         # SerDesError and subclasses pin to the base discriminator so replay
         # always reconstructs them as SerDesError.
         if isinstance(exception, SerDesError):

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

@@ -217,11 +217,12 @@ def test_callback_result_failed():
     callback = Callback("callback5", "op5", mock_state)
 
     # A FAILED callback (external SendDurableExecutionCallbackFailure) surfaces
-    # as CallbackExternalError. The external error's own type stays on the
-    # callback operation, and there is no synthetic __cause__.
+    # as CallbackExternalError carrying the sender's type as error_type, and
+    # there is no synthetic __cause__.
     with pytest.raises(CallbackExternalError) as exc_info:
         callback.result()
     assert exc_info.value.message == "Callback failed"
+    assert exc_info.value.error_type == "CallbackError"
     assert isinstance(exc_info.value, CallbackError)
     assert exc_info.value.__cause__ is None
 
@@ -1582,6 +1583,10 @@ def test_wait_for_callback_reraises_unrelated_child_context_error():
     """A non-callback, non-step failure re-raises the original ChildContextError."""
     mock_state = Mock(spec=ExecutionState)
     mock_state.durable_execution_arn = "test_arn"
+    # Neither the callback nor the submitter step failed.
+    mock_state.get_checkpoint_result.return_value = (
+        CheckpointedResult.create_not_found()
+    )
     context = create_test_context(state=mock_state)
 
     inner = InvokeError("nested invoke failed")