Log Durable API errors during state fetch, and save partially fetched state.

Author: denzyll

src/aws_durable_execution_sdk_python/state.py

@@ -16,6 +16,7 @@
     BackgroundThreadError,
     CallableRuntimeError,
     DurableExecutionsError,
+    GetExecutionStateError,
     OrphanedChildException,
 )
 from aws_durable_execution_sdk_python.lambda_service import (
@@ -275,20 +276,38 @@ def fetch_paginated_operations(
             initial_operations: initial operations to be added to ExecutionState
             checkpoint_token: checkpoint token used to call Durable Functions API.
             next_marker: a marker indicates that there are paginated operations.
+
+        Raises:
+            GetExecutionStateError: If the API call fails. The error is logged
+                with structured extras before re-raising. Callers are responsible
+                for deciding whether to fail the execution or allow Lambda retry
+                based on is_retryable().
         """
         all_operations: list[Operation] = (
             initial_operations.copy() if initial_operations else []
         )
-        while next_marker:
-            output: StateOutput = self._service_client.get_execution_state(
-                durable_execution_arn=self.durable_execution_arn,
-                checkpoint_token=checkpoint_token,
-                next_marker=next_marker,
+        try:
+            while next_marker:
+                output: StateOutput = self._service_client.get_execution_state(
+                    durable_execution_arn=self.durable_execution_arn,
+                    checkpoint_token=checkpoint_token,
+                    next_marker=next_marker,
+                )
+                all_operations.extend(output.operations)
+                next_marker = output.next_marker
+        except GetExecutionStateError as e:
+            logger.exception(
+                "Durable API error during state fetch.",
+                extra=e.build_logger_extras(),
             )
-            all_operations.extend(output.operations)
-            next_marker = output.next_marker
-        with self._operations_lock:
-            self.operations.update({op.operation_id: op for op in all_operations})
+            raise
+        finally:
+            # Always store whatever operations we successfully fetched
+            if all_operations:
+                with self._operations_lock:
+                    self.operations.update(
+                        {op.operation_id: op for op in all_operations}
+                    )
 
     def get_input_payload(self) -> str | None:
         # It is possible that backend will not provide an execution operation

tests/state_test.py

@@ -16,6 +16,8 @@
 from aws_durable_execution_sdk_python.exceptions import (
     BackgroundThreadError,
     CallableRuntimeError,
+    DurableApiErrorCategory,
+    GetExecutionStateError,
     OrphanedChildException,
 )
 from aws_durable_execution_sdk_python.identifier import OperationIdentifier
@@ -724,6 +726,88 @@ def mock_get_execution_state(durable_execution_arn, checkpoint_token, next_marke
         assert operation.operation_id == expected_op.operation_id
 
 
+def test_fetch_paginated_operations_stores_partial_results_on_error():
+    """Test that operations from successful pages are stored even when a later page fails."""
+    mock_lambda_client = Mock(spec=LambdaClient)
+
+    non_retryable_error = GetExecutionStateError(
+        message="KMS access denied",
+        error_category=DurableApiErrorCategory.EXECUTION,
+        error={"Code": "KMSAccessDeniedException", "Message": "KMS access denied"},
+        response_metadata={"HTTPStatusCode": 502},
+    )
+
+    def mock_get_execution_state(durable_execution_arn, checkpoint_token, next_marker):
+        if next_marker == "marker1":
+            return StateOutput(
+                operations=[
+                    Operation(
+                        operation_id="1",
+                        operation_type=OperationType.STEP,
+                        status=OperationStatus.STARTED,
+                    )
+                ],
+                next_marker="marker2",
+            )
+        raise non_retryable_error
+
+    mock_lambda_client.get_execution_state.side_effect = mock_get_execution_state
+
+    state = ExecutionState(
+        durable_execution_arn="test_arn",
+        initial_checkpoint_token="token123",  # noqa: S106
+        operations={},
+        service_client=mock_lambda_client,
+    )
+
+    with pytest.raises(GetExecutionStateError):
+        state.fetch_paginated_operations(
+            initial_operations=[
+                Operation(
+                    operation_id="0",
+                    operation_type=OperationType.STEP,
+                    status=OperationStatus.STARTED,
+                )
+            ],
+            checkpoint_token="test_token",  # noqa: S106
+            next_marker="marker1",
+        )
+
+    # Initial operation + page 1 should be stored despite page 2 failing
+    assert "0" in state.operations
+    assert "1" in state.operations
+    assert len(state.operations) == 2
+
+
+def test_fetch_paginated_operations_logs_error(caplog):
+    """Test that GetExecutionStateError is logged with structured extras."""
+    mock_lambda_client = Mock(spec=LambdaClient)
+
+    error = GetExecutionStateError(
+        message="Service error",
+        error_category=DurableApiErrorCategory.INVOCATION,
+        error={"Code": "ServiceException", "Message": "Service error"},
+        response_metadata={"HTTPStatusCode": 500},
+    )
+    mock_lambda_client.get_execution_state.side_effect = error
+
+    state = ExecutionState(
+        durable_execution_arn="test_arn",
+        initial_checkpoint_token="token123",  # noqa: S106
+        operations={},
+        service_client=mock_lambda_client,
+    )
+
+    with pytest.raises(GetExecutionStateError):
+        state.fetch_paginated_operations(
+            initial_operations=[],
+            checkpoint_token="test_token",  # noqa: S106
+            next_marker="marker1",
+        )
+
+    assert "Durable API error during state fetch." in caplog.text
+
+
 # ============================================================================
 # Checkpoint Batching Tests
 # ============================================================================