feat: Classify non-retryable customer errors for Durable API calls (#358)

Non-retryable customer errors from Lambda (e.g., KMSAccessDeniedException, KMSDisabledException) arrive as HTTP 502 during CheckpointDurableExecution and GetDurableExecutionState API calls. Without this change, these 502s are classified as retryable invocation errors, causing the SDK to retry invocations that will never self-resolve.

Extract shared classification logic into BotoClientError with a _NON_RETRYABLE_CUSTOMER_ERROR_CODES set and _classify_error_category method. Both CheckpointError and GetExecutionStateError now inherit from_exception() with unified classification so that non-retryable errors (KMS 502s, 4xx client errors) return Status: FAILED immediately, while retryable errors (5xx, 429, network errors) continue to raise for Lambda retry.

Add is_retryable() to InvocationError hierarchy so execution.py handlers use a single interface instead of isinstance checks.

Add backward-compatible aliases for CheckpointErrorCategory and is_retriable().

Log Durable API errors during state fetch, and save partially fetched state.

Testing — parameterized classification tests across both error types for all four KMS codes, 4xx, 429, 5xx, and 502 scenarios. Integration tests for non-retryable/retryable errors across all three execution.py code paths: initial pagination, background thread, and user thread.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Author: denzyll

src/aws_durable_execution_sdk_python/exceptions.py

@@ -13,6 +13,22 @@
 BAD_REQUEST_ERROR: int = 400
 TOO_MANY_REQUESTS_ERROR: int = 429
 SERVICE_ERROR: int = 500
+INVALID_PARAMETER_VALUE_EXCEPTION: str = "InvalidParameterValueException"
+INVALID_CHECKPOINT_TOKEN_PREFIX: str = "Invalid Checkpoint Token"
+
+# Non-retryable customer error codes that arrive as non-4xx (e.g. HTTP 502) from Lambda.
+# Unlike typical 5xx errors, these require customer intervention (e.g., fixing
+# a KMS key configuration) and will never succeed on retry.
+# Add new non-retryable error codes here — they are automatically classified
+# as EXECUTION (non-retryable) by _classify_error_category().
+_NON_RETRYABLE_CUSTOMER_ERROR_CODES: frozenset[str] = frozenset(
+    {
+        "KMSAccessDeniedException",
+        "KMSDisabledException",
+        "KMSInvalidStateException",
+        "KMSNotFoundException",
+    }
+)
 
 if TYPE_CHECKING:
     import datetime
@@ -77,6 +93,14 @@ def __init__(
     ):
         super().__init__(message, termination_reason)
 
+    def is_retryable(self) -> bool:
+        """Whether this error is retryable. Returns True by default.
+
+        Subclasses override to implement classification logic based on
+        error codes and HTTP status codes.
+        """
+        return True
+
 
 class CallbackError(ExecutionError):
     """Error in callback handling."""
@@ -86,27 +110,98 @@ def __init__(self, message: str, callback_id: str | None = None):
         self.callback_id = callback_id
 
 
+class DurableApiErrorCategory(Enum):
+    INVOCATION = "INVOCATION"
+    EXECUTION = "EXECUTION"
+
+
+# Backward-compatible alias
+CheckpointErrorCategory = DurableApiErrorCategory
+
+
 class BotoClientError(InvocationError):
+    """Error from a Lambda API call (e.g., CheckpointDurableExecution, GetDurableExecutionState).
+
+    Extends InvocationError because the default behavior for API failures is to retry
+    the Lambda invocation. However, some errors are non-retryable (e.g., 4xx client errors,
+    KMS key misconfiguration) and should fail the execution instead. The error_category field
+    and is_retryable() method distinguish these cases at runtime.
+    """
+
     def __init__(
         self,
         message: str,
+        error_category: DurableApiErrorCategory = DurableApiErrorCategory.INVOCATION,
         error: AwsErrorObj | None = None,
         response_metadata: AwsErrorMetadata | None = None,
         termination_reason=TerminationReason.INVOCATION_ERROR,
     ):
         super().__init__(message=message, termination_reason=termination_reason)
         self.error: AwsErrorObj | None = error
         self.response_metadata: AwsErrorMetadata | None = response_metadata
+        self.error_category: DurableApiErrorCategory = error_category
 
     @classmethod
     def from_exception(cls, exception: Exception) -> Self:
         response = getattr(exception, "response", {})
         response_metadata = response.get("ResponseMetadata")
         error = response.get("Error")
+        error_category = BotoClientError._classify_error_category(
+            error, response_metadata
+        )
         return cls(
-            message=str(exception), error=error, response_metadata=response_metadata
+            message=str(exception),
+            error_category=error_category,
+            error=error,
+            response_metadata=response_metadata,
         )
 
+    @staticmethod
+    def _classify_error_category(
+        error: AwsErrorObj | None,
+        response_metadata: AwsErrorMetadata | None,
+    ) -> DurableApiErrorCategory:
+        """Classify a Durable API error as retryable (INVOCATION) or non-retryable (EXECUTION).
+
+        Classification rules:
+        - Non-retryable customer error codes (e.g., KMS key issues) → EXECUTION
+          These arrive as HTTP 502 but require customer intervention to fix.
+        - 4xx errors → EXECUTION, except:
+          - 429 (TooManyRequests) → INVOCATION (throttling is transient)
+          - InvalidParameterValueException with "Invalid Checkpoint Token" → INVOCATION
+            (stale token from a concurrent checkpoint; next invocation gets a fresh token)
+        - 5xx, network errors → INVOCATION
+        """
+        error_code: str | None = (error and error.get("Code")) or None
+        if error_code and error_code in _NON_RETRYABLE_CUSTOMER_ERROR_CODES:
+            return DurableApiErrorCategory.EXECUTION
+
+        status_code: int | None = (
+            response_metadata and response_metadata.get("HTTPStatusCode")
+        ) or None
+        if (
+            status_code
+            and BAD_REQUEST_ERROR <= status_code < SERVICE_ERROR
+            and status_code != TOO_MANY_REQUESTS_ERROR
+            and error
+            and not (
+                (error.get("Code") or "") == INVALID_PARAMETER_VALUE_EXCEPTION
+                and (error.get("Message") or "").startswith(
+                    INVALID_CHECKPOINT_TOKEN_PREFIX
+                )
+            )
+        ):
+            return DurableApiErrorCategory.EXECUTION
+
+        return DurableApiErrorCategory.INVOCATION
+
+    def is_retryable(self) -> bool:
+        """Whether this error is retryable based on error_category."""
+        return self.error_category == DurableApiErrorCategory.INVOCATION
+
+    # Backward-compatible alias
+    is_retriable = is_retryable
+
     def build_logger_extras(self) -> dict:
         extras: dict = {}
         # preserve PascalCase to be consistent with other langauges
@@ -125,55 +220,23 @@ def __init__(self, message: str, step_id: str | None = None):
         self.step_id = step_id
 
 
-class CheckpointErrorCategory(Enum):
-    INVOCATION = "INVOCATION"
-    EXECUTION = "EXECUTION"
-
-
 class CheckpointError(BotoClientError):
     """Failure to checkpoint. Will terminate the lambda."""
 
     def __init__(
         self,
         message: str,
-        error_category: CheckpointErrorCategory,
+        error_category: DurableApiErrorCategory = DurableApiErrorCategory.INVOCATION,
         error: AwsErrorObj | None = None,
         response_metadata: AwsErrorMetadata | None = None,
     ):
         super().__init__(
             message,
+            error_category,
             error,
             response_metadata,
             termination_reason=TerminationReason.CHECKPOINT_FAILED,
         )
-        self.error_category: CheckpointErrorCategory = error_category
-
-    @classmethod
-    def from_exception(cls, exception: Exception) -> CheckpointError:
-        base = BotoClientError.from_exception(exception)
-        metadata: AwsErrorMetadata | None = base.response_metadata
-        error: AwsErrorObj | None = base.error
-        error_category: CheckpointErrorCategory = CheckpointErrorCategory.INVOCATION
-
-        # 4xx errors (except 429) are permanent failures (EXECUTION), unless it's an
-        # InvalidParameterValueException with "Invalid Checkpoint Token" which is retriable (INVOCATION).
-        # 5xx, 429, and network errors are retriable (INVOCATION).
-        status_code: int | None = (metadata and metadata.get("HTTPStatusCode")) or None
-        if (
-            status_code
-            and BAD_REQUEST_ERROR <= status_code < SERVICE_ERROR
-            and status_code != TOO_MANY_REQUESTS_ERROR
-            and error
-            and not (
-                (error.get("Code") or "") == "InvalidParameterValueException"
-                and (error.get("Message") or "").startswith("Invalid Checkpoint Token")
-            )
-        ):
-            error_category = CheckpointErrorCategory.EXECUTION
-        return CheckpointError(str(exception), error_category, error, metadata)
-
-    def is_retriable(self):
-        return self.error_category == CheckpointErrorCategory.INVOCATION
 
 
 class ValidationError(DurableExecutionsError):
@@ -186,11 +249,13 @@ class GetExecutionStateError(BotoClientError):
     def __init__(
         self,
         message: str,
+        error_category: DurableApiErrorCategory = DurableApiErrorCategory.INVOCATION,
         error: AwsErrorObj | None = None,
         response_metadata: AwsErrorMetadata | None = None,
     ):
         super().__init__(
             message,
+            error_category,
             error,
             response_metadata,
             termination_reason=TerminationReason.INVOCATION_ERROR,

src/aws_durable_execution_sdk_python/execution.py

@@ -260,11 +260,26 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
             else ReplayStatus.NEW,
         )
 
-        execution_state.fetch_paginated_operations(
-            invocation_input.initial_execution_state.operations,
-            invocation_input.checkpoint_token,
-            invocation_input.initial_execution_state.next_marker,
-        )
+        try:
+            execution_state.fetch_paginated_operations(
+                invocation_input.initial_execution_state.operations,
+                invocation_input.checkpoint_token,
+                invocation_input.initial_execution_state.next_marker,
+            )
+        except BotoClientError as e:
+            # Non-retryable Durable API errors (e.g., customer configuration issues,
+            # 4xx client errors) will never succeed on retry — fail the execution immediately.
+            if not e.is_retryable():
+                logger.exception(
+                    "Non-retryable Durable API error during initial state fetch. Must fail execution "
+                    "without retry.",
+                    extra=e.build_logger_extras(),
+                )
+                return DurableExecutionInvocationOutput(
+                    status=InvocationStatus.FAILED,
+                    error=ErrorObject.from_exception(e),
+                ).to_dict()
+            raise
 
         raw_input_payload: str | None = execution_state.get_input_payload()
 
@@ -356,11 +371,20 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
                         "Checkpoint processing failed",
                         extra=bg_error.source_exception.build_logger_extras(),
                     )
+                    # Non-retryable Durable API errors (e.g., customer configuration issues,
+                    # 4xx client errors) will never succeed on retry — fail the execution immediately.
+                    if not bg_error.source_exception.is_retryable():
+                        logger.exception(
+                            "Non-retryable Durable API error from background thread. Must fail execution "
+                            "without retry.",
+                            extra=bg_error.source_exception.build_logger_extras(),
+                        )
+                        return DurableExecutionInvocationOutput(
+                            status=InvocationStatus.FAILED,
+                            error=ErrorObject.from_exception(bg_error.source_exception),
+                        ).to_dict()
                 else:
                     logger.exception("Checkpoint processing failed")
-                # handle the original exception
-                if isinstance(bg_error.source_exception, CheckpointError):
-                    return handle_checkpoint_error(bg_error.source_exception).to_dict()
                 raise bg_error.source_exception from bg_error
 
             except SuspendExecution:
@@ -377,12 +401,23 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
                     extra=e.build_logger_extras(),
                 )
                 return handle_checkpoint_error(e).to_dict()
-            except InvocationError:
+            except InvocationError as e:
+                # Non-retryable Durable API errors (e.g., customer configuration issues,
+                # 4xx client errors) will never succeed on retry — fail the execution immediately.
+                if not e.is_retryable():
+                    logger.exception(
+                        "Non-retryable Durable API error. Must fail execution without retry.",
+                        extra=e.build_logger_extras(),  # type: ignore[attr-defined]
+                    )
+                    return DurableExecutionInvocationOutput(
+                        status=InvocationStatus.FAILED,
+                        error=ErrorObject.from_exception(e),
+                    ).to_dict()
                 logger.exception("Invocation error. Must terminate.")
                 # Throw the error to trigger Lambda retry
                 raise
             except ExecutionError as e:
-                logger.exception("Execution error. Must terminate without retry.")
+                logger.exception("Execution error. Must fail execution without retry.")
                 return DurableExecutionInvocationOutput(
                     status=InvocationStatus.FAILED,
                     error=ErrorObject.from_exception(e),
@@ -428,7 +463,7 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
 
 
 def handle_checkpoint_error(error: CheckpointError) -> DurableExecutionInvocationOutput:
-    if error.is_retriable():
+    if error.is_retryable():
         raise error from None  # Terminate Lambda immediately and have it be retried
     return DurableExecutionInvocationOutput(
         status=InvocationStatus.FAILED, error=ErrorObject.from_exception(error)

src/aws_durable_execution_sdk_python/state.py

@@ -16,6 +16,7 @@
     BackgroundThreadError,
     CallableRuntimeError,
     DurableExecutionsError,
+    GetExecutionStateError,
     OrphanedChildException,
 )
 from aws_durable_execution_sdk_python.lambda_service import (
@@ -275,20 +276,38 @@ def fetch_paginated_operations(
             initial_operations: initial operations to be added to ExecutionState
             checkpoint_token: checkpoint token used to call Durable Functions API.
             next_marker: a marker indicates that there are paginated operations.
+
+        Raises:
+            GetExecutionStateError: If the API call fails. The error is logged
+                with structured extras before re-raising. Callers are responsible
+                for deciding whether to fail the execution or allow Lambda retry
+                based on is_retryable().
         """
         all_operations: list[Operation] = (
             initial_operations.copy() if initial_operations else []
         )
-        while next_marker:
-            output: StateOutput = self._service_client.get_execution_state(
-                durable_execution_arn=self.durable_execution_arn,
-                checkpoint_token=checkpoint_token,
-                next_marker=next_marker,
+        try:
+            while next_marker:
+                output: StateOutput = self._service_client.get_execution_state(
+                    durable_execution_arn=self.durable_execution_arn,
+                    checkpoint_token=checkpoint_token,
+                    next_marker=next_marker,
+                )
+                all_operations.extend(output.operations)
+                next_marker = output.next_marker
+        except GetExecutionStateError as e:
+            logger.exception(
+                "Durable API error during state fetch.",
+                extra=e.build_logger_extras(),
             )
-            all_operations.extend(output.operations)
-            next_marker = output.next_marker
-        with self._operations_lock:
-            self.operations.update({op.operation_id: op for op in all_operations})
+            raise
+        finally:
+            # Always store whatever operations we successfully fetched
+            if all_operations:
+                with self._operations_lock:
+                    self.operations.update(
+                        {op.operation_id: op for op in all_operations}
+                    )
 
     def get_input_payload(self) -> str | None:
         # It is possible that backend will not provide an execution operation