Fail fast when pagination callbacks raise

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

Author: cursoragent

hyperbrowser/client/polling.py

@@ -94,6 +94,19 @@ def _ensure_non_awaitable(
         )
 
 
+def _invoke_non_retryable_callback(
+    callback: Callable[..., T], *args: object, callback_name: str, operation_name: str
+) -> T:
+    try:
+        return callback(*args)
+    except _NonRetryablePollingError:
+        raise
+    except Exception as exc:
+        raise _NonRetryablePollingError(
+            f"{callback_name} failed for {operation_name}: {exc}"
+        ) from exc
+
+
 def _validate_retry_config(
     *,
     max_attempts: int,
@@ -390,19 +403,34 @@ def collect_paginated_results(
                 callback_name="get_next_page",
                 operation_name=operation_name,
             )
-            callback_result = on_page_success(page_response)
+            callback_result = _invoke_non_retryable_callback(
+                on_page_success,
+                page_response,
+                callback_name="on_page_success",
+                operation_name=operation_name,
+            )
             _ensure_non_awaitable(
                 callback_result,
                 callback_name="on_page_success",
                 operation_name=operation_name,
             )
-            current_page_batch = get_current_page_batch(page_response)
+            current_page_batch = _invoke_non_retryable_callback(
+                get_current_page_batch,
+                page_response,
+                callback_name="get_current_page_batch",
+                operation_name=operation_name,
+            )
             _ensure_non_awaitable(
                 current_page_batch,
                 callback_name="get_current_page_batch",
                 operation_name=operation_name,
             )
-            total_page_batches = get_total_page_batches(page_response)
+            total_page_batches = _invoke_non_retryable_callback(
+                get_total_page_batches,
+                page_response,
+                callback_name="get_total_page_batches",
+                operation_name=operation_name,
+            )
             _ensure_non_awaitable(
                 total_page_batches,
                 callback_name="get_total_page_batches",
@@ -480,19 +508,34 @@ async def collect_paginated_results_async(
                 operation_name=operation_name,
             )
             page_response = await page_awaitable
-            callback_result = on_page_success(page_response)
+            callback_result = _invoke_non_retryable_callback(
+                on_page_success,
+                page_response,
+                callback_name="on_page_success",
+                operation_name=operation_name,
+            )
             _ensure_non_awaitable(
                 callback_result,
                 callback_name="on_page_success",
                 operation_name=operation_name,
             )
-            current_page_batch = get_current_page_batch(page_response)
+            current_page_batch = _invoke_non_retryable_callback(
+                get_current_page_batch,
+                page_response,
+                callback_name="get_current_page_batch",
+                operation_name=operation_name,
+            )
             _ensure_non_awaitable(
                 current_page_batch,
                 callback_name="get_current_page_batch",
                 operation_name=operation_name,
             )
-            total_page_batches = get_total_page_batches(page_response)
+            total_page_batches = _invoke_non_retryable_callback(
+                get_total_page_batches,
+                page_response,
+                callback_name="get_total_page_batches",
+                operation_name=operation_name,
+            )
             _ensure_non_awaitable(
                 total_page_batches,
                 callback_name="get_total_page_batches",

tests/test_polling.py

@@ -390,6 +390,50 @@ def on_page_success(response: dict) -> object:
     assert callback_attempts["count"] == 1
 
 
+def test_collect_paginated_results_fails_fast_when_on_page_success_raises():
+    page_attempts = {"count": 0}
+
+    def get_next_page(page: int) -> dict:
+        page_attempts["count"] += 1
+        return {"current": 1, "total": 1, "items": []}
+
+    with pytest.raises(HyperbrowserError, match="on_page_success failed"):
+        collect_paginated_results(
+            operation_name="sync paginated callback exception",
+            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: (_ for _ in ()).throw(ValueError("boom")),
+            max_wait_seconds=1.0,
+            max_attempts=5,
+            retry_delay_seconds=0.0001,
+        )
+
+    assert page_attempts["count"] == 1
+
+
+def test_collect_paginated_results_fails_fast_when_page_batch_callback_raises():
+    page_attempts = {"count": 0}
+
+    def get_next_page(page: int) -> dict:
+        page_attempts["count"] += 1
+        return {"current": 1, "total": 1, "items": []}
+
+    with pytest.raises(HyperbrowserError, match="get_current_page_batch failed"):
+        collect_paginated_results(
+            operation_name="sync paginated page-batch callback exception",
+            get_next_page=get_next_page,
+            get_current_page_batch=lambda response: response["missing"],  # type: ignore[index]
+            get_total_page_batches=lambda response: response["total"],
+            on_page_success=lambda response: None,
+            max_wait_seconds=1.0,
+            max_attempts=5,
+            retry_delay_seconds=0.0001,
+        )
+
+    assert page_attempts["count"] == 1
+
+
 def test_collect_paginated_results_rejects_awaitable_current_page_callback_result():
     with pytest.raises(
         HyperbrowserError,
@@ -520,6 +564,58 @@ def on_page_success(response: dict) -> object:
     asyncio.run(run())
 
 
+def test_collect_paginated_results_async_fails_fast_when_on_page_success_raises():
+    async def run() -> None:
+        page_attempts = {"count": 0}
+
+        async def get_next_page(page: int) -> dict:
+            page_attempts["count"] += 1
+            return {"current": 1, "total": 1, "items": []}
+
+        with pytest.raises(HyperbrowserError, match="on_page_success failed"):
+            await collect_paginated_results_async(
+                operation_name="async paginated callback exception",
+                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: (_ for _ in ()).throw(
+                    ValueError("boom")
+                ),
+                max_wait_seconds=1.0,
+                max_attempts=5,
+                retry_delay_seconds=0.0001,
+            )
+
+        assert page_attempts["count"] == 1
+
+    asyncio.run(run())
+
+
+def test_collect_paginated_results_async_fails_fast_when_page_batch_callback_raises():
+    async def run() -> None:
+        page_attempts = {"count": 0}
+
+        async def get_next_page(page: int) -> dict:
+            page_attempts["count"] += 1
+            return {"current": 1, "total": 1, "items": []}
+
+        with pytest.raises(HyperbrowserError, match="get_total_page_batches failed"):
+            await collect_paginated_results_async(
+                operation_name="async paginated page-batch callback exception",
+                get_next_page=get_next_page,
+                get_current_page_batch=lambda response: response["current"],
+                get_total_page_batches=lambda response: response["missing"],  # type: ignore[index]
+                on_page_success=lambda response: None,
+                max_wait_seconds=1.0,
+                max_attempts=5,
+                retry_delay_seconds=0.0001,
+            )
+
+        assert page_attempts["count"] == 1
+
+    asyncio.run(run())
+
+
 def test_collect_paginated_results_async_rejects_awaitable_current_page_callback_result():
     async def run() -> None:
         with pytest.raises(