fix(waits): repair built-in strategy, fail on exhaustion

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

@@ -21,6 +21,7 @@
     DurableExecutionsError,
     InvocationError,
     ValidationError,
+    WaitForConditionError,
 )
 
 # Core decorator - used in every durable function
@@ -39,6 +40,7 @@
     "ParallelBranch",
     "StepContext",
     "ValidationError",
+    "WaitForConditionError",
     "WithRetryConfig",
     "__version__",
     "durable_execution",

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

@@ -266,6 +266,12 @@ class InvalidStateError(DurableExecutionsError):
     """Raised when an operation is attempted on an object in an invalid state."""
 
 
+class WaitForConditionError(DurableExecutionsError):
+    """Raised when a wait_for_condition exhausts its polling attempts before the
+    condition is met, matching the exhaustion semantics of the JS and Java SDKs.
+    """
+
+
 class UserlandError(DurableExecutionsError):
     """Failure in user-land - i.e code passed into durable executions from the caller."""
 

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

@@ -239,7 +239,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

@@ -11,6 +11,7 @@
     CallableRuntimeError,
     InvocationError,
     SuspendExecution,
+    WaitForConditionError,
 )
 from aws_durable_execution_sdk_python.identifier import OperationIdentifier
 from aws_durable_execution_sdk_python.lambda_service import (
@@ -30,6 +31,8 @@
 from aws_durable_execution_sdk_python.waits import (
     WaitForConditionConfig,
     WaitForConditionDecision,
+    WaitStrategyConfig,
+    create_wait_strategy,
 )
 from tests.serdes_test import CustomDictSerDes
 
@@ -1386,6 +1389,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.
+    Typed reconstruction lands with #120, so today it surfaces as a
+    CallableRuntimeError carrying the original error_type."""
+    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(CallableRuntimeError) 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.