fix(waits): repair strategy, fail on exhaustion

Author: AdityaAudi

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

@@ -242,7 +242,7 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T:
             if delay_seconds is not None and delay_seconds < 1:
                 logger.warning(
                     (
-                        "WaitDecision delay_seconds step for id: %s, name: %s,"
+                        "wait_for_condition delay_seconds for id: %s, name: %s,"
                         "is %d < 1. Setting to minimum of 1 seconds."
                     ),
                     self.operation_identifier.operation_id,

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

@@ -7,6 +7,7 @@
 from typing import TYPE_CHECKING, Generic
 
 from aws_durable_execution_sdk_python.config import Duration, JitterStrategy, T
+from aws_durable_execution_sdk_python.exceptions import WaitForConditionError
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -16,29 +17,6 @@
 Numeric = int | float
 
 
-@dataclass
-class WaitDecision:
-    """Decision about whether to wait a step and with what delay."""
-
-    should_wait: bool
-    delay: Duration
-
-    @property
-    def delay_seconds(self) -> int:
-        """Get delay in seconds."""
-        return self.delay.to_seconds()
-
-    @classmethod
-    def wait(cls, delay: Duration) -> WaitDecision:
-        """Create a wait decision."""
-        return cls(should_wait=True, delay=delay)
-
-    @classmethod
-    def no_wait(cls) -> WaitDecision:
-        """Create a no-wait decision."""
-        return cls(should_wait=False, delay=Duration())
-
-
 @dataclass
 class WaitStrategyConfig(Generic[T]):
     should_continue_polling: Callable[[T], bool]
@@ -69,35 +47,6 @@ def timeout_seconds(self) -> int | None:
         return self.timeout.to_seconds()
 
 
-def create_wait_strategy(
-    config: WaitStrategyConfig[T],
-) -> Callable[[T, int], WaitDecision]:
-    def wait_strategy(result: T, attempts_made: int) -> WaitDecision:
-        # Check if condition is met
-        if not config.should_continue_polling(result):
-            return WaitDecision.no_wait()
-
-        # Check if we've exceeded max attempts
-        if attempts_made >= config.max_attempts:
-            return WaitDecision.no_wait()
-
-        # Calculate delay with exponential backoff
-        base_delay: float = min(
-            config.initial_delay_seconds * (config.backoff_rate ** (attempts_made - 1)),
-            config.max_delay_seconds,
-        )
-
-        # Apply jitter to get final delay
-        delay_with_jitter: float = config.jitter_strategy.apply_jitter(base_delay)
-
-        # Round up and ensure minimum of 1 second
-        final_delay: int = max(1, math.ceil(delay_with_jitter))
-
-        return WaitDecision.wait(Duration(seconds=final_delay))
-
-    return wait_strategy
-
-
 @dataclass(frozen=True)
 class WaitForConditionDecision:
     """Decision about whether to continue waiting."""
@@ -121,6 +70,41 @@ def stop_polling(cls) -> WaitForConditionDecision:
         return cls(should_continue=False, delay=Duration())
 
 
+def create_wait_strategy(
+    config: WaitStrategyConfig[T],
+) -> Callable[[T, int], WaitForConditionDecision]:
+    def wait_strategy(result: T, attempts_made: int) -> WaitForConditionDecision:
+        # Condition satisfied wins over exhaustion, so a condition met on the final
+        # attempt still succeeds (matches the JS and Java SDKs).
+        if not config.should_continue_polling(result):
+            return WaitForConditionDecision.stop_polling()
+
+        # Out of attempts: fail rather than stop, otherwise the executor would treat
+        # this as success and return partial state.
+        if attempts_made >= config.max_attempts:
+            msg = (
+                f"wait_for_condition exhausted {config.max_attempts} attempts "
+                "before the condition was met"
+            )
+            raise WaitForConditionError(msg)
+
+        # Calculate delay with exponential backoff
+        base_delay: float = min(
+            config.initial_delay_seconds * (config.backoff_rate ** (attempts_made - 1)),
+            config.max_delay_seconds,
+        )
+
+        # Apply jitter to get final delay
+        delay_with_jitter: float = config.jitter_strategy.apply_jitter(base_delay)
+
+        # Round up and ensure minimum of 1 second
+        final_delay: int = max(1, math.ceil(delay_with_jitter))
+
+        return WaitForConditionDecision.continue_waiting(Duration(seconds=final_delay))
+
+    return wait_strategy
+
+
 @dataclass(frozen=True)
 class WaitForConditionConfig(Generic[T]):
     """Configuration for wait_for_condition."""

packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py

@@ -35,6 +35,8 @@
 from aws_durable_execution_sdk_python.waits import (
     WaitForConditionConfig,
     WaitForConditionDecision,
+    WaitStrategyConfig,
+    create_wait_strategy,
 )
 from tests.serdes_test import CustomDictSerDes
 
@@ -1571,6 +1573,95 @@ def mock_wait_strategy(state, attempt):
     assert mock_state.create_checkpoint.call_count == 2  # START + SUCCESS checkpoints
 
 
+def test_wait_for_condition_exhaustion_raises_and_checkpoints_fail():
+    """Live path: the built-in strategy runs out of attempts, so it raises
+    WaitForConditionError, which is checkpointed as a FAIL and propagated."""
+    mock_state = Mock(spec=ExecutionState)
+    mock_state.durable_execution_arn = "arn:aws:test"
+    mock_state.get_checkpoint_result.return_value = (
+        CheckpointedResult.create_not_found()
+    )
+
+    mock_logger = Mock(spec=Logger)
+    mock_logger.with_log_info.return_value = mock_logger
+
+    op_id = OperationIdentifier(
+        "op1", OperationSubType.WAIT_FOR_CONDITION, None, "test_wait"
+    )
+
+    def check_func(state, context):
+        return state + 1
+
+    mock_state.wrap_user_function.return_value = check_func
+
+    # max_attempts=1 means attempt 1 is already the last one.
+    config = WaitForConditionConfig(
+        initial_state=5,
+        wait_strategy=create_wait_strategy(
+            WaitStrategyConfig(should_continue_polling=lambda x: True, max_attempts=1)
+        ),
+    )
+
+    with pytest.raises(WaitForConditionError):
+        wait_for_condition_handler(
+            state=mock_state,
+            operation_identifier=op_id,
+            check=check_func,
+            config=config,
+            context_logger=mock_logger,
+        )
+
+    assert mock_state.create_checkpoint.call_count == 2  # START and FAIL
+    fail_operation = mock_state.create_checkpoint.call_args_list[1][1][
+        "operation_update"
+    ]
+    assert fail_operation.error.type == "WaitForConditionError"
+
+
+def test_wait_for_condition_exhaustion_surfaces_on_replay():
+    """Replay path: the FAILED checkpoint short-circuits on the next invocation
+    and is reconstructed as the typed WaitForConditionError carrying the original
+    error_type, without re-running the check."""
+    mock_state = Mock(spec=ExecutionState)
+    mock_state.durable_execution_arn = "test_arn"
+    operation = Operation(
+        operation_id="op1",
+        operation_type=OperationType.STEP,
+        status=OperationStatus.FAILED,
+        step_details=StepDetails(
+            error=ErrorObject("exhausted attempts", "WaitForConditionError", None, None)
+        ),
+    )
+    mock_result = CheckpointedResult.create_from_operation(operation)
+    mock_state.get_checkpoint_result.return_value = mock_result
+
+    mock_logger = Mock(spec=Logger)
+    op_id = OperationIdentifier(
+        "op1", OperationSubType.WAIT_FOR_CONDITION, None, "test_wait"
+    )
+
+    def check_func(state, context):
+        msg = "Check function should not be called on replay of a failure"
+        raise AssertionError(msg)
+
+    config = WaitForConditionConfig(
+        initial_state=5,
+        wait_strategy=lambda s, a: WaitForConditionDecision.stop_polling(),
+    )
+
+    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 exc_info.value.error_type == "WaitForConditionError"
+    assert mock_state.create_checkpoint.call_count == 0  # Nothing new on replay
+
+
 def test_wait_for_condition_executes_check_when_checkpoint_not_terminal_duplicate():
     """Test backward compatibility: when checkpoint is not terminal (STARTED),
     the wait_for_condition operation executes the check function normally.