722 feature exit gracefully with pending when a checkpoint response has no checkpointtoken

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]
 
@@ -137,12 +141,25 @@ def apply(
             invocation_id=execution.current_invocation_id,
         ).to_str()
 
+        # 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.
+        #
+        # The withheld token is recorded as the idempotency record's outbound
+        # token so a retry of this call replays the same tokenless response,
+        # even after a resume has moved the execution on.
+        outbound_token: str | None = new_token
+        if execution.paused:
+            outbound_token = None
+            execution.deferred_invocation = True
+
         execution.last_checkpoint = CheckpointIdempotencyRecord(
             client_token=client_token or "",
             inbound_checkpoint_token=checkpoint_token,
-            outbound_checkpoint_token=new_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

@@ -70,11 +70,15 @@ class CheckpointIdempotencyRecord:
     ``(client_token, inbound_checkpoint_token)`` pair is entitled to a
     byte-identical response; this record is what we compare
     against and replay from.
+
+    ``outbound_checkpoint_token`` is None when the response withheld the
+    token because the execution was paused, so the replay stays identical
+    after a resume.
     """
 
     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 +98,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 +167,17 @@ def __init__(
         self.result: DurableExecutionInvocationOutput | None = None
         self.consecutive_failed_invocation_attempts: int = 0
         self.close_status: ExecutionStatus | None = None
+        # While True, every checkpoint from pause_execution() until resume_execution()
+        # is answered without a checkpoint token.
+        # If responses have no token then no new invocation will start.
+        self.paused: bool = False
+        # Set while paused when progress was stopped and must be resumed later.
+        # This means either a new handler invocation was not started because the
+        # execution is paused, or the current handler was given no next checkpoint
+        # token and therefore must stop as PENDING.
+        #
+        # resume_execution() clears this flag and starts one new invocation.
+        self.deferred_invocation: bool = False
 
     def touch_operation(self, operation_id: str) -> None:
         """Record a state-affecting event on an operation.
@@ -248,6 +263,8 @@ 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,
+            "Paused": self.paused,
+            "DeferredInvocation": self.deferred_invocation,
         }
 
     @classmethod
@@ -316,6 +333,8 @@ def from_json_dict(cls, data: dict[str, Any]) -> Execution:
         execution.close_status = (
             ExecutionStatus(close_status_str) if close_status_str else None
         )
+        execution.paused = data.get("Paused", False)
+        execution.deferred_invocation = data.get("DeferredInvocation", False)
 
         return execution
 

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

@@ -5,6 +5,7 @@
 import asyncio
 import logging
 import threading
+import time
 import uuid
 from datetime import datetime
 from typing import TYPE_CHECKING, assert_never
@@ -1454,19 +1455,29 @@ 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)
+                # PENDING is valid while paused because pause can make a
+                # checkpoint response omit the next checkpoint token, forcing
+                # the current handler invocation to stop as PENDING.
+                #
+                # Otherwise, PENDING requires either pending durable operations
+                # or a change the handler has not seen yet. The unseen-change
+                # case can happen when an operation completes after this
+                # invocation's input was built but before this response is
+                # validated, so a follow-up invocation is needed.
+                if (
+                    not execution.paused
+                    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 = (
-                        "Cannot return PENDING status with no pending operations."
+                        "Cannot return PENDING status unless execution is paused, "
+                        "has pending durable operations, or has unseen changes "
+                        "after invocation input was built."
                     )
                     raise InvalidParameterValueException(msg_pending_ops)
                 logger.info("[%s] Execution pending async work", execution_arn)
@@ -1521,6 +1532,15 @@ def _begin_invocation(
             )
             return None
 
+        if execution.paused:
+            execution.deferred_invocation = True
+            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 +1796,68 @@ def _invoke_execution(self, execution_arn: str, delay: float = 0) -> None:
             completion_event=completion_event,
         )
 
+    def pause_execution(self, execution_arn: str) -> None:
+        """Make the local checkpoint server answer this execution's
+        checkpoints without a token, starting now.
+
+        Experimental; may change or be removed in a future release.
+
+        The invocation running now, if any, is answered without a token
+        on its next checkpoint. That checkpoint is accepted - its updates
+        stay durable - but the invocation reports PENDING, and no further
+        checkpoint of its is accepted. No new invocation starts until
+        resume_execution() is called.
+
+        Idempotent; a no-op once the execution has finished.
+        Resolves once no invocation of this execution is running.
+        """
+        self._validate_execution_arn(execution_arn)
+        self._registry.submit(
+            execution_arn,
+            CallableTask(lambda: self._set_paused(execution_arn)),
+        ).result()
+        self._wait_until_idle(execution_arn)
+
+    def resume_execution(self, execution_arn: str) -> None:
+        """Make the local checkpoint server answer this execution's
+        checkpoints with a token again.
+
+        Experimental; may change or be removed in a future release.
+
+        Starts the invocation that pause_execution() held back, if any.
+
+        Idempotent; a no-op once the execution has finished or if it was
+        not paused.
+        """
+        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._store.load(execution_arn)
+        if execution.is_complete or execution.paused:
+            return
+        execution.paused = True
+        self._store.save(execution)
+
+    def _resume_execution(self, execution_arn: str) -> None:
+        execution = self._store.load(execution_arn)
+        if execution.is_complete or not execution.paused:
+            return
+        execution.paused = False
+        deferred = execution.deferred_invocation
+        execution.deferred_invocation = False
+        self._store.save(execution)
+        if deferred:
+            self._invoke_execution(execution_arn)
+
+    def _wait_until_idle(self, execution_arn: str) -> None:
+        """Block until no invocation of ``execution_arn`` is running."""
+        while self._invocation_gate(execution_arn) is InvocationState.INVOKING:
+            time.sleep(0.005)
+
     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

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

@@ -714,6 +714,10 @@ def __init__(
         self._checkpoint_processor.set_chained_invoke_preflight(
             self._executor.preflight_chained_invoke
         )
+        # ARN of the most recent execution started by run_async(), including calls
+        # through run(). Used as the default target for pause_execution() and
+        # resume_execution() when no ARN is provided.
+        self._default_execution_arn: str | None = None
 
     def register_durable_function(
         self,
@@ -861,8 +865,59 @@ def run_async(
         if output.execution_arn is None:
             msg_arn: str = "Execution ARN must exist to run test."
             raise DurableFunctionsTestError(msg_arn)
+        self._default_execution_arn = output.execution_arn
         return output.execution_arn
 
+    def pause_execution(self, execution_arn: str | None = None) -> None:
+        """Make the local checkpoint server answer this execution's
+        checkpoints without a token, starting now.
+
+        Experimental; may change or be removed in a future release.
+
+        Defaults to the execution run_async() last started. The
+        invocation running now, if any, is answered without a token on
+        its next checkpoint: that checkpoint is accepted, but the
+        invocation reports PENDING, and no new invocation starts until
+        resume_execution().
+
+        Idempotent; a no-op once the execution has finished.
+        Resolves once no invocation of this execution is running.
+
+        Raises:
+            DurableFunctionsTestError: If no execution is in progress
+                (run_async() has not been called).
+        """
+        self._executor.pause_execution(self._require_execution_arn(execution_arn))
+
+    def resume_execution(self, execution_arn: str | None = None) -> None:
+        """Make the local checkpoint server answer this execution's
+        checkpoints with a token again.
+
+        Experimental; may change or be removed in a future release.
+
+        Defaults to the execution run_async() last started. Starts the
+        invocation pause_execution() held back, if any.
+
+        Idempotent; a no-op once the execution has finished or if it was not paused.
+
+        Raises:
+            DurableFunctionsTestError: If no execution is in progress
+                (run_async() has not been called).
+        """
+        self._executor.resume_execution(self._require_execution_arn(execution_arn))
+
+    def _require_execution_arn(self, execution_arn: str | None) -> str:
+        arn = (
+            execution_arn if execution_arn is not None else self._default_execution_arn
+        )
+        if arn is None:
+            msg = (
+                "No execution in progress. Call run_async() or run() first, "
+                "or pass execution_arn explicitly."
+            )
+            raise DurableFunctionsTestError(msg)
+        return arn
+
     def wait_for_result(
         self, execution_arn: str, timeout: int = 60
     ) -> DurableFunctionTestResult:

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

@@ -408,3 +408,52 @@ def test_process_checkpoint_delivers_due_wait_completion() -> None:
         op for op in persisted.operations if op.operation_id == "wait-1"
     )
     assert persisted_wait.status is OperationStatus.SUCCEEDED
+
+
+def test_paused_checkpoint_retries_without_a_token_even_after_resume():
+    """A retry of a checkpoint answered while paused replays the same
+    tokenless response, so the invocation cannot keep checkpointing. The
+    withholding is recorded on the idempotency entry, so a resume in
+    between does not hand the retry a live token."""
+    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="inv-paused-idem",
+    )
+    execution = Execution.new(start_input)
+    execution.start()
+    execution.paused = True
+    store.save(execution)
+
+    inbound = CheckpointToken(
+        execution_arn=execution.durable_execution_arn, token_sequence=0
+    ).to_str()
+    updates = [
+        OperationUpdate(
+            operation_id="step-A",
+            operation_type=OperationType.STEP,
+            action=OperationAction.START,
+            name="step-A",
+        )
+    ]
+
+    first = processor.process_checkpoint(inbound, updates, "c1")
+    assert first.checkpoint_token is None
+
+    retry = processor.process_checkpoint(inbound, updates, "c1")
+    assert retry.checkpoint_token is None
+
+    resumed = store.load(execution.durable_execution_arn)
+    resumed.paused = False
+    store.save(resumed)
+
+    retry_after_resume = processor.process_checkpoint(inbound, updates, "c1")
+    assert retry_after_resume.checkpoint_token is None