aws-durable-execution-sdk-python-testing 1.2.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

aws_durable_execution_sdk_python_testing/execution.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
-from dataclasses import replace
-from datetime import UTC, datetime
+import json
+import logging
+from dataclasses import dataclass, replace
+from datetime import datetime
 from enum import Enum
 from threading import Lock
 from typing import Any
@@ -12,6 +14,7 @@ from aws_durable_execution_sdk_python.execution import (
     InvocationStatus,
 )
 from aws_durable_execution_sdk_python.lambda_service import (
+    ChainedInvokeDetails,
     ErrorObject,
     ExecutionDetails,
     Operation,
@@ -20,6 +23,7 @@ from aws_durable_execution_sdk_python.lambda_service import (
     OperationUpdate,
 )
 
+from aws_durable_execution_sdk_python_testing.clock import real_now
 from aws_durable_execution_sdk_python_testing.exceptions import (
     IllegalStateException,
     InvalidParameterValueException,
@@ -27,6 +31,7 @@ from aws_durable_execution_sdk_python_testing.exceptions import (
 
 # Import AWS exceptions
 from aws_durable_execution_sdk_python_testing.model import (
+    executed_version,
     InvocationCompletedDetails,
     StartDurableExecutionInput,
 )
@@ -34,6 +39,18 @@ from aws_durable_execution_sdk_python_testing.token import (
     CheckpointToken,
 )
 
+logger = logging.getLogger(__name__)
+
+
+_CHAINED_INVOKE_TERMINAL_STATUSES: frozenset[OperationStatus] = frozenset(
+    {
+        OperationStatus.SUCCEEDED,
+        OperationStatus.FAILED,
+        OperationStatus.TIMED_OUT,
+        OperationStatus.STOPPED,
+    }
+)
+
 
 class ExecutionStatus(Enum):
     """Execution status for API responses."""
@@ -45,6 +62,60 @@ class ExecutionStatus(Enum):
     TIMED_OUT = "TIMED_OUT"
 
 
+class PauseState(Enum):
+    """Whether a test paused this execution, and whether resume must invoke."""
+
+    NOT_PAUSED = "NOT_PAUSED"
+    # Paused while idle: no handler was running and no invocation was due.
+    # Resume lifts the pause and the next trigger invokes as usual.
+    PAUSED = "PAUSED"
+    # Paused while progress was owed. Either the running handler checkpointed
+    # and was answered without a token, so it stopped as PENDING before it
+    # finished, or a scheduled invocation came due and was held back.
+    # Resume must start one new invocation to make up for it.
+    PAUSED_INVOCATION_DEFERRED = "PAUSED_INVOCATION_DEFERRED"
+
+
+@dataclass(frozen=True)
+class CheckpointIdempotencyRecord:
+    """Single-slot cache of the most recent accepted checkpoint response
+    with a checkpoint token.
+
+    A matching ``(client_token, inbound_checkpoint_token)`` pair replays
+    this response without applying updates again. Pause-interrupted
+    checkpoints leave this record unchanged so their retries are rejected.
+    """
+
+    client_token: str
+    inbound_checkpoint_token: str
+    outbound_checkpoint_token: str | None
+    operations: list[Operation]
+    next_marker: str | None
+
+    def to_json_dict(self) -> dict[str, Any]:
+        """Serialize to JSON-compatible dict for persistence."""
+        return {
+            "ClientToken": self.client_token,
+            "InboundCheckpointToken": self.inbound_checkpoint_token,
+            "OutboundCheckpointToken": self.outbound_checkpoint_token,
+            "Operations": [op.to_json_dict() for op in self.operations],
+            "NextMarker": self.next_marker,
+        }
+
+    @classmethod
+    def from_json_dict(cls, data: dict[str, Any]) -> CheckpointIdempotencyRecord:
+        """Reconstruct from a to_json_dict output."""
+        return cls(
+            client_token=data["ClientToken"],
+            inbound_checkpoint_token=data["InboundCheckpointToken"],
+            outbound_checkpoint_token=data.get("OutboundCheckpointToken"),
+            operations=[
+                Operation.from_json_dict(op_data) for op_data in data["Operations"]
+            ],
+            next_marker=data.get("NextMarker"),
+        )
+
+
 class Execution:
     """Execution state."""
 
@@ -59,20 +130,111 @@ class Execution:
         self.start_input: StartDurableExecutionInput = start_input
         self.operations: list[Operation] = operations
         self.updates: list[OperationUpdate] = []
+        self.update_timestamps: list[datetime] = []
         self.invocation_completions: list[InvocationCompletedDetails] = []
-        self.used_tokens: set[str] = set()
-        # TODO: this will need to persist/rehydrate depending on inmemory vs sqllite store
-        self._token_sequence: int = 0
+        self.updated_operation_ids: list[str] = []
+        # Two monotonic counters. `token_sequence` is the checkpoint-response
+        # version carried on the wire; `seq_counter` is the internal
+        # event counter bumped on every state-affecting mutation.
+        self.token_sequence: int = 0
+        self.seq_counter: int = 0
+        # Watermark of the highest `seq_counter` the handler has been told
+        # about (via InitialExecutionState or a successful checkpoint
+        # response). Filters the "unseen" delta on the next checkpoint.
+        self.handler_seen_seq: int = 0
+        # Per-op "last touched at `seq_counter`" bookkeeping. Sparse:
+        # ops not in the dict are treated as touch-seq 0.
+        self.operation_last_touched_seq: dict[str, int] = {}
+        # Per-op payload size, tracked as a sidecar dict because
+        # ``Operation`` is frozen upstream.
+        self.operation_size_bytes: dict[str, int] = {}
+        # Chained-invoke linkage: operation_id -> the invoked child
+        # execution's ARN. Populated when a durable child execution is
+        # started for a CHAINED_INVOKE operation; absent for targets
+        # that run as a single invocation without their own execution.
+        self.chained_invoke_children: dict[str, str] = {}
+        # The parent execution's ARN when this execution was started by
+        # a chained invoke. A child's output is capped at the chained
+        # invoke limit, so validation needs to know.
+        self.parent_execution_arn: str | None = None
+        # The runner's region at creation; None when unknown (for
+        # example an execution built directly in a test).
+        self.region: str | None = None
+        # Set when a trigger arrived while an invocation was already
+        # in flight; the gate-release path consults it to decide whether
+        # to schedule another invocation.
+        self.needs_reinvoke: bool = False
+        # Single-slot cache for idempotent checkpoint replays.
+        self.last_checkpoint: CheckpointIdempotencyRecord | None = None
+        # Identity of the handler invocation currently allowed to
+        # checkpoint. Regenerated each time a new invocation is
+        # dispatched, so a checkpoint carrying a superseded invocation's
+        # token is rejected. Held in memory semantics: it defaults empty
+        # for executions that never went through an invocation.
+        self.current_invocation_id: str = ""
         self._state_lock: Lock = Lock()
         self.is_complete: bool = False
         self.result: DurableExecutionInvocationOutput | None = None
         self.consecutive_failed_invocation_attempts: int = 0
         self.close_status: ExecutionStatus | None = None
+        self._pause_state: PauseState = PauseState.NOT_PAUSED
+
+    @property
+    def pause_state(self) -> PauseState:
+        return self._pause_state
+
+    @property
+    def is_paused(self) -> bool:
+        return self._pause_state is not PauseState.NOT_PAUSED
 
     @property
-    def token_sequence(self) -> int:
-        """Get current token sequence value."""
-        return self._token_sequence
+    def has_deferred_invocation(self) -> bool:
+        return self._pause_state is PauseState.PAUSED_INVOCATION_DEFERRED
+
+    def pause(self) -> None:
+        if self._pause_state is PauseState.NOT_PAUSED:
+            self._pause_state = PauseState.PAUSED
+
+    def defer_invocation(self) -> None:
+        if self._pause_state is PauseState.PAUSED:
+            self._pause_state = PauseState.PAUSED_INVOCATION_DEFERRED
+
+    def resume(self) -> bool:
+        """Clear the pause and return whether the caller must schedule an invocation."""
+        deferred: bool = self._pause_state is PauseState.PAUSED_INVOCATION_DEFERRED
+        self._pause_state = PauseState.NOT_PAUSED
+        return deferred
+
+    def touch_operation(self, operation_id: str) -> None:
+        """Record a state-affecting event on an operation.
+
+        Bumps ``seq_counter`` by 1 and records the new value as the
+        operation's "last touched" sequence. Called from the checkpoint
+        dispatcher (once per accepted update) and from every async
+        completion path (``complete_wait``, ``complete_retry``,
+        ``complete_callback_*``, terminal transitions).
+        """
+        self.seq_counter += 1
+        self.operation_last_touched_seq[operation_id] = self.seq_counter
+
+    def advance_token_sequence(self) -> int:
+        """Bump the checkpoint-response version counter.
+
+        Called exactly once per accepted non-idempotent checkpoint call.
+        Idempotent replays and async completions do not call this.
+        Returns the new value.
+        """
+        self.token_sequence += 1
+        return self.token_sequence
+
+    def begin_new_invocation(self) -> str:
+        """Assign a fresh identity to the next handler invocation and
+        return it. A checkpoint token minted for a prior invocation no
+        longer matches once this is called, so a superseded invocation's
+        checkpoints are rejected.
+        """
+        self.current_invocation_id = str(uuid4())
+        return self.current_invocation_id
 
     def current_status(self) -> ExecutionStatus:
         """Get execution status."""
@@ -105,15 +267,29 @@ class Execution:
             "StartInput": self.start_input.to_dict(),
             "Operations": [op.to_json_dict() for op in self.operations],
             "Updates": [update.to_dict() for update in self.updates],
+            "UpdateTimestamps": [ts.isoformat() for ts in self.update_timestamps],
             "InvocationCompletions": [
                 completion.to_json_dict() for completion in self.invocation_completions
             ],
-            "UsedTokens": list(self.used_tokens),
-            "TokenSequence": self._token_sequence,
+            "UpdatedOperationIds": self.updated_operation_ids,
+            "TokenSequence": self.token_sequence,
+            "SeqCounter": self.seq_counter,
+            "HandlerSeenSeq": self.handler_seen_seq,
+            "OperationLastTouchedSeq": dict(self.operation_last_touched_seq),
+            "OperationSizeBytes": dict(self.operation_size_bytes),
+            "ChainedInvokeChildren": dict(self.chained_invoke_children),
+            "ParentExecutionArn": self.parent_execution_arn,
+            "Region": self.region,
+            "NeedsReinvoke": self.needs_reinvoke,
+            "LastCheckpoint": (
+                self.last_checkpoint.to_json_dict() if self.last_checkpoint else None
+            ),
             "IsComplete": self.is_complete,
             "Result": self.result.to_dict() if self.result else None,
             "ConsecutiveFailedInvocationAttempts": self.consecutive_failed_invocation_attempts,
             "CloseStatus": self.close_status.value if self.close_status else None,
+            "CurrentInvocationId": self.current_invocation_id,
+            "PauseState": self._pause_state.value,
         }
 
     @classmethod
@@ -138,12 +314,37 @@ class Execution:
         execution.updates = [
             OperationUpdate.from_dict(update_data) for update_data in data["Updates"]
         ]
+        execution.update_timestamps = [
+            datetime.fromisoformat(ts) for ts in data.get("UpdateTimestamps", [])
+        ]
         execution.invocation_completions = [
             InvocationCompletedDetails.from_json_dict(item)
             for item in data.get("InvocationCompletions", [])
         ]
-        execution.used_tokens = set(data["UsedTokens"])
-        execution._token_sequence = data["TokenSequence"]  # noqa: SLF001
+        execution.updated_operation_ids = list(data.get("UpdatedOperationIds", []))
+        # NOTE: prior format included a "UsedTokens" set. Current
+        # removed it (the Executor no longer consults it — token
+        # validity is via _is_current_token against token_sequence).
+        # Old-format dicts are loaded by ignoring "UsedTokens".
+        execution.token_sequence = data["TokenSequence"]
+        # Safe defaults for fields added after the original schema.
+        execution.seq_counter = data.get("SeqCounter", 0)
+        execution.handler_seen_seq = data.get("HandlerSeenSeq", 0)
+        execution.operation_last_touched_seq = dict(
+            data.get("OperationLastTouchedSeq", {})
+        )
+        execution.operation_size_bytes = dict(data.get("OperationSizeBytes", {}))
+        execution.chained_invoke_children = dict(data.get("ChainedInvokeChildren", {}))
+        execution.parent_execution_arn = data.get("ParentExecutionArn")
+        execution.region = data.get("Region")
+        execution.needs_reinvoke = data.get("NeedsReinvoke", False)
+        execution.current_invocation_id = data.get("CurrentInvocationId", "")
+        last_checkpoint_data = data.get("LastCheckpoint")
+        execution.last_checkpoint = (
+            CheckpointIdempotencyRecord.from_json_dict(last_checkpoint_data)
+            if last_checkpoint_data
+            else None
+        )
         execution.is_complete = data["IsComplete"]
         execution.result = (
             DurableExecutionInvocationOutput.from_dict(data["Result"])
@@ -157,10 +358,13 @@ class Execution:
         execution.close_status = (
             ExecutionStatus(close_status_str) if close_status_str else None
         )
+        execution._pause_state = PauseState(
+            data.get("PauseState", PauseState.NOT_PAUSED.value)
+        )
 
         return execution
 
-    def start(self) -> None:
+    def start(self, now: datetime | None = None) -> None:
         if self.start_input.invocation_id is None:
             msg: str = "invocation_id is required"
             raise InvalidParameterValueException(msg)
@@ -170,7 +374,7 @@ class Execution:
                     operation_id=self.start_input.invocation_id,
                     parent_id=None,
                     name=self.start_input.execution_name,
-                    start_timestamp=datetime.now(UTC),
+                    start_timestamp=now if now is not None else real_now(),
                     operation_type=OperationType.EXECUTION,
                     status=OperationStatus.STARTED,
                     execution_details=ExecutionDetails(
@@ -188,26 +392,49 @@ class Execution:
         return self.operations[0]
 
     def get_new_checkpoint_token(self) -> str:
-        """Generate a new checkpoint token with incremented sequence"""
+        """Serialise the current ``token_sequence`` as a checkpoint token.
+
+        Does NOT bump ``token_sequence``. Only
+        ``Executor.checkpoint_execution`` (via
+        :meth:`advance_token_sequence`) is allowed to advance the
+        counter. This method retains its name for backward
+        compatibility but is now semantically
+        "get_current_checkpoint_token".
+        """
         with self._state_lock:
-            self._token_sequence += 1
-            new_token_sequence = self._token_sequence
             token = CheckpointToken(
                 execution_arn=self.durable_execution_arn,
-                token_sequence=new_token_sequence,
+                token_sequence=self.token_sequence,
+                invocation_id=self.current_invocation_id,
             )
-            token_str = token.to_str()
-            self.used_tokens.add(token_str)
-            return token_str
+            return token.to_str()
 
     def get_navigable_operations(self) -> list[Operation]:
-        """Get list of operations, but exclude child operations where the parent has already completed."""
-        return self.operations
+        """Return a snapshot copy of the operation list.
+
+        The copy is taken under ``_state_lock`` so callers iterating the
+        result are not exposed to concurrent operation mutations. Operations
+        are frozen dataclasses, so a shallow copy is a stable snapshot.
+        """
+        with self._state_lock:
+            return list(self.operations)
 
     def get_assertable_operations(self) -> list[Operation]:
-        """Get list of operations, but exclude the EXECUTION operations"""
-        # TODO: this excludes EXECUTION at start, but can there be an EXECUTION at the end if there was a checkpoint with large payload?
-        return self.operations[1:]
+        """Return a snapshot of operations, excluding EXECUTION operations.
+
+        Filters by ``operation_type`` rather than assuming the EXECUTION
+        operation is always the first entry, so a second EXECUTION-type
+        entry (e.g. a future large-payload checkpoint recorded at the end
+        of the list) is excluded too. Snapshot taken under ``_state_lock``,
+        matching ``get_navigable_operations``, so a concurrent checkpoint
+        can't produce a torn read.
+        """
+        with self._state_lock:
+            return [
+                operation
+                for operation in self.operations
+                if operation.operation_type != OperationType.EXECUTION
+            ]
 
     def has_pending_operations(self, execution: Execution) -> bool:
         """True if execution has pending operations."""
@@ -228,6 +455,17 @@ class Execution:
                 return True
         return False
 
+    def has_unseen_changes(self) -> bool:
+        """True if any operation changed after the handler's last
+        observed sequence watermark."""
+        return self.has_changes_after(self.handler_seen_seq)
+
+    def has_changes_after(self, seq: int) -> bool:
+        """True if any operation changed after ``seq_counter`` was ``seq``."""
+        return any(
+            touched > seq for touched in self.operation_last_touched_seq.values()
+        )
+
     def record_invocation_completion(
         self, start_timestamp: datetime, end_timestamp: datetime, request_id: str
     ) -> None:
@@ -240,41 +478,62 @@ class Execution:
             )
         )
 
-    def complete_success(self, result: str | None) -> None:
+    def mark_state_delivered(self) -> None:
+        """Reset the list of changed operations once an invocation's input
+        carries them.
+
+        The service reports, on each invocation, the operations that
+        changed after the last state the handler observed. The runner
+        delivers state in the invocation input. So the list is reset when
+        that input is built, not when the invocation completes: an
+        operation that completes while the handler is still running is
+        reported on the next invocation.
+        """
+        self.updated_operation_ids = []
+
+    def _record_updated_operation(self, operation_id: str) -> None:
+        """Remember an operation that changed since the handler last saw
+        the state."""
+        if operation_id not in self.updated_operation_ids:
+            self.updated_operation_ids.append(operation_id)
+
+    def complete_success(self, result: str | None, now: datetime | None = None) -> None:
         """Complete execution successfully (DecisionType.COMPLETE_WORKFLOW_EXECUTION)."""
         self.result = DurableExecutionInvocationOutput(
             status=InvocationStatus.SUCCEEDED, result=result
         )
         self.is_complete = True
         self.close_status = ExecutionStatus.SUCCEEDED
-        self._end_execution(OperationStatus.SUCCEEDED)
+        self._end_execution(OperationStatus.SUCCEEDED, now)
 
-    def complete_fail(self, error: ErrorObject) -> None:
+    def complete_fail(
+        self, error: ErrorObject | None, now: datetime | None = None
+    ) -> None:
         """Complete execution with failure (DecisionType.FAIL_WORKFLOW_EXECUTION)."""
         self.result = DurableExecutionInvocationOutput(
             status=InvocationStatus.FAILED, error=error
         )
         self.is_complete = True
         self.close_status = ExecutionStatus.FAILED
-        self._end_execution(OperationStatus.FAILED)
+        self._end_execution(OperationStatus.FAILED, now)
 
-    def complete_timeout(self, error: ErrorObject) -> None:
+    def complete_timeout(self, error: ErrorObject, now: datetime | None = None) -> None:
         """Complete execution with timeout."""
         self.result = DurableExecutionInvocationOutput(
             status=InvocationStatus.FAILED, error=error
         )
         self.is_complete = True
         self.close_status = ExecutionStatus.TIMED_OUT
-        self._end_execution(OperationStatus.TIMED_OUT)
+        self._end_execution(OperationStatus.TIMED_OUT, now)
 
-    def complete_stopped(self, error: ErrorObject) -> None:
+    def complete_stopped(self, error: ErrorObject, now: datetime | None = None) -> None:
         """Complete execution as terminated (TerminateWorkflowExecutionV2Request)."""
         self.result = DurableExecutionInvocationOutput(
             status=InvocationStatus.FAILED, error=error
         )
         self.is_complete = True
         self.close_status = ExecutionStatus.STOPPED
-        self._end_execution(OperationStatus.STOPPED)
+        self._end_execution(OperationStatus.STOPPED, now)
 
     def find_operation(self, operation_id: str) -> tuple[int, Operation]:
         """Find operation by ID, return index and operation."""
@@ -296,7 +555,9 @@ class Execution:
         msg: str = f"Callback operation with callback_id [{callback_id}] not found"
         raise IllegalStateException(msg)
 
-    def complete_wait(self, operation_id: str) -> Operation:
+    def complete_wait(
+        self, operation_id: str, now: datetime | None = None
+    ) -> Operation:
         """Complete WAIT operation when timer fires."""
         index, operation = self.find_operation(operation_id)
 
@@ -312,13 +573,14 @@ class Execution:
 
         # Thread-safe increment sequence and operation update
         with self._state_lock:
-            self._token_sequence += 1
+            self.touch_operation(operation_id)
             # Build and assign updated operation
             self.operations[index] = replace(
                 operation,
                 status=OperationStatus.SUCCEEDED,
-                end_timestamp=datetime.now(UTC),
+                end_timestamp=now if now is not None else real_now(),
             )
+            self._record_updated_operation(operation_id)
             return self.operations[index]
 
     def complete_retry(self, operation_id: str) -> Operation:
@@ -337,7 +599,7 @@ class Execution:
 
         # Thread-safe increment sequence and operation update
         with self._state_lock:
-            self._token_sequence += 1
+            self.touch_operation(operation_id)
             # Build updated step_details with cleared next_attempt_timestamp
             new_step_details = None
             if operation.step_details:
@@ -352,10 +614,62 @@ class Execution:
 
             # Assign
             self.operations[index] = updated_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, result: bytes | None = None
+        self,
+        callback_id: str,
+        result: bytes | None = None,
+        now: datetime | None = None,
     ) -> Operation:
         """Complete CALLBACK operation with success."""
         index, operation = self.find_callback_operation(callback_id)
@@ -364,7 +678,7 @@ class Execution:
             raise IllegalStateException(msg)
 
         with self._state_lock:
-            self._token_sequence += 1
+            self.touch_operation(operation.operation_id)
             updated_callback_details = None
             if operation.callback_details:
                 updated_callback_details = replace(
@@ -375,13 +689,13 @@ class Execution:
             self.operations[index] = replace(
                 operation,
                 status=OperationStatus.SUCCEEDED,
-                end_timestamp=datetime.now(UTC),
+                end_timestamp=now if now is not None else real_now(),
                 callback_details=updated_callback_details,
             )
             return self.operations[index]
 
     def complete_callback_failure(
-        self, callback_id: str, error: ErrorObject
+        self, callback_id: str, error: ErrorObject, now: datetime | None = None
     ) -> Operation:
         """Complete CALLBACK operation with failure."""
         index, operation = self.find_callback_operation(callback_id)
@@ -391,7 +705,7 @@ class Execution:
             raise IllegalStateException(msg)
 
         with self._state_lock:
-            self._token_sequence += 1
+            self.touch_operation(operation.operation_id)
             updated_callback_details = None
             if operation.callback_details:
                 updated_callback_details = replace(
@@ -401,13 +715,13 @@ class Execution:
             self.operations[index] = replace(
                 operation,
                 status=OperationStatus.FAILED,
-                end_timestamp=datetime.now(UTC),
+                end_timestamp=now if now is not None else real_now(),
                 callback_details=updated_callback_details,
             )
             return self.operations[index]
 
     def complete_callback_timeout(
-        self, callback_id: str, error: ErrorObject
+        self, callback_id: str, error: ErrorObject, now: datetime | None = None
     ) -> Operation:
         """Complete CALLBACK operation with timeout."""
         index, operation = self.find_callback_operation(callback_id)
@@ -417,7 +731,7 @@ class Execution:
             raise IllegalStateException(msg)
 
         with self._state_lock:
-            self._token_sequence += 1
+            self.touch_operation(operation.operation_id)
             updated_callback_details = None
             if operation.callback_details:
                 updated_callback_details = replace(
@@ -427,18 +741,238 @@ class Execution:
             self.operations[index] = replace(
                 operation,
                 status=OperationStatus.TIMED_OUT,
-                end_timestamp=datetime.now(UTC),
+                end_timestamp=now if now is not None else real_now(),
                 callback_details=updated_callback_details,
             )
             return self.operations[index]
 
-    def _end_execution(self, status: OperationStatus) -> None:
+    def complete_chained_invoke(
+        self,
+        operation_id: str,
+        status: OperationStatus,
+        result: str | None = None,
+        error: ErrorObject | None = None,
+        now: datetime | None = None,
+    ) -> Operation:
+        """Transition a CHAINED_INVOKE operation to a terminal status.
+
+        ``status`` must be SUCCEEDED, FAILED, TIMED_OUT, or STOPPED.
+        SUCCEEDED carries ``result``; the failure statuses carry
+        ``error``.
+        """
+        index, operation = self.find_operation(operation_id)
+        self._require_chained_invoke(operation)
+
+        if status not in _CHAINED_INVOKE_TERMINAL_STATUSES:
+            msg_bad_status: str = (
+                f"Invalid terminal status for chained invoke: {status}"
+            )
+            raise IllegalStateException(msg_bad_status)
+        if operation.status is not OperationStatus.STARTED:
+            msg_not_active: str = (
+                f"Chained invoke operation [{operation_id}] is not active"
+            )
+            raise IllegalStateException(msg_not_active)
+        if status is OperationStatus.SUCCEEDED and error is not None:
+            msg_success_error: str = (
+                "Cannot provide an Error for a SUCCEEDED chained invoke."
+            )
+            raise IllegalStateException(msg_success_error)
+
+        with self._state_lock:
+            self.touch_operation(operation_id)
+            self.operations[index] = replace(
+                operation,
+                status=status,
+                end_timestamp=now if now is not None else real_now(),
+                chained_invoke_details=ChainedInvokeDetails(result=result, error=error),
+            )
+            # State paging sizes a completed operation by its result or
+            # error, not by its START payload.
+            self.operation_size_bytes[operation_id] = len(
+                (result if result is not None else "").encode()
+            ) + (len(json.dumps(error.to_dict())) if error is not None else 0)
+            self._record_updated_operation(operation_id)
+            return self.operations[index]
+
+    def executed_version(self) -> str:
+        """The version this execution runs, as the service reports it."""
+        return executed_version(self.start_input.function_qualifier)
+
+    def function_arn(self, default_region: str) -> str:
+        """The qualified function ARN of this execution.
+
+        The service qualifies the ARN with the executed version, and
+        ``Version`` carries the same value. An execution stored before the
+        runner recorded regions has none; it reports ``default_region``.
+        """
+        region: str = self.region if self.region is not None else default_region
+        return (
+            f"arn:aws:lambda:{region}:{self.start_input.account_id}"
+            f":function:{self.start_input.function_name}:{self.executed_version()}"
+        )
+
+    def record_chained_invoke_child(
+        self, operation_id: str, child_execution_arn: str
+    ) -> None:
+        """Record the child execution ARN for a CHAINED_INVOKE operation."""
+        with self._state_lock:
+            self.chained_invoke_children[operation_id] = child_execution_arn
+
+    @staticmethod
+    def _require_chained_invoke(operation: Operation) -> None:
+        if operation.operation_type != OperationType.CHAINED_INVOKE:
+            msg: str = (
+                f"Expected CHAINED_INVOKE operation, got {operation.operation_type}"
+            )
+            raise IllegalStateException(msg)
+
+    def _end_execution(
+        self, status: OperationStatus, now: datetime | None = None
+    ) -> None:
         """Set the end_timestamp on the main EXECUTION operation when execution completes."""
         execution_op: Operation = self.get_operation_execution_started()
         if execution_op.operation_type == OperationType.EXECUTION:
             with self._state_lock:
+                # Terminal transition of the EXECUTION op is a real
+                # state change — record it via touch_operation so
+                # introspection / GetDurableExecutionState reflects
+                # it. The handler has returned by this point, so the
+                # touch cannot change any checkpoint delta.
+                self.touch_operation(execution_op.operation_id)
                 self.operations[0] = replace(
                     execution_op,
                     status=status,
-                    end_timestamp=datetime.now(UTC),
+                    end_timestamp=now if now is not None else real_now(),
                 )
+
+
+@dataclass(frozen=True)
+class OperationPaginatorState:
+    """Pinned context for serving operation state to the handler.
+
+    One class does three jobs that share a pinned snapshot of the
+    execution's operation graph:
+
+    * :meth:`page` — bytes-bounded pagination across the snapshot,
+      used by invocation-input construction and
+      ``GetDurableExecutionState``.
+    * :meth:`unseen_operations` — strict ``> handler_seen_seq``
+      delta filter, used by ``checkpoint_execution`` to decide which
+      ops go back to the handler.
+    * :meth:`advance_handler_seen` — monotonic forward update of the
+      delivery watermark on the wrapped :class:`Execution`, called
+      after a checkpoint response is constructed and its ops are
+      known.
+
+
+    Pinning ``token_sequence`` and ``operations`` at construction
+    guarantees that paginated reads and delta computation see the same
+    snapshot, even if the underlying :class:`Execution` is mutated
+    between method calls.
+    """
+
+    execution: Execution
+    pinned_token_sequence: int
+    snapshot_operations: list[Operation]
+
+    @classmethod
+    def pin(cls, execution: Execution) -> OperationPaginatorState:
+        """Capture the execution's current token_sequence and operation
+        list. Subsequent mutations on the Execution do not affect this
+        instance's ``page`` / ``unseen_operations`` results."""
+        return cls(
+            execution=execution,
+            pinned_token_sequence=execution.token_sequence,
+            snapshot_operations=list(execution.operations),
+        )
+
+    def page(
+        self,
+        marker: str | None,
+        max_size_bytes: int,
+        max_items: int | None = None,
+    ) -> tuple[list[Operation], str | None]:
+        """Return a page of ops starting after ``marker``, bounded by
+        ``max_size_bytes`` and, when given, ``max_items``. Second element
+        of the tuple is a marker for the next page, or ``None`` when the
+        page fits everything."""
+        start_idx = self._resolve_marker(marker)
+        return self._walk_page(
+            self.snapshot_operations, start_idx, max_size_bytes, max_items
+        )
+
+    def unseen_operations(self) -> list[Operation]:
+        """Operations whose ``operation_last_touched_seq`` is strictly
+        greater than the wrapped execution's ``handler_seen_seq``,
+        in creation order.
+
+        Ops absent from ``operation_last_touched_seq`` are treated as
+        touch-seq 0 and never appear in a delta once the
+        watermark has moved.
+        """
+        touched = self.execution.operation_last_touched_seq
+        cutoff = self.execution.handler_seen_seq
+        return [
+            op
+            for op in self.snapshot_operations
+            if touched.get(op.operation_id, 0) > cutoff
+        ]
+
+    def advance_handler_seen(self, seq: int) -> None:
+        """Advance the wrapped execution's ``handler_seen_seq`` to
+        ``seq`` if that represents forward progress. Monotonic:
+        smaller or equal values are ignored."""
+        if seq > self.execution.handler_seen_seq:
+            self.execution.handler_seen_seq = seq
+
+    # --- internals -------------------------------------------------
+
+    def _walk_page(
+        self,
+        ops: list[Operation],
+        start_idx: int,
+        max_size_bytes: int,
+        max_items: int | None = None,
+    ) -> tuple[list[Operation], str | None]:
+        selected: list[Operation] = []
+        total = 0
+        for i in range(start_idx, len(ops)):
+            op = ops[i]
+            size = self._size_for(op)
+            over_bytes: bool = total + size > max_size_bytes
+            over_count: bool = max_items is not None and len(selected) >= max_items
+            if selected and (over_bytes or over_count):
+                return selected, self._encode_marker(i)
+            selected.append(op)
+            total += size
+        return selected, None
+
+    def _size_for(self, op: Operation) -> int:
+        # Zero-size ops are possible (no payload / no error / no input);
+        # floor at 1 to guarantee pagination always advances.
+        return max(self.execution.operation_size_bytes.get(op.operation_id, 0), 1)
+
+    def _resolve_marker(self, marker: str | None) -> int:
+        if marker is None:
+            return 0
+        seq, idx = self._decode_marker(marker)
+        if seq != self.pinned_token_sequence:
+            msg = "Invalid marker"
+            raise InvalidParameterValueException(msg)
+        if idx < 0 or idx > len(self.snapshot_operations):
+            msg = "Invalid marker"
+            raise InvalidParameterValueException(msg)
+        return idx
+
+    def _encode_marker(self, next_idx: int) -> str:
+        return f"{self.pinned_token_sequence}:{next_idx}"
+
+    @staticmethod
+    def _decode_marker(marker: str) -> tuple[int, int]:
+        try:
+            seq_part, idx_part = marker.split(":", 1)
+            return int(seq_part), int(idx_part)
+        except (ValueError, AttributeError) as exc:
+            msg = "Invalid marker"
+            raise InvalidParameterValueException(msg) from exc