Avoid unnecessary retry sleeps after final page batch

Co-authored-by: Shri Sukhani <shrisukhani@users.noreply.github.com>

Author: cursoragent

hyperbrowser/client/polling.py

@@ -201,20 +201,23 @@ def collect_paginated_results(
             raise HyperbrowserTimeoutError(
                 f"Timed out fetching paginated results for {operation_name} after {max_wait_seconds} seconds"
             )
+        should_sleep = True
         try:
             page_response = get_next_page(current_page_batch + 1)
             on_page_success(page_response)
             current_page_batch = get_current_page_batch(page_response)
             total_page_batches = get_total_page_batches(page_response)
             failures = 0
             first_check = False
+            should_sleep = current_page_batch < total_page_batches
         except Exception as exc:
             failures += 1
             if failures >= max_attempts:
                 raise HyperbrowserError(
                     f"Failed to fetch page batch {current_page_batch + 1} for {operation_name} after {max_attempts} attempts: {exc}"
                 ) from exc
-        time.sleep(retry_delay_seconds)
+        if should_sleep:
+            time.sleep(retry_delay_seconds)
 
 
 async def collect_paginated_results_async(
@@ -244,20 +247,23 @@ async def collect_paginated_results_async(
             raise HyperbrowserTimeoutError(
                 f"Timed out fetching paginated results for {operation_name} after {max_wait_seconds} seconds"
             )
+        should_sleep = True
         try:
             page_response = await get_next_page(current_page_batch + 1)
             on_page_success(page_response)
             current_page_batch = get_current_page_batch(page_response)
             total_page_batches = get_total_page_batches(page_response)
             failures = 0
             first_check = False
+            should_sleep = current_page_batch < total_page_batches
         except Exception as exc:
             failures += 1
             if failures >= max_attempts:
                 raise HyperbrowserError(
                     f"Failed to fetch page batch {current_page_batch + 1} for {operation_name} after {max_attempts} attempts: {exc}"
                 ) from exc
-        await asyncio.sleep(retry_delay_seconds)
+        if should_sleep:
+            await asyncio.sleep(retry_delay_seconds)
 
 
 def wait_for_job_result(

tests/test_polling.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+import hyperbrowser.client.polling as polling_helpers
 from hyperbrowser.client.polling import (
     collect_paginated_results,
     collect_paginated_results_async,
@@ -225,6 +226,55 @@ async def run() -> None:
     asyncio.run(run())
 
 
+def test_collect_paginated_results_does_not_sleep_after_last_page(monkeypatch):
+    sleep_calls = []
+
+    monkeypatch.setattr(
+        polling_helpers.time, "sleep", lambda delay: sleep_calls.append(delay)
+    )
+
+    collect_paginated_results(
+        operation_name="sync single-page",
+        get_next_page=lambda page: {"current": 1, "total": 1, "items": ["a"]},
+        get_current_page_batch=lambda response: response["current"],
+        get_total_page_batches=lambda response: response["total"],
+        on_page_success=lambda response: None,
+        max_wait_seconds=1.0,
+        max_attempts=2,
+        retry_delay_seconds=0.5,
+    )
+
+    assert sleep_calls == []
+
+
+def test_collect_paginated_results_async_does_not_sleep_after_last_page(monkeypatch):
+    sleep_calls = []
+
+    async def fake_sleep(delay: float) -> None:
+        sleep_calls.append(delay)
+
+    monkeypatch.setattr(polling_helpers.asyncio, "sleep", fake_sleep)
+
+    async def run() -> None:
+        async def get_next_page(page: int) -> dict:
+            return {"current": 1, "total": 1, "items": ["a"]}
+
+        await collect_paginated_results_async(
+            operation_name="async single-page",
+            get_next_page=get_next_page,
+            get_current_page_batch=lambda response: response["current"],
+            get_total_page_batches=lambda response: response["total"],
+            on_page_success=lambda response: None,
+            max_wait_seconds=1.0,
+            max_attempts=2,
+            retry_delay_seconds=0.5,
+        )
+
+    asyncio.run(run())
+
+    assert sleep_calls == []
+
+
 def test_collect_paginated_results_raises_after_page_failures():
     with pytest.raises(HyperbrowserError, match="Failed to fetch page batch 1"):
         collect_paginated_results(