diff --git a/strix/tools/browser/__init__.py b/strix/tools/browser/__init__.py index d81ed365..763a3a43 100644 --- a/strix/tools/browser/__init__.py +++ b/strix/tools/browser/__init__.py @@ -1,5 +1,4 @@ from .browser_actions import browser_actions -from .browser_manager import cleanup_agent -__all__ = ["browser_actions", "cleanup_agent"] +__all__ = ["browser_actions"] diff --git a/strix/tools/browser/browser_actions.py b/strix/tools/browser/browser_actions.py index 205ff1a6..dba89bad 100644 --- a/strix/tools/browser/browser_actions.py +++ b/strix/tools/browser/browser_actions.py @@ -13,11 +13,9 @@ from strix.tools.registry import register_tool from .browser_manager import ( BrowserSession, _close_session, - _ensure_healthy_session, _get_session, _launch_browser, _launch_local_browser, - llm_supports_vision, ) @@ -50,22 +48,11 @@ BrowserUseLocalAction = Literal[ ] _TASK_TIMEOUT = 300 -_WS_ERRORS = ( - "websocket", - "cdp", - "not initialized", - "not connected", - "connection closed", - "disconnected", -) -def _is_ws_error(exc: BaseException) -> bool: - msg = (type(exc).__name__ + str(exc)).lower() - return any(kw in msg for kw in _WS_ERRORS) +def _build_llm(metadata: dict[str, Any] | None = None) -> tuple[Any, bool]: + import litellm - -def _build_llm(metadata: dict[str, Any] | None = None) -> Any: from strix.config.config import resolve_llm_config from .litellm.chat import ChatLiteLLM @@ -79,7 +66,7 @@ def _build_llm(metadata: dict[str, Any] | None = None) -> Any: api_key=api_key, api_base=api_base, metadata=metadata, - ) + ), litellm.supports_vision(model) def _resolve_cdp_url(agent_state: Any) -> tuple[str, str]: @@ -93,39 +80,28 @@ def _resolve_cdp_url(agent_state: Any) -> tuple[str, str]: async def _execute_task(session: BrowserSession, operation: Any, desc: str) -> dict[str, Any]: session.task_count += 1 - task_num = session.task_count - logger.info("Task #%d: %s", task_num, desc[:200]) + logger.info("Task #%d: %s", session.task_count, desc[:200]) - if err := await _ensure_healthy_session(session, task_num): - session.consecutive_failures += 1 - return {"error": err, "is_running": False} - - for attempt in range(2): - try: - result = await asyncio.wait_for(operation(), timeout=_TASK_TIMEOUT) - session.consecutive_failures = 0 - return result if isinstance(result, dict) else {"result": result, "is_running": False} - except TimeoutError: - session.invalidated = True - session.consecutive_failures += 1 - - return {"error": f"Timeout after {_TASK_TIMEOUT}s", "is_running": False} - - except Exception as exc: # noqa: BLE001 - if _is_ws_error(exc) and attempt == 0: - session.invalidated = True - - if err := await _ensure_healthy_session(session, task_num): - session.consecutive_failures += 1 - return {"error": err, "is_running": False} - - continue - session.consecutive_failures += 1 - if _is_ws_error(exc): - session.invalidated = True - return {"error": f"Failed: {exc}", "is_running": False} - - return {"error": "Failed after 2 attempts", "is_running": False} + try: + result = await asyncio.wait_for(operation(), timeout=_TASK_TIMEOUT) + return ( + result + if isinstance(result, dict) + else { + "result": result, + "is_running": False, + } + ) + except TimeoutError: + return { + "error": f"Timeout after {_TASK_TIMEOUT}s", + "is_running": False, + } + except Exception as exc: # noqa: BLE001 + return { + "error": f"Failed: {exc}", + "is_running": False, + } async def _run_browser_agent( @@ -134,10 +110,9 @@ async def _run_browser_agent( return_fields: list[str] | None, metadata: dict[str, Any] | None = None, ) -> dict[str, Any]: - llm = _build_llm(metadata=metadata) + llm, vision = _build_llm(metadata=metadata) - # [fix] prevent browseruse from killing the cdp - # connection after execution (go figure) + # [monkeypatch] prevent browseruse from killing the cdp connection after execution session.browser.browser_profile.keep_alive = True agent: Any = Agent( @@ -145,7 +120,7 @@ async def _run_browser_agent( llm=llm, browser=session.browser, flash_mode=True, - use_vision=llm_supports_vision(), + use_vision=vision, ) async def log_step(step: Any) -> None: @@ -157,14 +132,21 @@ async def _run_browser_agent( final_result = ( result.final_result() if callable(result.final_result) else result.final_result ) - return {"error": final_result or "Agent failed", "is_running": False} + return { + "error": final_result or "Agent failed", + "is_running": False, + } final_result = ( result.final_result() if hasattr(result, "final_result") and callable(result.final_result) else getattr(result, "final_result", str(result)) ) - out = {"message": "Task completed", "result": final_result, "is_running": False} + out = { + "message": "Task completed", + "result": final_result, + "is_running": False, + } if return_fields: fields = {f: getattr(result, f, None) for f in return_fields} @@ -197,7 +179,7 @@ async def _run_browser_tool( if not session.browser.is_cdp_connected: await session.browser.start() - llm = _build_llm(metadata=metadata) + llm, _ = _build_llm(metadata=metadata) if session.local: from pathlib import Path @@ -241,13 +223,11 @@ async def populate_response(session: BrowserSession, response: dict[str, Any]) - url = await session.browser.get_current_page_url() all_tabs = await session.browser.get_tabs() except Exception as e: # noqa: BLE001 - # TODO: Consider a fallback way to retrieve the url? url = f"URL retrieval failed: {e}" title = "Title retrieval failed with same error" all_tabs = [] try: - # this can fail if the browser disconnects during the tool completion vp = getattr(session.browser.browser_profile, "viewport", None) viewport = { "width": vp.width if vp else None, @@ -304,20 +284,23 @@ async def browser_actions( "is_running": True, } - if not llm_supports_vision(): - result["warning"] = "Model does not support vision" - return result if action == "close_browser": await _close_session(agent_id) - return {"message": "Browser closed", "is_running": False} + return { + "message": "Browser closed", + "is_running": False, + } session = _get_session(agent_id) if action == "run": if not task: - return {"error": "task required for run action", "is_running": False} + return { + "error": "task required for run action", + "is_running": False, + } runner = partial(_run_browser_agent, session, task, return_fields, metadata) desc = task @@ -335,4 +318,7 @@ async def browser_actions( except Exception as error: logger.exception("browser_actions error: %s", action) - return {"error": str(error), "is_running": False} + return { + "error": str(error), + "is_running": False, + } diff --git a/strix/tools/browser/browser_manager.py b/strix/tools/browser/browser_manager.py index 04cf2da1..b1297b55 100644 --- a/strix/tools/browser/browser_manager.py +++ b/strix/tools/browser/browser_manager.py @@ -2,19 +2,16 @@ import asyncio import atexit import contextlib import logging -import time from typing import Any +from urllib.parse import urlparse +import httpx from browser_use import Browser +from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_fixed logger = logging.getLogger(__name__) -_MAX_CONSECUTIVE_FAILURES = 3 -_MAX_SESSION_AGE = 1800 -_CDP_RECOVERY_TIMEOUT = 60 -_CDP_RECOVERY_INTERVAL = 2 - class BrowserSessionManager: def __init__(self) -> None: @@ -50,18 +47,13 @@ class BrowserSessionManager: def close_all(self) -> None: for session in list(self.sessions.values()): - session.invalidated = True browser = session.browser session.browser = None - if browser is None: - continue - for method_name in ("close", "stop"): - fn = getattr(browser, method_name, None) - if callable(fn): - with contextlib.suppress(Exception): - if asyncio.iscoroutine(coro := fn()): - coro.close() - break + if browser is not None: + with contextlib.suppress(Exception): + coro = browser.stop() + if asyncio.iscoroutine(coro): + coro.close() self.sessions.clear() @@ -70,9 +62,6 @@ class BrowserSession: "auth_token", "browser", "cdp_url", - "consecutive_failures", - "created_at", - "invalidated", "local", "profile_directory", "task_count", @@ -95,210 +84,95 @@ class BrowserSession: self.auth_token = auth_token self.local = local self.profile_directory = profile_directory - self.created_at = time.monotonic() self.task_count = 0 - self.consecutive_failures = 0 - self.invalidated = False - - @property - def age(self) -> float: - return time.monotonic() - self.created_at - - @property - def needs_refresh(self) -> bool: - return ( - self.invalidated - or self.consecutive_failures >= _MAX_CONSECUTIVE_FAILURES - or self.age > _MAX_SESSION_AGE - ) async def close(self) -> None: - self.invalidated = True - await _close_browser(self.browser, label="shutdown") + await _close_browser(self.browser) self.browser = None - async def refresh(self) -> None: - old_browser = self.browser - self.browser = None - await _close_browser(old_browser, label="stale") - - from browser_use import Browser - - if self.local: - kwargs: dict[str, Any] = {} - if self.profile_directory: - kwargs["profile_directory"] = self.profile_directory - self.browser = Browser.from_system_chrome(**kwargs) - else: - ws_url, _ = await asyncio.to_thread( - _wait_for_cdp, self.cdp_url, self.auth_token, max_attempts=10, interval=2.0 - ) - self.browser = Browser(cdp_url=ws_url) - self.ws_url = ws_url - - self.invalidated = False - self.consecutive_failures = 0 - - async def ensure_healthy(self, task_num: int) -> str | None: - if self.local: - if self.needs_refresh: - try: - await self.refresh() - except Exception as exc: # noqa: BLE001 - return f"Failed to refresh local browser session: {exc}" - return None - - cdp_alive = await asyncio.to_thread(_check_cdp_alive, self.cdp_url, self.auth_token) - - if not self.needs_refresh and cdp_alive: - return None - - if not cdp_alive and not await _wait_for_cdp_recovery(self, task_num): - if not self.needs_refresh: - self.invalidated = True - return ( - f"Chromium CDP at {self.cdp_url} is not responding after {_CDP_RECOVERY_TIMEOUT}s" - ) - - try: - await self.refresh() - except Exception as exc: # noqa: BLE001 - if not self.needs_refresh: - self.invalidated = True - return f"Failed to refresh browser session: {exc}" - return None - _manager = BrowserSessionManager() -async def _close_browser(browser: Any, label: str = "") -> None: +async def _close_browser(browser: Any) -> None: if browser is None: return + # suppress "Client is stopping" noise that CDP fires during intentional teardown loop = asyncio.get_running_loop() - original_handler = loop.get_exception_handler() - - # [debug] ignore client is stopping exceptions since we... know - def squelch_cdp_error(loop: asyncio.AbstractEventLoop, context: dict[str, Any]) -> None: - exc = context.get("exception") - if isinstance(exc, ConnectionError) and "Client is stopping" in str(exc): - return - loop.default_exception_handler(context) - - loop.set_exception_handler(squelch_cdp_error) + prev_handler = loop.get_exception_handler() + loop.set_exception_handler( + lambda loop_, ctx: ( + None + if isinstance(ctx.get("exception"), ConnectionError) + and "Client is stopping" in str(ctx["exception"]) + else loop_.default_exception_handler(ctx) + ) + ) try: - for method_name in ("close", "stop"): - if ( - (fn := getattr(browser, method_name, None)) - and callable(fn) - and asyncio.iscoroutine(coro := fn()) - ): - try: - await asyncio.wait_for(coro, timeout=10) - except TimeoutError: - logger.warning("Browser (%s) %s() timed out", label, method_name) - except Exception: # noqa: BLE001,S110 - pass - else: - return + await asyncio.wait_for(browser.stop(), timeout=10) + except Exception: # noqa: BLE001,S110 + pass finally: await asyncio.sleep(0.1) - loop.set_exception_handler(original_handler) + loop.set_exception_handler(prev_handler) -def _wait_for_cdp( +class _CDPNotReadyError(Exception): + pass + + +@retry( # type: ignore[misc] + stop=stop_after_attempt(30), + wait=wait_fixed(1), + retry=retry_if_exception_type(_CDPNotReadyError), + reraise=True, +) +async def _wait_for_cdp( cdp_url: str, auth_token: str = "", # nosec B107 - max_attempts: int = 30, - interval: float = 1.0, ) -> tuple[str, dict[str, Any]]: - from urllib.parse import urlparse - - import httpx - version_url = cdp_url.rstrip("/") + "/json/version" - headers: dict[str, str] = {} - if auth_token: - headers["Authorization"] = f"Bearer {auth_token}" - last_error: str = "" + headers = {"Authorization": f"Bearer {auth_token}"} if auth_token else {} - for attempt in range(1, max_attempts + 1): + async with httpx.AsyncClient(trust_env=False, timeout=5) as client: try: - with httpx.Client(trust_env=False, timeout=5) as client: - resp = client.get(version_url, headers=headers) - if resp.status_code == 200 and "webSocketDebuggerUrl" in resp.text: - version_info = resp.json() - raw_ws = version_info.get("webSocketDebuggerUrl", "") + resp = await client.get(version_url, headers=headers) + except httpx.HTTPError as e: + raise _CDPNotReadyError(f"{type(e).__name__}: {e}") from e - # Rewrite the WebSocket URL: Chromium reports the container- - # internal address (e.g. ws://127.0.0.1:19222/devtools/...) - # but we need it routed through the tool server's CDP proxy. - parsed_cdp = urlparse(cdp_url) - parsed_ws = urlparse(raw_ws) - ws_url = parsed_ws._replace( - netloc=parsed_cdp.netloc, - path=parsed_cdp.path.rstrip("/") + parsed_ws.path, - ).geturl() + if resp.status_code != 200 or "webSocketDebuggerUrl" not in resp.text: + raise _CDPNotReadyError(f"HTTP {resp.status_code}") - # Append auth token so browser-use's WebSocket upgrade - # request passes through the CDP auth proxy. - if auth_token: - sep = "&" if "?" in ws_url else "?" - ws_url = f"{ws_url}{sep}token={auth_token}" + info = resp.json() + ws_url = _rewrite_ws_url(cdp_url, info.get("webSocketDebuggerUrl", ""), auth_token) + logger.info("CDP ready: %s", info.get("Browser", "?")) + return ws_url, info - # Log the WS URL with the token redacted to avoid - # leaking credentials into log files. - import re - safe_ws = re.sub(r"([?&])token=[^&]+", r"\1token=REDACTED", ws_url) - logger.info( - "CDP ready at %s (attempt %d): browser=%s, ws_raw=%s, ws_rewritten=%s", - cdp_url, - attempt, - version_info.get("Browser", "unknown"), - raw_ws, - safe_ws, - ) - return ws_url, version_info - last_error = f"HTTP {resp.status_code}, body={resp.text[:200]}" - logger.debug("CDP not ready at %s (attempt %d): %s", cdp_url, attempt, last_error) - except httpx.ConnectError as e: - last_error = f"ConnectError: {e}" - logger.debug("CDP connect failed (attempt %d): %s", attempt, last_error) - except httpx.TimeoutException as e: - last_error = f"Timeout: {e}" - logger.debug("CDP timeout (attempt %d): %s", attempt, last_error) - except httpx.RequestError as e: - last_error = f"RequestError({type(e).__name__}): {e}" - logger.debug("CDP request error (attempt %d): %s", attempt, last_error) - time.sleep(interval) - - raise ConnectionError( - f"Chromium CDP at {cdp_url} did not become ready after {max_attempts}s. " - f"Last error: {last_error}. " - "The sandbox browser may have failed to start — check container logs." - ) +def _rewrite_ws_url(cdp_url: str, raw_ws: str, auth_token: str) -> str: + parsed_cdp = urlparse(cdp_url) + parsed_ws = urlparse(raw_ws) + ws_url = parsed_ws._replace( + netloc=parsed_cdp.netloc, + path=parsed_cdp.path.rstrip("/") + parsed_ws.path, + ).geturl() + if auth_token: + sep = "&" if "?" in ws_url else "?" + ws_url = f"{ws_url}{sep}token={auth_token}" + return ws_url async def _launch_browser(cdp_url: str, agent_id: str, auth_token: str = "") -> BrowserSession: # nosec B107 if session := _manager.get(agent_id): return session - ws_url, version_info = await asyncio.to_thread(_wait_for_cdp, cdp_url, auth_token) - logger.info( - "CDP ready: %s, protocol: %s", - version_info.get("Browser", "?"), - version_info.get("Protocol-Version", "?"), - ) - - from browser_use import Browser - + ws_url, _ = await _wait_for_cdp(cdp_url, auth_token) browser = Browser(cdp_url=ws_url) if session := _manager.get(agent_id): - task = asyncio.ensure_future(_close_browser(browser, label="race-discarded")) + task = asyncio.ensure_future(_close_browser(browser)) _manager.background_tasks.add(task) task.add_done_callback(_manager.background_tasks.discard) return session @@ -312,8 +186,6 @@ async def _launch_local_browser( if session := _manager.get(agent_id): return session - from browser_use import Browser - kwargs: dict[str, Any] = {} if profile_directory: kwargs["profile_directory"] = profile_directory @@ -328,11 +200,9 @@ async def _launch_local_browser( def _get_session(agent_id: str) -> BrowserSession: if session := _manager.get(agent_id): return session - msg = ( - f"Browser not launched. Active sessions: {list(_manager.sessions.keys())}, " - f"requested: {agent_id}" + raise ValueError( + f"Browser not launched. Active: {list(_manager.sessions.keys())}, requested: {agent_id}" ) - raise ValueError(msg) async def _close_session(agent_id: str) -> None: @@ -340,60 +210,4 @@ async def _close_session(agent_id: str) -> None: await session.close() -def cleanup_agent(agent_id: str) -> None: - if not (session := _manager.remove(agent_id)): - return - try: - loop = asyncio.get_event_loop() - if loop.is_running(): - asyncio.run_coroutine_threadsafe(session.close(), loop) - else: - loop.run_until_complete(session.close()) - except Exception: # noqa: BLE001,S110 - pass - - atexit.register(_manager.close_all) - - -def _check_cdp_alive(cdp_url: str, auth_token: str = "") -> bool: # nosec B107 - import httpx - - version_url = cdp_url.rstrip("/") + "/json/version" - headers: dict[str, str] = {} - if auth_token: - headers["Authorization"] = f"Bearer {auth_token}" - try: - with httpx.Client(trust_env=False, timeout=5) as client: - resp = client.get(version_url, headers=headers) - return resp.status_code == 200 and "webSocketDebuggerUrl" in resp.text - except Exception: # noqa: BLE001 - return False - - -async def _wait_for_cdp_recovery(session: BrowserSession, task_num: int) -> bool: - max_attempts = int(_CDP_RECOVERY_TIMEOUT / _CDP_RECOVERY_INTERVAL) - for attempt in range(1, max_attempts + 1): - await asyncio.sleep(_CDP_RECOVERY_INTERVAL) - if await asyncio.to_thread(_check_cdp_alive, session.cdp_url, session.auth_token): - logger.info( - "Task #%d: CDP recovered after %ds", task_num, attempt * _CDP_RECOVERY_INTERVAL - ) - return True - return False - - -async def _ensure_healthy_session(session: BrowserSession, task_num: int) -> str | None: - return await session.ensure_healthy(task_num) - - -def llm_supports_vision() -> bool: - try: - import litellm - - from strix.config.config import resolve_llm_config - - model, _, _ = resolve_llm_config() - return bool(model and litellm.supports_vision(model)) - except Exception: # noqa: BLE001 - return False