From 742d3edc2027db1eceb854088e9cbd7a60601fea Mon Sep 17 00:00:00 2001
From: MagMueller <mamagnus00@gmail.com>
Date: Wed, 19 Aug 2026 07:27:36 -0700
Subject: [PATCH] fix: allow slow cloud screenshot responses

---
 src/browser_harness/helpers.py | 38 ++++++++++++++++++++++++-----
 tests/unit/test_helpers.py     | 44 ++++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 6 deletions(-)

diff --git a/src/browser_harness/helpers.py b/src/browser_harness/helpers.py
index 95a6dc31..e912b0ae 100644
--- a/src/browser_harness/helpers.py
+++ b/src/browser_harness/helpers.py
@@ -38,21 +38,37 @@ def _load_env_file(p):
 NAME = os.environ.get("BU_NAME", "default")
 SOCK = ipc.sock_addr(NAME)
 INTERNAL = ("chrome://", "chrome-untrusted://", "devtools://", "chrome-extension://", "about:")
+IPC_CONNECT_TIMEOUT_SECONDS = 5.0
+DEFAULT_IPC_RESPONSE_TIMEOUT_SECONDS = 5.0
+# Cloud screenshots routinely take longer than ordinary CDP round trips. Keep
+# their IPC socket alive within the caller's existing 90-second process budget.
+SCREENSHOT_IPC_RESPONSE_TIMEOUT_SECONDS = 60.0
 
 
-def _send(req):
-    c, token = ipc.connect(NAME, timeout=5.0)
+class _IPCResponseTimeout(TimeoutError):
+    pass
+
+
+def _send(req, response_timeout=DEFAULT_IPC_RESPONSE_TIMEOUT_SECONDS):
+    c, token = ipc.connect(NAME, timeout=IPC_CONNECT_TIMEOUT_SECONDS)
     try:
-        r = ipc.request(c, token, req)
+        c.settimeout(response_timeout)
+        try:
+            r = ipc.request(c, token, req)
+        except TimeoutError as e:
+            raise _IPCResponseTimeout from e
     finally:
         c.close()
     if "error" in r: raise RuntimeError(r["error"])
     return r
 
 
-def cdp(method, session_id=None, **params):
+def cdp(method, session_id=None, _response_timeout=DEFAULT_IPC_RESPONSE_TIMEOUT_SECONDS, **params):
     """Raw CDP. cdp('Page.navigate', url='...'), cdp('DOM.getDocument', depth=-1)."""
-    return _send({"method": method, "params": params, "session_id": session_id}).get("result", {})
+    return _send(
+        {"method": method, "params": params, "session_id": session_id},
+        response_timeout=_response_timeout,
+    ).get("result", {})
 
 
 def drain_events():  return _send({"meta": "drain_events"})["events"]
@@ -243,7 +259,17 @@ def capture_screenshot(path=None, full=False, max_dim=None):
     """Save a PNG of the current viewport. Set max_dim=1800 on a 2× display to
     keep the file under the 2000px-per-side limit some image-aware LLMs enforce."""
     path = path or str(ipc._TMP / "shot.png")
-    r = cdp("Page.captureScreenshot", format="png", captureBeyondViewport=full)
+    try:
+        r = cdp(
+            "Page.captureScreenshot",
+            _response_timeout=SCREENSHOT_IPC_RESPONSE_TIMEOUT_SECONDS,
+            format="png",
+            captureBeyondViewport=full,
+        )
+    except _IPCResponseTimeout as e:
+        raise RuntimeError(
+            f"Page.captureScreenshot timed out after {SCREENSHOT_IPC_RESPONSE_TIMEOUT_SECONDS:g}s"
+        ) from e
     open(path, "wb").write(base64.b64decode(r["data"]))
     if max_dim:
         from PIL import Image
diff --git a/tests/unit/test_helpers.py b/tests/unit/test_helpers.py
index 9a099030..26318f99 100644
--- a/tests/unit/test_helpers.py
+++ b/tests/unit/test_helpers.py
@@ -29,6 +29,50 @@ def test_max_dim_default_is_no_resize(fake_png):
     assert _run(fake_png, 4592, 2286) == (4592, 2286)
 
 
+def test_send_keeps_connect_timeout_short_and_sets_response_budget():
+    class FakeSocket:
+        def __init__(self):
+            self.timeouts = []
+
+        def settimeout(self, value):
+            self.timeouts.append(value)
+
+        def close(self):
+            pass
+
+    socket = FakeSocket()
+    with patch("browser_harness.helpers.ipc.connect", return_value=(socket, None)) as connect, \
+         patch("browser_harness.helpers.ipc.request", return_value={}):
+        helpers._send({"meta": "ping"}, response_timeout=60.0)
+
+    connect.assert_called_once_with(helpers.NAME, timeout=helpers.IPC_CONNECT_TIMEOUT_SECONDS)
+    assert socket.timeouts == [60.0]
+
+
+def test_screenshot_uses_long_response_timeout_without_forwarding_it_to_cdp(fake_png, tmp_path):
+    with patch(
+        "browser_harness.helpers._send",
+        return_value={"result": {"data": fake_png(800, 400)}},
+    ) as send:
+        helpers.capture_screenshot(str(tmp_path / "shot.png"))
+
+    request = send.call_args.args[0]
+    assert request == {
+        "method": "Page.captureScreenshot",
+        "params": {"format": "png", "captureBeyondViewport": False},
+        "session_id": None,
+    }
+    assert send.call_args.kwargs == {
+        "response_timeout": helpers.SCREENSHOT_IPC_RESPONSE_TIMEOUT_SECONDS
+    }
+
+
+def test_screenshot_timeout_has_context(tmp_path):
+    with patch("browser_harness.helpers._send", side_effect=helpers._IPCResponseTimeout):
+        with pytest.raises(RuntimeError, match="Page.captureScreenshot timed out after 60s"):
+            helpers.capture_screenshot(str(tmp_path / "shot.png"))
+
+
 def _seed_skill(tmp_path):
     site = tmp_path / "domain-skills" / "example"
     site.mkdir(parents=True)