feat(testing): invoker state and checkpoint stamping

packages/aws-durable-execution-sdk-python-examples/test/conftest.py

@@ -106,14 +106,23 @@ def run(
         input: str | None = None,  # noqa: A002
         timeout: int = 60,
     ) -> DurableFunctionTestResult:
-        """Execute the durable function and return results."""
+        """Execute the durable function and return results.
+
+        The caller's ``timeout`` is the execution deadline. Local runs
+        pass it as ``execution_timeout``; cloud runs pass it as the
+        runner's ``timeout``.
+        """
+        if isinstance(self._runner, DurableFunctionTestRunner):
+            return self._runner.run(input=input, execution_timeout=timeout)
         return self._runner.run(input=input, timeout=timeout)
 
     def run_async(
         self,
         input: str | None = None,  # noqa: A002
         timeout: int = 60,
     ) -> str:
+        if isinstance(self._runner, DurableFunctionTestRunner):
+            return self._runner.run_async(input=input, execution_timeout=timeout)
         return self._runner.run_async(input=input, timeout=timeout)
 
     def send_callback_success(
@@ -222,7 +231,7 @@ def test_hello_world(durable_runner):
         if not handler:
             pytest.fail("handler is required for local mode tests")
         # Create local runner (needs cleanup via context manager)
-        runner = DurableFunctionTestRunner(handler=handler)
+        runner = DurableFunctionTestRunner(handler=handler, execution_timeout=60)
 
     # Wrap in adapter and use context manager for proper cleanup
     with TestRunnerAdapter(runner, runner_mode) as adapter:

packages/aws-durable-execution-sdk-python-examples/test/wait_for_callback/test_wait_for_callback_timeout.py

@@ -29,4 +29,4 @@ def test_handle_wait_for_callback_timeout_scenarios(durable_runner):
     assert result_data["success"] is False
     assert isinstance(result_data["error"], str)
     assert len(result_data["error"]) > 0
-    assert "Callback timed out: Callback.Timeout" == result_data["error"]
+    assert "Callback timed out" == result_data["error"]

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

@@ -0,0 +1,45 @@
+"""The shared checkpoint write-transaction core.
+
+Both checkpoint entry points need to detect an idempotent replay before
+doing real work. This class holds that shared logic so neither caller
+duplicates it.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python_testing.execution import (
+        CheckpointIdempotencyRecord,
+        Execution,
+    )
+
+
+class CheckpointCore:
+    """Shared checkpoint logic used by both the web and in-process paths."""
+
+    @staticmethod
+    def match_cached(
+        execution: Execution,
+        checkpoint_token: str,
+        client_token: str | None,
+    ) -> CheckpointIdempotencyRecord | None:
+        """Return the cached idempotency record when the incoming
+        ``(client_token, checkpoint_token)`` matches the last checkpoint.
+
+        Returns ``None`` when there is no match. A missing or empty
+        ``client_token`` cannot replay because idempotency requires an
+        explicit client identifier.
+        """
+        if not client_token:
+            return None
+        cached: CheckpointIdempotencyRecord | None = execution.last_checkpoint
+        if cached is None:
+            return None
+        if (
+            cached.client_token != client_token
+            or cached.inbound_checkpoint_token != checkpoint_token
+        ):
+            return None
+        return cached

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

@@ -1,4 +1,10 @@
-"""Main checkpoint processor that orchestrates operation transformations."""
+"""In-process checkpoint flow.
+
+Orchestrates the full checkpoint request for in-process callers
+(``InMemoryServiceClient``). Mirrors the flow inside
+``Executor.checkpoint_execution`` used by the HTTP path, so both
+entry points share identical delta + watermark semantics.
+"""
 
 from __future__ import annotations
 
@@ -7,37 +13,57 @@
 from aws_durable_execution_sdk_python.lambda_service import (
     CheckpointOutput,
     CheckpointUpdatedExecutionState,
-    OperationUpdate,
     StateOutput,
 )
 
+from aws_durable_execution_sdk_python_testing.checkpoint.core import CheckpointCore
 from aws_durable_execution_sdk_python_testing.checkpoint.transformer import (
-    OperationTransformer,
+    CheckpointRequestDispatcher,
 )
 from aws_durable_execution_sdk_python_testing.checkpoint.validators.checkpoint import (
     CheckpointValidator,
 )
 from aws_durable_execution_sdk_python_testing.exceptions import (
     InvalidParameterValueException,
 )
+from aws_durable_execution_sdk_python_testing.execution import (
+    CheckpointIdempotencyRecord,
+    OperationPaginatorState,
+)
 from aws_durable_execution_sdk_python_testing.observer import ExecutionNotifier
 from aws_durable_execution_sdk_python_testing.token import CheckpointToken
 
 
 if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python.lambda_service import OperationUpdate
+
     from aws_durable_execution_sdk_python_testing.execution import Execution
     from aws_durable_execution_sdk_python_testing.scheduler import Scheduler
     from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore
 
 
+# Matches the Executor default; see design.md § Data model.
+DEFAULT_MAX_INVOCATION_PAGE_BYTES = 5 * 1024 * 1024
+
+
 class CheckpointProcessor:
-    """Handle OperationUpdate transformations and execution state updates."""
+    """In-process checkpoint flow used by ``InMemoryServiceClient``.
+
+    Uses the same :class:`CheckpointRequestDispatcher` + paginator
+    primitives as ``Executor.checkpoint_execution`` so observable
+    behaviour is identical across the two entry points.
+    """
 
-    def __init__(self, store: ExecutionStore, scheduler: Scheduler):
+    def __init__(
+        self,
+        store: ExecutionStore,
+        scheduler: Scheduler,  # noqa: ARG002 — kept for backward-compatible signature
+        max_page_bytes: int = DEFAULT_MAX_INVOCATION_PAGE_BYTES,
+    ):
         self._store = store
-        self._scheduler = scheduler
         self._notifier = ExecutionNotifier()
-        self._transformer = OperationTransformer()
+        self._dispatcher = CheckpointRequestDispatcher()
+        self._max_page_bytes = max_page_bytes
 
     def add_execution_observer(self, observer) -> None:
         """Add observer for execution events."""
@@ -47,41 +73,75 @@ def process_checkpoint(
         self,
         checkpoint_token: str,
         updates: list[OperationUpdate],
-        client_token: str | None,  # noqa: ARG002
+        client_token: str | None,
     ) -> CheckpointOutput:
-        """Process checkpoint updates and return result with updated execution state."""
-        # 1. Get current execution state
-        token: CheckpointToken = CheckpointToken.from_str(checkpoint_token)
+        """Apply ``updates`` and return the delta since the handler's
+        last observation.
+
+        Advances ``token_sequence`` exactly once; bumps
+        ``seq_counter`` once per accepted update via
+        ``touch_operation``; advances ``handler_seen_seq`` only for
+        operations actually returned in the response.
+        """
+        token = CheckpointToken.from_str(checkpoint_token)
         execution: Execution = self._store.load(token.execution_arn)
 
-        # 2. Validate checkpoint token
-        if execution.is_complete or token.token_sequence != execution.token_sequence:
-            msg: str = "Invalid checkpoint token"
-
+        # the relevant invariant: idempotency first, before the token-sequence check.
+        cached = _maybe_replay_cached(execution, checkpoint_token, client_token)
+        if cached is not None:
+            return cached
+
+        if (
+            execution.is_complete
+            or token.token_sequence != execution.token_sequence
+            or token.invocation_id != execution.current_invocation_id
+        ):
+            msg = "Invalid checkpoint token"
             raise InvalidParameterValueException(msg)
 
-        # 3. Validate all updates, state transitions are valid, sizes etc.
-        CheckpointValidator.validate_input(updates, execution)
-
-        # 4. Transform OperationUpdate -> Operation and schedule future replays
-        updated_operations, all_updates = self._transformer.process_updates(
-            updates=updates,
-            current_operations=execution.operations,
-            notifier=self._notifier,
-            execution_arn=token.execution_arn,
+        if updates:
+            CheckpointValidator.validate_input(updates, execution)
+            self._dispatcher.apply_updates(
+                execution=execution,
+                updates=updates,
+                client_token=client_token,
+                notifier=self._notifier,
+                touch=execution.touch_operation,
+            )
+
+        new_token_sequence = execution.advance_token_sequence()
+
+        paginator = OperationPaginatorState.pin(execution)
+        delta = paginator.unseen_operations()
+        response_ops, _ = paginator.page_subset(delta, self._max_page_bytes)
+        if response_ops:
+            highest_delivered_seq = max(
+                execution.operation_last_touched_seq[op.operation_id]
+                for op in response_ops
+            )
+            paginator.advance_handler_seen(highest_delivered_seq)
+
+        new_token = CheckpointToken(
+            execution_arn=execution.durable_execution_arn,
+            token_sequence=new_token_sequence,
+            invocation_id=execution.current_invocation_id,
+        ).to_str()
+
+        execution.last_checkpoint = CheckpointIdempotencyRecord(
+            client_token=client_token or "",
+            inbound_checkpoint_token=checkpoint_token,
+            outbound_checkpoint_token=new_token,
+            operations=list(response_ops),
+            next_marker=None,
         )
 
-        # 5. Generate a new checkpoint token and save updated operations
-        new_checkpoint_token = execution.get_new_checkpoint_token()
-        execution.operations = updated_operations
-        execution.updates.extend(all_updates)
         self._store.update(execution)
 
-        # 6. Return checkpoint result
         return CheckpointOutput(
-            checkpoint_token=new_checkpoint_token,
+            checkpoint_token=new_token,
             new_execution_state=CheckpointUpdatedExecutionState(
-                operations=execution.get_navigable_operations(), next_marker=None
+                operations=response_ops,
+                next_marker=None,
             ),
         )
 
@@ -91,11 +151,35 @@ def get_execution_state(
         next_marker: str,  # noqa: ARG002
         max_items: int = 1000,  # noqa: ARG002
     ) -> StateOutput:
-        """Get current execution state."""
-        token: CheckpointToken = CheckpointToken.from_str(checkpoint_token)
-        execution: Execution = self._store.load(token.execution_arn)
-
-        # TODO: paging when size or max
+        """Get current execution state.
+
+        Returns the full navigable operation list with a null marker.
+        Marker round-tripping and the invocation-state gate are
+        enforced on the HTTP path in :class:`Executor`.
+        """
+        token = CheckpointToken.from_str(checkpoint_token)
+        execution = self._store.load(token.execution_arn)
         return StateOutput(
-            operations=execution.get_navigable_operations(), next_marker=None
+            operations=execution.get_navigable_operations(),
+            next_marker=None,
         )
+
+
+def _maybe_replay_cached(
+    execution: Execution,
+    checkpoint_token: str,
+    client_token: str | None,
+) -> CheckpointOutput | None:
+    """Replay cached CheckpointOutput when the incoming call matches
+    the last checkpoint. Returns ``None`` when there is no match.
+    """
+    cached = CheckpointCore.match_cached(execution, checkpoint_token, client_token)
+    if cached is None:
+        return None
+    return CheckpointOutput(
+        checkpoint_token=cached.outbound_checkpoint_token,
+        new_execution_state=CheckpointUpdatedExecutionState(
+            operations=list(cached.operations),
+            next_marker=cached.next_marker,
+        ),
+    )

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

@@ -61,16 +61,8 @@ def process(
                 new_step_details = StepDetails(
                     attempt=current_attempt + 1,
                     next_attempt_timestamp=next_attempt_time,
-                    result=(
-                        current_op.step_details.result
-                        if current_op and current_op.step_details
-                        else None
-                    ),
-                    error=(
-                        current_op.step_details.error
-                        if current_op and current_op.step_details
-                        else None
-                    ),
+                    result=update.payload,
+                    error=update.error,
                 )
 
                 # Create new operation with updated step_details
@@ -99,12 +91,12 @@ def process(
                     else None,
                 )
 
-                # Schedule step retry timer to fire after delay
-                notifier.notify_step_retry_scheduled(
-                    execution_arn=execution_arn,
-                    operation_id=update.operation_id,
-                    delay=delay,
-                )
+                # Timer scheduling is handled centrally by
+                # Executor._schedule_earliest_pending. The
+                # processor only needs to set
+                # step_details.next_attempt_timestamp; the executor
+                # finds the earliest pending wake across all ops and
+                # arms a single timer.
                 return retry_operation
             case OperationAction.SUCCEED:
                 return self._translate_update_to_operation(

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

@@ -74,12 +74,11 @@ def process(
                     chained_invoke_details=None,
                 )
 
-                # Schedule wait timer to complete after delay
-                notifier.notify_wait_timer_scheduled(
-                    execution_arn=execution_arn,
-                    operation_id=update.operation_id,
-                    delay=scaled_wait_seconds,
-                )
+                # Timer scheduling is handled centrally by
+                # Executor._schedule_earliest_pending. The processor
+                # only needs to set wait_details.scheduled_end_timestamp;
+                # the executor finds the earliest pending wake across
+                # all ops and arms a single timer.
                 return wait_operation
             case OperationAction.CANCEL:
                 # TODO: need to cancel the WAIT in the executor