From 45de84178d4e907cf3c83f4118bd6dd3185bc0d8 Mon Sep 17 00:00:00 2001
From: Sanjay Santhanam <51058514+Sanjays2402@users.noreply.github.com>
Date: Wed, 22 Jul 2026 22:35:08 -0700
Subject: [PATCH] fix: fail invalid polling state

Propagate checkpoint deserialization failures through the existing terminal operation error path instead of restarting from the initial state.\n\nUpdate the regression test to require a FAIL checkpoint and typed wait-for-condition error.
---
 .../operation/wait_for_condition.py           | 30 ++++++++-----------
 .../operation/wait_for_condition_test.py      | 26 +++++++++-------
 2 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/wait_for_condition.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/wait_for_condition.py
index 1a830541..3e7bd900 100644
--- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/wait_for_condition.py
+++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/wait_for_condition.py
@@ -152,33 +152,27 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T:
             Suspends if condition not met
             Raises error if check function fails
         """
-        # Determine current state from checkpoint
-        if checkpointed_result.is_started_or_ready() and checkpointed_result.result:
-            try:
+        try:
+            # Determine current state from checkpoint
+            if checkpointed_result.is_started_or_ready() and checkpointed_result.result:
                 current_state = deserialize(
                     serdes=self.config.serdes,
                     data=checkpointed_result.result,
                     operation_id=self.operation_identifier.operation_id,
                     durable_execution_arn=self.state.durable_execution_arn,
                 )
-            except Exception:
-                # Default to initial state if there's an error getting checkpointed state
-                logger.exception(
-                    "⚠️ wait_for_condition failed to deserialize state for id: %s, name: %s. Using initial state.",
-                    self.operation_identifier.operation_id,
-                    self.operation_identifier.name,
-                )
+            else:
                 current_state = self.config.initial_state
-        else:
-            current_state = self.config.initial_state
 
-        # Get attempt number - current attempt is checkpointed attempts + 1
-        # The checkpoint stores completed attempts, so the current attempt being executed is one more
-        attempt: int = 1
-        if checkpointed_result.operation and checkpointed_result.operation.step_details:
-            attempt = checkpointed_result.operation.step_details.attempt + 1
+            # Get attempt number - current attempt is checkpointed attempts + 1
+            # The checkpoint stores completed attempts, so the current attempt being executed is one more
+            attempt: int = 1
+            if (
+                checkpointed_result.operation
+                and checkpointed_result.operation.step_details
+            ):
+                attempt = checkpointed_result.operation.step_details.attempt + 1
 
-        try:
             # Execute the check function with the injected logger
             check_context = WaitForConditionCheckContext(
                 logger=self.context_logger.with_log_info(
diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py
index 363be21a..40f97a0d 100644
--- a/packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py
+++ b/packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py
@@ -332,8 +332,8 @@ def check_func(state, context):
     assert result == 6  # 5 (initial) + 1
 
 
-def test_wait_for_condition_retry_invalid_json_state():
-    """Test wait_for_condition on retry with invalid JSON state."""
+def test_wait_for_condition_retry_invalid_json_state_fails():
+    """Test invalid checkpointed state fails instead of restarting polling."""
     mock_state = Mock(spec=ExecutionState)
     mock_state.durable_execution_arn = "arn:aws:test"
     operation = Operation(
@@ -362,15 +362,21 @@ def check_func(state, context):
         wait_strategy=lambda s, a: WaitForConditionDecision.stop_polling(),
     )
 
-    result = wait_for_condition_handler(
-        state=mock_state,
-        operation_identifier=op_id,
-        check=check_func,
-        config=config,
-        context_logger=mock_logger,
-    )
+    with pytest.raises(WaitForConditionError) as exc_info:
+        wait_for_condition_handler(
+            state=mock_state,
+            operation_identifier=op_id,
+            check=check_func,
+            config=config,
+            context_logger=mock_logger,
+        )
 
-    assert result == 6  # Falls back to initial state
+    assert exc_info.value.error_type == "ExecutionError"
+    mock_state.create_checkpoint.assert_called_once()
+    assert (
+        mock_state.create_checkpoint.call_args.kwargs["operation_update"].action
+        == OperationAction.FAIL
+    )
 
 
 def test_wait_for_condition_check_function_exception():