diff --git a/.env.example b/.env.example index 0e1e7579..1915342a 100644 --- a/.env.example +++ b/.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 diff --git a/interaction-skills/connection.md b/interaction-skills/connection.md index 85e264c2..31445db5 100644 --- a/interaction-skills/connection.md +++ b/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()` diff --git a/pyproject.toml b/pyproject.toml index f812a6ab..384143dc 100644 --- a/pyproject.toml +++ b/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"} diff --git a/src/browser_harness/daemon.py b/src/browser_harness/daemon.py index 0f0f2555..f357d9ef 100644 --- a/src/browser_harness/daemon.py +++ b/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} diff --git a/src/browser_harness/helpers.py b/src/browser_harness/helpers.py index 2014887b..0e52ab44 100644 --- a/src/browser_harness/helpers.py +++ b/src/browser_harness/helpers.py @@ -38,10 +38,36 @@ 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() @@ -49,9 +75,12 @@ def _send(req): 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} diff --git a/tests/unit/test_helpers.py b/tests/unit/test_helpers.py index 4a45ee07..bf9e5312 100644 --- a/tests/unit/test_helpers.py +++ b/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 diff --git a/tests/unit/test_reconnect.py b/tests/unit/test_reconnect.py new file mode 100644 index 00000000..55bc867a --- /dev/null +++ b/tests/unit/test_reconnect.py @@ -0,0 +1,121 @@ +"""Fix #2: daemon rebuilds the CDP websocket and retries once when it drops. + +These exercise Daemon.handle()'s recovery branch and _conn_dead() without a real +browser: the CDP client is a stub, and _connect_cdp is monkeypatched to swap in a +fresh stub the way the real reconnect swaps in a live client. +""" +import asyncio + +from browser_harness import daemon as d + + +class FakeCDP: + def __init__(self, results=None, fail_times=0, exc=None): + self.results = results or {} + self.fail_times = fail_times + self.exc = exc or ConnectionError("WebSocket connection closed") + self.calls = [] + self.stopped = False + + async def send_raw(self, method, params=None, session_id=None): + self.calls.append((method, params, session_id)) + if self.fail_times > 0: + self.fail_times -= 1 + raise self.exc + return self.results.get(method, {"ok": method}) + + async def stop(self): + self.stopped = True + + +def _daemon_with(dead): + dm = d.Daemon() + dm.cdp = dead + dm.session = "sess-old" + dm.target_id = "tgt-old" + return dm + + +def test_conn_dead_matches_all_drop_signatures(): + assert d._conn_dead(ConnectionError("WebSocket connection closed")) + assert d._conn_dead(ConnectionError("Client is stopping")) + assert d._conn_dead(RuntimeError("Client is not started. Call start() first")) + assert d._conn_dead(RuntimeError("no close frame received or sent")) + + class ConnectionClosedError(Exception): # matched by type name, no websockets import + pass + + assert d._conn_dead(ConnectionClosedError("1006")) + # A normal CDP protocol error must NOT look like a dead socket. + assert not d._conn_dead(RuntimeError("Cannot find context with specified id")) + + +def test_handle_reconnects_and_retries_on_dead_ws(): + dead = FakeCDP(fail_times=1) # first call drops, would keep dropping if reused + dm = _daemon_with(dead) + good = FakeCDP(results={"Runtime.evaluate": {"value": 42}}) + + async def fake_connect(): + dm.cdp = good + dm.session = "sess-new" + dm.target_id = "tgt-new" + + dm._connect_cdp = fake_connect + + resp = asyncio.run(dm.handle({"method": "Runtime.evaluate", "params": {"expression": "1"}})) + + assert resp == {"result": {"value": 42}}, resp + assert dead.stopped is True # old client torn down + assert good.calls[-1][2] == "sess-new" # retry ran on the fresh session + + +def test_handle_reports_error_when_rebuild_fails(): + dead = FakeCDP(fail_times=1) + dm = _daemon_with(dead) + + async def fake_connect(): + raise RuntimeError("Chrome gone") + + dm._connect_cdp = fake_connect + + resp = asyncio.run(dm.handle({"method": "Runtime.evaluate", "params": {}})) + + assert "error" in resp + assert "CDP reconnect failed" in resp["error"] + + +def test_handle_passes_through_non_socket_errors_without_reconnect(): + boom = FakeCDP(fail_times=1, exc=RuntimeError("Cannot find context with specified id")) + dm = _daemon_with(boom) + + called = {"connect": False} + + async def fake_connect(): + called["connect"] = True + + dm._connect_cdp = fake_connect + + resp = asyncio.run(dm.handle({"method": "Runtime.evaluate", "params": {}})) + + assert resp == {"error": "Cannot find context with specified id"} + assert called["connect"] is False # a plain CDP error never reconnects + assert boom.stopped is False + + +def test_reconnect_is_noop_when_client_already_swapped(): + """Second handler to see the same drop must reuse the rebuilt client, not + tear down the live one.""" + dead = FakeCDP() + dm = _daemon_with(dead) + live = FakeCDP() + dm.cdp = live # simulate: another coroutine already reconnected + + async def fake_connect(): + raise AssertionError("_connect_cdp must not run when client already swapped") + + dm._connect_cdp = fake_connect + + asyncio.run(dm._reconnect(dead)) # `dead` is the stale one + + assert dm.cdp is live + assert dead.stopped is False diff --git a/tests/unit/test_run.py b/tests/unit/test_run.py index 20783e03..9e7ef0cf 100644 --- a/tests/unit/test_run.py +++ b/tests/unit/test_run.py @@ -7,6 +7,20 @@ from browser_harness import run +@pytest.fixture(autouse=True) +def _isolate_browser_env(monkeypatch): + """run.main()'s cloud-bootstrap guard reads process env directly + (BROWSER_USE_API_KEY, BU_AUTOSPAWN, BU_CDP_URL, BU_CDP_WS). A dev box or CI + runner that exports any of these leaks it into every test — e.g. an ambient + BU_CDP_URL (a running local Chrome debug endpoint) makes + _explicit_cdp_configured() true and silently blocks the bootstrap these + tests assert on, and an ambient BU_AUTOSPAWN + key would fire the real + start_remote_daemon in tests that don't mock it. Start every test from a + clean slate; tests that need a var set it themselves.""" + for var in ("BROWSER_USE_API_KEY", "BU_AUTOSPAWN", "BU_CDP_URL", "BU_CDP_WS"): + monkeypatch.delenv(var, raising=False) + + def test_stdin_executes_code(): stdout = StringIO() fake_stdin = StringIO("print('hello from stdin')")