fix(sdk): answer PENDING on revoked token

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

@@ -38,9 +38,13 @@
 
 class CheckpointResult(NamedTuple):
     """Outcome of applying a checkpoint: the new token, the operations to
-    return to the handler this round, and the lifecycle effects raised."""
+    return to the handler this round, and the lifecycle effects raised.
 
-    checkpoint_token: str
+    ``checkpoint_token`` is None when the execution is paused and the token
+    is withheld, telling the SDK this invocation may checkpoint no further.
+    """
+
+    checkpoint_token: str | None
     operations: list[Operation]
     effects: list[CheckpointEffect]
 
@@ -86,9 +90,11 @@ def apply(
         """Apply ``updates`` to ``execution`` and compute the response delta.
 
         Advances ``token_sequence`` exactly once, returns the full set of
-        operations the handler has not yet seen, advances
-        ``handler_seen_seq`` to cover them, and records the
-        idempotency entry for a byte-identical replay of a retried call.
+        operations the handler has not yet seen unless paused, advances
+        ``handler_seen_seq`` only for returned operations, and records an
+        idempotency entry for a byte-identical replay of a retried call
+        only when returning a checkpoint token.
+
         The caller is responsible for the invocation gate, locking,
         persistence, and applying the returned effects.
 
@@ -123,7 +129,10 @@ def apply(
         # The checkpoint response returns the full unseen delta in a single
         # response. Advance handler_seen_seq to cover every returned op so
         # the next delta carries only operations touched after this response.
-        response_ops: list[Operation] = paginator.unseen_operations()
+        # Paused checkpoints persist updates without delivering any state.
+        response_ops: list[Operation] = (
+            [] if execution.is_paused else paginator.unseen_operations()
+        )
         if response_ops:
             highest_delivered_seq: int = max(
                 execution.operation_last_touched_seq[op.operation_id]
@@ -137,12 +146,26 @@ def apply(
             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,
-        )
+        # A paused execution registers this checkpoint's updates but withholds
+        # the token: a response without one tells the SDK this invocation may
+        # checkpoint no further, so it reports PENDING at its next checkpoint
+        # rather than continuing, and owes a re-invoke once resumed.
+        #
+        # Checkpoints that were interrupted by a pause are not added to the
+        # idempotency record. A retry of this checkpoint will fail rather than be
+        # answered again.
+
+        outbound_token: str | None = new_token
+        if execution.is_paused:
+            outbound_token = None
+            execution.defer_invocation()
+        else:
+            execution.last_checkpoint = CheckpointIdempotencyRecord(
+                client_token=client_token or "",
+                inbound_checkpoint_token=checkpoint_token,
+                outbound_checkpoint_token=outbound_token,
+                operations=list(response_ops),
+                next_marker=None,
+            )
 
-        return CheckpointResult(new_token, response_ops, effects)
+        return CheckpointResult(outbound_token, response_ops, effects)

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

@@ -62,19 +62,33 @@ 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.
+    """Single-slot cache of the most recent accepted checkpoint response
+    with a checkpoint token.
 
-    Single-slot cache of the most recent accepted checkpoint response.
-    ``(client_token, inbound_checkpoint_token)`` pair is entitled to a
-    byte-identical response; this record is what we compare
-    against and replay from.
+    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
+    outbound_checkpoint_token: str | None
     operations: list[Operation]
     next_marker: str | None
 
@@ -94,7 +108,7 @@ def from_json_dict(cls, data: dict[str, Any]) -> CheckpointIdempotencyRecord:
         return cls(
             client_token=data["ClientToken"],
             inbound_checkpoint_token=data["InboundCheckpointToken"],
-            outbound_checkpoint_token=data["OutboundCheckpointToken"],
+            outbound_checkpoint_token=data.get("OutboundCheckpointToken"),
             operations=[
                 Operation.from_json_dict(op_data) for op_data in data["Operations"]
             ],
@@ -163,6 +177,33 @@ def __init__(
         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 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.
@@ -248,6 +289,7 @@ def to_json_dict(self) -> dict[str, Any]:
             "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
@@ -316,6 +358,9 @@ def from_json_dict(cls, data: dict[str, Any]) -> 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
 

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

@@ -4,8 +4,11 @@
 
 import asyncio
 import logging
+import math
 import threading
+import time
 import uuid
+from concurrent.futures import wait
 from datetime import datetime
 from typing import TYPE_CHECKING, assert_never
 
@@ -351,7 +354,7 @@ def get_execution(self, execution_arn: str) -> Execution:
         try:
             return self._store.load(execution_arn)
         except KeyError as e:
-            msg: str = f"Execution {execution_arn} not found"
+            msg: str = "Durable Execution does not exist"
             raise ResourceNotFoundException(msg) from e
 
     def get_execution_details(self, execution_arn: str) -> GetDurableExecutionResponse:
@@ -1454,15 +1457,22 @@ def _validate_invocation_response_and_store(
                 )
 
             case InvocationStatus.PENDING:
-                # An operation the handler waited on may complete between
-                # the handler's return and this check. A change the
-                # handler has not seen, after the invocation's input was
-                # built, earns a re-invoke, so PENDING is valid; only a
-                # handler that waited on nothing is in error.
-                if not execution.has_pending_operations(execution) and not (
-                    invocation_seq is not None
-                    and execution.has_changes_after(
-                        max(invocation_seq, execution.handler_seen_seq)
+                # A paused execution answers the running invocation's
+                # checkpoint without a token and defers its next invocation;
+                # that invocation must stop as PENDING even with nothing
+                # pending, and resume re-invokes it. Any other PENDING needs
+                # pending operations or a change the handler has not seen
+                # since its input was built. The unseen-change case happens
+                # when an operation completes after this invocation's input
+                # was built but before this response is validated.
+                if (
+                    not execution.has_deferred_invocation
+                    and not execution.has_pending_operations(execution)
+                    and not (
+                        invocation_seq is not None
+                        and execution.has_changes_after(
+                            max(invocation_seq, execution.handler_seen_seq)
+                        )
                     )
                 ):
                     msg_pending_ops: str = (
@@ -1521,6 +1531,15 @@ def _begin_invocation(
             )
             return None
 
+        if execution.is_paused:
+            execution.defer_invocation()
+            self._store.save(execution)
+            logger.debug(
+                "[%s] Holding back scheduled invocation while paused",
+                execution_arn,
+            )
+            return None
+
         # Claim the gate: at most one handler invocation per
         # execution in flight.
         self._set_invocation_gate(execution_arn, InvocationState.INVOKING)
@@ -1776,6 +1795,107 @@ def _invoke_execution(self, execution_arn: str, delay: float = 0) -> None:
             completion_event=completion_event,
         )
 
+    def pause_execution(self, execution_arn: str, timeout: float = 60) -> None:
+        """Experimental; may change or be removed in a future release.
+
+        Pause an execution and wait for its current invocation to stop.
+
+        If a running invocation sends another checkpoint after the pause is
+        applied, its updates are accepted and saved. The response omits the
+        checkpoint token, causing the SDK to end the invocation with PENDING.
+        Further checkpoints from that invocation are rejected.
+
+        Returns once the pause request is applied and no invocation is running.
+        A running step is not interrupted.
+
+        Repeated calls leave an existing pause unchanged and wait again for
+        the running invocation to stop. Completed executions are unchanged.
+
+        If this call times out, the pause request is not cancelled.
+
+        Args:
+            execution_arn: Execution to pause.
+            timeout: Maximum wait in real seconds. Must be finite and positive.
+
+        Raises:
+            InvalidParameterValueException: If the ARN is blank.
+            ResourceNotFoundException: If the execution does not exist.
+            ValueError: If timeout is not finite and positive.
+            TimeoutError: If the pause wait times out.
+        """
+
+        if not math.isfinite(timeout) or timeout <= 0:
+            msg = "timeout must be finite and positive"
+            raise ValueError(msg)
+        deadline = time.monotonic() + timeout
+
+        self._validate_execution_arn(execution_arn)
+        pause = self._registry.submit(
+            execution_arn,
+            CallableTask(lambda: self._set_paused(execution_arn)),
+        )
+
+        timeout_message = (
+            f"Execution {execution_arn} did not become idle within {timeout:g} seconds "
+            "of requesting pause. The pause request has not been cancelled."
+        )
+
+        done, _ = wait([pause], timeout=max(0, deadline - time.monotonic()))
+        if not done:
+            raise TimeoutError(timeout_message)
+        pause.result()
+
+        self._wait_until_idle(execution_arn, deadline, timeout_message)
+
+    def resume_execution(self, execution_arn: str) -> None:
+        """Experimental; may change or be removed in a future release.
+
+        Make the local checkpoint server answer this execution's
+        checkpoints with a token again.
+
+        Starts the invocation that pause_execution(arn) held back, if any.
+
+        Repeated calls leave an execution that is not paused unchanged.
+        Completed executions are unchanged.
+
+        Raises:
+            InvalidParameterValueException: If the ARN is blank.
+            ResourceNotFoundException: If the execution does not exist.
+        """
+        self._validate_execution_arn(execution_arn)
+        self._registry.submit(
+            execution_arn,
+            CallableTask(lambda: self._resume_execution(execution_arn)),
+        ).result()
+
+    def _set_paused(self, execution_arn: str) -> None:
+        execution = self.get_execution(execution_arn)
+        if execution.is_complete or execution.is_paused:
+            return
+        execution.pause()
+        self._store.save(execution)
+
+    def _resume_execution(self, execution_arn: str) -> None:
+        execution = self.get_execution(execution_arn)
+        if execution.is_complete or not execution.is_paused:
+            return
+
+        deferred = execution.resume()
+
+        self._store.save(execution)
+        if deferred:
+            self._invoke_execution(execution_arn)
+
+    def _wait_until_idle(
+        self, execution_arn: str, deadline: float, timeout_message: str
+    ) -> None:
+        """Wait for no invocation to be running. Raise at the deadline."""
+        while self._invocation_gate(execution_arn) is InvocationState.INVOKING:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise TimeoutError(timeout_message)
+            time.sleep(min(0.005, remaining))
+
     def _complete_workflow(
         self, execution_arn: str, result: str | None, error: ErrorObject | None
     ):

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

@@ -3222,9 +3222,12 @@ def to_dict(self) -> dict[str, Any]:
 
 @dataclass(frozen=True)
 class CheckpointDurableExecutionResponse:
-    """Response from checkpointing a durable execution."""
+    """Response from checkpointing a durable execution.
 
-    checkpoint_token: str
+    ``checkpoint_token`` is None when this invocation may checkpoint no further
+    """
+
+    checkpoint_token: str | None
     new_execution_state: CheckpointUpdatedExecutionState | None = None
 
     @classmethod
@@ -3234,12 +3237,14 @@ def from_dict(cls, data: dict) -> CheckpointDurableExecutionResponse:
             new_execution_state = CheckpointUpdatedExecutionState.from_dict(state_data)
 
         return cls(
-            checkpoint_token=data["CheckpointToken"],
+            checkpoint_token=data.get("CheckpointToken"),
             new_execution_state=new_execution_state,
         )
 
     def to_dict(self) -> dict[str, Any]:
-        result: dict[str, Any] = {"CheckpointToken": self.checkpoint_token}
+        result: dict[str, Any] = {}
+        if self.checkpoint_token is not None:
+            result["CheckpointToken"] = self.checkpoint_token
         if self.new_execution_state is not None:
             result["NewExecutionState"] = self.new_execution_state.to_dict()
         return result