fix: handle missing checkpoint token from service

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py

@@ -252,6 +252,13 @@ def resubmitter(ready: list[ExecutableWithState]) -> None:
             """
             try:
                 execution_state.create_checkpoint()
+            except OrphanedChildException:
+                # The execution has already completed, so there is nothing left
+                # to resume. Stop cleanly instead of recording a resume error.
+                # OrphanedChildException is a BaseException and so is not caught
+                # by the except below.
+                self._completion_event.set()
+                return
             except Exception as exc:  # noqa: BLE001
                 # resubmitter runs only on the single timer thread, so this
                 # check-then-set needs no lock. First error wins: keep the

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py

@@ -1086,7 +1086,8 @@ def from_dict(
 class CheckpointOutput:
     """Representation of the CheckpointDurableExecutionOutput structure of the DEX CheckpointDurableExecution API."""
 
-    checkpoint_token: str
+    # None on the terminal checkpoint that ends the execution.
+    checkpoint_token: str | None
     new_execution_state: CheckpointUpdatedExecutionState
 
     @classmethod
@@ -1109,8 +1110,7 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> CheckpointOutput:
             new_execution_state = CheckpointUpdatedExecutionState()
 
         return cls(
-            # TODO: maybe should throw if empty?
-            checkpoint_token=data.get("CheckpointToken", ""),
+            checkpoint_token=data.get("CheckpointToken"),
             new_execution_state=new_execution_state,
         )
 
@@ -1201,6 +1201,11 @@ def checkpoint(
         updates: list[OperationUpdate],
         client_token: str | None,
     ) -> CheckpointOutput:
+        # A checkpoint token is required. Raise a clear, retryable error (so the
+        # invocation re-drives) rather than letting the client reject an empty
+        # value with an opaque validation error.
+        if not checkpoint_token:
+            raise CheckpointError("Cannot checkpoint without a checkpoint token.")
         try:
             optional_params: dict[str, str] = {}
             if client_token is not None:
@@ -1230,6 +1235,13 @@ def get_execution_state(
         next_marker: str,
         max_items: int = 1000,
     ) -> StateOutput:
+        # A checkpoint token is required. Raise a clear, retryable error (so the
+        # invocation re-drives) rather than letting the client reject an empty
+        # value with an opaque validation error.
+        if not checkpoint_token:
+            raise GetExecutionStateError(
+                "Cannot get execution state without a checkpoint token."
+            )
         try:
             result: GetDurableExecutionStateResponseTypeDef = (
                 self.client.get_durable_execution_state(

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py

@@ -15,6 +15,7 @@
 
 from aws_durable_execution_sdk_python.exceptions import (
     BackgroundThreadError,
+    CheckpointError,
     DurableExecutionsError,
     DurableOperationError,
     GetExecutionStateError,
@@ -294,6 +295,12 @@ def __init__(
         self._overflow_queue: queue.Queue[QueuedOperation] = queue.Queue()
         self._checkpointing_stopped: threading.Event = threading.Event()
         self._checkpointing_failed: CompletionEvent = CompletionEvent()
+        # Set once the service confirms the execution has completed (a checkpoint
+        # response with no token). No further checkpoint can succeed afterward.
+        self._execution_completed: threading.Event = threading.Event()
+        # Serializes the execution-completed check with enqueueing so a checkpoint
+        # is never enqueued after _settle_after_execution_completed drains the queue.
+        self._completion_lock: Lock = Lock()
 
         # Concurrency management for parallel operations: parent_id -> {child_operation_ids}
         self._parent_to_children: dict[str, set[str]] = {}
@@ -608,8 +615,27 @@ def create_checkpoint(
         # Create wrapper object for queue
         queued_op = QueuedOperation(operation_update, completion_event)
 
-        # Enqueue the wrapper object (operation_update can be None for empty checkpoints)
-        self._checkpoint_queue.put(queued_op)
+        # Enqueue under the same lock the background loop holds while it stops and
+        # drains the queue - on completion via _settle_after_execution_completed, or
+        # on failure in the exception handler. Re-check both terminal conditions
+        # inside the lock so a checkpoint is never enqueued after a drain and left
+        # with a waiter that blocks forever.
+        with self._completion_lock:
+            if self._checkpointing_failed.is_set():
+                # Raises the stored BackgroundThreadError.
+                self._checkpointing_failed.wait()
+            if self._execution_completed.is_set():
+                operation_id = (
+                    operation_update.operation_id
+                    if operation_update is not None
+                    else ""
+                )
+                raise OrphanedChildException(
+                    "Execution already completed; checkpoint will not be processed.",
+                    operation_id=operation_id,
+                )
+            # Enqueue the wrapper object (operation_update can be None for empty checkpoints)
+            self._checkpoint_queue.put(queued_op)
 
         # Conditionally wait for completion based on is_sync parameter
         if is_sync:
@@ -756,15 +782,34 @@ def checkpoint_batches_forever(self) -> None:
 
                     logger.debug("Checkpoint batch processed successfully")
 
-                    # Update local token for next iteration
-                    current_checkpoint_token = output.checkpoint_token
+                    # The service omits the token only when it completes the
+                    # execution. Advance on a returned token; otherwise treat a
+                    # terminal batch as completion and any other omission as an
+                    # invalid response.
+                    execution_completed: bool = False
+                    if output.checkpoint_token:
+                        current_checkpoint_token = output.checkpoint_token
+                    elif any(
+                        update.operation_type is OperationType.EXECUTION
+                        and (
+                            update.action is OperationAction.SUCCEED
+                            or update.action is OperationAction.FAIL
+                        )
+                        for update in updates
+                    ):
+                        execution_completed = True
+                    else:
+                        raise CheckpointError(
+                            "Checkpoint response omitted the token outside of "
+                            "execution completion."
+                        )
 
                     previous_operations = self.operations
 
                     # Fetch new operations from the API before unblocking sync waiters
                     updated_operations = self.fetch_paginated_operations(
                         output.new_execution_state.operations,
-                        output.checkpoint_token,
+                        current_checkpoint_token,
                         output.new_execution_state.next_marker,
                     )
                     self._plugin_executor.on_operation_update(
@@ -777,6 +822,15 @@ def checkpoint_batches_forever(self) -> None:
                     for queued_op in batch:
                         if queued_op.completion_event is not None:
                             queued_op.completion_event.set()
+
+                    if execution_completed:
+                        # The execution is complete; no further checkpoint can
+                        # succeed. Stop the loop and settle any still-queued
+                        # operations (e.g. from orphaned concurrent branches) so
+                        # their waiters do not block and no further checkpoint
+                        # request is issued with the consumed token.
+                        self._settle_after_execution_completed()
+                        break
                 except Exception as e:
                     # Checkpoint failed - wake all blocked threads so they can raise error
                     # Drain both queues and signal all completion events
@@ -785,38 +839,74 @@ def checkpoint_batches_forever(self) -> None:
                         "Checkpoint creation failed", e
                     )
 
-                    # FIFO: although at this point order not really import any anymore
-                    # Signal completion events for the failed batch
+                    # Signal completion events for the failed batch (already dequeued,
+                    # so no producer can race these).
                     for queued_op in batch:
                         if queued_op.completion_event is not None:
                             queued_op.completion_event.set(bg_error)
 
-                    # overflow 1st: although at this point order not really import any anymore
-                    while not self._overflow_queue.empty():
-                        try:
-                            item = self._overflow_queue.get_nowait()
-                            if item.completion_event:
-                                item.completion_event.set(bg_error)
-                        except queue.Empty:
-                            break
-
-                    # finally Wake all blocked threads in main queue
-                    while not self._checkpoint_queue.empty():
-                        try:
-                            item = self._checkpoint_queue.get_nowait()
-                            if item.completion_event:
-                                item.completion_event.set(bg_error)
-                        except queue.Empty:
-                            break
-
-                    # Set the failure event so future checkpoint attempts fail immediately
-                    self._checkpointing_failed.set(bg_error)
+                    # Drain the queues and set the failure flag under the completion
+                    # lock so a concurrent create_checkpoint either observes the
+                    # failure and raises, or has its operation drained here - never
+                    # enqueued after the drain and left blocked.
+                    with self._completion_lock:
+                        while not self._overflow_queue.empty():
+                            try:
+                                item = self._overflow_queue.get_nowait()
+                                if item.completion_event:
+                                    item.completion_event.set(bg_error)
+                            except queue.Empty:
+                                break
+
+                        while not self._checkpoint_queue.empty():
+                            try:
+                                item = self._checkpoint_queue.get_nowait()
+                                if item.completion_event:
+                                    item.completion_event.set(bg_error)
+                            except queue.Empty:
+                                break
+
+                        # Future checkpoint attempts fail immediately.
+                        self._checkpointing_failed.set(bg_error)
 
                     # Exit the loop - error has been signaled to main thread via completion events
                     break
 
         logger.debug("Background checkpoint processing stopped")
 
+    def _settle_after_execution_completed(self) -> None:
+        """Stop checkpointing and settle queued operations after the execution ends.
+
+        Called when a checkpoint response omits the token on a terminal batch,
+        which the service does only when it completes the execution. Any operation
+        still queued belongs to work that can no longer be checkpointed (typically
+        an orphaned concurrent branch), so its waiter is settled with
+        OrphanedChildException rather than left blocking or sent with the consumed
+        token.
+        """
+        with self._completion_lock:
+            self._execution_completed.set()
+            self._checkpointing_stopped.set()
+
+            for pending_queue in (self._overflow_queue, self._checkpoint_queue):
+                while not pending_queue.empty():
+                    try:
+                        queued_op: QueuedOperation = pending_queue.get_nowait()
+                    except queue.Empty:
+                        break
+                    if queued_op.completion_event is not None:
+                        operation_id: str = (
+                            queued_op.operation_update.operation_id
+                            if queued_op.operation_update is not None
+                            else ""
+                        )
+                        queued_op.completion_event.set(
+                            OrphanedChildException(
+                                "Execution already completed; checkpoint will not be processed.",
+                                operation_id=operation_id,
+                            )
+                        )
+
     def stop_checkpointing(self) -> None:
         """Signal background thread to stop checkpointing.
 

packages/aws-durable-execution-sdk-python/tests/concurrency_test.py

@@ -1,5 +1,6 @@
 """Tests for the concurrency module."""
 
+import contextlib
 import json
 import random
 import threading
@@ -38,6 +39,7 @@
 from aws_durable_execution_sdk_python.exceptions import (
     ChildContextError,
     InvalidStateError,
+    OrphanedChildException,
     SuspendExecution,
     TimedSuspendExecution,
 )
@@ -1574,6 +1576,77 @@ def checkpoint(*args, **kwargs):
     executor.long_runner_release.set()
 
 
+def test_concurrent_executor_resume_orphaned_stops_cleanly():
+    """An orphaned resume refresh must be handled, not propagated.
+
+    When the execution completes while a concurrent operation is orphaned, the
+    timer resubmit's checkpoint refresh raises OrphanedChildException. That is a
+    BaseException, so the resubmitter's generic `except Exception` does not catch
+    it; it must be handled explicitly so it neither crashes the timer thread nor
+    escapes execute(). Contrast with the RuntimeError case above, which does
+    propagate.
+    """
+
+    class TestExecutor(ConcurrentExecutor):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.long_runner_release = threading.Event()
+
+        def execute_item(self, child_context, executable):
+            if executable.index == 0:
+                # Long-runner keeps the map alive so task 1 resumes in-process.
+                self.long_runner_release.wait(timeout=5)
+                return "result_A"
+            # Task 1 suspends with a past timestamp -> immediate in-process resume.
+            msg = "resume-me"
+            raise TimedSuspendExecution(msg, time.time() - 1)
+
+    executables = [Executable(0, lambda: "task_A"), Executable(1, lambda: "task_B")]
+    completion_config = CompletionConfig(
+        min_successful=2,
+        tolerated_failure_count=None,
+        tolerated_failure_percentage=None,
+    )
+
+    executor = TestExecutor(
+        executables=executables,
+        max_concurrency=2,
+        completion_config=completion_config,
+        sub_type_top="TOP",
+        sub_type_iteration="ITER",
+        name_prefix="test_",
+        serdes=None,
+    )
+
+    execution_state = Mock()
+
+    def checkpoint(*args, **kwargs):
+        # The resume refresh calls create_checkpoint() with no arguments; the
+        # execution has completed, so it is orphaned.
+        if not args and not kwargs:
+            raise OrphanedChildException(
+                "execution already completed", operation_id="op"
+            )
+
+    execution_state.create_checkpoint = Mock(side_effect=checkpoint)
+
+    executor_context = Mock()
+    executor_context._create_step_id_for_logical_step = lambda *args: "1"  # noqa: SLF001
+    child_context = Mock()
+    child_context.state.wrap_user_function = lambda func, *args, **kwargs: func
+    executor_context.create_child_context = lambda *args, **kwargs: child_context
+
+    try:
+        # The orphaned resume must not surface as an error; the map may suspend
+        # or complete, but OrphanedChildException must never escape.
+        with contextlib.suppress(SuspendExecution):
+            executor.execute(execution_state, executor_context)
+    except OrphanedChildException:
+        pytest.fail("OrphanedChildException from the resume refresh must be handled")
+    finally:
+        executor.long_runner_release.set()
+
+
 def test_concurrent_executor_with_timed_resubmit_while_other_task_running():
     """Test timed resubmission while other tasks are still running."""