Improve polling resilience on transient status errors

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

Author: cursoragent

hyperbrowser/client/polling.py

@@ -20,16 +20,29 @@ def poll_until_terminal_status(
     is_terminal_status: Callable[[str], bool],
     poll_interval_seconds: float,
     max_wait_seconds: Optional[float],
+    max_status_failures: int = 5,
 ) -> str:
     start_time = time.monotonic()
+    failures = 0
 
     while True:
         if has_exceeded_max_wait(start_time, max_wait_seconds):
             raise HyperbrowserError(
                 f"Timed out waiting for {operation_name} after {max_wait_seconds} seconds"
             )
 
-        status = get_status()
+        try:
+            status = get_status()
+            failures = 0
+        except Exception as exc:
+            failures += 1
+            if failures >= max_status_failures:
+                raise HyperbrowserError(
+                    f"Failed to poll {operation_name} after {max_status_failures} attempts: {exc}"
+                ) from exc
+            time.sleep(poll_interval_seconds)
+            continue
+
         if is_terminal_status(status):
             return status
         time.sleep(poll_interval_seconds)
@@ -62,16 +75,29 @@ async def poll_until_terminal_status_async(
     is_terminal_status: Callable[[str], bool],
     poll_interval_seconds: float,
     max_wait_seconds: Optional[float],
+    max_status_failures: int = 5,
 ) -> str:
     start_time = time.monotonic()
+    failures = 0
 
     while True:
         if has_exceeded_max_wait(start_time, max_wait_seconds):
             raise HyperbrowserError(
                 f"Timed out waiting for {operation_name} after {max_wait_seconds} seconds"
             )
 
-        status = await get_status()
+        try:
+            status = await get_status()
+            failures = 0
+        except Exception as exc:
+            failures += 1
+            if failures >= max_status_failures:
+                raise HyperbrowserError(
+                    f"Failed to poll {operation_name} after {max_status_failures} attempts: {exc}"
+                ) from exc
+            await asyncio.sleep(poll_interval_seconds)
+            continue
+
         if is_terminal_status(status):
             return status
         await asyncio.sleep(poll_interval_seconds)

tests/test_polling.py

@@ -36,6 +36,38 @@ def test_poll_until_terminal_status_times_out():
         )
 
 
+def test_poll_until_terminal_status_retries_transient_status_errors():
+    attempts = {"count": 0}
+
+    def get_status() -> str:
+        attempts["count"] += 1
+        if attempts["count"] < 3:
+            raise ValueError("temporary")
+        return "completed"
+
+    status = poll_until_terminal_status(
+        operation_name="sync poll retries",
+        get_status=get_status,
+        is_terminal_status=lambda value: value in {"completed", "failed"},
+        poll_interval_seconds=0.0001,
+        max_wait_seconds=1.0,
+    )
+
+    assert status == "completed"
+
+
+def test_poll_until_terminal_status_raises_after_status_failures():
+    with pytest.raises(HyperbrowserError, match="Failed to poll sync poll failure"):
+        poll_until_terminal_status(
+            operation_name="sync poll failure",
+            get_status=lambda: (_ for _ in ()).throw(ValueError("always")),
+            is_terminal_status=lambda value: value in {"completed", "failed"},
+            poll_interval_seconds=0.0001,
+            max_wait_seconds=1.0,
+            max_status_failures=2,
+        )
+
+
 def test_retry_operation_retries_and_returns_value():
     attempts = {"count": 0}
 
@@ -95,3 +127,25 @@ async def operation() -> str:
         assert result == "ok"
 
     asyncio.run(run())
+
+
+def test_async_poll_until_terminal_status_retries_transient_status_errors():
+    async def run() -> None:
+        attempts = {"count": 0}
+
+        async def get_status() -> str:
+            attempts["count"] += 1
+            if attempts["count"] < 3:
+                raise ValueError("temporary")
+            return "completed"
+
+        status = await poll_until_terminal_status_async(
+            operation_name="async poll retries",
+            get_status=get_status,
+            is_terminal_status=lambda value: value in {"completed", "failed"},
+            poll_interval_seconds=0.0001,
+            max_wait_seconds=1.0,
+        )
+        assert status == "completed"
+
+    asyncio.run(run())