fix(testing): fail fast on invalid handler output

packages/aws-durable-execution-sdk-python-testing/docs/architecture.md

@@ -430,15 +430,24 @@ retry budget cannot be bypassed.
 
 ### Retry budget
 
-A handler exception or an output validation failure increases the
-execution's `consecutive_failed_invocation_attempts`. At
+A handler exception increases the execution's
+`consecutive_failed_invocation_attempts`. At
 `MAX_CONSECUTIVE_FAILED_ATTEMPTS` the execution is failed with the last
 observed error. Below that, a retry is scheduled through
 `_invoke_execution` after `RETRY_BACKOFF_SECONDS`. Any clean invocation
 resets the counter. A clean invocation is one that succeeded or returned
 PENDING. So a transient failure followed by a success does not count
 against future budget.
 
+An invalid output is handled as the service handles it. A FAILED output
+with a Result, a SUCCEEDED output with an Error, a missing or unknown
+Status, or an output that cannot be parsed fails the execution at once.
+The error has type `InvalidParameterValueException` and the validation
+message. The only invalid output that is retried is PENDING with no
+pending operations, because an SDK race that resolves on replay can
+produce it. It counts against the same counter, with a cap of
+`MAX_CONSECUTIVE_INVALID_PENDING_ATTEMPTS`.
+
 ### Invocation timeout
 
 `execution_timeout` and `invocation_timeout` are set on the runner

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

@@ -109,6 +109,20 @@ class DurableFunctionsTestError(Exception):
     """Base class for testing errors."""
 
 
+class InvalidInvocationOutputError(DurableFunctionsTestError):
+    """The handler's output cannot be parsed as an invocation output."""
+
+    def __init__(self, message: str, request_id: str | None = None) -> None:
+        """Initialize with the message and the invocation's request ID.
+
+        Args:
+            message: Why the output cannot be parsed
+            request_id: Request ID of the invocation that returned the output
+        """
+        self.request_id = request_id
+        super().__init__(message)
+
+
 # endregion Testing
 
 

packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py

@@ -54,6 +54,7 @@
 )
 from aws_durable_execution_sdk_python_testing.exceptions import (
     IllegalStateException,
+    InvalidInvocationOutputError,
     InvalidParameterValueException,
     ResourceNotFoundException,
 )
@@ -126,6 +127,11 @@ class Executor(ExecutionObserver):
     # target and creating a durable child. Each takes milliseconds.
     QUICK_DISPATCH_WORKERS: int = 8
     MAX_CONSECUTIVE_FAILED_ATTEMPTS: int = 5
+    # The service fails an execution on an invalid invocation output at
+    # once, except PENDING with no pending operations. An SDK race that
+    # resolves on replay can produce that one, so it gets this many
+    # consecutive attempts.
+    MAX_CONSECUTIVE_INVALID_PENDING_ATTEMPTS: int = 3
     RETRY_BACKOFF_SECONDS: int = 5
     # GetDurableExecutionState page-count bounds, mirroring the service
     # (default 100, hard cap 1000 operations per page).
@@ -1540,8 +1546,8 @@ def _finish_invocation(
     ) -> None:
         """Apply a completed handler invocation.
 
-        Records the attempt, validates the response (completing or
-        retrying), resets the retry counter on a clean run, and releases
+        Records the attempt, validates the response (completing, failing
+        or retrying), resets the retry counter on a clean run, and releases
         the gate to PRE_INVOKE (re-invoking if a trigger was deferred) or
         COMPLETED.
         """
@@ -1574,11 +1580,20 @@ def _finish_invocation(
             logger.warning(
                 "[%s] Lambda output validation failure: %s", execution_arn, e
             )
-            error_obj = ErrorObject.from_exception(e)
+            error_obj = self._invalid_output_error(str(e))
             # Release the gate before retry scheduling so the retry's
             # invoke finds PRE_INVOKE.
             self._set_invocation_gate(execution_arn, InvocationState.PRE_INVOKE)
-            self._retry_invocation(execution, error_obj)
+            # Only PENDING with no pending operations is retried; every
+            # other invalid output fails the execution at once.
+            if response.status is InvocationStatus.PENDING:
+                self._retry_invocation(
+                    execution,
+                    error_obj,
+                    max_attempts=self.MAX_CONSECUTIVE_INVALID_PENDING_ATTEMPTS,
+                )
+            else:
+                self._fail_workflow(execution_arn, error_obj)
             return
 
         # A clean invocation resets the retry counter.
@@ -1616,6 +1631,37 @@ def _fail_invocation_not_found(
         self._set_invocation_gate(execution_arn, InvocationState.PRE_INVOKE)
         self._fail_workflow(execution_arn, error)
 
+    @staticmethod
+    def _invalid_output_error(message: str) -> ErrorObject:
+        """The error the service fails an execution with for an invalid output."""
+        return ErrorObject(
+            message=message,
+            type="InvalidParameterValueException",
+            data=None,
+            stack_trace=None,
+        )
+
+    def _fail_after_invalid_output(
+        self,
+        execution_arn: str,
+        error: ErrorObject,
+        invocation_start: datetime,
+        invocation_end: datetime,
+        request_id: str | None,
+    ) -> None:
+        """Fail the execution for an output the invoker could not parse."""
+        execution = self._store.load(execution_arn)
+        execution.record_invocation_completion(
+            invocation_start, invocation_end, request_id or str(uuid.uuid4())
+        )
+        self._store.save(execution)
+        if execution.is_complete:
+            # A stop or timeout landed during the invocation and stands.
+            self._set_invocation_gate(execution_arn, InvocationState.COMPLETED)
+            return
+        self._set_invocation_gate(execution_arn, InvocationState.PRE_INVOKE)
+        self._fail_workflow(execution_arn, error)
+
     def _retry_after_timeout(
         self,
         execution_arn: str,
@@ -1721,6 +1767,32 @@ async def invoke() -> None:
                     )
                 )
 
+            except InvalidInvocationOutputError as err:
+                if self._closing:
+                    return
+                # The handler ran, but the service cannot parse its output,
+                # so it fails the execution without invoking it again.
+                invocation_end = self._clock.now()
+                logger.warning(
+                    "[%s] Lambda output validation failure: %s", execution_arn, err
+                )
+                error_obj = self._invalid_output_error(str(err))
+                request_id = err.request_id
+                await asyncio.wrap_future(
+                    self._registry.submit(
+                        execution_arn,
+                        CallableTask(
+                            lambda: self._fail_after_invalid_output(
+                                execution_arn,
+                                error_obj,
+                                invocation_start,
+                                invocation_end,
+                                request_id,
+                            )
+                        ),
+                    )
+                )
+
             except asyncio.TimeoutError:
                 if self._closing:
                     return
@@ -1803,22 +1875,26 @@ def _fail_workflow(self, execution_arn: str, error: ErrorObject | None):
 
         self.fail_execution(execution_arn, error)
 
-    def _retry_invocation(self, execution: Execution, error: ErrorObject):
+    def _retry_invocation(
+        self,
+        execution: Execution,
+        error: ErrorObject,
+        max_attempts: int | None = None,
+    ) -> None:
         """Handle retry logic or fail execution if retries exhausted.
 
-        Budget: ``MAX_CONSECUTIVE_FAILED_ATTEMPTS`` attempts
-        per ``RETRY_BACKOFF_SECONDS`` backoff. Increments the counter
-        BEFORE the threshold check so an attempt-N failure that
-        reaches the ceiling fails the execution rather than scheduling
-        a pointless retry.
+        Budget: ``max_attempts`` attempts, by default
+        ``MAX_CONSECUTIVE_FAILED_ATTEMPTS``, per ``RETRY_BACKOFF_SECONDS``
+        backoff. Increments the counter BEFORE the threshold check so an
+        attempt-N failure that reaches the ceiling fails the execution
+        rather than scheduling a pointless retry.
         """
+        if max_attempts is None:
+            max_attempts = self.MAX_CONSECUTIVE_FAILED_ATTEMPTS
         execution.consecutive_failed_invocation_attempts += 1
         self._store.save(execution)
 
-        if (
-            execution.consecutive_failed_invocation_attempts
-            >= self.MAX_CONSECUTIVE_FAILED_ATTEMPTS
-        ):
+        if execution.consecutive_failed_invocation_attempts >= max_attempts:
             # Budget exhausted — fail the execution with the last
             # observed error.
             self._fail_workflow(

packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/invoker.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import json
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from threading import Lock
 from typing import TYPE_CHECKING, Any, Protocol
 from uuid import uuid4
@@ -15,12 +15,14 @@
     DurableExecutionInvocationOutput,
     InitialExecutionState,
 )
+from aws_durable_execution_sdk_python.lambda_service import ErrorObject
 
 from aws_durable_execution_sdk_python_testing.checkpoint.processor import (
     DEFAULT_MAX_INVOCATION_PAGE_BYTES,
 )
 from aws_durable_execution_sdk_python_testing.exceptions import (
     DurableFunctionsTestError,
+    InvalidInvocationOutputError,
     InvalidParameterValueException,
     ResourceNotFoundException,
 )
@@ -107,6 +109,45 @@ class InvokeResponse:
     invocation_output: DurableExecutionInvocationOutput
     request_id: str
 
+    @classmethod
+    def from_output(cls, output: Any, request_id: str) -> InvokeResponse:
+        """Create from the value a handler returned.
+
+        Raises:
+            InvalidInvocationOutputError: If the value is not an object with
+                a known ``Status`` and, if present, an object ``Error``.
+        """
+        try:
+            invocation_output = DurableExecutionInvocationOutput.from_dict(output)
+            # from_dict treats any falsey Error as absent, but only a missing
+            # or null Error is. An empty object is kept, so SUCCEEDED with it
+            # fails validation; a value that is not an object raises here.
+            error: Any = output.get("Error")
+            if error is not None and invocation_output.error is None:
+                invocation_output = replace(
+                    invocation_output, error=ErrorObject.from_dict(error)
+                )
+        except (AttributeError, TypeError, ValueError) as e:
+            msg: str = f"Invalid invocation output: {e}"
+            raise InvalidInvocationOutputError(msg, request_id) from e
+        return cls(invocation_output=invocation_output, request_id=request_id)
+
+    @classmethod
+    def from_payload(cls, payload: bytes, request_id: str) -> InvokeResponse:
+        """Create from the payload a Lambda invocation returned.
+
+        Raises:
+            InvalidInvocationOutputError: If the payload is not UTF-8 JSON or
+                not a valid invocation output.
+        """
+        try:
+            # UnicodeDecodeError is a ValueError.
+            output: Any = json.loads(payload.decode("utf-8"))
+        except ValueError as e:
+            msg: str = f"Invalid invocation output: {e}"
+            raise InvalidInvocationOutputError(msg, request_id) from e
+        return cls.from_output(output, request_id)
+
 
 # Values the in-process Lambda context reports when the caller gives none.
 DEFAULT_TEST_REGION = "us-west-2"
@@ -261,10 +302,7 @@ def invoke(
         )
         handler: Callable = self._resolve_handler(function_name)
         response_dict = handler(input_with_client, context)
-        output = DurableExecutionInvocationOutput.from_dict(response_dict)
-        return InvokeResponse(
-            invocation_output=output, request_id=context.aws_request_id
-        )
+        return InvokeResponse.from_output(response_dict, context.aws_request_id)
 
     def update_endpoint(self, endpoint_url: str, region_name: str) -> None:
         """No-op for in-process invoker."""
@@ -459,6 +497,7 @@ def invoke(
         Raises:
             ResourceNotFoundException: If function does not exist
             InvalidParameterValueException: If parameters are invalid
+            InvalidInvocationOutputError: If the function's output cannot be parsed
             DurableFunctionsTestError: For other invocation failures
         """
 
@@ -496,22 +535,23 @@ def invoke(
                 msg = f"Lambda invocation failed with status {status_code}: {error_payload}"
                 raise DurableFunctionsTestError(msg)
 
-            # Parse response payload
-            response_payload = response["Payload"].read().decode("utf-8")
-            response_dict = json.loads(response_payload)
-
-            # Extract request ID from response headers (x-amzn-RequestId or x-amzn-request-id)
-            headers = response.get("ResponseMetadata", {}).get("HTTPHeaders", {})
-            request_id = (
-                headers.get("x-amzn-RequestId")
+            # Read response payload; from_payload decodes it, below
+            response_payload: bytes = response["Payload"].read()
+
+            # botocore puts the request ID in ResponseMetadata and keeps the
+            # response headers with lowercased names.
+            metadata: dict[str, Any] = response.get("ResponseMetadata", {})
+            headers: dict[str, str] = {
+                name.lower(): value
+                for name, value in metadata.get("HTTPHeaders", {}).items()
+            }
+            request_id: str = (
+                metadata.get("RequestId")
+                or headers.get("x-amzn-requestid")
                 or headers.get("x-amzn-request-id")
                 or f"local-{uuid4()}"
             )
 
-            # Convert to DurableExecutionInvocationOutput
-            output = DurableExecutionInvocationOutput.from_dict(response_dict)
-            return InvokeResponse(invocation_output=output, request_id=request_id)
-
         except client.exceptions.ResourceNotFoundException as e:
             msg = f"Function not found: {function_name}"
             raise ResourceNotFoundException(msg) from e
@@ -565,3 +605,7 @@ def invoke(
                 raise DurableFunctionsTestError(msg) from e
             msg = f"Unexpected error during Lambda invocation: {e}"
             raise DurableFunctionsTestError(msg) from e
+
+        # Parsed outside the try: an output the service cannot parse fails
+        # the execution, so it must not be reported as a failed invocation.
+        return InvokeResponse.from_payload(response_payload, request_id)