fix(testing): deliver due waits on checkpoint

packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/core.py

@@ -95,6 +95,11 @@ def apply(
         ``now`` is the checkpoint's single source of "now", resolved from
         the execution's clock, so all timestamps stamped by this apply
         advance with modeled time under a skip clock.
+
+        Before computing the delta, completes every wait and step retry
+        whose scheduled moment has passed, so those transitions are
+        delivered in this response and a suspended branch resumes within
+        the running invocation instead of waiting for the next one.
         """
         effects: list[CheckpointEffect] = []
         if updates:
@@ -109,6 +114,11 @@ def apply(
 
         new_token_sequence: int = execution.advance_token_sequence()
 
+        # Pick up scheduler-due completions before pinning the paginator:
+        # each transition touches its operation, so it lands in this
+        # response's unseen delta and the handler sees it immediately.
+        execution.complete_due_operations(now)
+
         paginator: OperationPaginatorState = OperationPaginatorState.pin(execution)
         # The checkpoint response returns the full unseen delta in a single
         # response. Advance handler_seen_seq to cover every returned op so

packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from dataclasses import dataclass, replace
 from datetime import datetime
 from enum import Enum
@@ -35,6 +36,8 @@
     CheckpointToken,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class ExecutionStatus(Enum):
     """Execution status for API responses."""
@@ -500,6 +503,54 @@ def complete_retry(self, operation_id: str) -> Operation:
             self._record_updated_operation(operation_id)
             return updated_operation
 
+    def complete_due_operations(self, now: datetime) -> bool:
+        """Complete every operation whose scheduled moment has passed.
+
+        Transitions due WAIT operations (``STARTED`` with a
+        ``scheduled_end_timestamp`` at or before ``now``) via
+        :meth:`complete_wait` and due STEP retries (``PENDING`` with a
+        ``next_attempt_timestamp`` at or before ``now``) via
+        :meth:`complete_retry`. A transition failure is logged and
+        skipped so one bad operation cannot block the rest.
+
+        Returns True when at least one operation completed.
+        """
+        completed_any: bool = False
+        for op in list(self.operations):
+            if (
+                op.operation_type is OperationType.WAIT
+                and op.status is OperationStatus.STARTED
+                and op.wait_details is not None
+                and op.wait_details.scheduled_end_timestamp is not None
+                and op.wait_details.scheduled_end_timestamp <= now
+            ):
+                try:
+                    self.complete_wait(op.operation_id, now=now)
+                    completed_any = True
+                except Exception:  # noqa: BLE001 — skip one bad op, keep the rest
+                    logger.exception(
+                        "[%s] complete_wait failed for due operation %s",
+                        self.durable_execution_arn,
+                        op.operation_id,
+                    )
+            elif (
+                op.operation_type is OperationType.STEP
+                and op.status is OperationStatus.PENDING
+                and op.step_details is not None
+                and op.step_details.next_attempt_timestamp is not None
+                and op.step_details.next_attempt_timestamp <= now
+            ):
+                try:
+                    self.complete_retry(op.operation_id)
+                    completed_any = True
+                except Exception:  # noqa: BLE001 — skip one bad op, keep the rest
+                    logger.exception(
+                        "[%s] complete_retry failed for due operation %s",
+                        self.durable_execution_arn,
+                        op.operation_id,
+                    )
+        return completed_any
+
     def complete_callback_success(
         self,
         callback_id: str,

packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py

@@ -1036,41 +1036,7 @@ def _fire_due_operations(self, execution_arn: str) -> bool:
         if execution.is_complete:
             return False
 
-        now = self._clock.now()
-        completed_any = False
-        for op in list(execution.operations):
-            if (
-                op.operation_type == OperationType.WAIT
-                and op.status == OperationStatus.STARTED
-                and op.wait_details is not None
-                and op.wait_details.scheduled_end_timestamp is not None
-                and op.wait_details.scheduled_end_timestamp <= now
-            ):
-                try:
-                    execution.complete_wait(op.operation_id, now=now)
-                    completed_any = True
-                except Exception:  # noqa: BLE001
-                    logger.exception(
-                        "[%s] earliest-pending: complete_wait failed for %s",
-                        execution_arn,
-                        op.operation_id,
-                    )
-            elif (
-                op.operation_type == OperationType.STEP
-                and op.status == OperationStatus.PENDING
-                and op.step_details is not None
-                and op.step_details.next_attempt_timestamp is not None
-                and op.step_details.next_attempt_timestamp <= now
-            ):
-                try:
-                    execution.complete_retry(op.operation_id)
-                    completed_any = True
-                except Exception:  # noqa: BLE001
-                    logger.exception(
-                        "[%s] earliest-pending: complete_retry failed for %s",
-                        execution_arn,
-                        op.operation_id,
-                    )
+        completed_any = execution.complete_due_operations(self._clock.now())
         if completed_any:
             self._store.update(execution)
         return completed_any

packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/processor_test.py

@@ -1,15 +1,19 @@
 """Unit tests for CheckpointProcessor."""
 
+from datetime import UTC, datetime, timedelta
 from unittest.mock import Mock, patch
 
 import pytest
 from aws_durable_execution_sdk_python.lambda_service import (
     CheckpointOutput,
     CheckpointUpdatedExecutionState,
+    Operation,
     OperationAction,
+    OperationStatus,
     OperationType,
     OperationUpdate,
     StateOutput,
+    WaitDetails,
 )
 
 from aws_durable_execution_sdk_python_testing.checkpoint.processor import (
@@ -354,3 +358,53 @@ def test_checkpoint_from_superseded_invocation_is_rejected():
     # The current invocation's checkpoint is accepted.
     result = processor.process_checkpoint(current_token, updates, "client-token")
     assert isinstance(result, CheckpointOutput)
+
+
+def test_process_checkpoint_delivers_due_wait_completion() -> None:
+    """A wait whose scheduled end has passed is completed by the
+    checkpoint and returned in the same response delta."""
+    store = InMemoryExecutionStore()
+    scheduler = Mock(spec=Scheduler)
+    processor = CheckpointProcessor(store, scheduler)
+
+    start_input = StartDurableExecutionInput(
+        account_id="123456789012",
+        function_name="test-function",
+        function_qualifier="$LATEST",
+        execution_name="test-execution",
+        execution_timeout_seconds=300,
+        execution_retention_period_days=7,
+        invocation_id="test-inv-id",
+    )
+    execution = Execution.new(start_input)
+    execution.start()
+
+    past: datetime = datetime.now(UTC) - timedelta(seconds=5)
+    execution.operations.append(
+        Operation(
+            operation_id="wait-1",
+            parent_id=None,
+            name="due-wait",
+            start_timestamp=past,
+            operation_type=OperationType.WAIT,
+            status=OperationStatus.STARTED,
+            wait_details=WaitDetails(scheduled_end_timestamp=past),
+        )
+    )
+    store.save(execution)
+
+    token: str = CheckpointToken(
+        execution_arn=execution.durable_execution_arn, token_sequence=0
+    ).to_str()
+
+    result: CheckpointOutput = processor.process_checkpoint(token, [], None)
+
+    returned = {op.operation_id: op for op in result.new_execution_state.operations}
+    assert "wait-1" in returned
+    assert returned["wait-1"].status is OperationStatus.SUCCEEDED
+
+    persisted = store.load(execution.durable_execution_arn)
+    persisted_wait = next(
+        op for op in persisted.operations if op.operation_id == "wait-1"
+    )
+    assert persisted_wait.status is OperationStatus.SUCCEEDED

packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_inline_completion_test.py

@@ -0,0 +1,123 @@
+"""End-to-end test proving a wait inside a parallel branch completes inline.
+
+Demonstrates that the local runner matches the real service: a wait
+started in one parallel branch completes while its sibling is still
+working, within a single invocation. Without the fix, the wait would
+only complete after the invocation returns and a second invocation
+re-checks it.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from typing import Any
+
+import pytest
+from aws_durable_execution_sdk_python.config import Duration
+from aws_durable_execution_sdk_python.context import DurableContext, durable_step
+from aws_durable_execution_sdk_python.execution import (
+    InvocationStatus,
+    durable_execution,
+)
+from aws_durable_execution_sdk_python.types import StepContext
+
+from aws_durable_execution_sdk_python_testing.runner import (
+    DurableFunctionTestResult,
+    DurableFunctionTestRunner,
+)
+
+
+# Wall-clock seconds branch B keeps working. Must exceed the modeled
+# wait plus the SDK's branch-retry overhead so branch A resumes while
+# B is still running.
+_SIBLING_WORK_SECONDS: float = 3.0
+
+# Modeled wait duration for branch A (minimum 1 second).
+_WAIT_SECONDS: int = 1
+
+
+@durable_step
+def do_work(step_context: StepContext, seconds: float) -> str:  # noqa: ARG001
+    """Simulate long-running work by sleeping."""
+    time.sleep(seconds)
+    return "work-done"
+
+
+# Timeline collected during the execution for post-hoc assertions.
+_timeline: list[tuple[float, str]] = []
+
+
+def _record(start: float, msg: str) -> None:
+    _timeline.append((time.time() - start, msg))
+
+
+@durable_execution
+def parallel_wait_and_work(event: Any, context: DurableContext) -> list[str]:  # noqa: ARG001
+    """Two parallel branches: A waits, B does work."""
+    start: float = time.time()
+    _record(start, "INVOCATION START")
+
+    def branch_a(ctx: DurableContext) -> str:
+        _record(start, "A: before wait")
+        ctx.wait(Duration.from_seconds(_WAIT_SECONDS))
+        _record(start, "A: after wait")
+        return "a-done"
+
+    def branch_b(ctx: DurableContext) -> str:
+        _record(start, "B: start work")
+        ctx.step(do_work(_SIBLING_WORK_SECONDS))
+        _record(start, "B: finished work")
+        return "b-done"
+
+    batch = context.parallel(functions=[branch_a, branch_b], name="test-parallel")
+    results: list[str] = batch.get_results()
+    _record(start, "parallel returned")
+    return results
+
+
+@pytest.mark.parametrize("skip_time", [True, False])
+def test_wait_completes_inline_during_parallel_invocation(skip_time: bool) -> None:  # noqa: FBT001
+    """A wait in one parallel branch completes while the sibling is still
+    working, within a single invocation.
+
+    The wait's completion is delivered through a checkpoint response
+    while the invocation is in flight, so the waiting branch resumes
+    without an extra invocation. Runs on both the skip clock and the
+    wall clock.
+    """
+    _timeline.clear()
+
+    with DurableFunctionTestRunner(
+        handler=parallel_wait_and_work,
+        execution_timeout=15,
+        skip_time=skip_time,
+    ) as runner:
+        result: DurableFunctionTestResult = runner.run(input="x")
+
+    assert result.status is InvocationStatus.SUCCEEDED
+    assert result.result == json.dumps(["a-done", "b-done"])
+
+    # Branch A must resume from its wait BEFORE branch B finishes its
+    # work. This proves the wait completed inline while the sibling was
+    # still working.
+    a_after_wait_times = [t for t, msg in _timeline if msg == "A: after wait"]
+    b_finished_times = [t for t, msg in _timeline if msg == "B: finished work"]
+
+    assert len(a_after_wait_times) >= 1, "Branch A never resumed after wait"
+    assert len(b_finished_times) >= 1, "Branch B never finished"
+
+    a_after_wait: float = a_after_wait_times[0]
+    b_finished: float = b_finished_times[0]
+
+    assert a_after_wait < b_finished, (
+        f"Branch A resumed at {a_after_wait:.2f}s but B finished at "
+        f"{b_finished:.2f}s — wait did not complete inline"
+    )
+
+    # Only ONE invocation should be recorded. Look for exactly one
+    # "INVOCATION START" event.
+    invocation_starts = [msg for _, msg in _timeline if msg == "INVOCATION START"]
+    assert len(invocation_starts) == 1, (
+        f"Expected 1 invocation, got {len(invocation_starts)}"
+    )