diff --git a/containers/Dockerfile b/containers/Dockerfile index 721ec8e7..c29c20c8 100644 --- a/containers/Dockerfile +++ b/containers/Dockerfile @@ -216,8 +216,8 @@ RUN echo 'export PATH="/home/pentester/go/bin:/home/pentester/.local/bin:/home/p USER root COPY containers/docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh COPY containers/healthcheck.sh /usr/local/bin/healthcheck.sh -COPY containers/cdp-auth-proxy.py /usr/local/bin/cdp-auth-proxy.py -RUN chmod +x /usr/local/bin/docker-entrypoint.sh /usr/local/bin/healthcheck.sh /usr/local/bin/cdp-auth-proxy.py +COPY containers/proxy.py /usr/local/bin/proxy.py +RUN chmod +x /usr/local/bin/docker-entrypoint.sh /usr/local/bin/healthcheck.sh /usr/local/bin/proxy.py HEALTHCHECK --interval=15s --timeout=5s --start-period=60s --retries=3 \ CMD healthcheck.sh diff --git a/containers/cdp-auth-proxy.py b/containers/cdp-auth-proxy.py deleted file mode 100644 index f11e6327..00000000 --- a/containers/cdp-auth-proxy.py +++ /dev/null @@ -1,193 +0,0 @@ -#!/usr/bin/env python3 -"""CDP authentication proxy. - -Replaces the raw socat forwarder with an auth-aware TCP proxy. Every -connection (HTTP *and* WebSocket upgrade) must carry a valid token — -either as an ``Authorization: Bearer `` header or a -``?token=`` query parameter. The latter is needed for WebSocket -clients (e.g. Playwright) that don't support custom headers on the -upgrade request. - -Auth credentials are stripped before the request is forwarded to -Chromium so the upstream sees a clean, standard CDP request. - -Environment variables: - TOOL_SERVER_TOKEN — required, shared secret - CDP_PORT — listen port (default 9222) - CDP_INTERNAL_PORT — upstream Chromium port (default 19222) -""" - -from __future__ import annotations - -import asyncio -import os -import re -import sys - - -TOKEN: str = os.environ["TOOL_SERVER_TOKEN"] -LISTEN_PORT: int = int(os.environ.get("CDP_PORT", "9222")) -UPSTREAM_PORT: int = int(os.environ.get("CDP_INTERNAL_PORT", "19222")) - -# Maximum bytes to buffer while looking for the end-of-headers marker. -# Prevents memory exhaustion from slow-loris / oversized-header attacks. -_MAX_HEADER_SIZE: int = 16 * 1024 # 16 KiB — plenty for CDP requests - - -# --------------------------------------------------------------------------- -# Auth helpers -# --------------------------------------------------------------------------- - - -def _check_auth(header_bytes: bytes) -> bool: - text = header_bytes.decode("latin-1", errors="replace") - - # 1) Authorization: Bearer - m = re.search(r"(?i)Authorization:\s*Bearer\s+(\S+)", text) - if m and m.group(1) == TOKEN: - return True - - # 2) ?token= or &token= query param - m = re.search(r"[?&]token=([^&\s]+)", text) - return bool(m and m.group(1) == TOKEN) - - -def _sanitize(header_bytes: bytes) -> bytes: - """Strip auth credentials before forwarding upstream.""" - text = header_bytes.decode("latin-1", errors="replace") - - # Remove Authorization header line - text = re.sub(r"(?im)^Authorization:[^\r\n]*\r\n", "", text) - - # Remove token query param — handle all positions: - # ?token=val&rest → ?rest (first of many) - # ?token=val → (nothing) (sole param) - # &token=val → (nothing) (not first) - text = re.sub(r"\?token=[^&\s]+&", "?", text) - text = re.sub(r"\?token=[^&\s]+", "", text) - text = re.sub(r"&token=[^&\s]+", "", text) - - return text.encode("latin-1") - - -_REJECT = ( - b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 12\r\nConnection: close\r\n\r\nUnauthorized" -) - - -# --------------------------------------------------------------------------- -# TCP proxy -# --------------------------------------------------------------------------- - - -async def _pipe( - reader: asyncio.StreamReader, - writer: asyncio.StreamWriter, -) -> None: - """Forward bytes until EOF or error.""" - try: - while True: - data = await reader.read(65536) - if not data: - break - writer.write(data) - await writer.drain() - except (ConnectionResetError, BrokenPipeError, OSError): - pass - finally: - try: - if writer.can_write_eof(): - writer.write_eof() - except OSError: - pass - - -async def _handle( - client_reader: asyncio.StreamReader, - client_writer: asyncio.StreamWriter, -) -> None: - up_writer: asyncio.StreamWriter | None = None - try: - # --- Read HTTP headers (up to the blank line) --- - buf = b"" - while b"\r\n\r\n" not in buf: - chunk = await asyncio.wait_for(client_reader.read(8192), timeout=10) - if not chunk: - return - buf += chunk - if len(buf) > _MAX_HEADER_SIZE: - client_writer.write( - b"HTTP/1.1 431 Request Header Fields Too Large\r\nConnection: close\r\n\r\n" - ) - await client_writer.drain() - return - - sep = buf.index(b"\r\n\r\n") + 4 - headers = buf[:sep] - remainder = buf[sep:] - - # --- Authenticate --- - if not _check_auth(headers): - client_writer.write(_REJECT) - await client_writer.drain() - return - - # --- Strip credentials & forward --- - headers = _sanitize(headers) - - up_reader, up_writer = await asyncio.open_connection("127.0.0.1", UPSTREAM_PORT) - up_writer.write(headers) - if remainder: - up_writer.write(remainder) - await up_writer.drain() - - # --- Bidirectional pipe --- - done, pending = await asyncio.wait( - [ - asyncio.create_task(_pipe(client_reader, up_writer)), - asyncio.create_task(_pipe(up_reader, client_writer)), - ], - return_when=asyncio.FIRST_COMPLETED, - ) - for t in pending: - t.cancel() - - except (ConnectionResetError, BrokenPipeError, OSError, TimeoutError): - pass - finally: - for w in (client_writer, up_writer): - if w is not None: - with _suppress_os(): - w.close() - - -class _suppress_os: # noqa: N801 - """Tiny context manager — cheaper than contextlib.suppress(OSError).""" - - def __enter__(self) -> None: - pass - - def __exit__(self, *exc: object) -> bool: - return isinstance(exc[1], OSError) if exc[1] is not None else False - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - - -async def _main() -> None: - server = await asyncio.start_server(_handle, "0.0.0.0", LISTEN_PORT) # nosec B104 - print( - f"CDP auth proxy: 0.0.0.0:{LISTEN_PORT} -> 127.0.0.1:{UPSTREAM_PORT}", - flush=True, - ) - async with server: - await server.serve_forever() - - -if __name__ == "__main__": - try: - asyncio.run(_main()) - except KeyboardInterrupt: - sys.exit(0) diff --git a/containers/docker-entrypoint.sh b/containers/docker-entrypoint.sh index 548e5560..fb86e357 100644 --- a/containers/docker-entrypoint.sh +++ b/containers/docker-entrypoint.sh @@ -222,7 +222,7 @@ start_chromium() { fi # Expose CDP on 0.0.0.0 via an authenticated proxy (replaces raw socat). - python3 /usr/local/bin/cdp-auth-proxy.py & + python3 /usr/local/bin/proxy.py & CDP_PROXY_PID=$! echo "Started CDP auth proxy (PID $CDP_PROXY_PID): 0.0.0.0:${CDP_PORT} -> 127.0.0.1:${CDP_INTERNAL_PORT}" diff --git a/containers/proxy.py b/containers/proxy.py new file mode 100644 index 00000000..2f83fe19 --- /dev/null +++ b/containers/proxy.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +import asyncio +import os +import re +from typing import Any + +import aiohttp +from aiohttp import WSMsgType, web + + +TOKEN = os.environ["TOOL_SERVER_TOKEN"] +LISTEN_PORT = int(os.environ.get("CDP_PORT", "9222")) +UPSTREAM = f"http://127.0.0.1:{os.environ.get('CDP_INTERNAL_PORT', '19222')}" + + +def _is_authorized(req: web.Request) -> bool: + auth_header: str = req.headers.get("Authorization", "") + token_param: str | None = req.query.get("token") + return auth_header == f"Bearer {TOKEN}" or token_param == TOKEN + + +def _strip_token(url: str) -> str: + return re.sub(r"[?&]token=[^&]*", "", url).replace("?&", "?").rstrip("?") + + +async def _proxy_ws(req: web.Request, url: str, headers: dict[str, Any]) -> web.WebSocketResponse: + client_ws = web.WebSocketResponse() + await client_ws.prepare(req) + + async with ( + aiohttp.ClientSession() as session, + session.ws_connect(url.replace("http", "ws", 1), headers=headers) as upstream_ws, + ): + + async def relay(src: Any, dst: Any) -> None: + async for msg in src: + if msg.type == WSMsgType.TEXT: + await dst.send_str(msg.data) + elif msg.type == WSMsgType.BINARY: + await dst.send_bytes(msg.data) + else: + print(str.format("Unexpected WebSocket message type: %s", msg.type)) + break + + await asyncio.gather( + relay(upstream_ws, client_ws), + relay(client_ws, upstream_ws), + return_exceptions=True, + ) + return client_ws + + +async def _proxy_http(req: web.Request, url: str, headers: dict[str, Any]) -> web.StreamResponse: + async with ( + aiohttp.ClientSession() as session, + session.request(req.method, url, headers=headers, data=await req.read()) as resp, + ): + out = web.StreamResponse( + status=resp.status, + headers={k: v for k, v in resp.headers.items() if k.lower() != "transfer-encoding"}, + ) + await out.prepare(req) + async for chunk in resp.content.iter_any(): + await out.write(chunk) + return out + + +async def _handle(req: web.Request) -> web.StreamResponse: + if not _is_authorized(req): + return web.Response(status=401, text="Unauthorized") + + url = _strip_token(f"{UPSTREAM}{req.path_qs}") + headers = {k: v for k, v in req.headers.items() if k.lower() not in ("host", "authorization")} + + if req.headers.get("Upgrade", "").lower() == "websocket": + return await _proxy_ws(req, url, headers) + return await _proxy_http(req, url, headers) + + +app = web.Application() +app.router.add_route("*", "/{path:.*}", _handle) + +if __name__ == "__main__": + print(f"CDP auth proxy: 0.0.0.0:{LISTEN_PORT} -> {UPSTREAM}", flush=True) + web.run_app(app, host="0.0.0.0", port=LISTEN_PORT, print=None) diff --git a/pyproject.toml b/pyproject.toml index 46117db5..fba826a2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -140,10 +140,10 @@ module = [ "traceloop.*", "browser_use", "browser_use.*", - "langchain_community", - "langchain_community.*", "cdp_use", "cdp_use.*", + "aiohttp", + "aiohttp.*", ] ignore_missing_imports = true diff --git a/strix/llm/__init__.py b/strix/llm/__init__.py index ed742d88..f35d6dd9 100644 --- a/strix/llm/__init__.py +++ b/strix/llm/__init__.py @@ -17,4 +17,9 @@ litellm.suppress_debug_info = True litellm._logging._disable_debugging() logging.getLogger("asyncio").setLevel(logging.CRITICAL) logging.getLogger("asyncio").propagate = False +logging.getLogger("browser_use").setLevel(logging.CRITICAL) +logging.getLogger("browser_use").propagate = False warnings.filterwarnings("ignore", category=RuntimeWarning, module="asyncio") +warnings.filterwarnings( + "ignore", category=RuntimeWarning, module="litellm.llms.custom_httpx.async_client_cleanup" +) diff --git a/strix/tools/browser/browser_actions.py b/strix/tools/browser/browser_actions.py index a6e9c7d4..183fc71f 100644 --- a/strix/tools/browser/browser_actions.py +++ b/strix/tools/browser/browser_actions.py @@ -1,20 +1,22 @@ import asyncio -import contextlib +import json import logging import os -import time +import re from typing import Any, Literal +from browser_use import Agent +from browser_use.tools.service import Tools + from strix.tools.registry import register_tool from .browser_manager import ( - _BrowserSession, + BrowserSession, _close_session, _ensure_healthy_session, _get_session, _launch_browser, _launch_local_browser, - _reinitialize_after_agent, llm_supports_vision, ) @@ -24,464 +26,117 @@ logger = logging.getLogger(__name__) BrowserUseLocalAction = Literal[ "launch", "run", - "close", - # Granular browser commands - "open", + "close_browser", + "search", + "navigate", + "go_back", + "wait", "click", - "type", "input", + "upload_file", "scroll", - "back", + "find_text", + "send_keys", + "extract", + "search_page", + "find_elements", "screenshot", - "state", + "save_as_pdf", + "dropdown_options", + "select_dropdown", + "evaluate", "switch", "close_tab", - "keys", - "select", - "eval", - "extract", - "hover", - "dblclick", - "rightclick", - "cookies", - "wait", - "get", + "write_file", + "read_file", + "replace_file", + "read_long_content", + "done", ] -# Actions that are handled by browser_commands (granular CDP control). -_GRANULAR_ACTIONS = frozenset( - { - "open", - "click", - "type", - "input", - "scroll", - "back", - "screenshot", - "state", - "switch", - "close_tab", - "keys", - "select", - "eval", - "extract", - "hover", - "dblclick", - "rightclick", - "cookies", - "wait", - "get", - } +_TASK_TIMEOUT = 300 +_WS_ERRORS = ( + "websocket", + "cdp", + "not initialized", + "not connected", + "connection closed", + "disconnected", ) -# Hard timeout for a single browser task (seconds). -_TASK_TIMEOUT = 300 + +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() -> Any: - """Build a browser-use compatible LLM from the strix LLM config. - - Returns a ``ChatLiteLLM`` instance that routes to any provider via litellm. - """ - from langchain_community.chat_models import ChatLiteLLM - from strix.config.config import resolve_llm_config + from .litellm.chat import ChatLiteLLM + model, api_key, api_base = resolve_llm_config() if not model: raise ValueError("STRIX_LLM environment variable must be set") return ChatLiteLLM( model=model, - api_key=api_key or None, - api_base=api_base or None, + api_key=api_key, + api_base=api_base, ) -# --------------------------------------------------------------------------- -# Task execution helpers -# --------------------------------------------------------------------------- - - -def _is_websocket_error(exc: BaseException) -> bool: - """Return True if the exception looks like a WebSocket / CDP disconnect.""" - # Check exception type hierarchy (websockets lib, ConnectionError, etc.) - exc_type = type(exc).__name__.lower() - if any(kw in exc_type for kw in ("connectionclosed", "websocket", "connectionerror")): - return True - - msg = str(exc).lower() - ws_keywords = ( - "websocket", - "cdp", - "not initialized", - "not connected", - "connection closed", - "disconnected", - "broken pipe", - "connection reset", - "eof occurred", - "protocol error", - "target closed", - "session closed", - "page closed", - "browser has been closed", - "browser was closed", - "client is stopping", - "no close frame", - "close frame", - "connection was closed", - "invalid state", - ) - return any(kw in msg for kw in ws_keywords) - - -async def _cleanup_agent(agent: Any) -> None: - """Best-effort cleanup of a browser-use Agent after failure. - - browser-use's Agent may leave dangling asyncio tasks (CDP listeners, - DOM observers) that produce "Future exception was never retrieved" - warnings when the WebSocket is already dead. This drains them. - """ - # Agent may expose a close/stop/cleanup method. - for method_name in ("close", "stop", "cleanup"): - fn = getattr(agent, method_name, None) - if not callable(fn): - continue - try: - coro = fn() - if asyncio.iscoroutine(coro): - await asyncio.wait_for(coro, timeout=5) - except Exception: # noqa: BLE001 - logger.debug("Agent %s() cleanup failed", method_name, exc_info=True) - else: - return - - # Fallback: try to close the browser context the agent was using, - # which cancels its internal CDP subscriptions. - browser_ctx = getattr(agent, "browser_context", None) - if browser_ctx is not None: - close_fn = getattr(browser_ctx, "close", None) - if callable(close_fn): - with contextlib.suppress(Exception): - coro = close_fn() - if asyncio.iscoroutine(coro): - await asyncio.wait_for(coro, timeout=5) - - -async def _run_agent_task( - task: str, - session: _BrowserSession, - *, - return_fields: list[str] | None = None, -) -> dict[str, Any]: - """Run a browser-use task on the session's browser with resilient reconnection. - - Lifecycle: - 1. Ensure the session is healthy (CDP alive, Browser fresh). - 2. Run the task. - 3. On WebSocket/CDP errors, invalidate the session and retry once. - 4. On any terminal failure, mark the session so the next call refreshes. - """ - from browser_use import Agent - - session.task_count += 1 - task_num = session.task_count - - logger.info( - "Task #%d for session (cdp=%s, ws=%s, age=%.1fs, failures=%d): %.200s", - task_num, - session.cdp_url, - session.ws_url, - session.age, - session.consecutive_failures, - task, - ) - - # --- Pre-flight health check (with recovery) --- - preflight_err = await _ensure_healthy_session(session, task_num) - if preflight_err: - session.consecutive_failures += 1 - return {"error": preflight_err, "is_running": False} - - # --- Execute (with one automatic retry on WebSocket disconnect) --- - max_attempts = 2 - last_error: str = "" - - for attempt in range(1, max_attempts + 1): - llm = _build_llm() - agent = Agent( - task=task, - llm=llm, - browser=session.browser, - flash_mode=True, - use_vision=llm.supports_vision, - ) - - try: - result = await asyncio.wait_for(agent.run(), timeout=_TASK_TIMEOUT) - except TimeoutError: - logger.exception( - "Task #%d timed out after %ds (attempt %d)", - task_num, - _TASK_TIMEOUT, - attempt, - ) - await _cleanup_agent(agent) - # Timeout leaves the Browser in an unknown state — invalidate. - session.invalidated = True - session.consecutive_failures += 1 - return { - "error": f"Browser task timed out after {_TASK_TIMEOUT}s", - "is_running": False, - } - except Exception as exc: - logger.exception( - "Task #%d raised exception (attempt %d)", - task_num, - attempt, - ) - await _cleanup_agent(agent) - last_error = str(exc) - - if _is_websocket_error(exc) and attempt < max_attempts: - logger.warning( - "Task #%d: WebSocket/CDP error — refreshing and retrying...", - task_num, - ) - # Force-refresh the session and retry the task. - session.invalidated = True - refresh_err = await _ensure_healthy_session(session, task_num) - if refresh_err: - session.consecutive_failures += 1 - return {"error": refresh_err, "is_running": False} - continue # retry - - # Non-retryable error or final attempt. - session.consecutive_failures += 1 - # If it smells like a connection issue, invalidate for next call. - if _is_websocket_error(exc): - session.invalidated = True - return {"error": f"Browser task failed: {exc}", "is_running": False} - - # --- Success path --- - interpreted = _interpret_agent_result(result, return_fields=return_fields) - - if "error" in interpreted: - error_msg = str(interpreted["error"]) - logger.warning("Task #%d failed: %s", task_num, error_msg) - - # If the agent reported a CDP/WS error, retry once. - if _is_websocket_error(Exception(error_msg)) and attempt < max_attempts: - logger.warning( - "Task #%d: agent result contains WebSocket error — refreshing and retrying...", - task_num, - ) - session.invalidated = True - refresh_err = await _ensure_healthy_session(session, task_num) - if refresh_err: - session.consecutive_failures += 1 - return {"error": refresh_err, "is_running": False} - continue # retry - - session.consecutive_failures += 1 - if _is_websocket_error(Exception(error_msg)): - session.invalidated = True - return interpreted - - # Task succeeded — reset failure counter. - session.consecutive_failures = 0 - result_preview = str(interpreted.get("result", ""))[:200] - logger.info("Task #%d succeeded: %s", task_num, result_preview) - return interpreted - - # Should not reach here, but just in case: - session.consecutive_failures += 1 - return { - "error": f"Browser task failed after {max_attempts} attempts: {last_error}", - "is_running": False, - } - - -_JSON_PRIMITIVES = str | int | float | bool | None - - -def _json_safe(value: Any) -> Any: - """Coerce a value into a JSON-serialisable form.""" - if isinstance(value, _JSON_PRIMITIVES): - return value - if isinstance(value, list): - return [v if isinstance(v, _JSON_PRIMITIVES) else str(v) for v in value] - return str(value) - - -_HISTORY_ACCESSORS: dict[str, str] = { - "urls": "urls", - "screenshot_paths": "screenshot_paths", - "screenshots": "screenshots", - "action_names": "action_names", - "extracted_content": "extracted_content", - "errors": "errors", - "model_actions": "model_actions", - "model_outputs": "model_outputs", - "last_action": "last_action", - "final_result": "final_result", - "is_done": "is_done", - "has_errors": "has_errors", - "model_thoughts": "model_thoughts", - "action_results": "action_results", - "action_history": "action_history", - "number_of_steps": "number_of_steps", - "total_duration_seconds": "total_duration_seconds", -} - - -def _interpret_agent_result( - result: Any, - return_fields: list[str] | None = None, -) -> dict[str, Any]: - """Inspect an ``AgentHistoryList`` and return a success or error dict. - - If the browser-use agent encountered errors (CDP failures, navigation - errors, etc.) this returns an ``{"error": ...}`` dict so the calling - strix agent treats the task as failed rather than succeeded. - - When *return_fields* is provided, the requested history accessors are - included in the response under a ``"fields"`` key. - """ - # Collect errors reported by the browser-use agent. - errors: list[str] = [] - try: - if hasattr(result, "has_errors") and result.has_errors(): - errors = [e for e in result.errors() if e is not None] - except Exception as exc: # noqa: BLE001 - logger.debug("Error checking agent errors: %s", exc) - - # Determine whether the agent considers the task successful. - success: bool | None = None - try: - if hasattr(result, "is_successful"): - success = result.is_successful() - except Exception as exc: # noqa: BLE001 - logger.debug("Error checking agent success: %s", exc) - - # Extract the final content the agent produced. - final_result: str | None = None - try: - if hasattr(result, "final_result"): - final_result = result.final_result() - except Exception as exc: # noqa: BLE001 - logger.debug("Error extracting final_result: %s", exc) - - logger.debug( - "Agent result interpretation: errors=%d, success=%s, has_final_result=%s, result_type=%s", - len(errors), - success, - final_result is not None, - type(result).__name__, - ) - - # If there are errors AND the agent didn't explicitly mark itself as - # successful, treat as failure and surface the errors. - if errors and success is not True: - error_summary = "; ".join(errors) - return {"error": error_summary, "is_running": False} - - # If the agent explicitly reported failure (is_done=True, success=False). - if success is False: - return { - "error": final_result or "Browser task failed (agent reported failure)", - "is_running": False, - } - - out: dict[str, Any] = { - "message": "Task completed", - "result": final_result or str(result), - "is_running": False, - } - - # Attach optional fields requested by the caller. - if return_fields: - fields: dict[str, Any] = {} - available = ", ".join(sorted(_HISTORY_ACCESSORS)) - for field in return_fields: - accessor = _HISTORY_ACCESSORS.get(field) - if accessor is None: - fields[field] = f"unknown field (available: {available})" - continue - try: - attr = getattr(result, accessor, None) - value = attr() if callable(attr) else attr - fields[field] = _json_safe(value) - except Exception as exc: # noqa: BLE001 - logger.debug("Error extracting field %s: %s", field, exc) - fields[field] = f"error: {exc}" - out["fields"] = fields - - return out - - -# --------------------------------------------------------------------------- -# CDP URL resolution -# --------------------------------------------------------------------------- - - def _resolve_cdp_url(agent_state: Any) -> tuple[str, str]: - """Extract the CDP URL and auth token from agent_state's sandbox_info. - - The sandbox container runs Chromium with ``--remote-debugging-port`` - and the Docker runtime maps it to an ephemeral host port stored in - ``sandbox_info["browser_cdp_port"]``. The auth token (shared with - the tool server) is used to authenticate against the CDP auth proxy. - - Returns ``(cdp_url, auth_token)``. - """ - if not hasattr(agent_state, "sandbox_info") or not agent_state.sandbox_info: - raise ValueError( - "agent_state must have sandbox_info with browser_cdp_port. " - "Ensure the sandbox is initialized before launching the browser. " - f"agent_state type={type(agent_state).__name__}, " - f"has sandbox_info={hasattr(agent_state, 'sandbox_info')}" - ) - - sandbox_info = agent_state.sandbox_info - cdp_port = sandbox_info.get("browser_cdp_port") + info = agent_state.sandbox_info + cdp_port = info.get("browser_cdp_port") if not cdp_port: - raise ValueError( - "sandbox_info is missing browser_cdp_port. " - "The sandbox container may not have Chromium CDP enabled. " - f"Available keys: {list(sandbox_info.keys())}" - ) + raise ValueError("Missing browser_cdp_port in sandbox_info") - auth_token: str = sandbox_info.get("auth_token", "") - - # Resolve the Docker host (same logic used by the tool server URL). - docker_host = os.getenv("DOCKER_HOST", "") - if docker_host: + host = "127.0.0.1" + if docker_host := os.getenv("DOCKER_HOST"): from urllib.parse import urlparse - parsed = urlparse(docker_host) - if parsed.scheme in ("tcp", "http", "https") and parsed.hostname: + if (parsed := urlparse(docker_host)).hostname: host = parsed.hostname - else: - host = "127.0.0.1" - else: - host = "127.0.0.1" - cdp_url = f"http://{host}:{cdp_port}" - logger.info( - "Resolved CDP URL: %s (port=%s, host=%s, sandbox_id=%s)", - cdp_url, - cdp_port, - host, - sandbox_info.get("workspace_id", "?")[:12], - ) - return cdp_url, auth_token + return f"http://{host}:{cdp_port}", info.get("auth_token", "") -# --------------------------------------------------------------------------- -# Tool entry-point -# --------------------------------------------------------------------------- +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]) + + 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} @register_tool(sandbox_execution=False) @@ -490,198 +145,168 @@ async def browser_actions( task: str | None = None, use_local: bool = False, profile_directory: str | None = None, - # --- Granular command parameters --- - url: str | None = None, - index: int | None = None, - text: str | None = None, - value: str | None = None, - selector: str | None = None, - keys: str | None = None, - js: str | None = None, - direction: str | None = None, - amount: int | None = None, - tab: int | None = None, - x: float | None = None, - y: float | None = None, - full: bool = False, - path: str | None = None, - subcommand: str | None = None, - query: str | None = None, - name: str | None = None, - domain: str | None = None, - file: str | None = None, - timeout: int | None = None, - state: str | None = None, - secure: bool = False, - http_only: bool = False, - same_site: str | None = None, - expires: float | None = None, return_fields: list[str] | None = None, *, agent_state: Any = None, + **kwargs: Any, ) -> dict[str, Any]: - """Use a browser to perform web tasks via natural language or granular commands. - - This is a **blocking** tool — the calling agent will wait for the browser - task to fully complete before receiving results. The browser session - persists across calls so state (cookies, auth, tabs) carries over. - - Two modes: - - **Sandboxed** (default): The browser runs inside the sandbox container - (with Caido proxy for traffic interception) and is controlled remotely - via CDP. Requires ``agent_state`` with ``sandbox_info``. - - **Local** (``use_local=True``): Uses the system Chrome installation - directly via ``Browser.from_system_chrome()``. No sandbox required. - Preserves login sessions, cookies, and extensions. You may need to - fully close Chrome before launching. Optionally pass - ``profile_directory`` (e.g. ``"Profile 1"``, ``"Default"``) to select - a specific Chrome profile. - - Actions: - launch - Connect to the browser. MUST be called first. - run - Execute a natural-language browser task (requires ``task``). - close - Shut down the browser session. - open - Navigate to a URL (requires ``url``). - click - Click an element by ``index`` or coordinates (``x``, ``y``). - type - Insert text at the cursor (requires ``text``). - input - Click an element and type into it (``index`` + ``text``). - scroll - Scroll the page (``direction``, ``amount``). - back - Navigate back. - screenshot - Capture a screenshot (``full``, ``path``). - state - Get the current page DOM state. - switch - Switch to a tab by index (``tab``). - close_tab - Close a tab (``tab``). - keys - Send keyboard keys (``keys``). - select - Select a dropdown option (``index`` + ``value``). - eval - Execute JavaScript (``js``). - extract - Extract information (requires agent mode — use ``run``). - hover - Hover over an element (``index``). - dblclick - Double-click an element (``index``). - rightclick - Right-click an element (``index``). - cookies - Cookie operations (``subcommand``: get/set/clear/export/import). - wait - Wait for element/text (``subcommand``: selector/text). - get - Get page info (``subcommand``: title/html/text/value/attributes/bbox). - """ try: from strix.tools.context import get_current_agent_id agent_id = get_current_agent_id() - logger.info( - "browser_actions called: action=%s, agent_id=%s, has_task=%s, " - "has_agent_state=%s, use_local=%s, profile_directory=%s", - action, - agent_id, - task is not None, - agent_state is not None, - use_local, - profile_directory, - ) + # Launch if action == "launch": if use_local: session = await _launch_local_browser(agent_id, profile_directory) - result: dict[str, Any] = { - "message": "Local browser launched and ready", + + result = { + "message": "Local browser ready", "mode": "local", - "profile_directory": session.profile_directory or "auto", + "profile": session.profile_directory or "auto", "is_running": True, } else: cdp_url, auth_token = _resolve_cdp_url(agent_state) session = await _launch_browser(cdp_url, agent_id, auth_token) - # Strip auth token from the ws_url before returning — the - # token is a secret and must not leak to the calling agent. - import re + ws_url = re.sub(r"[?&]token=[^&]+", "", session.ws_url) - safe_ws = re.sub(r"[?&]token=[^&]+", "", session.ws_url) result = { - "message": "Browser launched and ready", + "message": "Browser ready", "mode": "sandboxed", - "cdp_url": session.cdp_url, - "ws_url": safe_ws, + "ws_url": ws_url, "is_running": True, } + if not llm_supports_vision(): - result["warning"] = ( - "The current model does not support vision — " - "the browser agent will operate without screenshots." - ) + result["warning"] = "Model does not support vision" + return result - if action == "close": + # Close + if action == "close_browser": await _close_session(agent_id) return {"message": "Browser closed", "is_running": False} + session = _get_session(agent_id) + + # Agent mode if action == "run": if not task: - raise ValueError("task parameter is required for run action") # noqa: TRY301 - session = _get_session(agent_id) - logger.info( - "Starting task for agent %s (uptime=%.1fs, tasks=%d, local=%s): %.120s", - agent_id, - time.monotonic() - session.created_at, - session.task_count, - session.local, - task, - ) - result = await _run_agent_task(task, session, return_fields=return_fields) - logger.info("Browser task completed for agent %s", agent_id) + msg = "task required for run action" + return {"error": msg, "is_running": False} - # Agent.run() calls browser_session.kill() which destroys the - # CDP client and event bus. Re-create the Browser so subsequent - # granular commands (state, eval, click, …) have a working session. - try: - await _reinitialize_after_agent(session) - except Exception: # noqa: BLE001 - logger.warning( - "Failed to reinitialize session after agent run for %s", - agent_id, - exc_info=True, + async def run_agent() -> dict[str, Any]: + llm = _build_llm() + + agent: Any = Agent( + task=task, + llm=llm, + browser=session.browser, + flash_mode=True, + use_vision=llm_supports_vision(), ) + async def log_step(step: Any) -> None: + logger.info("Agent step completed: %s", step) + + try: + result = await agent.run(on_step_end=log_step) + + # Extract result + if hasattr(result, "is_successful") and not result.is_successful(): + final_result = ( + result.final_result() + if callable(result.final_result) + else result.final_result + ) + 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} + + # Return fields + if return_fields: + fields = {f: getattr(result, f, None) for f in return_fields} + out["fields"] = fields + + return out + finally: + # Cleanup + for name in ("close", "stop"): + if fn := getattr(agent, name, None): + try: + if asyncio.iscoroutine(coro := fn()): + await asyncio.wait_for(coro, timeout=5) + except Exception: # noqa: BLE001,S110 + pass + + return await _execute_task(session, run_agent, task) + + # [fix] parse json nested in xml + def fix_json_in_xml(kws: dict[str, Any]) -> dict[str, Any]: + result = {} + for k, v in kws.items(): + if v is None: + continue + if isinstance(v, str) and v.startswith(("[", "{")): + try: + result[k] = json.loads(v) + except (json.JSONDecodeError, ValueError): + result[k] = v + else: + result[k] = v return result - if action in _GRANULAR_ACTIONS: - session = _get_session(agent_id) - from .browser_commands import handle_command + params = fix_json_in_xml(kwargs) - return await handle_command( + async def run_tool() -> Any: + # [resiliency] this happens really randomly. Better to be proactive + if not session.browser.is_cdp_connected: + await session.browser.start() + + llm = _build_llm() + + if session.local: + from pathlib import Path + + from browser_use.filesystem.file_system import FileSystem + + base_dir = Path.cwd() / "browser_files" + base_dir.mkdir(parents=True, exist_ok=True) + file_system = FileSystem(base_dir=str(base_dir), create_default_files=False) + else: + # [monkeypatch] this is to make the screenshot tool work + class StubFileSystem: + def __getattr__(self, name: str) -> Any: + def soft_error(*args: Any, **kwargs: Any) -> dict[str, str]: + error_msg = ( + f"File operation '{name}' not available in sandboxed environment" + ) + logger.warning(error_msg) + return {"error": error_msg} + + return soft_error + + file_system = StubFileSystem() + + return await Tools().registry.execute_action( action, - session.browser, - url=url, - index=index, - text=text, - value=value, - selector=selector, - keys=keys, - js=js, - direction=direction, - amount=amount, - tab=tab, - x=x, - y=y, - full=full, - path=path, - subcommand=subcommand, - query=query, - name=name, - domain=domain, - file=file, - timeout=timeout, - state=state, - secure=secure, - http_only=http_only, - same_site=same_site, - expires=expires, + params=params, + browser_session=session.browser, + page_extraction_llm=llm, + file_system=file_system, ) - raise ValueError(f"Unknown action: {action}") # noqa: TRY301 + return await _execute_task(session, run_tool, f"{action}({list(params.keys())[:3]})") except Exception as error: - logger.exception( - "browser_actions error: action=%s", - action, - ) + logger.exception("browser_actions error: %s", action) return {"error": str(error), "is_running": False} diff --git a/strix/tools/browser/browser_actions_schema.xml b/strix/tools/browser/browser_actions_schema.xml index 26410a06..898ade2a 100644 --- a/strix/tools/browser/browser_actions_schema.xml +++ b/strix/tools/browser/browser_actions_schema.xml @@ -14,7 +14,7 @@ Two usage styles: - **Agent mode** (action="run"): Give a natural-language task and the browser-use agent autonomously navigates, clicks, fills forms, and extracts information. - - **Granular mode**: Use specific actions (open, click, type, scroll, etc.) for precise, + - **Granular mode**: Use specific actions matching browser-use's built-in tools for precise, deterministic control of the browser via CDP. @@ -22,42 +22,54 @@ **Lifecycle:** - launch: Start the browser. MUST be called before any other actions. Pass use_local=true for local Chrome. Optionally pass profile_directory. - - close: Shut down the browser session. + - close_browser: Shut down the browser session. **Agent mode:** - run: Execute a natural-language browser task. Requires 'task' parameter. Optionally pass 'return_fields' to select which history data to include in the response. - **Navigation:** - - open: Navigate to a URL. Requires 'url'. - - back: Navigate back in history. + **Navigation & Browser Control:** + - search: Search queries (DuckDuckGo, Google, Bing). Requires 'query', optional 'engine'. + - navigate: Navigate to URLs. Requires 'url', optional 'new_tab'. + - go_back: Go back in browser history. + - wait: Wait for specified seconds. Optional 'seconds'. - **Interaction:** - - click: Click an element by 'index' or coordinates ('x', 'y'). - - type: Insert text at the current cursor position. Requires 'text'. - - input: Click an element then type into it. Requires 'index' and 'text'. - - hover: Hover over an element. Requires 'index'. - - dblclick: Double-click an element. Requires 'index'. - - rightclick: Right-click an element. Requires 'index'. - - keys: Send keyboard keys (e.g. "Enter", "Ctrl+a"). Requires 'keys'. - - select: Select a dropdown option. Requires 'index' and 'value'. - - scroll: Scroll the page. Optional 'direction' (up/down/left/right) and 'amount'. + **Page Interaction:** + - click: Click element by index. Requires 'index'. + - input: Input text into form fields. Requires 'index' and 'text', optional 'clear'. + - upload_file: Upload files. Requires 'index' and 'path'. + - scroll: Scroll the page. Optional 'down' (bool), 'pages' (float), 'index'. + - find_text: Scroll to specific text. Requires 'text'. + - send_keys: Send special keys (Enter, Escape, etc.). Requires 'keys'. - **Information:** - - screenshot: Capture a screenshot. Optional 'full' (full page) and 'path' (save to file). - - state: Get the current page DOM state with element indices. - - eval: Execute JavaScript. Requires 'js'. - - get: Get page info. Requires 'subcommand': title, html, text, value, attributes, bbox. + **Content Extraction:** + - extract: Extract data from webpages using LLM. Requires 'query', optional 'extract_links', 'start_from_char', 'output_schema'. + - search_page: Search page text (like grep). Requires 'pattern', optional 'regex', 'case_sensitive', 'context_chars', 'css_scope', 'max_results'. + - find_elements: Query DOM by CSS selector. Requires 'selector', optional 'attributes', 'max_results', 'include_text'. - **Cookies:** - - cookies: Cookie operations. Requires 'subcommand': get, set, clear, export, import. + **Visual Analysis:** [Requires to be running in non-sandbox environment] + - screenshot: Capture screenshot. Optional 'file_name'. + - save_as_pdf: Save page as PDF. Optional 'file_name', 'print_background', 'landscape', 'scale', 'paper_format'. - **Waiting:** - - wait: Wait for conditions. Requires 'subcommand': selector, text. + **Form Controls:** + - dropdown_options: Get dropdown options. Requires 'index'. + - select_dropdown: Select dropdown option. Requires 'index' and 'text'. - **Tabs:** - - switch: Switch to a tab by index. Requires 'tab'. - - close_tab: Close a tab. Optional 'tab' (defaults to focused tab). + **JavaScript Execution:** + - evaluate: Execute custom JavaScript code. Requires 'code'. + + **Tab Management:** + - switch: Switch between tabs. Requires 'tab_id'. + - close_tab: Close browser tab. Requires 'tab_id'. + + **File Operations:** [Requires to be running in non-sandbox environment] + - write_file: Write content to files. Requires 'file_name' and 'content', optional 'append', 'trailing_newline', 'leading_newline'. + - read_file: Read file contents. Requires 'file_name'. + - replace_file: Replace text in files. Requires 'file_name', 'old_str', 'new_str'. + - read_long_content: Read long content intelligently. Requires 'goal', optional 'source', 'context'. + + **Task Completion:** + - done: Complete the task. Requires 'text', optional 'success', 'files_to_display'. Required for 'run' action. A natural-language description of what to do in @@ -71,84 +83,131 @@ Only used with action='launch' and use_local=true. The Chrome profile directory to use, e.g. "Default", "Profile 1". If not specified, auto-selected. + + Search query. Required for 'search' and 'extract' actions. + + + Search engine for 'search' action. Options: "google", "duckduckgo", "bing". + - URL to navigate to. Required for 'open' action. Also used by cookies get/clear/export. + URL to navigate to. Required for 'navigate' action. + + + Open URL in new tab for 'navigate' action. Default: false. + + + Seconds to wait for 'wait' action. - Element index from the DOM state. Used by: click, input, select, hover, dblclick, - rightclick, and get (text/value/attributes/bbox). + Element index from the DOM state. Used by: click, input, upload_file, dropdown_options, select_dropdown, scroll (for scrollable elements). - Text content. Required for 'type' and 'input' actions. Also used by wait text. + Text content. Required for 'input', 'select_dropdown', 'find_text', and 'done' actions. - - Value to select. Required for 'select' action. Also used for cookie value in cookies set. - - - CSS selector. Used by wait selector and get html. - - - Keyboard keys to send. Required for 'keys' action. Examples: "Enter", "Ctrl+a", "Escape". - - - JavaScript code to execute. Required for 'eval' action. - - - Scroll direction: "up", "down", "left", "right". Default: "down". - - - Scroll amount in pixels. Default: 500. - - - Tab index. Required for 'switch', optional for 'close_tab'. - - - X coordinate for click action (use with 'y' for coordinate-based clicking). - - - Y coordinate for click action (use with 'x' for coordinate-based clicking). - - - For 'screenshot': capture the full page. Default: false. + + Clear existing text before input. Used by 'input' action. Default: false. - For 'screenshot': save to this file path instead of returning base64. + File path for 'upload_file' action. - - Sub-action for compound commands: - - cookies: "get", "set", "clear", "export", "import" - - wait: "selector", "text" - - get: "title", "html", "text", "value", "attributes", "bbox" + + Scroll direction for 'scroll' action. True=down, False=up. Default: true. - - Query for 'extract' action (requires agent mode). + + Number of pages to scroll for 'scroll' action. Range: 0.5-10.0. Default: 1.0. - - Cookie name. Used by cookies set. + + Keyboard keys to send. Required for 'send_keys' action. Examples: "Enter", "Escape", "Tab". - - Cookie domain. Used by cookies set. + + Extract URLs from page for 'extract' action. Default: false. - - File path for cookies export/import. + + Character position to start extraction from for 'extract' action (for truncated content). - - Timeout in milliseconds for wait actions. Default: 30000. + + Pydantic model schema for structured extraction in 'extract' action. - - Wait state for wait selector: "visible", "hidden", "attached", "detached". Default: "visible". + + Search pattern for 'search_page' action. Can be regex if regex=true. - - Cookie secure flag. Used by cookies set. Default: false. + + Treat pattern as regex for 'search_page' action. Default: false. - - Cookie httpOnly flag. Used by cookies set. Default: false. + + Case-sensitive search for 'search_page' action. Default: false. - - Cookie SameSite policy. Used by cookies set. Values: "Strict", "Lax", "None". + + Number of context characters around matches for 'search_page' action. - - Cookie expiration timestamp. Used by cookies set. + + CSS selector to limit search scope for 'search_page' action. + + + Maximum results for 'search_page' and 'find_elements' actions. + + + CSS selector for 'find_elements' action. + + + Attribute names to extract for 'find_elements' action. E.g. ["href", "src"]. + + + Include element text content for 'find_elements' action. Default: true. + + + File name for 'screenshot', 'save_as_pdf', 'write_file', 'read_file', 'replace_file' actions. + + + Include background graphics in PDF for 'save_as_pdf' action. Default: false. + + + Landscape orientation for 'save_as_pdf' action. Default: false. + + + Scale factor for 'save_as_pdf' action. Range: 0.1-2.0. + + + Paper format for 'save_as_pdf' action. E.g. "A4", "Letter". + + + JavaScript code to execute. Required for 'evaluate' action. + + + Tab ID for 'switch' and 'close_tab' actions. Last 4 chars of target_id from browser state. + + + File content for 'write_file' action. + + + Append to file instead of overwriting for 'write_file' action. Default: false. + + + Add trailing newline for 'write_file' action. + + + Add leading newline for 'write_file' action. + + + String to replace for 'replace_file' action. + + + Replacement string for 'replace_file' action. + + + Information goal for 'read_long_content' action. + + + Content source for 'read_long_content' action. Options: "page", file path. + + + Additional context for 'read_long_content' action. + + + Task success status for 'done' action. + + + Files to display for 'done' action. Only used with action='run'. A list of history fields to include in the @@ -208,100 +267,171 @@ true - # Navigate to a URL - - open - https://example.com - - - # Get current page state with element indices - - state - - - # Click element by index (from state output) - - click - 5 - - - # Click by coordinates - - click - 100 - 200 - - - # Type into an input field (click + type) - - input - 3 - admin@example.com - - - # Send keyboard keys - - keys - Enter - - - # Scroll down - - scroll - down - 500 - - - # Take a screenshot - - screenshot - true - /tmp/page.png - - - # Get page title - - get - title - - - # Wait for an element to appear - - wait - selector - #login-form - 10000 - - - # Get cookies for a URL - - cookies - get - https://example.com - - - # Execute JavaScript - - eval - document.title - - # Run a natural-language browser task run Go to https://example.com/login, fill in username "admin" and password "secret", then click the login button - # Run a task and request specific output fields + # Search with DuckDuckGo - run - Navigate through the site and collect all page titles - ["urls", "extracted_content", "number_of_steps"] + search + Python web scraping - # Close the browser when done + # Navigate to a URL - close + navigate + https://example.com + + + # Go back in history + + go_back + + + # Wait 2 seconds + + wait + 2 + + + # Click element by index + + click + 5 + + + # Type into an input field + + input + 3 + admin@example.com + + + # Upload a file + + upload_file + 7 + /path/to/file.pdf + + + # Scroll down + + scroll + true + 1.5 + + + # Scroll to text + + find_text + Contact Us + + + # Send keyboard keys + + send_keys + Enter + + + # Extract structured data from page + + extract + Extract all product names and prices + true + + + # Search within the page + + search_page + error|warning + true + false + + + # Find elements by CSS selector + + find_elements + a.product-link + ["href", "title"] + + + # Take a screenshot + + screenshot + page_screenshot.png + + + # Save page as PDF + + save_as_pdf + report.pdf + true + + + # Get dropdown options + + dropdown_options + 10 + + + # Select dropdown option + + select_dropdown + 10 + Option 2 + + + # Execute JavaScript + + evaluate + (function(){return document.title})() + + + # Switch to another tab + + switch + a3f2 + + + # Close a tab + + close_tab + a3f2 + + + # Write content to file + + write_file + results.txt + Extracted data here + + + # Read a file + + read_file + data.json + + + # Replace text in file + + replace_file + config.txt + old_value + new_value + + + # Complete the task + + done + Successfully extracted all product data + true + + + # Close the browser when completely done + + close_browser diff --git a/strix/tools/browser/browser_commands.py b/strix/tools/browser/browser_commands.py deleted file mode 100644 index 423d707e..00000000 --- a/strix/tools/browser/browser_commands.py +++ /dev/null @@ -1,726 +0,0 @@ -from __future__ import annotations - -import asyncio -import base64 -import json -import logging -from pathlib import Path -from typing import Any - - -logger = logging.getLogger(__name__) - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -async def _ensure_browser_started(bs: Any) -> None: - """Ensure the BrowserSession has an active CDP connection and watchdogs. - - Checks both the root CDP client and session manager, which are required - by ``get_or_create_cdp_session()``. Calls ``start()`` if either is - missing — ``start()`` is idempotent and safe to call when already connected. - """ - has_cdp = getattr(bs, "_cdp_client_root", None) is not None - has_manager = getattr(bs, "session_manager", None) is not None - if has_cdp and has_manager: - return - await bs.start() - - -async def _execute_js(bs: Any, js: str) -> Any: - """Execute JavaScript in the browser via CDP.""" - cdp_session = await bs.get_or_create_cdp_session(target_id=None, focus=False) - if not cdp_session: - raise RuntimeError("No active browser session") - result = await cdp_session.cdp_client.send.Runtime.evaluate( - params={"expression": js, "returnByValue": True}, - session_id=cdp_session.session_id, - ) - return result.get("result", {}).get("value") - - -async def _get_element_center(bs: Any, node: Any) -> tuple[float, float] | None: - """Get the center coordinates of an element.""" - try: - cdp_session = await bs.cdp_client_for_node(node) - session_id = cdp_session.session_id - backend_node_id = node.backend_node_id - - try: - await cdp_session.cdp_client.send.DOM.scrollIntoViewIfNeeded( - params={"backendNodeId": backend_node_id}, session_id=session_id - ) - await asyncio.sleep(0.05) - except Exception: # noqa: BLE001 - logger.debug("scrollIntoViewIfNeeded failed", exc_info=True) - - element_rect = await bs.get_element_coordinates(backend_node_id, cdp_session) - except Exception: - logger.exception("Failed to get element center") - return None - else: - if element_rect: - center_x = element_rect.x + element_rect.width / 2 - center_y = element_rect.y + element_rect.height / 2 - return center_x, center_y - return None - - -async def _get_node_or_error( - bs: Any, index: int | None -) -> tuple[Any | None, dict[str, Any] | None]: - """Look up a DOM node by index, returning (node, None) or (None, error_dict).""" - if index is None: - return None, {"error": "index parameter is required"} - node = await bs.get_element_by_index(index) - if node is None: - return None, {"error": f"Element index {index} not found - page may have changed"} - return node, None - - -# --------------------------------------------------------------------------- -# Command dispatcher -# --------------------------------------------------------------------------- - - -async def handle_command(action: str, bs: Any, **params: Any) -> dict[str, Any]: - """Dispatch a granular browser command. - - Parameters are passed as keyword args — each action uses only the keys it - needs and ignores the rest. - """ - await _ensure_browser_started(bs) - - if action == "open": - return await _cmd_open(bs, **params) - if action == "click": - return await _cmd_click(bs, **params) - if action == "type": - return await _cmd_type(bs, **params) - if action == "input": - return await _cmd_input(bs, **params) - if action == "scroll": - return await _cmd_scroll(bs, **params) - if action == "back": - return await _cmd_back(bs) - if action == "screenshot": - return await _cmd_screenshot(bs, **params) - if action == "state": - return await _cmd_state(bs) - if action == "switch": - return await _cmd_switch(bs, **params) - if action in ("close_tab", "close-tab"): - return await _cmd_close_tab(bs, **params) - if action == "keys": - return await _cmd_keys(bs, **params) - if action == "select": - return await _cmd_select(bs, **params) - if action == "eval": - return await _cmd_eval(bs, **params) - if action == "extract": - return await _cmd_extract(**params) - if action == "hover": - return await _cmd_hover(bs, **params) - if action == "dblclick": - return await _cmd_dblclick(bs, **params) - if action == "rightclick": - return await _cmd_rightclick(bs, **params) - if action == "cookies": - return await _cmd_cookies(bs, **params) - if action == "wait": - return await _cmd_wait(bs, **params) - if action == "get": - return await _cmd_get(bs, **params) - raise ValueError(f"Unknown browser command: {action}") - - -# --------------------------------------------------------------------------- -# Individual commands -# --------------------------------------------------------------------------- - - -async def _cmd_open(bs: Any, *, url: str | None = None, **_: Any) -> dict[str, Any]: - if not url: - return {"error": "url parameter is required for open action"} - if not url.startswith(("http://", "https://", "file://")): - url = "https://" + url - from browser_use.browser.events import NavigateToUrlEvent - - await bs.event_bus.dispatch(NavigateToUrlEvent(url=url)) - result: dict[str, Any] = {"url": url} - if getattr(bs, "browser_profile", None) and bs.browser_profile.use_cloud and bs.cdp_url: - from urllib.parse import quote - - result["live_url"] = f"https://live.browser-use.com/?wss={quote(bs.cdp_url, safe='')}" - return result - - -async def _cmd_click( - bs: Any, - *, - index: int | None = None, - x: float | None = None, - y: float | None = None, - **_: Any, -) -> dict[str, Any]: - if x is not None and y is not None: - from browser_use.browser.events import ClickCoordinateEvent - - await bs.event_bus.dispatch(ClickCoordinateEvent(coordinate_x=x, coordinate_y=y)) - return {"clicked_coordinate": {"x": x, "y": y}} - if index is not None: - from browser_use.browser.events import ClickElementEvent - - node, err = await _get_node_or_error(bs, index) - if err: - return err - await bs.event_bus.dispatch(ClickElementEvent(node=node)) - return {"clicked": index} - return {"error": "Provide index or (x, y) coordinates"} - - -async def _cmd_type(bs: Any, *, text: str | None = None, **_: Any) -> dict[str, Any]: - if not text: - return {"error": "text parameter is required for type action"} - cdp_session = await bs.get_or_create_cdp_session(target_id=None, focus=False) - if not cdp_session: - return {"error": "No active browser session"} - await cdp_session.cdp_client.send.Input.insertText( - params={"text": text}, session_id=cdp_session.session_id - ) - return {"typed": text} - - -async def _cmd_input( - bs: Any, *, index: int | None = None, text: str | None = None, **_: Any -) -> dict[str, Any]: - if not text: - return {"error": "text parameter is required for input action"} - from browser_use.browser.events import ClickElementEvent, TypeTextEvent - - node, err = await _get_node_or_error(bs, index) - if err: - return err - await bs.event_bus.dispatch(ClickElementEvent(node=node)) - await bs.event_bus.dispatch(TypeTextEvent(node=node, text=text)) - return {"input": text, "element": index} - - -async def _cmd_scroll( - bs: Any, - *, - direction: str | None = None, - amount: int | None = None, - **_: Any, -) -> dict[str, Any]: - from browser_use.browser.events import ScrollEvent - - d = direction or "down" - a = amount or 500 - await bs.event_bus.dispatch(ScrollEvent(direction=d, amount=a)) - return {"scrolled": d, "amount": a} - - -async def _cmd_back(bs: Any) -> dict[str, Any]: - from browser_use.browser.events import GoBackEvent - - await bs.event_bus.dispatch(GoBackEvent()) - return {"back": True} - - -async def _cmd_screenshot( - bs: Any, *, full: bool = False, path: str | None = None, **_: Any -) -> dict[str, Any]: - data = await bs.take_screenshot(full_page=full) - if path: - p = Path(path) - p.write_bytes(data) - return {"saved": str(p), "size": len(data)} - return {"screenshot": base64.b64encode(data).decode(), "size": len(data)} - - -async def _cmd_state(bs: Any) -> dict[str, Any]: - state = await bs.get_browser_state_summary() - if state.dom_state is None: - return {"error": "Browser state has no DOM — the page may still be loading"} - state_text = state.dom_state.llm_representation() - if state.page_info: - pi = state.page_info - header = ( - f"viewport: {pi.viewport_width}x{pi.viewport_height}\n" - f"page: {pi.page_width}x{pi.page_height}\n" - f"scroll: ({pi.scroll_x}, {pi.scroll_y})\n" - ) - state_text = header + state_text - return {"_raw_text": state_text} - - -async def _cmd_switch(bs: Any, *, tab: int | None = None, **_: Any) -> dict[str, Any]: - if tab is None: - return {"error": "tab parameter is required for switch action"} - from browser_use.browser.events import SwitchTabEvent - - page_targets = bs.session_manager.get_all_page_targets() if bs.session_manager else [] - if tab < 0 or tab >= len(page_targets): - return {"error": f"Invalid tab index {tab}. Available: 0-{len(page_targets) - 1}"} - target_id = page_targets[tab].target_id - await bs.event_bus.dispatch(SwitchTabEvent(target_id=target_id)) - return {"switched": tab} - - -async def _cmd_close_tab(bs: Any, *, tab: int | None = None, **_: Any) -> dict[str, Any]: - from browser_use.browser.events import CloseTabEvent - - page_targets = bs.session_manager.get_all_page_targets() if bs.session_manager else [] - if tab is not None: - if tab < 0 or tab >= len(page_targets): - return {"error": f"Invalid tab index {tab}. Available: 0-{len(page_targets) - 1}"} - target_id = page_targets[tab].target_id - else: - focused = bs.session_manager.get_focused_target() if bs.session_manager else None - if not focused: - return {"error": "No focused tab to close"} - target_id = focused.target_id - await bs.event_bus.dispatch(CloseTabEvent(target_id=target_id)) - return {"closed": tab} - - -async def _cmd_keys(bs: Any, *, keys: str | None = None, **_: Any) -> dict[str, Any]: - if not keys: - return {"error": "keys parameter is required"} - from browser_use.browser.events import SendKeysEvent - - await bs.event_bus.dispatch(SendKeysEvent(keys=keys)) - return {"sent": keys} - - -async def _cmd_select( - bs: Any, *, index: int | None = None, value: str | None = None, **_: Any -) -> dict[str, Any]: - if not value: - return {"error": "value parameter is required for select action"} - from browser_use.browser.events import SelectDropdownOptionEvent - - node, err = await _get_node_or_error(bs, index) - if err: - return err - await bs.event_bus.dispatch(SelectDropdownOptionEvent(node=node, text=value)) - return {"selected": value, "element": index} - - -async def _cmd_eval(bs: Any, *, js: str | None = None, **_: Any) -> dict[str, Any]: - if not js: - return {"error": "js parameter is required for eval action"} - result = await _execute_js(bs, js) - return {"result": result} - - -async def _cmd_extract(**params: Any) -> dict[str, Any]: - query = params.get("query") - if not query: - return {"error": "query parameter is required for extract action"} - return {"query": query, "error": "extract requires agent mode — use action='run'"} - - -async def _cmd_hover(bs: Any, *, index: int | None = None, **_: Any) -> dict[str, Any]: - node, err = await _get_node_or_error(bs, index) - if err: - return err - coords = await _get_element_center(bs, node) - if not coords: - return {"error": "Could not get element coordinates for hover"} - cx, cy = coords - cdp_session = await bs.cdp_client_for_node(node) - await cdp_session.cdp_client.send.Input.dispatchMouseEvent( - params={"type": "mouseMoved", "x": cx, "y": cy}, - session_id=cdp_session.session_id, - ) - return {"hovered": index} - - -async def _cmd_dblclick(bs: Any, *, index: int | None = None, **_: Any) -> dict[str, Any]: - node, err = await _get_node_or_error(bs, index) - if err: - return err - coords = await _get_element_center(bs, node) - if not coords: - return {"error": "Could not get element coordinates for double-click"} - cx, cy = coords - cdp_session = await bs.cdp_client_for_node(node) - sid = cdp_session.session_id - send = cdp_session.cdp_client.send.Input.dispatchMouseEvent - await send(params={"type": "mouseMoved", "x": cx, "y": cy}, session_id=sid) - await asyncio.sleep(0.05) - await send( - params={"type": "mousePressed", "x": cx, "y": cy, "button": "left", "clickCount": 2}, - session_id=sid, - ) - await asyncio.sleep(0.05) - await send( - params={"type": "mouseReleased", "x": cx, "y": cy, "button": "left", "clickCount": 2}, - session_id=sid, - ) - return {"double_clicked": index} - - -async def _cmd_rightclick(bs: Any, *, index: int | None = None, **_: Any) -> dict[str, Any]: - node, err = await _get_node_or_error(bs, index) - if err: - return err - coords = await _get_element_center(bs, node) - if not coords: - return {"error": "Could not get element coordinates for right-click"} - cx, cy = coords - cdp_session = await bs.cdp_client_for_node(node) - sid = cdp_session.session_id - send = cdp_session.cdp_client.send.Input.dispatchMouseEvent - await send(params={"type": "mouseMoved", "x": cx, "y": cy}, session_id=sid) - await asyncio.sleep(0.05) - await send( - params={"type": "mousePressed", "x": cx, "y": cy, "button": "right", "clickCount": 1}, - session_id=sid, - ) - await asyncio.sleep(0.05) - await send( - params={"type": "mouseReleased", "x": cx, "y": cy, "button": "right", "clickCount": 1}, - session_id=sid, - ) - return {"right_clicked": index} - - -# --------------------------------------------------------------------------- -# Cookies -# --------------------------------------------------------------------------- - - -async def _cmd_cookies(bs: Any, **params: Any) -> dict[str, Any]: - sub = params.get("subcommand") - if sub == "get": - return await _cookies_get(bs, **params) - if sub == "set": - return await _cookies_set(bs, **params) - if sub == "clear": - return await _cookies_clear(bs, **params) - if sub == "export": - return await _cookies_export(bs, **params) - if sub == "import": - return await _cookies_import(bs, **params) - return {"error": "subcommand required: get, set, clear, export, import"} - - -def _filter_cookies_by_url(cookie_list: list[dict[str, Any]], url: str) -> list[dict[str, Any]]: - from urllib.parse import urlparse - - domain = urlparse(url).netloc - return [ - c - for c in cookie_list - if domain.endswith(str(c.get("domain", "")).lstrip(".")) - or str(c.get("domain", "")).lstrip(".").endswith(domain) - ] - - -def _cookie_to_dict(c: Any) -> dict[str, Any]: - d: dict[str, Any] = { - "name": c.get("name", ""), - "value": c.get("value", ""), - "domain": c.get("domain", ""), - "path": c.get("path", "/"), - "secure": c.get("secure", False), - "httpOnly": c.get("httpOnly", False), - } - if "sameSite" in c: - d["sameSite"] = c.get("sameSite") - if "expires" in c: - d["expires"] = c.get("expires") - return d - - -async def _cookies_get(bs: Any, **params: Any) -> dict[str, Any]: - cookies = await bs._cdp_get_cookies() - cookie_list = [_cookie_to_dict(c) for c in cookies] - url = params.get("url") - if url: - cookie_list = _filter_cookies_by_url(cookie_list, url) - return {"cookies": cookie_list} - - -async def _cookies_set(bs: Any, **params: Any) -> dict[str, Any]: - from cdp_use.cdp.network import Cookie - - cookie_dict: dict[str, Any] = { - "name": params.get("name", ""), - "value": params.get("value", ""), - "path": params.get("path", "/"), - "secure": params.get("secure", False), - "httpOnly": params.get("http_only", False), - } - if params.get("domain"): - cookie_dict["domain"] = params["domain"] - if params.get("same_site"): - cookie_dict["sameSite"] = params["same_site"] - if params.get("expires"): - cookie_dict["expires"] = params["expires"] - if not params.get("domain"): - hostname = await _execute_js(bs, "window.location.hostname") - if hostname: - cookie_dict["domain"] = hostname - try: - cookie_obj = Cookie(**cookie_dict) - await bs._cdp_set_cookies([cookie_obj]) - return {"set": params.get("name"), "success": True} - except Exception as e: - logger.exception("Failed to set cookie") - return {"set": params.get("name"), "success": False, "error": str(e)} - - -async def _cookies_clear(bs: Any, **params: Any) -> dict[str, Any]: - url = params.get("url") - if url: - from urllib.parse import urlparse - - cookies = await bs._cdp_get_cookies() - domain = urlparse(url).netloc - cdp_session = await bs.get_or_create_cdp_session(target_id=None, focus=False) - if cdp_session: - for cookie in cookies: - cookie_domain = str(cookie.get("domain", "")).lstrip(".") - if domain.endswith(cookie_domain) or cookie_domain.endswith(domain): - await cdp_session.cdp_client.send.Network.deleteCookies( - params={ - "name": cookie.get("name", ""), - "domain": cookie.get("domain"), - "path": cookie.get("path", "/"), - }, - session_id=cdp_session.session_id, - ) - else: - await bs._cdp_clear_cookies() - return {"cleared": True, "url": url} - - -async def _cookies_export(bs: Any, **params: Any) -> dict[str, Any]: - file_path = params.get("file") - if not file_path: - return {"error": "file parameter is required for cookies export"} - cookies = await bs._cdp_get_cookies() - cookie_list = [_cookie_to_dict(c) for c in cookies] - url = params.get("url") - if url: - cookie_list = _filter_cookies_by_url(cookie_list, url) - p = Path(file_path) - p.write_text(json.dumps(cookie_list, indent=2, ensure_ascii=False), encoding="utf-8") - return {"exported": len(cookie_list), "file": str(p)} - - -async def _cookies_import(bs: Any, **params: Any) -> dict[str, Any]: - file_path = params.get("file") - if not file_path: - return {"error": "file parameter is required for cookies import"} - p = Path(file_path) - if not p.exists(): - return {"error": f"File not found: {p}"} - cookies = json.loads(p.read_text()) - cdp_session = await bs.get_or_create_cdp_session(target_id=None, focus=False) - if not cdp_session: - return {"error": "No active browser session"} - cookie_list = [] - for c in cookies: - cookie_params: dict[str, Any] = { - "name": c["name"], - "value": c["value"], - "domain": c.get("domain"), - "path": c.get("path", "/"), - "secure": c.get("secure", False), - "httpOnly": c.get("httpOnly", False), - } - if c.get("sameSite"): - cookie_params["sameSite"] = c["sameSite"] - if c.get("expires"): - cookie_params["expires"] = c["expires"] - cookie_list.append(cookie_params) - try: - await cdp_session.cdp_client.send.Network.setCookies( - params={"cookies": cookie_list}, - session_id=cdp_session.session_id, - ) - return {"imported": len(cookie_list), "file": str(p)} - except Exception as e: # noqa: BLE001 - return {"error": f"Failed to import cookies: {e}"} - - -# --------------------------------------------------------------------------- -# Wait -# --------------------------------------------------------------------------- - - -async def _cmd_wait(bs: Any, **params: Any) -> dict[str, Any]: - sub = params.get("subcommand") - if sub == "selector": - return await _wait_selector(bs, **params) - if sub == "text": - return await _wait_text(bs, **params) - return {"error": "subcommand required: selector, text"} - - -async def _wait_selector(bs: Any, **params: Any) -> dict[str, Any]: - selector = params.get("selector") - if not selector: - return {"error": "selector parameter is required"} - timeout_ms = params.get("timeout", 30000) - timeout_s = timeout_ms / 1000.0 - wait_state = params.get("state", "visible") - poll = 0.1 - elapsed = 0.0 - - js_map = { - "attached": f"document.querySelector({json.dumps(selector)}) !== null", - "detached": f"document.querySelector({json.dumps(selector)}) === null", - "visible": ( - f"(function(){{ const el = document.querySelector({json.dumps(selector)});" - " if (!el) return false; const s = window.getComputedStyle(el);" - " const r = el.getBoundingClientRect();" - " return s.display !== 'none' && s.visibility !== 'hidden'" - " && s.opacity !== '0' && r.width > 0 && r.height > 0; }})()" - ), - "hidden": ( - f"(function(){{ const el = document.querySelector({json.dumps(selector)});" - " if (!el) return true; const s = window.getComputedStyle(el);" - " const r = el.getBoundingClientRect();" - " return s.display === 'none' || s.visibility === 'hidden'" - " || s.opacity === '0' || r.width === 0 || r.height === 0; }})()" - ), - } - check_js = js_map.get(wait_state, js_map["attached"]) - - while elapsed < timeout_s: - if await _execute_js(bs, check_js): - return {"selector": selector, "found": True} - await asyncio.sleep(poll) - elapsed += poll - return {"selector": selector, "found": False} - - -async def _wait_text(bs: Any, **params: Any) -> dict[str, Any]: - text = params.get("text") - if not text: - return {"error": "text parameter is required"} - timeout_s = params.get("timeout", 30000) / 1000.0 - poll = 0.1 - elapsed = 0.0 - check_js = f"document.body.innerText.includes({json.dumps(text)})" - - while elapsed < timeout_s: - if await _execute_js(bs, check_js): - return {"text": text, "found": True} - await asyncio.sleep(poll) - elapsed += poll - return {"text": text, "found": False} - - -# --------------------------------------------------------------------------- -# Get -# --------------------------------------------------------------------------- - - -async def _cmd_get(bs: Any, **params: Any) -> dict[str, Any]: - sub = params.get("subcommand") - if sub == "title": - title = await _execute_js(bs, "document.title") - return {"title": title or ""} - if sub == "html": - selector = params.get("selector") - if selector: - js = ( - f"(function(){{ const el = document.querySelector({json.dumps(selector)});" - " return el ? el.outerHTML : null; }})()" - ) - else: - js = "document.documentElement.outerHTML" - html = await _execute_js(bs, js) - return {"html": html or ""} - if sub == "text": - return await _get_text(bs, **params) - if sub == "value": - return await _get_value(bs, **params) - if sub == "attributes": - return await _get_attributes(bs, **params) - if sub == "bbox": - return await _get_bbox(bs, **params) - return {"error": "subcommand required: title, html, text, value, attributes, bbox"} - - -async def _get_text(bs: Any, **params: Any) -> dict[str, Any]: - index = params.get("index") - node, err = await _get_node_or_error(bs, index) - if err: - return err - text = node.get_all_children_text(max_depth=10) if node else "" - return {"index": index, "text": text} - - -async def _get_value(bs: Any, **params: Any) -> dict[str, Any]: - index = params.get("index") - node, err = await _get_node_or_error(bs, index) - if err: - return err - assert node is not None - try: - cdp_session = await bs.cdp_client_for_node(node) - resolve_result = await cdp_session.cdp_client.send.DOM.resolveNode( - params={"backendNodeId": node.backend_node_id}, - session_id=cdp_session.session_id, - ) - object_id = resolve_result["object"].get("objectId") - if object_id: - value_result = await cdp_session.cdp_client.send.Runtime.callFunctionOn( - params={ - "objectId": object_id, - "functionDeclaration": "function() { return this.value; }", - "returnByValue": True, - }, - session_id=cdp_session.session_id, - ) - val = value_result.get("result", {}).get("value") - return {"index": index, "value": val or ""} - except Exception: - logger.exception("Failed to get element value") - return {"index": index, "value": ""} - - -async def _get_attributes(bs: Any, **params: Any) -> dict[str, Any]: - index = params.get("index") - node, err = await _get_node_or_error(bs, index) - if err: - return err - assert node is not None - attrs = node.attributes or {} - return {"index": index, "attributes": dict(attrs)} - - -async def _get_bbox(bs: Any, **params: Any) -> dict[str, Any]: - index = params.get("index") - node, err = await _get_node_or_error(bs, index) - if err: - return err - assert node is not None - try: - cdp_session = await bs.cdp_client_for_node(node) - box_result = await cdp_session.cdp_client.send.DOM.getBoxModel( - params={"backendNodeId": node.backend_node_id}, - session_id=cdp_session.session_id, - ) - model = box_result["model"] - content = model.get("content", []) - if len(content) >= 8: - x = min(content[0], content[2], content[4], content[6]) - y = min(content[1], content[3], content[5], content[7]) - w = max(content[0], content[2], content[4], content[6]) - x - h = max(content[1], content[3], content[5], content[7]) - y - return {"index": index, "bbox": {"x": x, "y": y, "width": w, "height": h}} - except Exception: - logger.exception("Failed to get element bbox") - return {"index": index, "bbox": {}} diff --git a/strix/tools/browser/browser_manager.py b/strix/tools/browser/browser_manager.py index 92b63086..f0d807c7 100644 --- a/strix/tools/browser/browser_manager.py +++ b/strix/tools/browser/browser_manager.py @@ -2,39 +2,68 @@ import asyncio import atexit import contextlib import logging -import threading import time from typing import Any logger = logging.getLogger(__name__) -# Maximum consecutive task failures before the session is forcibly invalidated. _MAX_CONSECUTIVE_FAILURES = 3 - -# Maximum session age (seconds) before forcing a fresh Browser object on next task. -_MAX_SESSION_AGE = 1800 # 30 minutes - -# CDP recovery: how long to wait for watchdog restart (seconds) and poll interval. +_MAX_SESSION_AGE = 1800 _CDP_RECOVERY_TIMEOUT = 60 _CDP_RECOVERY_INTERVAL = 2 -# Strong refs to fire-and-forget background tasks so they aren't GC'd. -_background_tasks: set[asyncio.Task[None]] = set() + +class BrowserSessionManager: + def __init__(self) -> None: + self.sessions: dict[str, BrowserSession] = {} + self.background_tasks: set[asyncio.Task[None]] = set() + + def get(self, agent_id: str) -> "BrowserSession | None": + return self.sessions.get(agent_id) + + def create( + self, + agent_id: str, + browser: Any, + cdp_url: str = "", + ws_url: str = "", + auth_token: str = "", + local: bool = False, + profile_directory: str | None = None, + ) -> "BrowserSession": + session = BrowserSession( + browser, + cdp_url, + ws_url, + auth_token=auth_token, + local=local, + profile_directory=profile_directory, + ) + self.sessions[agent_id] = session + return session + + def remove(self, agent_id: str) -> "BrowserSession | None": + return self.sessions.pop(agent_id, None) + + 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 + self.sessions.clear() -# --------------------------------------------------------------------------- -# Persistent browser session -# --------------------------------------------------------------------------- -# browser-use's Agent.run() calls browser_session.start() internally, which -# is idempotent — it only connects to CDP if _cdp_client_root is None. -# We therefore do NOT call browser.start() during launch; we just create the -# Browser object with the CDP URL and let Agent.run() manage the connection -# on the caller's event loop. -# --------------------------------------------------------------------------- - - -class _BrowserSession: +class BrowserSession: __slots__ = ( "auth_token", "browser", @@ -54,20 +83,20 @@ class _BrowserSession: cdp_url: str, ws_url: str, *, - auth_token: str = "", # nosec B107 + auth_token: str = "", local: bool = False, profile_directory: str | None = None, - ) -> None: - self.auth_token = auth_token + ): self.browser = browser self.cdp_url = cdp_url self.ws_url = ws_url + 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 - self.local = local - self.profile_directory = profile_directory @property def age(self) -> float: @@ -75,135 +104,123 @@ class _BrowserSession: @property def needs_refresh(self) -> bool: - """Session should get a fresh Browser object before the next task.""" 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") + self.browser = None -_sessions: dict[str, _BrowserSession] = {} -_lock = threading.Lock() + 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 + + if self.needs_refresh: + cdp_alive = await asyncio.to_thread(_check_cdp_alive, self.cdp_url, self.auth_token) + if not cdp_alive and not await _wait_for_cdp_recovery(self, task_num): + return ( + f"Chromium CDP at {self.cdp_url} is not responding " + f"after {_CDP_RECOVERY_TIMEOUT}s" + ) + try: + await self.refresh() + except Exception as exc: # noqa: BLE001 + return f"Failed to refresh browser session: {exc}" + return None + + if await asyncio.to_thread(_check_cdp_alive, self.cdp_url, self.auth_token): + return None + + if not await _wait_for_cdp_recovery(self, task_num): + 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 + self.invalidated = True + return f"Chromium restarted but reconnection failed: {exc}" + + return None -def _squelch_connection_error( - loop: asyncio.AbstractEventLoop, - context: dict[str, Any], -) -> None: - """Swallow 'Client is stopping' futures left behind by the CDP client. - - During browser shutdown the CDP WebSocket layer may reject pending - futures with ``ConnectionError('Client is stopping')``. Because no - caller ever awaits them, asyncio logs noisy "Future exception was - never retrieved" tracebacks. This handler silences that specific - case and delegates everything else to the default handler. - """ - exc = context.get("exception") - if isinstance(exc, ConnectionError) and "Client is stopping" in str(exc): - logger.debug("Suppressed CDP teardown noise: %s", exc) - return - loop.default_exception_handler(context) +_manager = BrowserSessionManager() -async def _safe_close_browser(browser: Any, label: str = "") -> None: - """Best-effort close/stop of a Browser object, swallowing all errors.""" +async def _close_browser(browser: Any, label: str = "") -> None: if browser is None: return loop = asyncio.get_running_loop() original_handler = loop.get_exception_handler() - loop.set_exception_handler(_squelch_connection_error) - tag = f" ({label})" if label else "" + # [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) + try: for method_name in ("close", "stop"): - fn = getattr(browser, method_name, None) - if not callable(fn): - continue - try: - coro = fn() - if asyncio.iscoroutine(coro): + 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 after 10s", tag, method_name) - except Exception as exc: # noqa: BLE001 - logger.debug("Browser%s %s() failed: %s", tag, method_name, exc) - else: - logger.debug("Browser%s closed via %s()", tag, method_name) - return + except TimeoutError: + logger.warning("Browser (%s) %s() timed out", label, method_name) + except Exception: # noqa: BLE001,S110 + pass + else: + return finally: - # Give straggler futures a moment to settle before restoring. await asyncio.sleep(0.1) loop.set_exception_handler(original_handler) -async def _refresh_browser(session: _BrowserSession) -> None: - """Close the existing Browser and create a fresh one. - - For sandboxed sessions: reconnects via CDP (the primary recovery path when - the WebSocket is stale after a crash, timeout, or too many failures). - - For local sessions: creates a fresh ``Browser.from_system_chrome()`` with - the same profile directory. - """ - old_browser = session.browser - session.browser = None - - # Close the old browser in the background — don't let it block recovery. - await _safe_close_browser(old_browser, label="stale") - - from browser_use import Browser - - if session.local: - kwargs: dict[str, Any] = {} - if session.profile_directory: - kwargs["profile_directory"] = session.profile_directory - session.browser = Browser.from_system_chrome(**kwargs) - session.invalidated = False - session.consecutive_failures = 0 - logger.info( - "Local session refreshed: new Browser.from_system_chrome(profile=%s)", - session.profile_directory or "auto", - ) - return - - ws_url, _ = await asyncio.to_thread( - _wait_for_cdp, - session.cdp_url, - session.auth_token, - max_attempts=10, - interval=2.0, - ) - - session.browser = Browser(cdp_url=ws_url) - session.ws_url = ws_url - session.invalidated = False - session.consecutive_failures = 0 - logger.info( - "Session refreshed: new Browser with ws=%s (old ws was stale)", - ws_url, - ) - - 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]]: - """Block until the CDP endpoint at *cdp_url* responds to ``/json/version``. - - Returns ``(ws_url, version_info)`` where *ws_url* is the WebSocket - debugger URL rewritten to be reachable from the host. Chromium inside - Docker reports ``ws://127.0.0.1:/...`` which is not - reachable from the host — we replace the host:port with the values from - *cdp_url* (the Docker-mapped endpoint). - - When *auth_token* is provided it is sent as a Bearer header on the - HTTP probe and appended as a ``?token=`` query parameter on the - rewritten WebSocket URL (for the CDP auth proxy). - """ from urllib.parse import urlparse import httpx @@ -272,228 +289,81 @@ def _wait_for_cdp( ) -async def _launch_browser(cdp_url: str, agent_id: str, auth_token: str = "") -> _BrowserSession: # nosec B107 - """Create a new browser session connected to the sandbox browser via CDP. +async def _launch_browser(cdp_url: str, agent_id: str, auth_token: str = "") -> BrowserSession: # nosec B107 + if session := _manager.get(agent_id): + return session - The *cdp_url* points to the Chromium instance running inside the Docker - sandbox (exposed via port mapping). browser-use's ``Browser`` connects - over the Chrome DevTools Protocol rather than launching a local process. - - We do NOT call ``browser.start()`` here — ``Agent.run()`` does that - internally (it's idempotent), which ensures the CDP connection is - established on the correct event loop. - """ - with _lock: - if agent_id in _sessions: - logger.info("Reusing existing browser session for agent %s", agent_id) - return _sessions[agent_id] - - logger.info("Launching browser for agent %s: cdp_url=%s", agent_id, cdp_url) - - # Wait for the container's Chromium to be ready and get the rewritten WS URL. - # We pass the ws:// URL directly so browser-use skips its own /json/version - # fetch (which would get the container-internal address). ws_url, version_info = await asyncio.to_thread(_wait_for_cdp, cdp_url, auth_token) logger.info( - "Chromium version: %s, protocol: %s, user-agent: %s", + "CDP ready: %s, protocol: %s", version_info.get("Browser", "?"), version_info.get("Protocol-Version", "?"), - version_info.get("User-Agent", "?")[:80], ) from browser_use import Browser browser = Browser(cdp_url=ws_url) - logger.info( - "Browser object created for agent %s (ws=%s) — CDP connect deferred to Agent.run()", - agent_id, - ws_url, - ) - session = _BrowserSession( - browser=browser, cdp_url=cdp_url, ws_url=ws_url, auth_token=auth_token - ) + if session := _manager.get(agent_id): + task = asyncio.ensure_future(_close_browser(browser, label="race-discarded")) + _manager.background_tasks.add(task) + task.add_done_callback(_manager.background_tasks.discard) + return session - with _lock: - if agent_id in _sessions: - # Another coroutine raced us; close ours and return theirs. - logger.info("Race: another session appeared for agent %s, discarding ours", agent_id) - task = asyncio.ensure_future(_safe_close_browser(browser, label="race-discarded")) - _background_tasks.add(task) - task.add_done_callback(_background_tasks.discard) - return _sessions[agent_id] - _sessions[agent_id] = session - - return session + return _manager.create(agent_id, browser, cdp_url, ws_url, auth_token=auth_token) async def _launch_local_browser( agent_id: str, profile_directory: str | None = None -) -> _BrowserSession: - """Create a browser session using the local system Chrome. - - Uses ``Browser.from_system_chrome()`` which auto-detects the Chrome - executable and user data directory. An optional *profile_directory* - (e.g. ``"Profile 1"``, ``"Default"``) selects a specific Chrome profile. - """ - with _lock: - if agent_id in _sessions: - logger.info("Reusing existing browser session for agent %s", agent_id) - return _sessions[agent_id] - - logger.info( - "Launching local browser for agent %s (profile=%s)", - agent_id, - profile_directory or "auto", - ) +) -> BrowserSession: + 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 - browser = Browser.from_system_chrome(headless=False, **kwargs) - logger.info( - "Local Browser object created for agent %s (profile=%s)", - agent_id, - profile_directory or "auto", + + if session := _manager.get(agent_id): + return session + + return _manager.create(agent_id, browser, local=True, profile_directory=profile_directory) + + +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}" ) - - session = _BrowserSession( - browser=browser, - cdp_url="", - ws_url="", - local=True, - profile_directory=profile_directory, - ) - - with _lock: - if agent_id in _sessions: - logger.info("Race: another session appeared for agent %s, discarding ours", agent_id) - return _sessions[agent_id] - _sessions[agent_id] = session - - return session - - -def _get_session(agent_id: str) -> _BrowserSession: - """Return the existing session for *agent_id*. - - Raises ``ValueError`` if no session exists (i.e. ``launch`` was not called). - """ - with _lock: - session = _sessions.get(agent_id) - if session is None: - with _lock: - active = list(_sessions.keys()) - raise ValueError( - "Browser not launched. You must call browser_use_local_action " - f"with action='launch' before running tasks. " - f"Active sessions: {active}, requested agent_id: {agent_id}" - ) - return session - - -async def _shutdown_session(session: _BrowserSession) -> None: - """Tear down a session's browser.""" - logger.info("Shutting down browser session (cdp_url=%s)", session.cdp_url) - session.invalidated = True - await _safe_close_browser(session.browser, label="shutdown") - session.browser = None - logger.info("Browser session shut down") + raise ValueError(msg) async def _close_session(agent_id: str) -> None: - """Tear down the browser session for *agent_id*.""" - with _lock: - session = _sessions.pop(agent_id, None) - if session is None: - logger.debug("close_session called but no session for agent %s", agent_id) - return - await _shutdown_session(session) + if session := _manager.remove(agent_id): + await session.close() def cleanup_agent(agent_id: str) -> None: - """Best-effort sync cleanup, called from non-async contexts.""" - with _lock: - session = _sessions.pop(agent_id, None) - if session is None: + if not (session := _manager.remove(agent_id)): return try: loop = asyncio.get_event_loop() if loop.is_running(): - asyncio.run_coroutine_threadsafe(_shutdown_session(session), loop) + asyncio.run_coroutine_threadsafe(session.close(), loop) else: - loop.run_until_complete(_shutdown_session(session)) - except Exception: # noqa: BLE001 - logger.debug("cleanup_agent: best-effort shutdown failed for %s", agent_id) + loop.run_until_complete(session.close()) + except Exception: # noqa: BLE001,S110 + pass -def _close_all() -> None: - with _lock: - sessions = [_sessions.pop(aid) for aid in list(_sessions) if aid in _sessions] - if not sessions: - return - for session in sessions: - session.invalidated = True - browser = session.browser - session.browser = None - if browser is None: - continue - # Attempt a synchronous best-effort close without spinning up a - # full asyncio.run(), which hangs when a signal handler fires - # inside the new loop's select(). - for method_name in ("close", "stop"): - fn = getattr(browser, method_name, None) - if not callable(fn): - continue - with contextlib.suppress(Exception): - coro = fn() - if asyncio.iscoroutine(coro): - coro.close() # discard; can't await at atexit - break - - -atexit.register(_close_all) - - -async def _reinitialize_after_agent(session: _BrowserSession) -> None: - """Re-create and start the Browser after Agent.run() kills it. - - browser-use's ``Agent.run()`` calls ``browser_session.kill()`` on - completion, which destroys the CDP client (``_cdp_client_root = None``), - clears the event bus (removing all handlers), and creates a fresh empty - ``EventBus``. Core handlers registered in ``model_post_init()`` are lost - and ``start()`` on the dead session would hang. - - The fix: create a brand-new ``Browser`` object (which runs - ``model_post_init()`` → registers fresh handlers) pointing at the same - CDP endpoint, then call ``start()`` to connect and attach watchdogs. - This ensures subsequent granular commands (state, eval, click, etc.) - have a fully functional session. - """ - from browser_use import Browser - - if session.local: - kwargs: dict[str, Any] = {} - if session.profile_directory: - kwargs["profile_directory"] = session.profile_directory - session.browser = Browser.from_system_chrome(**kwargs) - else: - session.browser = Browser(cdp_url=session.ws_url) - - await session.browser.start() - logger.info("Session reinitialized after agent run (local=%s)", session.local) - - -# --------------------------------------------------------------------------- -# CDP health checks & recovery -# --------------------------------------------------------------------------- +atexit.register(_manager.close_all) def _check_cdp_alive(cdp_url: str, auth_token: str = "") -> bool: # nosec B107 - """Quick check that the CDP endpoint is reachable (non-blocking from sync).""" import httpx version_url = cdp_url.rstrip("/") + "/json/version" @@ -508,115 +378,23 @@ def _check_cdp_alive(cdp_url: str, auth_token: str = "") -> bool: # nosec B107 return False -async def _wait_for_cdp_recovery(session: _BrowserSession, task_num: int) -> bool: - """Wait for the CDP endpoint to come back after a crash/restart. - - Returns True if CDP recovered, False if it's still dead after the timeout. - """ +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, + "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: - """Ensure the session has a working Browser before running a task. - - Returns ``None`` on success or an error message string on failure. - """ - # --- Local sessions: only refresh on invalidation / too many failures / age --- - if session.local: - if session.needs_refresh: - reason = ( - "invalidated" - if session.invalidated - else f"{session.consecutive_failures} consecutive failures" - if session.consecutive_failures >= _MAX_CONSECUTIVE_FAILURES - else f"age {session.age:.0f}s > {_MAX_SESSION_AGE}s" - ) - logger.warning( - "Task #%d: local session needs refresh (%s) — rebuilding Browser", - task_num, - reason, - ) - try: - await _refresh_browser(session) - except Exception as exc: - logger.exception("Task #%d: local session refresh failed", task_num) - return f"Failed to refresh local browser session: {exc}" - return None - - # --- Sandboxed sessions: CDP health checks --- - - # 1) If the session was invalidated (prior crash/timeout/too many failures) - # or is too old, proactively refresh the Browser object. - if session.needs_refresh: - reason = ( - "invalidated" - if session.invalidated - else f"{session.consecutive_failures} consecutive failures" - if session.consecutive_failures >= _MAX_CONSECUTIVE_FAILURES - else f"age {session.age:.0f}s > {_MAX_SESSION_AGE}s" - ) - logger.warning( - "Task #%d: session needs refresh (%s) — rebuilding Browser", - task_num, - reason, - ) - # Make sure CDP is alive first (may need watchdog restart). - cdp_alive = await asyncio.to_thread(_check_cdp_alive, session.cdp_url, session.auth_token) - if not cdp_alive and not await _wait_for_cdp_recovery(session, task_num): - return ( - f"Chromium CDP at {session.cdp_url} is not responding after " - f"{_CDP_RECOVERY_TIMEOUT}s. The sandbox browser may have " - "crashed. Check container logs." - ) - try: - await _refresh_browser(session) - except Exception as exc: - logger.exception("Task #%d: session refresh failed", task_num) - return f"Failed to refresh browser session: {exc}" - return None - - # 2) Normal pre-flight: verify CDP is alive. - if await asyncio.to_thread(_check_cdp_alive, session.cdp_url, session.auth_token): - return None # all good - - # 3) CDP down — wait for watchdog, then refresh. - logger.warning( - "Task #%d: CDP not responding at %s — waiting for watchdog restart...", - task_num, - session.cdp_url, - ) - if not await _wait_for_cdp_recovery(session, task_num): - session.invalidated = True - return ( - f"Chromium CDP at {session.cdp_url} is not responding after " - f"{_CDP_RECOVERY_TIMEOUT}s. The sandbox browser may have crashed. " - "Check container logs." - ) - try: - await _refresh_browser(session) - except Exception as exc: - session.invalidated = True - logger.exception( - "Task #%d: reconnection after CDP recovery failed", - task_num, - ) - return f"Chromium restarted but reconnection failed: {exc}" - - return None +async def _ensure_healthy_session(session: BrowserSession, task_num: int) -> str | None: + return await session.ensure_healthy(task_num) def llm_supports_vision() -> bool: - """Check whether the configured LLM supports vision/image input.""" try: import litellm diff --git a/strix/tools/browser/litellm/__init__.py b/strix/tools/browser/litellm/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/strix/tools/browser/litellm/chat.py b/strix/tools/browser/litellm/chat.py new file mode 100644 index 00000000..0ab82d8f --- /dev/null +++ b/strix/tools/browser/litellm/chat.py @@ -0,0 +1,248 @@ +import logging +from dataclasses import dataclass, field +from typing import Any, TypeVar, overload + +# mypy: disable-error-code="attr-defined" +from browser_use.llm.base import BaseChatModel +from browser_use.llm.exceptions import ModelProviderError, ModelRateLimitError +from browser_use.llm.messages import BaseMessage +from browser_use.llm.schema import SchemaOptimizer +from browser_use.llm.views import ChatInvokeCompletion, ChatInvokeUsage +from pydantic import BaseModel + +from .serializer import LiteLLMMessageSerializer + + +logger = logging.getLogger(__name__) + +T = TypeVar("T", bound=BaseModel) + + +@dataclass +class ChatLiteLLM(BaseChatModel): + """Chat model that routes to any provider via LiteLLM. + + Uses litellm's unified ``acompletion`` API to support all providers + (OpenAI, Anthropic, Google, Ollama, OpenRouter, DeepSeek, etc.) + through a single interface. + + The ``model`` parameter uses litellm's model format, e.g.:: + + "gpt-4o" + "anthropic/claude-sonnet-4-20250514" + "openrouter/google/gemini-2.0-flash-001" + "ollama/llama3" + + Structured output (``output_format``) is handled via litellm's + ``response_format`` parameter which translates across providers. + """ + + model: str + api_key: str | None = None + api_base: str | None = None + temperature: float | None = 0.0 + max_tokens: int | None = 4096 + max_retries: int = 3 + + _provider_name: str = field( + default="", + init=False, + repr=False, + ) + _clean_model: str = field( + default="", + init=False, + repr=False, + ) + + def __post_init__(self) -> None: + """Resolve provider info from the model string via litellm.""" + try: + from litellm import get_llm_provider + + self._clean_model, self._provider_name, _, _ = get_llm_provider(self.model) + except Exception: # noqa: BLE001 + if "/" in self.model: + self._provider_name, self._clean_model = self.model.split("/", 1) + else: + self._provider_name = "openai" + self._clean_model = self.model + + logger.debug( + "ChatLiteLLM initialized: model=%s, provider=%s, clean=%s, api_base=%s", + self.model, + self._provider_name, + self._clean_model, + self.api_base or "(default)", + ) + + @property + def provider(self) -> str: + return self._provider_name or "litellm" + + @property + def name(self) -> str: + return self._clean_model or self.model + + @staticmethod + def _parse_usage(response: Any) -> ChatInvokeUsage | None: + """Extract token usage from a litellm response.""" + usage = getattr(response, "usage", None) + if usage is None: + return None + + prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0 + completion_tokens = getattr(usage, "completion_tokens", 0) or 0 + + prompt_cached = getattr(usage, "cache_read_input_tokens", None) + cache_creation = getattr(usage, "cache_creation_input_tokens", None) + + if prompt_cached is None: + details = getattr(usage, "prompt_tokens_details", None) + if details: + prompt_cached = getattr(details, "cached_tokens", None) + + return ChatInvokeUsage( + prompt_tokens=prompt_tokens, + prompt_cached_tokens=int(prompt_cached) if prompt_cached is not None else None, + prompt_cache_creation_tokens=int(cache_creation) + if cache_creation is not None + else None, + prompt_image_tokens=None, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + + @overload + async def ainvoke( + self, + messages: list[BaseMessage], + output_format: None = None, + **kwargs: Any, + ) -> ChatInvokeCompletion[str]: ... + + @overload + async def ainvoke( + self, + messages: list[BaseMessage], + output_format: type[T], + **kwargs: Any, + ) -> ChatInvokeCompletion[T]: ... + + async def ainvoke( + self, + messages: list[BaseMessage], + output_format: type[T] | None = None, + **kwargs: Any, + ) -> ChatInvokeCompletion[T] | ChatInvokeCompletion[str]: + """Invoke the model via litellm. + + Args: + messages: List of browser-use chat messages. + output_format: Optional Pydantic model class for structured output. + **kwargs: Extra keyword args (``session_id`` etc.) — ignored. + + Returns: + ``ChatInvokeCompletion`` with either a string or parsed Pydantic model. + """ + import litellm + + litellm_messages = LiteLLMMessageSerializer.serialize(messages) + + params: dict[str, Any] = { + "model": self.model, + "messages": litellm_messages, + "num_retries": self.max_retries, + } + + if self.temperature is not None: + params["temperature"] = self.temperature + if self.max_tokens is not None: + params["max_tokens"] = self.max_tokens + if self.api_key: + params["api_key"] = self.api_key + if self.api_base: + params["api_base"] = self.api_base + + if output_format is not None: + schema = SchemaOptimizer.create_optimized_json_schema(output_format) + params["response_format"] = { + "type": "json_schema", + "json_schema": { + "name": "agent_output", + "strict": True, + "schema": schema, + }, + } + + try: + response = await litellm.acompletion(**params) + except litellm.RateLimitError as e: + raise ModelRateLimitError( + message=str(e), + model=self.name, + ) from e + except litellm.Timeout as e: + raise ModelProviderError( + message=f"Request timed out: {e}", + model=self.name, + ) from e + except litellm.APIConnectionError as e: + raise ModelProviderError( + message=str(e), + model=self.name, + ) from e + except litellm.APIError as e: + status = getattr(e, "status_code", 502) or 502 + raise ModelProviderError( + message=str(e), + status_code=status, + model=self.name, + ) from e + except ModelProviderError: + raise + except Exception as e: + raise ModelProviderError( + message=str(e), + model=self.name, + ) from e + + choice = response.choices[0] if response.choices else None + if choice is None: + raise ModelProviderError( + message="Empty response: no choices returned by the model", + status_code=502, + model=self.name, + ) + + content = choice.message.content or "" + usage = self._parse_usage(response) + stop_reason = choice.finish_reason + + thinking: str | None = None + msg_obj = choice.message + reasoning = getattr(msg_obj, "reasoning_content", None) + if reasoning: + thinking = str(reasoning) + + if output_format is not None: + if not content: + raise ModelProviderError( + message="Model returned empty content for structured output request", + status_code=500, + model=self.name, + ) + parsed = output_format.model_validate_json(content) + return ChatInvokeCompletion( + completion=parsed, + thinking=thinking, + usage=usage, + stop_reason=stop_reason, + ) + + return ChatInvokeCompletion( + completion=content, + thinking=thinking, + usage=usage, + stop_reason=stop_reason, + ) diff --git a/strix/tools/browser/litellm/serializer.py b/strix/tools/browser/litellm/serializer.py new file mode 100644 index 00000000..f802b4c0 --- /dev/null +++ b/strix/tools/browser/litellm/serializer.py @@ -0,0 +1,130 @@ +from typing import Any + +from browser_use.llm.messages import ( + AssistantMessage, + BaseMessage, + ContentPartImageParam, + ContentPartTextParam, + SystemMessage, + UserMessage, +) + + +class LiteLLMMessageSerializer: + """Serializer for converting browser-use message types to LiteLLM format.""" + + @staticmethod + def _serialize_user_content( + content: str | list[ContentPartTextParam | ContentPartImageParam], + ) -> str | list[dict[str, Any]]: + """Convert user message content for LiteLLM compatibility.""" + if isinstance(content, str): + return content + + parts: list[dict[str, Any]] = [] + for part in content: + if part.type == "text": + parts.append( + { + "type": "text", + "text": part.text, + } + ) + elif part.type == "image_url": + parts.append( + { + "type": "image_url", + "image_url": { + "url": part.image_url.url, + "detail": part.image_url.detail, + }, + } + ) + return parts + + @staticmethod + def _serialize_system_content( + content: str | list[ContentPartTextParam], + ) -> str | list[dict[str, Any]]: + """Convert system message content for LiteLLM compatibility.""" + if isinstance(content, str): + return content + + return [ + { + "type": "text", + "text": p.text, + } + for p in content + ] + + @staticmethod + def _serialize_assistant_content( + content: str | list[Any] | None, + ) -> str | list[dict[str, Any]] | None: + """Convert assistant message content for LiteLLM compatibility.""" + if content is None: + return None + if isinstance(content, str): + return content + + parts = [] + for part in content: + if part.type == "text": + parts.append( + { + "type": "text", + "text": part.text, + } + ) + elif part.type == "refusal": + parts.append( + { + "type": "text", + "text": f"[Refusal] {part.refusal}", + } + ) + return parts + + @staticmethod + def serialize(messages: list[BaseMessage]) -> list[dict[str, Any]]: + """Convert browser-use messages to litellm-compatible dicts (OpenAI format). + + LiteLLM accepts OpenAI-format message dicts for all providers, handling + provider-specific conversion (e.g. image blocks for Anthropic) internally. + """ + result: list[dict[str, Any]] = [] + for msg in messages: + if isinstance(msg, UserMessage): + d: dict[str, Any] = {"role": "user"} + d["content"] = LiteLLMMessageSerializer._serialize_user_content(msg.content) + if msg.name is not None: + d["name"] = msg.name + result.append(d) + + elif isinstance(msg, SystemMessage): + d = {"role": "system"} + d["content"] = LiteLLMMessageSerializer._serialize_system_content(msg.content) + if msg.name is not None: + d["name"] = msg.name + result.append(d) + + elif isinstance(msg, AssistantMessage): + d = {"role": "assistant"} + d["content"] = LiteLLMMessageSerializer._serialize_assistant_content(msg.content) + if msg.name is not None: + d["name"] = msg.name + if msg.tool_calls: + d["tool_calls"] = [ + { + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments, + }, + } + for tc in msg.tool_calls + ] + result.append(d) + return result