Add checkpoint error handling when final checkpointing fails

FullyTyped · FullyTyped · commit 1c513c071ebe · 2025-11-08T14:05:56.000Z
diff --git a/src/aws_durable_execution_sdk_python/exceptions.py b/src/aws_durable_execution_sdk_python/exceptions.py
@@ -8,7 +8,7 @@
 import time
 from dataclasses import dataclass
 from enum import Enum
-from typing import TYPE_CHECKING, Literal, Self, TypedDict
+from typing import TYPE_CHECKING, Self, TypedDict
 
 BAD_REQUEST_ERROR: int = 400
 SERVICE_ERROR: int = 500
@@ -124,7 +124,9 @@ def __init__(self, message: str, step_id: str | None = None):
         self.step_id = step_id
 
 
-CheckpointErrorKind = Literal["Execution", "Invocation"]
+class CheckpointErrorCategory(Enum):
+    INVOCATION = "INVOCATION"
+    EXECUTION = "EXECUTION"
 
 
 class CheckpointError(BotoClientError):
@@ -133,7 +135,7 @@ class CheckpointError(BotoClientError):
     def __init__(
         self,
         message: str,
-        error_kind: CheckpointErrorKind,
+        error_category: CheckpointErrorCategory,
         error: AwsErrorObj | None = None,
         response_metadata: AwsErrorMetadata | None = None,
     ):
@@ -143,14 +145,14 @@ def __init__(
             response_metadata,
             termination_reason=TerminationReason.CHECKPOINT_FAILED,
         )
-        self.error_kind: CheckpointErrorKind = error_kind
+        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_kind: CheckpointErrorKind = "Invocation"
+        error_category: CheckpointErrorCategory = CheckpointErrorCategory.INVOCATION
 
         # InvalidParameterValueException and error message starts with "Invalid Checkpoint Token" is an InvocationError
         # all other 4xx errors are Execution Errors and should be retried
@@ -172,11 +174,11 @@ def from_exception(cls, exception: Exception) -> CheckpointError:
                 )
             )
         ):
-            error_kind = "Execution"
-        return CheckpointError(str(exception), error_kind, error, metadata)
+            error_category = CheckpointErrorCategory.EXECUTION
+        return CheckpointError(str(exception), error_category, error, metadata)
 
-    def should_be_retried(self):
-        return self.error_kind == "Execution"
+    def is_retriable(self):
+        return self.error_category == CheckpointErrorCategory.EXECUTION
 
 
 class ValidationError(DurableExecutionsError):
diff --git a/src/aws_durable_execution_sdk_python/execution.py b/src/aws_durable_execution_sdk_python/execution.py
@@ -281,7 +281,6 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
                     invocation_input.durable_execution_arn,
                 )
                 serialized_result = json.dumps(result)
-
                 # large response handling here. Remember if checkpointing to complete, NOT to include
                 # payload in response
                 if (
@@ -300,8 +299,12 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
                     # Must ensure the result is persisted before returning to Lambda.
                     # Large results exceed Lambda response limits and must be stored durably
                     # before the execution completes.
-                    execution_state.create_checkpoint(success_operation, is_sync=True)
-
+                    try:
+                        execution_state.create_checkpoint(
+                            success_operation, is_sync=True
+                        )
+                    except CheckpointError as e:
+                        return handle_checkpoint_error(e).to_dict()
                     return DurableExecutionInvocationOutput.create_succeeded(
                         result=""
                     ).to_dict()
@@ -320,17 +323,9 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
                     )
                 else:
                     logger.exception("Checkpoint processing failed")
-                # Raise the original exception
-                if (
-                    isinstance(bg_error.source_exception, CheckpointError)
-                    and bg_error.source_exception.should_be_retried()
-                ):
-                    raise bg_error.source_exception from None  # Terminate Lambda immediately and have it be retried
+                # handle the original exception
                 if isinstance(bg_error.source_exception, CheckpointError):
-                    return DurableExecutionInvocationOutput(
-                        status=InvocationStatus.FAILED,
-                        error=ErrorObject.from_exception(bg_error.source_exception),
-                    ).to_dict()
+                    return handle_checkpoint_error(bg_error.source_exception).to_dict()
                 raise bg_error.source_exception from bg_error
 
             except SuspendExecution:
@@ -346,11 +341,7 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
                     "Checkpoint system failed",
                     extra=e.build_logger_extras(),
                 )
-                if e.should_be_retried():
-                    raise  # Terminate Lambda immediately and have it be retried
-                return DurableExecutionInvocationOutput(
-                    status=InvocationStatus.FAILED, error=ErrorObject.from_exception(e)
-                ).to_dict()
+                return handle_checkpoint_error(e).to_dict()
             except InvocationError:
                 logger.exception("Invocation error. Must terminate.")
                 # Throw the error to trigger Lambda retry
@@ -388,12 +379,22 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
                     # Must ensure the result is persisted before returning to Lambda.
                     # Large results exceed Lambda response limits and must be stored durably
                     # before the execution completes.
-                    execution_state.create_checkpoint_sync(failed_operation)
-
+                    try:
+                        execution_state.create_checkpoint_sync(failed_operation)
+                    except CheckpointError as e:
+                        return handle_checkpoint_error(e).to_dict()
                     return DurableExecutionInvocationOutput(
                         status=InvocationStatus.FAILED
                     ).to_dict()
 
                 return result
 
     return wrapper
+
+
+def handle_checkpoint_error(error: CheckpointError) -> DurableExecutionInvocationOutput:
+    if error.is_retriable():
+        raise error from None  # Terminate Lambda immediately and have it be retried
+    return DurableExecutionInvocationOutput(
+        status=InvocationStatus.FAILED, error=ErrorObject.from_exception(error)
+    )
diff --git a/tests/exceptions_test.py b/tests/exceptions_test.py
@@ -10,6 +10,7 @@
     CallableRuntimeError,
     CallableRuntimeErrorSerializableDetails,
     CheckpointError,
+    CheckpointErrorCategory,
     DurableExecutionsError,
     ExecutionError,
     InvocationError,
@@ -42,7 +43,9 @@ def test_invocation_error():
 
 def test_checkpoint_error():
     """Test CheckpointError exception."""
-    error = CheckpointError("checkpoint failed", error_kind="Execution")
+    error = CheckpointError(
+        "checkpoint failed", error_category=CheckpointErrorCategory.EXECUTION
+    )
     assert str(error) == "checkpoint failed"
     assert isinstance(error, InvocationError)
     assert isinstance(error, UnrecoverableError)
@@ -62,8 +65,8 @@ def test_checkpoint_error_classification_invalid_token_invocation():
 
     result = CheckpointError.from_exception(client_error)
 
-    assert result.error_kind == "Invocation"
-    assert not result.should_be_retried()
+    assert result.error_category == CheckpointErrorCategory.INVOCATION
+    assert not result.is_retriable()
 
 
 def test_checkpoint_error_classification_other_4xx_execution():
@@ -76,8 +79,8 @@ def test_checkpoint_error_classification_other_4xx_execution():
 
     result = CheckpointError.from_exception(client_error)
 
-    assert result.error_kind == "Execution"
-    assert result.should_be_retried()
+    assert result.error_category == CheckpointErrorCategory.EXECUTION
+    assert result.is_retriable()
 
 
 def test_checkpoint_error_classification_invalid_param_without_token_execution():
@@ -93,8 +96,8 @@ def test_checkpoint_error_classification_invalid_param_without_token_execution()
 
     result = CheckpointError.from_exception(client_error)
 
-    assert result.error_kind == "Execution"
-    assert result.should_be_retried()
+    assert result.error_category == CheckpointErrorCategory.EXECUTION
+    assert result.is_retriable()
 
 
 def test_checkpoint_error_classification_5xx_invocation():
@@ -107,8 +110,8 @@ def test_checkpoint_error_classification_5xx_invocation():
 
     result = CheckpointError.from_exception(client_error)
 
-    assert result.error_kind == "Invocation"
-    assert not result.should_be_retried()
+    assert result.error_category == CheckpointErrorCategory.INVOCATION
+    assert not result.is_retriable()
 
 
 def test_checkpoint_error_classification_unknown_invocation():
@@ -117,8 +120,8 @@ def test_checkpoint_error_classification_unknown_invocation():
 
     result = CheckpointError.from_exception(unknown_error)
 
-    assert result.error_kind == "Invocation"
-    assert not result.should_be_retried()
+    assert result.error_category == CheckpointErrorCategory.INVOCATION
+    assert not result.is_retriable()
 
 
 def test_validation_error():
diff --git a/tests/execution_test.py b/tests/execution_test.py
