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

@@ -1098,7 +1098,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
@@ -1121,8 +1122,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,
         )
 
@@ -1213,6 +1213,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:
@@ -1242,6 +1247,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,9 @@ 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()
 
         # Concurrency management for parallel operations: parent_id -> {child_operation_ids}
         self._parent_to_children: dict[str, set[str]] = {}
@@ -587,6 +591,18 @@ def create_checkpoint(
                         operation_id=operation_update.operation_id,
                     )
 
+        # The execution has completed; no further checkpoint can succeed. Reject
+        # late attempts (e.g. from an orphaned concurrent branch) rather than
+        # enqueueing work that would block or be sent with the consumed token.
+        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,
+            )
+
         # Check if background checkpointing has failed
         if self._checkpointing_failed.is_set():
             # This will raise the stored BackgroundThreadError
@@ -748,15 +764,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,
                     )
                     for update in updates:
@@ -776,6 +811,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
@@ -816,6 +860,38 @@ def checkpoint_batches_forever(self) -> None:
 
         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.
+        """
+        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."""
 

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

@@ -1641,10 +1641,14 @@ def test_checkpoint_output_from_dict():
 
 
 def test_checkpoint_output_from_dict_empty():
-    """Test CheckpointOutput.from_dict with empty data."""
+    """Test CheckpointOutput.from_dict with empty data.
+
+    The service omits CheckpointToken on the terminal checkpoint; from_dict
+    surfaces that as None rather than an empty string.
+    """
     data = {}
     output = CheckpointOutput.from_dict(data)
-    assert not output.checkpoint_token
+    assert output.checkpoint_token is None
     assert len(output.new_execution_state.operations) == 0
     assert output.new_execution_state.next_marker is None
 
@@ -1825,6 +1829,37 @@ def test_lambda_client_checkpoint_with_explicit_none_client_token():
     assert result.checkpoint_token == "new_token"  # noqa: S105
 
 
+@pytest.mark.parametrize("token", ["", None])
+def test_lambda_client_checkpoint_empty_token_raises(token):
+    """An empty/None checkpoint token raises a retryable CheckpointError, not a client error."""
+    mock_client = Mock()
+    lambda_client = LambdaClient(mock_client)
+    update = OperationUpdate(
+        operation_id="op1",
+        operation_type=OperationType.STEP,
+        action=OperationAction.START,
+    )
+
+    with pytest.raises(CheckpointError) as exc_info:
+        lambda_client.checkpoint("arn123", token, [update], None)
+
+    assert exc_info.value.is_retryable()
+    mock_client.checkpoint_durable_execution.assert_not_called()
+
+
+@pytest.mark.parametrize("token", ["", None])
+def test_lambda_client_get_execution_state_empty_token_raises(token):
+    """An empty/None checkpoint token raises a retryable GetExecutionStateError, not a client error."""
+    mock_client = Mock()
+    lambda_client = LambdaClient(mock_client)
+
+    with pytest.raises(GetExecutionStateError) as exc_info:
+        lambda_client.get_execution_state("arn123", token, "marker")
+
+    assert exc_info.value.is_retryable()
+    mock_client.get_durable_execution_state.assert_not_called()
+
+
 def test_lambda_client_checkpoint_with_empty_string_client_token():
     """Test LambdaClient.checkpoint method with empty string client_token - should pass empty string."""
     mock_client = Mock()