slim down browser-manager

This commit is contained in:
STJ 2026-03-13 17:01:33 -07:00
parent e461905ab9
commit 1ffcbcf9d3
3 changed files with 111 additions and 312 deletions

View file

@ -1,5 +1,4 @@
from .browser_actions import browser_actions
from .browser_manager import cleanup_agent
__all__ = ["browser_actions", "cleanup_agent"]
__all__ = ["browser_actions"]

View file

@ -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,
}

View file

@ -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