fix: replay status after paginated state (#377)

* fix: replay status after paginated state

* fix: move replay status check into state

---------

Co-authored-by: Frank Chen <65260095+zhongkechen@users.noreply.github.com>

Author: tejakusireddy

src/aws_durable_execution_sdk_python/execution.py

@@ -281,11 +281,7 @@ def execute(self):
             initial_checkpoint_token=self.invocation_input.checkpoint_token,
             operations={},
             service_client=self.service_client,
-            # If there are operations other than the initial EXECUTION one, current state is in replay mode
-            # todo: replay status will be wrong if initial_execution_state contains only one operation and more in next pages
-            replay_status=ReplayStatus.REPLAY
-            if len(self.invocation_input.initial_execution_state.operations) > 1
-            else ReplayStatus.NEW,
+            replay_status=ReplayStatus.NEW,
         )
 
         try:
@@ -307,6 +303,8 @@ def execute(self):
                 exception=e, retryable=e.is_retryable()
             )
 
+        execution_state.mark_replaying_if_prior_operations_exist()
+
         raw_input_payload: str | None = execution_state.get_input_payload()
 
         # Python RIC LambdaMarshaller just uses standard json deserialization for event

src/aws_durable_execution_sdk_python/state.py

@@ -379,6 +379,18 @@ def is_replaying(self) -> bool:
         with self._replay_status_lock:
             return self._replay_status is ReplayStatus.REPLAY
 
+    def mark_replaying_if_prior_operations_exist(self) -> None:
+        """Mark execution state as replaying when non-execution operations exist."""
+        with self._operations_lock:
+            has_prior_operations: bool = any(
+                op.operation_type is not OperationType.EXECUTION
+                for op in self.operations.values()
+            )
+
+        if has_prior_operations:
+            with self._replay_status_lock:
+                self._replay_status = ReplayStatus.REPLAY
+
     def get_checkpoint_result(self, checkpoint_id: str) -> CheckpointedResult:
         """Get checkpoint result.
 

tests/execution_test.py

@@ -43,6 +43,7 @@
     OperationStatus,
     OperationType,
     OperationUpdate,
+    StateOutput,
     StepDetails,
     WaitDetails,
 )
@@ -2686,6 +2687,36 @@ def _make_lambda_context():
     return ctx
 
 
+def test_durable_execution_replays_when_paginated_state_has_prior_operations():
+    """Test paginated execution state starts in replay mode when prior operations exist."""
+    mock_client = Mock(spec=DurableServiceClient)
+    step_operation = Operation(
+        operation_id="step1",
+        operation_type=OperationType.STEP,
+        status=OperationStatus.SUCCEEDED,
+    )
+    mock_client.get_execution_state.return_value = StateOutput(
+        operations=[step_operation],
+        next_marker=None,
+    )
+
+    invocation_input = _make_invocation_input(mock_client, next_marker="page2")
+
+    @durable_execution
+    def test_handler(event: Any, context: DurableContext) -> dict:
+        return {"is_replaying": context.state.is_replaying()}
+
+    result = test_handler(invocation_input, _make_lambda_context())
+
+    assert result["Status"] == InvocationStatus.SUCCEEDED.value
+    assert json.loads(result["Result"]) == {"is_replaying": True}
+    mock_client.get_execution_state.assert_called_once_with(
+        durable_execution_arn="arn:test:execution",
+        checkpoint_token="token123",
+        next_marker="page2",
+    )
+
+
 def test_durable_execution_non_retryable_invocation_error_returns_failed():
     """Test that non-retryable InvocationError returns FAILED instead of retrying."""
     mock_client = Mock(spec=DurableServiceClient)