Fix/timeout and reconnect

.env.example

@@ -1,2 +1,8 @@
 # Copy to .env — auto-loaded. Only needed for remote browsers.
 BROWSER_USE_API_KEY=bu_your_key_here
+
+# Optional. Per-call read budget (seconds) for every CDP command the helpers
+# send to the daemon. Default 30 (screenshots 60). Raise it for pages with very
+# slow navigations or long awaited promises; lower it to fail fast. Individual
+# calls can also override via cdp(..., _timeout=N) / js(..., timeout=N).
+# BH_CMD_TIMEOUT=30

interaction-skills/connection.md

@@ -6,6 +6,24 @@ When Chrome opens fresh, the only CDP `type: "page"` targets are `chrome://inspe
 
 The daemon's `attach_first_page()` handles this by creating an `about:blank` tab when no real pages exist. If you still end up on an invisible tab, use `switch_tab()` which calls `Target.activateTarget` to bring the tab to front.
 
+## Auto-reconnect
+
+If the CDP websocket drops mid-session (Chrome closed/restarted, a remote
+endpoint hiccup), the daemon rebuilds the connection and retries the failing
+call once — for a local Chrome it re-resolves the live websocket via
+`get_ws_url()`, so it also recovers from a Chrome restart on the same port. The
+call the caller made succeeds transparently; a follow-up call re-attaches any
+tab-specific session. Only if the rebuild itself fails (e.g. a remote
+`BU_CDP_WS` whose endpoint is gone) do you get an error back. No manual
+`--reload` needed for a plain drop.
+
+## Command timeout
+
+Every helper→daemon CDP call has a read budget (default 30s, screenshots 60s),
+tunable via `BH_CMD_TIMEOUT` or per-call `cdp(..., _timeout=N)` /
+`js(..., timeout=N)`. This is separate from the short socket-connect budget, so
+a slow-but-valid navigation or awaited promise no longer fails at connect time.
+
 ## Startup sequence
 
 1. Check if a daemon is already running with `daemon_alive()`

pyproject.toml

@@ -17,6 +17,14 @@ dependencies = [
 [project.scripts]
 browser-harness = "browser_harness.run:main"
 
+# Dev tooling lives in the default `dev` group so `uv run pytest` resolves it
+# from the persistent project env. Running `uv run --with pytest pytest` instead
+# builds a fresh ephemeral overlay venv on every cold cache (~2.8s vs ~0.5s).
+[dependency-groups]
+dev = [
+    "pytest>=8",
+]
+
 [tool.setuptools]
 package-dir = {"" = "src"}
 

src/browser_harness/daemon.py

@@ -179,6 +179,29 @@ def is_real_page(t):
     return t["type"] == "page" and not t.get("url", "").startswith(INTERNAL)
 
 
+def _conn_dead(e):
+    """True when an exception from send_raw means the CDP websocket is gone.
+
+    cdp_use surfaces a drop three ways: ConnectionError('WebSocket connection
+    closed') / ('Client is stopping') when a pending future is failed, a bare
+    RuntimeError('Client is not started') when self.ws was already None, and a
+    websockets ConnectionClosed raised straight out of ws.send(). We match by
+    type where we can and fall back to message text so a websockets version
+    bump can't silently drop the ConnectionClosed case."""
+    if isinstance(e, ConnectionError):
+        return True
+    name = type(e).__name__
+    if "ConnectionClosed" in name:
+        return True
+    msg = str(e)
+    return (
+        "WebSocket connection closed" in msg
+        or "Client is not started" in msg
+        or "Client is stopping" in msg
+        or "no close frame received" in msg
+    )
+
+
 class Daemon:
     def __init__(self):
         self.cdp = None
@@ -187,6 +210,9 @@ def __init__(self):
         self.events = deque(maxlen=BUF)
         self.dialog = None
         self.stop = None  # asyncio.Event, set inside start()
+        # Serializes CDP-client rebuilds so concurrent handlers that all see the
+        # same dropped websocket don't each spin up a fresh connection.
+        self._reconnect_lock = asyncio.Lock()
 
     async def attach_first_page(self):
         """Attach to a real page (or any page). Sets self.session. Returns attached target or None."""
@@ -229,8 +255,27 @@ async def enable_one(d):
                 log(f"enable {d} on {session_id}: {e}")
         await asyncio.gather(*(enable_one(d) for d in ("Page", "DOM", "Runtime", "Network")))
 
-    async def start(self):
-        self.stop = asyncio.Event()
+    def _install_tap(self):
+        """Wrap the current CDP client's event dispatch so the daemon mirrors
+        events into its buffer, tracks dialogs, and re-marks tab titles. Must be
+        re-run against every fresh CDP client (each has its own event registry),
+        so it lives here rather than inline in start()."""
+        orig = self.cdp._event_registry.handle_event
+        mark_js = "if(!document.title.startsWith('\U0001F434'))document.title='\U0001F434 '+document.title"
+        async def tap(method, params, session_id=None):
+            self.events.append({"method": method, "params": params, "session_id": session_id})
+            if method == "Page.javascriptDialogOpening":
+                self.dialog = params
+            elif method == "Page.javascriptDialogClosed":
+                self.dialog = None
+            elif method in ("Page.loadEventFired", "Page.domContentEventFired"):
+                asyncio.create_task(_silent(asyncio.wait_for(self.cdp.send_raw("Runtime.evaluate", {"expression": mark_js}, session_id=self.session), timeout=2)))
+            return await orig(method, params, session_id)
+        self.cdp._event_registry.handle_event = tap
+
+    async def _connect_cdp(self):
+        """Build the CDP client, connect the websocket, attach a page, and install
+        the event tap. Shared by initial start() and reconnect()."""
         url = get_ws_url()
         log(f"connecting to {url}")
         self.cdp = CDPClient(url)
@@ -245,18 +290,27 @@ async def start(self):
                 )
             raise RuntimeError(f"CDP WS handshake failed: {e} -- click Allow in Chrome if prompted, then retry")
         await self.attach_first_page()
-        orig = self.cdp._event_registry.handle_event
-        mark_js = "if(!document.title.startsWith('\U0001F434'))document.title='\U0001F434 '+document.title"
-        async def tap(method, params, session_id=None):
-            self.events.append({"method": method, "params": params, "session_id": session_id})
-            if method == "Page.javascriptDialogOpening":
-                self.dialog = params
-            elif method == "Page.javascriptDialogClosed":
-                self.dialog = None
-            elif method in ("Page.loadEventFired", "Page.domContentEventFired"):
-                asyncio.create_task(_silent(asyncio.wait_for(self.cdp.send_raw("Runtime.evaluate", {"expression": mark_js}, session_id=self.session), timeout=2)))
-            return await orig(method, params, session_id)
-        self.cdp._event_registry.handle_event = tap
+        self._install_tap()
+
+    async def _reconnect(self, failed_cdp):
+        """Rebuild the CDP connection after the websocket dropped.
+
+        `failed_cdp` is the client the caller saw fail. Under the lock we bail if
+        someone already swapped in a live client, so only the first failing
+        handler rebuilds and the rest reuse the result. For a local Chrome,
+        get_ws_url() re-resolves the live websocket, so this also recovers from a
+        Chrome restart. For a remote BU_CDP_WS the URL is fixed; if that endpoint
+        is gone the rebuild raises and the original error is returned."""
+        async with self._reconnect_lock:
+            if self.cdp is not failed_cdp:
+                return
+            log("CDP websocket dropped — reconnecting")
+            await _silent(failed_cdp.stop())
+            await self._connect_cdp()
+
+    async def start(self):
+        self.stop = asyncio.Event()
+        await self._connect_cdp()
 
     async def handle(self, req):
         # Token guard for Windows TCP loopback: any local process can otherwise
@@ -352,7 +406,27 @@ async def disable_old():
             if "Session with given id not found" in msg and sid == self.session and sid:
                 log(f"stale session {sid}, re-attaching")
                 if await self.attach_first_page():
-                    return {"result": await self.cdp.send_raw(method, params, session_id=self.session)}
+                    try:
+                        return {"result": await self.cdp.send_raw(method, params, session_id=self.session)}
+                    except Exception as e2:
+                        e, msg = e2, str(e2)
+            if _conn_dead(e):
+                # Websocket dropped (Chrome closed/restarted, remote hiccup).
+                # Rebuild the connection and retry once. After a rebuild every
+                # prior session id is void, so a non-Target call runs on the
+                # freshly attached default session rather than the stale id the
+                # caller computed client-side.
+                failed = self.cdp
+                try:
+                    await self._reconnect(failed)
+                except Exception as re:
+                    log(f"reconnect failed: {re}")
+                    return {"error": f"{msg} -- CDP reconnect failed: {re}"}
+                retry_sid = None if method.startswith("Target.") else self.session
+                try:
+                    return {"result": await self.cdp.send_raw(method, params, session_id=retry_sid)}
+                except Exception as e3:
+                    return {"error": str(e3)}
             return {"error": msg}
 
 

src/browser_harness/helpers.py

@@ -38,20 +38,49 @@ def _load_env_file(p):
 SOCK = ipc.sock_addr(NAME)
 INTERNAL = ("chrome://", "chrome-untrusted://", "devtools://", "chrome-extension://", "about:")
 
+# Establishing the AF_UNIX/loopback socket to the (local) daemon should be
+# instant, so the connect budget stays short. The command budget is the time we
+# allow the daemon to round-trip a CDP call to Chrome and answer — screenshots,
+# heavy DOM reads, and awaited promises routinely need more than a couple of
+# seconds. These were previously fused into a single 5s value (connect() sets
+# both the connect AND read timeout), which hard-capped every operation at 5s
+# and turned any slow-but-valid call into a failure.
+_CONNECT_TIMEOUT = 5.0
+_DEFAULT_CMD_TIMEOUT = 30.0
+_SCREENSHOT_CMD_TIMEOUT = 60.0
+
+
+def _cmd_timeout(default=_DEFAULT_CMD_TIMEOUT):
+    """Per-call read budget in seconds. BH_CMD_TIMEOUT overrides the default."""
+    v = os.environ.get("BH_CMD_TIMEOUT")
+    if v:
+        try:
+            return float(v)
+        except ValueError:
+            pass
+    return default
 
-def _send(req):
-    c, token = ipc.connect(NAME, timeout=5.0)
+
+def _send(req, timeout=None):
+    # Connect on the short budget, then raise the socket's read timeout to the
+    # command budget before waiting on the daemon's reply. Without the second
+    # settimeout, the reply had to arrive within _CONNECT_TIMEOUT.
+    c, token = ipc.connect(NAME, timeout=_CONNECT_TIMEOUT)
     try:
+        c.settimeout(_cmd_timeout() if timeout is None else timeout)
         r = ipc.request(c, token, req)
     finally:
         c.close()
     if "error" in r: raise RuntimeError(r["error"])
     return r
 
 
-def cdp(method, session_id=None, **params):
-    """Raw CDP. cdp('Page.navigate', url='...'), cdp('DOM.getDocument', depth=-1)."""
-    return _send({"method": method, "params": params, "session_id": session_id}).get("result", {})
+def cdp(method, session_id=None, _timeout=None, **params):
+    """Raw CDP. cdp('Page.navigate', url='...'), cdp('DOM.getDocument', depth=-1).
+
+    _timeout overrides the per-call read budget in seconds (env BH_CMD_TIMEOUT
+    sets the default). Underscored so it can't collide with a CDP param name."""
+    return _send({"method": method, "params": params, "session_id": session_id}, timeout=_timeout).get("result", {})
 
 
 def drain_events():  return _send({"meta": "drain_events"})["events"]
@@ -109,19 +138,51 @@ def _runtime_value(response, expression):
     return None
 
 
-def _runtime_evaluate(expression, session_id=None, await_promise=False):
+def _runtime_evaluate(expression, session_id=None, await_promise=False, timeout=None):
     try:
-        r = cdp("Runtime.evaluate", session_id=session_id, expression=expression, returnByValue=True, awaitPromise=await_promise)
+        r = cdp("Runtime.evaluate", session_id=session_id, expression=expression, returnByValue=True, awaitPromise=await_promise, _timeout=timeout)
     except TimeoutError as e:
         raise RuntimeError(f"Runtime.evaluate timed out; expression: {_js_snippet(expression)}") from e
     return _runtime_value(r, expression)
 
 
+_CONTROL_KEYWORDS = frozenset(("if", "for", "while", "switch", "catch", "with"))
+
+
+def _word_before(expression, i):
+    """The identifier immediately preceding index i, skipping whitespace."""
+    j = i - 1
+    while j >= 0 and expression[j].isspace():
+        j -= 1
+    end = j + 1
+    while j >= 0 and (expression[j] == "_" or expression[j] == "$" or expression[j].isalnum()):
+        j -= 1
+    return expression[j + 1:end]
+
+
+def _opens_function_body(sig, sig_ctl):
+    """Whether the `{` about to be consumed opens a function/arrow body rather than a block."""
+    if len(sig) >= 2 and sig[-1] == ">" and sig[-2] == "=":
+        return True                      # `=> {`
+    if sig and sig[-1] == ")":
+        return not sig_ctl[-1]           # `) {` is a function body unless it closed `if(...)` etc
+    return False                         # object literal, bare block, `else {`, `try {`, `do {`
+
+
 def _has_return_statement(expression):
+    """True only when `expression` has a `return` at the snippet's own top level.
+
+    A `return` inside a nested function or arrow body belongs to that callback, so
+    wrapping the snippet in an IIFE would discard the value the caller wanted.
+    """
     i = 0
     n = len(expression)
     state = "code"
     quote = ""
+    sig = []        # significant (non-whitespace) code characters, in order
+    sig_ctl = []    # parallel to sig: True when a ')' closed a control-statement paren
+    paren_ctl = []  # stack: was each open '(' introduced by a control keyword
+    brace_fn = []   # stack: does each open '{' delimit a function body
     while i < n:
         ch = expression[i]
         nxt = expression[i + 1] if i + 1 < n else ""
@@ -136,7 +197,20 @@ def _has_return_statement(expression):
                 before = expression[i - 1] if i > 0 else ""
                 after = expression[i + 6] if i + 6 < n else ""
                 if not (before == "_" or before.isalnum()) and not (after == "_" or after.isalnum()):
-                    return True
+                    if not any(brace_fn):
+                        return True
+                    i += 6; continue
+            if ch == "(":
+                paren_ctl.append(_word_before(expression, i) in _CONTROL_KEYWORDS)
+            elif ch == ")":
+                ctl = paren_ctl.pop() if paren_ctl else False
+                sig.append(ch); sig_ctl.append(ctl); i += 1; continue
+            elif ch == "{":
+                brace_fn.append(_opens_function_body(sig, sig_ctl))
+            elif ch == "}":
+                if brace_fn: brace_fn.pop()
+            if not ch.isspace():
+                sig.append(ch); sig_ctl.append(False)
             i += 1; continue
         if state == "line_comment":
             if ch == "\n":
@@ -270,7 +344,10 @@ 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)
+    # Screenshots serialize a base64 PNG back over the socket — on a large or
+    # retina page that can take longer than the default command budget, so give
+    # it a wider one (still overridable via BH_CMD_TIMEOUT).
+    r = cdp("Page.captureScreenshot", format="png", captureBeyondViewport=full, _timeout=_cmd_timeout(_SCREENSHOT_CMD_TIMEOUT))
     open(path, "wb").write(base64.b64decode(r["data"]))
     if max_dim:
         from PIL import Image
@@ -432,16 +509,19 @@ def wait_for_network_idle(timeout=10.0, idle_ms=500):
         time.sleep(0.1)
     return False
 
-def js(expression, target_id=None):
+def js(expression, target_id=None, timeout=None):
     """Run JS in the attached tab (default) or inside an iframe target (via iframe_target()).
 
     Expressions with top-level `return` are automatically wrapped in an IIFE, so both
     `document.title` and `const x = 1; return x` are valid inputs.
+
+    Pass timeout (seconds) to widen the read budget for a deliberately slow awaited
+    promise; env BH_CMD_TIMEOUT sets the default for all calls.
     """
     sid = cdp("Target.attachToTarget", targetId=target_id, flatten=True)["sessionId"] if target_id else None
-    if _has_return_statement(expression) and not expression.strip().startswith("("):
+    if _has_return_statement(expression):
         expression = f"(function(){{{expression}}})()"
-    return _runtime_evaluate(expression, session_id=sid, await_promise=True)
+    return _runtime_evaluate(expression, session_id=sid, await_promise=True, timeout=timeout)
 
 
 _KC = {"Enter": 13, "Tab": 9, "Escape": 27, "Backspace": 8, " ": 32, "ArrowLeft": 37, "ArrowUp": 38, "ArrowRight": 39, "ArrowDown": 40}

tests/unit/test_helpers.py

@@ -350,3 +350,42 @@ def fake_send(req):
         "session filter, the background rWS/lF pair would have updated "
         "last_activity and prevented the idle window from elapsing."
     )
+
+
+# --- js() IIFE wrapping: only a snippet's OWN top-level return should trigger it ---
+
+@pytest.mark.parametrize("snippet", [
+    "[1,2,3].map(function(x){return x*2})",
+    "[1,2,3].map(x=>{return x*2})",
+    "Array.from(document.querySelectorAll('a')).map(a=>{const h=a.href; return h})",
+    "(()=>{const x=1; return x})()",
+    "(function(){return 1})()",
+    "function f(){ return 1 }; f()",
+    "class C { m(){ return 1 } }",
+    "const o = { run(){ return 1 } }",
+    "myreturn + returnValue",
+    "'return 1'",
+    "// return 1",
+    "/* return 1 */ 42",
+    "`return ${x}`",
+])
+def test_nested_return_is_not_a_top_level_return(snippet):
+    assert helpers._has_return_statement(snippet) is False, (
+        "A `return` inside a nested function/arrow body belongs to that callback. "
+        "Wrapping the snippet in an IIFE would make it evaluate to undefined, which "
+        "surfaces to the caller as a silent None."
+    )
+
+
+@pytest.mark.parametrize("snippet", [
+    "const x = 1; return x",
+    "if (a) { return 1 } return 2",
+    "for (const a of b) { return a }",
+    "while (x) { return 1 }",
+    "switch(x){case 1: return 2}",
+    "try { foo() } catch (e) { return e }",
+    "const o = {a:1}; return o",
+    "els.forEach(e=>{ if(e){ return } }); return 5",
+])
+def test_top_level_return_still_wraps(snippet):
+    assert helpers._has_return_statement(snippet) is True