Harden Kernel replay finalization

Author: rgarcia

src/clawbench/runner/batch.py

@@ -49,6 +49,7 @@ def detect_engine() -> str:
     "claw-eval": "test-cases/claw-eval",
 }
 DEFAULT_CASES_SUITE = "v2"
+MANAGED_BROWSER_RUNTIMES = frozenset({"browserbase", "kernel"})
 
 
 def load_models_yaml() -> dict:
@@ -289,7 +290,7 @@ async def run_job(
                     cmd_parts += ["--harness", harness]
                 if browser_runtime:
                     cmd_parts += ["--browser-runtime", browser_runtime]
-                    if browser_runtime == "kernel":
+                    if browser_runtime in MANAGED_BROWSER_RUNTIMES:
                         cmd_parts.append("--hide-browser-viewer")
                 if browser_cdp_url:
                     cmd_parts += ["--browser-cdp-url", browser_cdp_url]
@@ -591,11 +592,10 @@ async def async_main(args: argparse.Namespace) -> int:
     shutdown_event = asyncio.Event()
     running_procs.clear()
     browser_runtime = getattr(args, "browser_runtime", None) or "local"
-    managed_browser_runtimes = {"browserbase", "kernel"}
     if getattr(args, "max_concurrent", None) is None:
-        args.max_concurrent = 1 if browser_runtime in managed_browser_runtimes else 2
+        args.max_concurrent = 1 if browser_runtime in MANAGED_BROWSER_RUNTIMES else 2
     if (
-        browser_runtime in managed_browser_runtimes
+        browser_runtime in MANAGED_BROWSER_RUNTIMES
         and args.harness == "claude-code-chrome-extension"
     ):
         print(

src/clawbench/runner/run_support/browser_runtime/providers.py

@@ -442,7 +442,7 @@ def __init__(
         options: dict[str, Any],
         api_url: str = _KERNEL_API_URL,
         replay_poll_interval_s: float = 1,
-        replay_poll_timeout_s: float = 30,
+        replay_poll_timeout_s: float = 300,
     ) -> None:
         unknown = sorted(set(options) - _KERNEL_ALLOWED_OPTIONS)
         if unknown:
@@ -471,14 +471,14 @@ def __init__(
         self.replay_poll_interval_s = replay_poll_interval_s
         self.replay_poll_timeout_s = replay_poll_timeout_s
 
-    def _request(
+    def _request_response(
         self,
         method: str,
         path: str,
         payload: dict[str, Any] | None = None,
         *,
         accept: str = "application/json",
-    ) -> bytes:
+    ) -> tuple[int, Any, bytes]:
         assert self.api_key is not None
         data = (
             json.dumps(payload, separators=(",", ":")).encode()
@@ -497,7 +497,11 @@ def _request(
         )
         try:
             with urllib.request.urlopen(request, timeout=30) as response:
-                return response.read()
+                return (
+                    int(getattr(response, "status", 200)),
+                    getattr(response, "headers", {}),
+                    response.read(),
+                )
         except urllib.error.HTTPError as e:
             if e.code in {401, 403}:
                 message = "Kernel authentication failed"
@@ -517,6 +521,22 @@ def _request(
                     reason = reason.replace(secret, "[REDACTED]")
             raise _KernelApiError(f"Kernel API request failed: {reason}") from None
 
+    def _request(
+        self,
+        method: str,
+        path: str,
+        payload: dict[str, Any] | None = None,
+        *,
+        accept: str = "application/json",
+    ) -> bytes:
+        _status, _headers, raw = self._request_response(
+            method,
+            path,
+            payload,
+            accept=accept,
+        )
+        return raw
+
     def _request_json(
         self,
         method: str,
@@ -669,11 +689,31 @@ def finalize(self, session: BrowserSession, output_dir: Path) -> None:
                     )
                 time.sleep(self.replay_poll_interval_s)
 
-            recording = self._request(
-                "GET",
-                f"/browsers/{session.session_id}/replays/{replay_id}",
-                accept="video/mp4",
-            )
+            while True:
+                status, headers, recording = self._request_response(
+                    "GET",
+                    f"/browsers/{session.session_id}/replays/{replay_id}",
+                    accept="video/mp4",
+                )
+                if status != 202:
+                    break
+                if time.monotonic() >= deadline:
+                    raise _KernelApiError(
+                        "Kernel replay did not finish downloading before timeout"
+                    )
+                retry_after: str | None = (
+                    headers.get("Retry-After") if headers else None
+                )
+                if retry_after is None:
+                    delay = self.replay_poll_interval_s
+                else:
+                    try:
+                        delay = float(retry_after)
+                    except ValueError:
+                        delay = self.replay_poll_interval_s
+                time.sleep(min(max(0, delay), max(0, deadline - time.monotonic())))
+            if status != 200:
+                raise _KernelApiError(f"Kernel replay download returned HTTP {status}")
             if not recording:
                 raise _KernelApiError("Kernel replay download was empty")
             recording_path = output_dir / "data" / "recording.mp4"

tests/test_browser_runtime.py

@@ -26,8 +26,16 @@
 
 
 class _FakeResponse:
-    def __init__(self, payload: object) -> None:
+    def __init__(
+        self,
+        payload: object,
+        *,
+        status: int = 200,
+        headers: dict[str, str] | None = None,
+    ) -> None:
         self.payload = payload
+        self.status = status
+        self.headers = headers or {}
 
     def __enter__(self) -> _FakeResponse:
         return self
@@ -322,11 +330,13 @@ def test_kernel_session_replay_and_cleanup(
     tmp_path: Path,
 ) -> None:
     calls: list[tuple[str, str, object, str | None]] = []
+    replay_downloads = 0
 
     def fake_urlopen(
         request: urllib.request.Request,
         timeout: int,
     ) -> _FakeResponse:
+        nonlocal replay_downloads
         assert timeout == 30
         if request.data:
             assert isinstance(request.data, bytes)
@@ -373,6 +383,13 @@ def fake_urlopen(
         if request.full_url.endswith("/replays/replay_123/stop"):
             return _FakeResponse(b"")
         if request.full_url.endswith("/replays/replay_123"):
+            replay_downloads += 1
+            if replay_downloads == 1:
+                return _FakeResponse(
+                    b"not-ready",
+                    status=202,
+                    headers={"Retry-After": "0"},
+                )
             return _FakeResponse(b"mp4-data")
         if request.full_url.endswith("/browsers/browser_123"):
             return _FakeResponse(b"")
@@ -383,7 +400,7 @@ def fake_urlopen(
         api_key="kernel-secret",
         options={"stealth": True, "region": "us-east"},
         replay_poll_interval_s=0,
-        replay_poll_timeout_s=0,
+        replay_poll_timeout_s=1,
     )
 
     session = provider.start({}, 1800)
@@ -401,6 +418,7 @@ def fake_urlopen(
     assert session.provider == "kernel"
     assert session.recording_mode == "provider-download"
     assert session.cleanup_status == "deleted"
+    assert replay_downloads == 2
     assert (tmp_path / "data" / "recording.mp4").read_bytes() == b"mp4-data"
     metadata = json.dumps(session.to_metadata())
     assert "cdp-secret" not in metadata