feat: execution state pagination and token validation

- Add paging logic in checkpoint processor with next_marker support
- Implement checkpoint token validation
- Add token expiration checking with error responses
- Handle missing token cases with context-appropriate validation
- Add pagination metadata to responses with configurable max_items
- Add test coverage for all validation scenarios

Author: Rares Polenciuc

src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py

@@ -88,14 +88,35 @@ def process_checkpoint(
     def get_execution_state(
         self,
         checkpoint_token: str,
-        next_marker: str,  # noqa: ARG002
-        max_items: int = 1000,  # noqa: ARG002
+        next_marker: str | None = None,
+        max_items: int = 1000,
     ) -> StateOutput:
-        """Get current execution state."""
+        """Get current execution state with batched checkpoint token validation and pagination."""
+        if not checkpoint_token:
+            msg: str = "Checkpoint token is required"
+            raise InvalidParameterValueException(msg)
+
         token: CheckpointToken = CheckpointToken.from_str(checkpoint_token)
         execution: Execution = self._store.load(token.execution_arn)
+        token.validate_for_execution(execution)
+
+        # Get all operations
+        all_operations = execution.get_navigable_operations()
+
+        # Apply pagination
+        start_index = 0
+        if next_marker:
+            try:
+                start_index = int(next_marker)
+            except ValueError:
+                start_index = 0
+
+        end_index = start_index + max_items
+        paginated_operations = all_operations[start_index:end_index]
+
+        # Determine next marker
+        next_marker_result = str(end_index) if end_index < len(all_operations) else None
 
-        # TODO: paging when size or max
         return StateOutput(
-            operations=execution.get_navigable_operations(), next_marker=None
+            operations=paginated_operations, next_marker=next_marker_result
         )

src/aws_durable_execution_sdk_python_testing/executor.py

@@ -41,6 +41,7 @@
     Execution as ExecutionSummary,
 )
 from aws_durable_execution_sdk_python_testing.observer import ExecutionObserver
+from aws_durable_execution_sdk_python_testing.token import CheckpointToken
 
 
 if TYPE_CHECKING:
@@ -359,9 +360,12 @@ def get_execution_state(
         """
         execution = self.get_execution(execution_arn)
 
-        # TODO: Validate checkpoint token if provided
-        if checkpoint_token and checkpoint_token not in execution.used_tokens:
-            msg: str = f"Invalid checkpoint token: {checkpoint_token}"
+        # Checkpoint token validation
+        if checkpoint_token:
+            token = CheckpointToken.from_str(checkpoint_token)
+            token.validate_for_execution(execution)
+        elif not execution.is_complete and marker is not None:
+            msg = "Checkpoint token is required for paginated requests on active executions"
             raise InvalidParameterValueException(msg)
 
         # Get operations (excluding the initial EXECUTION operation for state)
@@ -371,7 +375,6 @@ def get_execution_state(
         if max_items is None:
             max_items = 100
 
-        # Simple pagination - in real implementation would need proper marker handling
         start_index = 0
         if marker:
             try:
@@ -468,11 +471,16 @@ def checkpoint_execution(
         """
         execution = self.get_execution(execution_arn)
 
-        # Validate checkpoint token
-        if checkpoint_token not in execution.used_tokens:
-            msg: str = f"Invalid checkpoint token: {checkpoint_token}"
+        # Comprehensive checkpoint token validation
+        from aws_durable_execution_sdk_python_testing.token import CheckpointToken
+
+        if not checkpoint_token:
+            msg = "Checkpoint token is required for checkpoint operations"
             raise InvalidParameterValueException(msg)
 
+        token = CheckpointToken.from_str(checkpoint_token)
+        token.validate_for_execution(execution)
+
         # TODO: Process operation updates using the checkpoint processor
         # This would integrate with the existing checkpoint processing pipeline
 

src/aws_durable_execution_sdk_python_testing/token.py

@@ -6,6 +6,10 @@
 import json
 from dataclasses import dataclass
 
+from aws_durable_execution_sdk_python_testing.exceptions import (
+    InvalidParameterValueException,
+)
+
 
 @dataclass(frozen=True)
 class CheckpointToken:
@@ -27,6 +31,23 @@ def from_str(cls, token: str) -> CheckpointToken:
         data = json.loads(decoded)
         return cls(execution_arn=data["arn"], token_sequence=data["seq"])
 
+    def validate_for_execution(self, execution) -> None:
+        """Validate token against execution"""
+
+        if self.execution_arn != execution.durable_execution_arn:
+            msg = "Checkpoint token does not match execution ARN"
+            raise InvalidParameterValueException(msg)
+
+        if execution.is_complete or self.token_sequence > execution.token_sequence:
+            msg = "Invalid or expired checkpoint token"
+            raise InvalidParameterValueException(msg)
+
+        # Check if token has been used
+        token_str = self.to_str()
+        if token_str not in execution.used_tokens:
+            msg = f"Invalid checkpoint token: {token_str}"
+            raise InvalidParameterValueException(msg)
+
 
 @dataclass(frozen=True)
 class CallbackToken:

tests/executor_test.py

@@ -33,6 +33,7 @@
     StartDurableExecutionInput,
 )
 from aws_durable_execution_sdk_python_testing.observer import ExecutionObserver
+from aws_durable_execution_sdk_python_testing.token import CheckpointToken
 
 
 class MockExecutionObserver(ExecutionObserver):
@@ -1889,9 +1890,15 @@ def test_stop_execution_with_custom_error(executor, mock_store):
 
 def test_get_execution_state(executor, mock_store):
     """Test get_execution_state method."""
-
     mock_execution = Mock()
-    mock_execution.used_tokens = {"token1", "token2"}
+    mock_execution.durable_execution_arn = "test-arn"
+    mock_execution.token_sequence = 5
+    mock_execution.is_complete = False
+
+    # Create valid token and add to used_tokens
+    token = CheckpointToken("test-arn", 3)
+    valid_token = token.to_str()
+    mock_execution.used_tokens = {valid_token}
 
     # Create mock operations
     operations = [
@@ -1916,7 +1923,7 @@ def test_get_execution_state(executor, mock_store):
 
     mock_store.load.return_value = mock_execution
 
-    result = executor.get_execution_state("test-arn", checkpoint_token="token1")  # noqa: S106
+    result = executor.get_execution_state("test-arn", checkpoint_token=valid_token)
 
     assert len(result.operations) == 2
     assert result.next_marker is None
@@ -1928,11 +1935,13 @@ def test_get_execution_state_invalid_token(executor, mock_store):
     mock_execution = Mock()
     mock_execution.used_tokens = {"token1", "token2"}
     mock_store.load.return_value = mock_execution
-
+    token = CheckpointToken("invalid-token", 3)
+    invalid_token = token.to_str()
     with pytest.raises(
-        InvalidParameterValueException, match="Invalid checkpoint token"
+        InvalidParameterValueException,
+        match="Checkpoint token does not match execution ARN",
     ):
-        executor.get_execution_state("test-arn", checkpoint_token="invalid-token")  # noqa: S106
+        executor.get_execution_state("test-arn", checkpoint_token=invalid_token)
 
 
 def test_get_execution_history(executor, mock_store):
@@ -1950,13 +1959,22 @@ def test_get_execution_history(executor, mock_store):
 def test_checkpoint_execution(executor, mock_store):
     """Test checkpoint_execution method."""
     mock_execution = Mock()
-    mock_execution.used_tokens = {"token1", "token2"}
-    mock_execution.get_new_checkpoint_token.return_value = "new-token"
+    mock_execution.durable_execution_arn = "test-arn"
+    mock_execution.token_sequence = 5
+    mock_execution.is_complete = False
+    new_token = "new-token"  # noqa:S105
+    mock_execution.get_new_checkpoint_token.return_value = new_token
+
+    # Create valid token and add to used_tokens
+    token = CheckpointToken("test-arn", 3)
+    valid_token = token.to_str()
+    mock_execution.used_tokens = {valid_token}
+
     mock_store.load.return_value = mock_execution
 
-    result = executor.checkpoint_execution("test-arn", "token1")
+    result = executor.checkpoint_execution("test-arn", valid_token)
 
-    assert result.checkpoint_token == "new-token"  # noqa: S105
+    assert result.checkpoint_token == new_token
     assert result.new_execution_state is None
     mock_store.load.assert_called_once_with("test-arn")
     mock_execution.get_new_checkpoint_token.assert_called_once()
@@ -1967,11 +1985,13 @@ def test_checkpoint_execution_invalid_token(executor, mock_store):
     mock_execution = Mock()
     mock_execution.used_tokens = {"token1", "token2"}
     mock_store.load.return_value = mock_execution
-
+    token = CheckpointToken("invalid-token", 3)
+    invalid_token = token.to_str()
     with pytest.raises(
-        InvalidParameterValueException, match="Invalid checkpoint token"
+        InvalidParameterValueException,
+        match="Checkpoint token does not match execution ARN",
     ):
-        executor.checkpoint_execution("test-arn", "invalid-token")
+        executor.checkpoint_execution("test-arn", invalid_token)
 
 
 # Callback method tests
@@ -2050,3 +2070,35 @@ def test_send_callback_heartbeat_none_callback_id(executor):
     """Test send_callback_heartbeat with None callback_id."""
     with pytest.raises(InvalidParameterValueException, match="callback_id is required"):
         executor.send_callback_heartbeat(None)
+
+
+def test_get_execution_state_no_token_with_marker_active_execution(
+    mock_store, mock_scheduler, mock_invoker
+):
+    """Test get_execution_state fails when no token provided with marker on active execution."""
+    executor = Executor(mock_store, mock_scheduler, mock_invoker)
+    execution_arn = "test-arn"
+
+    # Create an active execution
+    execution = Execution(execution_arn, "test-function", {})
+    execution.is_complete = False
+    mock_store.load.return_value = execution
+
+    with pytest.raises(
+        InvalidParameterValueException, match="Checkpoint token is required"
+    ):
+        executor.get_execution_state(execution_arn, marker="some-marker")
+
+
+def test_checkpoint_execution_no_token(mock_store, mock_scheduler, mock_invoker):
+    """Test checkpoint_execution fails when no token provided."""
+    executor = Executor(mock_store, mock_scheduler, mock_invoker)
+    execution_arn = "test-arn"
+
+    execution = Execution(execution_arn, "test-function", {})
+    mock_store.load.return_value = execution
+
+    with pytest.raises(
+        InvalidParameterValueException, match="Checkpoint token is required"
+    ):
+        executor.checkpoint_execution(execution_arn, "", [], "client-token")

tests/token_test.py

@@ -2,9 +2,13 @@
 
 import base64
 import json
+from unittest.mock import Mock
 
 import pytest
 
+from aws_durable_execution_sdk_python_testing.exceptions import (
+    InvalidParameterValueException,
+)
 from aws_durable_execution_sdk_python_testing.token import (
     CallbackToken,
     CheckpointToken,
@@ -130,3 +134,94 @@ def test_callback_token_frozen_dataclass():
 
     with pytest.raises(AttributeError):
         token.operation_id = "new-op"
+
+
+def test_checkpoint_token_validate_for_execution_success():
+    """Test successful token validation."""
+    token = CheckpointToken("test-arn", 5)
+    execution = Mock()
+    execution.durable_execution_arn = "test-arn"
+    execution.token_sequence = 10
+    execution.is_complete = False
+    execution.used_tokens = {token.to_str()}
+
+    token.validate_for_execution(execution)
+
+
+def test_checkpoint_token_validate_for_execution_arn_mismatch():
+    """Test token validation fails when ARN doesn't match."""
+    token = CheckpointToken("test-arn", 5)
+    execution = Mock()
+    execution.durable_execution_arn = "different-arn"
+    execution.token_sequence = 10
+    execution.is_complete = False
+
+    with pytest.raises(
+        InvalidParameterValueException, match="does not match execution ARN"
+    ):
+        token.validate_for_execution(execution)
+
+
+def test_checkpoint_token_validate_for_execution_completed():
+    """Test token validation fails when execution is complete."""
+    token = CheckpointToken("test-arn", 5)
+    execution = Mock()
+    execution.durable_execution_arn = "test-arn"
+    execution.token_sequence = 10
+    execution.is_complete = True
+
+    with pytest.raises(InvalidParameterValueException, match="Invalid or expired"):
+        token.validate_for_execution(execution)
+
+
+def test_checkpoint_token_validate_for_execution_future_sequence():
+    """Test token validation fails when token sequence is from future."""
+    token = CheckpointToken("test-arn", 15)
+    execution = Mock()
+    execution.durable_execution_arn = "test-arn"
+    execution.token_sequence = 10
+    execution.is_complete = False
+
+    with pytest.raises(InvalidParameterValueException, match="Invalid or expired"):
+        token.validate_for_execution(execution)
+
+
+def test_checkpoint_token_validate_for_execution_equal_sequence():
+    """Test token validation succeeds when sequences are equal."""
+    token = CheckpointToken("test-arn", 10)
+    execution = Mock()
+    execution.durable_execution_arn = "test-arn"
+    execution.token_sequence = 10
+    execution.is_complete = False
+    execution.used_tokens = {token.to_str()}
+
+    token.validate_for_execution(execution)
+
+
+def test_checkpoint_token_validate_for_execution_not_in_used_tokens():
+    """Test token validation fails when token not in used_tokens."""
+    token = CheckpointToken("test-arn", 5)
+    execution = Mock()
+    execution.durable_execution_arn = "test-arn"
+    execution.token_sequence = 10
+    execution.is_complete = False
+    execution.used_tokens = {"other-token"}
+
+    with pytest.raises(
+        InvalidParameterValueException, match="Invalid checkpoint token"
+    ):
+        token.validate_for_execution(execution)
+
+
+def test_checkpoint_token_validate_for_execution_in_used_tokens():
+    """Test token validation succeeds when token is in used_tokens."""
+    token = CheckpointToken("test-arn", 5)
+    execution = Mock()
+    execution.durable_execution_arn = "test-arn"
+    execution.token_sequence = 10
+    execution.is_complete = False
+    # Mock the token string that would be generated
+    token_str = token.to_str()
+    execution.used_tokens = {token_str}
+
+    token.validate_for_execution(execution)