feat: exit early when pending

- suspend execution of step and wait_for_condition when the checkpointed result is pending
- test cases

Author: Alex Wang

src/aws_durable_execution_sdk_python/operation/step.py

@@ -14,9 +14,14 @@
 from aws_durable_execution_sdk_python.exceptions import (
     FatalError,
     StepInterruptedError,
+    SuspendExecution,
     TimedSuspendExecution,
 )
-from aws_durable_execution_sdk_python.lambda_service import ErrorObject, OperationUpdate
+from aws_durable_execution_sdk_python.lambda_service import (
+    ErrorObject,
+    Operation,
+    OperationUpdate,
+)
 from aws_durable_execution_sdk_python.logger import Logger, LogInfo
 from aws_durable_execution_sdk_python.retries import RetryPresets
 from aws_durable_execution_sdk_python.serdes import deserialize, serialize
@@ -52,7 +57,9 @@ def step_handler(
     if not config:
         config = StepConfig()
 
-    checkpointed_result = state.get_checkpoint_result(operation_identifier.operation_id)
+    checkpointed_result: CheckpointedResult = state.get_checkpoint_result(
+        operation_identifier.operation_id
+    )
     if checkpointed_result.is_succeeded():
         logger.debug(
             "Step already completed, skipping execution for id: %s, name: %s",
@@ -73,6 +80,12 @@ def step_handler(
         # have to throw the exact same error on replay as the checkpointed failure
         checkpointed_result.raise_callable_error()
 
+    if checkpointed_result.is_pending():
+        _suspend_with_operation(
+            operation_identifier=operation_identifier,
+            operation=checkpointed_result.operation,
+        )
+
     if checkpointed_result.is_started():
         # step was previously interrupted
         if config.step_semantics is StepSemantics.AT_MOST_ONCE_PER_RETRY:
@@ -193,7 +206,7 @@ def retry_handler(
 
         state.create_checkpoint(operation_update=retry_operation)
 
-        _suspend(operation_identifier, retry_decision)
+        _suspend_with_decision(operation_identifier, retry_decision)
 
     # no retry
     fail_operation: OperationUpdate = OperationUpdate.create_step_fail(
@@ -208,10 +221,30 @@ def retry_handler(
     raise error_object.to_callable_runtime_error()
 
 
-def _suspend(operation_identifier: OperationIdentifier, retry_decision: RetryDecision):
+def _suspend_with_decision(
+    operation_identifier: OperationIdentifier, retry_decision: RetryDecision
+):
     scheduled_timestamp = time.time() + retry_decision.delay_seconds
     msg = f"Retry scheduled for {operation_identifier.operation_id} in {retry_decision.delay_seconds} seconds"
     raise TimedSuspendExecution(
         msg,
         scheduled_timestamp=scheduled_timestamp,
     )
+
+
+def _suspend_with_operation(
+    operation_identifier: OperationIdentifier, operation: Operation | None
+) -> None:
+    if (
+        operation is None
+        or operation.step_details is None
+        or operation.step_details.next_attempt_timestamp is None
+    ):
+        msg = f"next_attempt_timestamp is None for {operation_identifier.operation_id}"
+        raise SuspendExecution(msg)
+    scheduled_timestamp = float(operation.step_details.next_attempt_timestamp)
+    msg = f"Retry scheduled for {operation_identifier.name or operation_identifier.operation_id} will retry at timestamp {scheduled_timestamp}"
+    raise TimedSuspendExecution(
+        msg,
+        scheduled_timestamp=scheduled_timestamp,
+    )

src/aws_durable_execution_sdk_python/operation/wait_for_condition.py

@@ -8,9 +8,14 @@
 
 from aws_durable_execution_sdk_python.exceptions import (
     FatalError,
+    SuspendExecution,
     TimedSuspendExecution,
 )
-from aws_durable_execution_sdk_python.lambda_service import ErrorObject, OperationUpdate
+from aws_durable_execution_sdk_python.lambda_service import (
+    ErrorObject,
+    Operation,
+    OperationUpdate,
+)
 from aws_durable_execution_sdk_python.logger import LogInfo
 from aws_durable_execution_sdk_python.serdes import deserialize, serialize
 from aws_durable_execution_sdk_python.types import WaitForConditionCheckContext
@@ -24,7 +29,10 @@
     )
     from aws_durable_execution_sdk_python.identifier import OperationIdentifier
     from aws_durable_execution_sdk_python.logger import Logger
-    from aws_durable_execution_sdk_python.state import ExecutionState
+    from aws_durable_execution_sdk_python.state import (
+        CheckpointedResult,
+        ExecutionState,
+    )
 
 
 T = TypeVar("T")
@@ -49,7 +57,9 @@ def wait_for_condition_handler(
         operation_identifier.name,
     )
 
-    checkpointed_result = state.get_checkpoint_result(operation_identifier.operation_id)
+    checkpointed_result: CheckpointedResult = state.get_checkpoint_result(
+        operation_identifier.operation_id
+    )
 
     # Check if already completed
     if checkpointed_result.is_succeeded():
@@ -70,6 +80,12 @@ def wait_for_condition_handler(
     if checkpointed_result.is_failed():
         checkpointed_result.raise_callable_error()
 
+    if checkpointed_result.is_pending():
+        _suspend_execution_with_operation(
+            operation_identifier=operation_identifier,
+            operation=checkpointed_result.operation,
+        )
+
     attempt: int = 1
     if checkpointed_result.is_started_or_ready():
         # This is a retry - get state from previous checkpoint
@@ -164,7 +180,7 @@ def wait_for_condition_handler(
 
         state.create_checkpoint(operation_update=retry_operation)
 
-        _suspend_execution(operation_identifier, decision)
+        _suspend_execution_with_decision(operation_identifier, decision)
 
     except Exception as e:
         # Mark as failed - waitForCondition doesn't have its own retry logic for errors
@@ -186,7 +202,7 @@ def wait_for_condition_handler(
     raise FatalError(msg)
 
 
-def _suspend_execution(
+def _suspend_execution_with_decision(
     operation_identifier: OperationIdentifier, decision: WaitForConditionDecision
 ) -> None:
     scheduled_timestamp = time.time() + (decision.delay_seconds or 0)
@@ -195,3 +211,21 @@ def _suspend_execution(
         msg,
         scheduled_timestamp=scheduled_timestamp,
     )
+
+
+def _suspend_execution_with_operation(
+    operation_identifier: OperationIdentifier, operation: Operation | None
+) -> None:
+    if (
+        operation is None
+        or operation.step_details is None
+        or operation.step_details.next_attempt_timestamp is None
+    ):
+        msg = f"next_attempt_timestamp is None for {operation_identifier.operation_id}"
+        raise SuspendExecution(msg)
+    scheduled_timestamp = float(operation.step_details.next_attempt_timestamp)
+    msg = f"wait_for_condition {operation_identifier.name or operation_identifier.operation_id} will retry at timestamp {scheduled_timestamp}"
+    raise TimedSuspendExecution(
+        msg,
+        scheduled_timestamp=scheduled_timestamp,
+    )

src/aws_durable_execution_sdk_python/state.py

@@ -106,6 +106,13 @@ def is_started_or_ready(self) -> bool:
             return False
         return op.status in (OperationStatus.STARTED, OperationStatus.READY)
 
+    def is_pending(self) -> bool:
+        """Return True if the checkpointed operation is PENDING."""
+        op = self.operation
+        if not op:
+            return False
+        return op.status is OperationStatus.PENDING
+
     def is_timed_out(self) -> bool:
         """Return True if the checkpointed operation is TIMED_OUT."""
         op = self.operation

tests/operation/step_test.py

@@ -375,7 +375,7 @@ def test_step_handler_retry_with_existing_attempts():
         operation_id="step12",
         operation_type=OperationType.STEP,
         status=OperationStatus.PENDING,
-        step_details=StepDetails(attempt=2),
+        step_details=StepDetails(attempt=2, next_attempt_timestamp="1764547200"),
     )
     mock_result = CheckpointedResult.create_from_operation(operation)
     mock_state.get_checkpoint_result.return_value = mock_result
@@ -398,10 +398,44 @@ def test_step_handler_retry_with_existing_attempts():
             mock_logger,
         )
 
-    # Verify retry strategy was called with correct attempt count (2 + 1 = 3)
-    mock_retry_strategy.assert_called_once()
-    call_args = mock_retry_strategy.call_args[0]
-    assert call_args[1] == 3  # retry_attempt + 1
+    # Verify retry strategy was not called because we already have attempt timestamp in the checkpointed location
+    mock_retry_strategy.assert_not_called()
+
+
+def test_step_handler_pending_without_existing_attempts():
+    """Test step_handler retry logic with existing attempt count."""
+    mock_state = Mock(spec=ExecutionState)
+
+    # Simulate a retry operation that was previously checkpointed
+    operation = Operation(
+        operation_id="step12",
+        operation_type=OperationType.STEP,
+        status=OperationStatus.PENDING,
+        step_details=StepDetails(attempt=2),
+    )
+    mock_result = CheckpointedResult.create_from_operation(operation)
+    mock_state.get_checkpoint_result.return_value = mock_result
+    mock_state.durable_execution_arn = "test_arn"
+
+    mock_retry_strategy = Mock(
+        return_value=RetryDecision(should_retry=True, delay_seconds=10)
+    )
+    config = StepConfig(retry_strategy=mock_retry_strategy)
+    mock_callable = Mock(side_effect=RuntimeError("Test error"))
+    mock_logger = Mock(spec=Logger)
+    mock_logger.with_log_info.return_value = mock_logger
+
+    with pytest.raises(SuspendExecution, match="next_attempt_timestamp is None"):
+        step_handler(
+            mock_callable,
+            mock_state,
+            OperationIdentifier("step12", None, "test_step"),
+            config,
+            mock_logger,
+        )
+
+    # Verify retry strategy was not called because we already have attempt timestamp in the checkpointed location
+    mock_retry_strategy.assert_not_called()
 
 
 @patch("aws_durable_execution_sdk_python.operation.step.retry_handler")

tests/operation/wait_for_condition_test.py

@@ -11,6 +11,7 @@
 )
 from aws_durable_execution_sdk_python.exceptions import (
     CallableRuntimeError,
+    FatalError,
     SuspendExecution,
 )
 from aws_durable_execution_sdk_python.identifier import OperationIdentifier
@@ -372,6 +373,7 @@ def test_wait_for_condition_no_operation_in_checkpoint():
     mock_result = Mock()
     mock_result.is_succeeded.return_value = False
     mock_result.is_failed.return_value = False
+    mock_result.is_pending.return_value = False
     mock_result.is_started_or_ready.return_value = True
     mock_result.is_existent.return_value = True
     mock_result.result = json.dumps(10)
@@ -416,6 +418,7 @@ def test_wait_for_condition_operation_no_step_details():
     mock_result = Mock()
     mock_result.is_succeeded.return_value = False
     mock_result.is_failed.return_value = False
+    mock_result.is_pending.return_value = False
     mock_result.is_started_or_ready.return_value = True
     mock_result.is_existent.return_value = True
     mock_result.result = json.dumps(10)
@@ -651,3 +654,72 @@ def check_func(state, context):
     )
 
     assert result == {"key": "value", "number": 42, "list": [1, 2, 3]}
+
+
+def test_wait_for_condition_pending():
+    mock_state = Mock(spec=ExecutionState)
+    mock_state.durable_execution_arn = "arn:aws:test"
+    operation = Operation(
+        operation_id="XXX",
+        operation_type=OperationType.STEP,
+        status=OperationStatus.PENDING,
+        step_details=StepDetails(
+            result='{"key": "VALUE", "number": "84", "list": [1, 2, 3]}',
+            next_attempt_timestamp="1764547200",
+        ),
+    )
+    mock_result = CheckpointedResult.create_from_operation(operation)
+    mock_state.get_checkpoint_result.return_value = mock_result
+
+    mock_logger = Mock(spec=Logger)
+    mock_logger.with_log_info.return_value = mock_logger
+
+    op_id = OperationIdentifier("op1", None, "test_wait")
+
+    def check_func(state, context):
+        msg = "Should not be called"
+        raise FatalError(msg)
+
+    config = WaitForConditionConfig(
+        initial_state=5,
+        wait_strategy=lambda s, a: WaitForConditionDecision.stop_polling(),
+        serdes=CustomDictSerDes(),
+    )
+
+    with pytest.raises(
+        SuspendExecution, match="wait_for_condition test_wait will retry at timestamp"
+    ):
+        wait_for_condition_handler(check_func, config, mock_state, op_id, mock_logger)
+
+
+def test_wait_for_condition_pending_without_next_attempt():
+    mock_state = Mock(spec=ExecutionState)
+    mock_state.durable_execution_arn = "arn:aws:test"
+    operation = Operation(
+        operation_id="XXX",
+        operation_type=OperationType.STEP,
+        status=OperationStatus.PENDING,
+        step_details=StepDetails(
+            result='{"key": "VALUE", "number": "84", "list": [1, 2, 3]}',
+        ),
+    )
+    mock_result = CheckpointedResult.create_from_operation(operation)
+    mock_state.get_checkpoint_result.return_value = mock_result
+
+    mock_logger = Mock(spec=Logger)
+    mock_logger.with_log_info.return_value = mock_logger
+
+    op_id = OperationIdentifier("op1", None, "test_wait")
+
+    def check_func(state, context):
+        msg = "Should not be called"
+        raise FatalError(msg)
+
+    config = WaitForConditionConfig(
+        initial_state=5,
+        wait_strategy=lambda s, a: WaitForConditionDecision.stop_polling(),
+        serdes=CustomDictSerDes(),
+    )
+
+    with pytest.raises(SuspendExecution, match="next_attempt_timestamp is None"):
+        wait_for_condition_handler(check_func, config, mock_state, op_id, mock_logger)

tests/state_test.py

@@ -214,6 +214,21 @@ def test_checkpointed_result_is_failed():
     assert result_no_op.is_failed() is False
 
 
+def test_checkpointerd_result_is_pending():
+    """Test CheckpointedResult.is_pending method."""
+    operation = Operation(
+        operation_id="op1",
+        operation_type=OperationType.STEP,
+        status=OperationStatus.PENDING,
+    )
+    result = CheckpointedResult.create_from_operation(operation)
+    assert result.is_pending() is True
+
+    # Test with no operation
+    result_no_op = CheckpointedResult.create_not_found()
+    assert result_no_op.is_pending() is False
+
+
 def test_checkpointed_result_is_started():
     """Test CheckpointedResult.is_started method."""
     operation = Operation(