import os import asyncio import json from typing import Any, Dict, List, Optional import websockets import websockets.exceptions from camel.toolkits.hybrid_browser_toolkit.hybrid_browser_toolkit_ts import ( HybridBrowserToolkit as BaseHybridBrowserToolkit, ) from camel.toolkits.hybrid_browser_toolkit.ws_wrapper import WebSocketBrowserWrapper as BaseWebSocketBrowserWrapper from app.component.command import bun, uv from app.component.environment import env from app.service.task import Agents from app.utils.listen.toolkit_listen import auto_listen_toolkit from app.utils.toolkit.abstract_toolkit import AbstractToolkit from utils import traceroot_wrapper as traceroot logger = traceroot.get_logger("hybrid_browser_toolkit") class WebSocketBrowserWrapper(BaseWebSocketBrowserWrapper): def __init__(self, config: Optional[Dict[str, Any]] = None): """Initialize wrapper.""" super().__init__(config) logger.info(f"WebSocketBrowserWrapper using ts_dir: {self.ts_dir}") def _ensure_local_no_proxy(self) -> None: local_hosts = ["localhost", "127.0.0.1", "::1"] for key in ("NO_PROXY", "no_proxy"): current = os.environ.get(key, "") if not current: os.environ[key] = ",".join(local_hosts) continue parts = [item.strip() for item in current.split(",") if item.strip()] updated = False for host in local_hosts: if host not in parts: parts.append(host) updated = True if updated: os.environ[key] = ",".join(parts) async def _receive_loop(self): """Background task to receive messages from WebSocket with enhanced logging.""" logger.debug("WebSocket receive loop started") disconnect_reason = None try: while self.websocket: try: response_data = await self.websocket.recv() response = json.loads(response_data) message_id = response.get("id") if message_id and message_id in self._pending_responses: # Set the result for the waiting coroutine future = self._pending_responses.pop(message_id) if not future.done(): future.set_result(response) logger.debug(f"Processed response for message {message_id}") else: message_summary = { "id": response.get("id"), "success": response.get("success"), "has_result": "result" in response, "result_type": type(response.get("result")).__name__ if "result" in response else None } logger.debug(f"Received unexpected message: {message_summary}") except asyncio.CancelledError: disconnect_reason = "Receive loop cancelled" logger.info(f"WebSocket disconnect: {disconnect_reason}") break except websockets.exceptions.ConnectionClosed as e: disconnect_reason = f"WebSocket closed: code={e.code}, reason={e.reason}" logger.warning(f"WebSocket disconnect: {disconnect_reason}") break except websockets.exceptions.WebSocketException as e: disconnect_reason = f"WebSocket error: {type(e).__name__}: {e}" logger.error(f"WebSocket disconnect: {disconnect_reason}") break except json.JSONDecodeError as e: logger.error(f"Failed to decode WebSocket message: {e}") continue # Try to continue on JSON errors except Exception as e: disconnect_reason = f"Unexpected error: {type(e).__name__}: {e}" logger.error(f"WebSocket disconnect: {disconnect_reason}", exc_info=True) # Notify all pending futures of the error for future in self._pending_responses.values(): if not future.done(): future.set_exception(e) self._pending_responses.clear() break finally: logger.info(f"WebSocket receive loop terminated. Reason: {disconnect_reason or 'Normal shutdown'}") # Mark the websocket as None to indicate disconnection self.websocket = None async def start(self): # Simply use the parent implementation which uses system npm/node self._ensure_local_no_proxy() logger.info("Starting WebSocket server using parent implementation (system npm/node)") await super().start() async def _send_command(self, command: str, params: Dict[str, Any]) -> Dict[str, Any]: """Send a command to the WebSocket server with enhanced error handling.""" try: # First ensure we have a valid connection if self.websocket is None: raise RuntimeError("WebSocket connection not established") # Check connection state before sending if hasattr(self.websocket, "state"): import websockets.protocol if self.websocket.state != websockets.protocol.State.OPEN: raise RuntimeError(f"WebSocket is in {self.websocket.state} state, not OPEN") logger.debug(f"Sending command '{command}' with params: {params}") # Call parent's _send_command result = await super()._send_command(command, params) logger.debug(f"Command '{command}' completed successfully") return result except RuntimeError as e: logger.error(f"Failed to send command '{command}': {e}") # Check if it's a connection issue if "WebSocket" in str(e) or "connection" in str(e).lower(): # Mark connection as dead self.websocket = None raise except Exception as e: logger.error(f"Unexpected error sending command '{command}': {type(e).__name__}: {e}") raise # WebSocket connection pool class WebSocketConnectionPool: """Manage WebSocket browser connections with session-based pooling.""" def __init__(self): self._connections: Dict[str, WebSocketBrowserWrapper] = {} self._lock = asyncio.Lock() async def get_connection(self, session_id: str, config: Dict[str, Any]) -> WebSocketBrowserWrapper: """Get or create a connection for the given session ID.""" async with self._lock: # Check if we have an existing connection for this session if session_id in self._connections: wrapper = self._connections[session_id] # Comprehensive connection health check is_healthy = False if wrapper.websocket: try: # Check WebSocket state based on available attributes if hasattr(wrapper.websocket, "state"): import websockets.protocol is_healthy = wrapper.websocket.state == websockets.protocol.State.OPEN if not is_healthy: logger.debug(f"Session {session_id} WebSocket state: {wrapper.websocket.state}") elif hasattr(wrapper.websocket, "open"): is_healthy = wrapper.websocket.open else: # Try ping as last resort try: await asyncio.wait_for(wrapper.websocket.ping(), timeout=1.0) is_healthy = True except: is_healthy = False except Exception as e: logger.debug(f"Health check failed for session {session_id}: {e}") is_healthy = False if is_healthy: logger.debug(f"Reusing healthy WebSocket connection for session {session_id}") return wrapper else: # Connection is unhealthy, clean it up logger.info(f"Removing unhealthy WebSocket connection for session {session_id}") try: await wrapper.stop() except Exception as e: logger.debug(f"Error stopping unhealthy wrapper: {e}") del self._connections[session_id] # Create a new connection logger.info(f"Creating new WebSocket connection for session {session_id}") wrapper = WebSocketBrowserWrapper(config) await wrapper.start() self._connections[session_id] = wrapper logger.info(f"Successfully created WebSocket connection for session {session_id}") return wrapper async def close_connection(self, session_id: str): """Close and remove a connection for the given session ID.""" async with self._lock: if session_id in self._connections: wrapper = self._connections[session_id] try: await wrapper.stop() except Exception as e: logger.error(f"Error closing WebSocket connection for session {session_id}: {e}") del self._connections[session_id] logger.info(f"Closed WebSocket connection for session {session_id}") async def _close_connection_unlocked(self, session_id: str): """Close connection without acquiring lock (for internal use).""" if session_id in self._connections: wrapper = self._connections[session_id] try: await wrapper.stop() except Exception as e: logger.error(f"Error closing WebSocket connection for session {session_id}: {e}") del self._connections[session_id] logger.info(f"Closed WebSocket connection for session {session_id}") async def close_all(self): """Close all connections in the pool.""" async with self._lock: for session_id in list(self._connections.keys()): await self._close_connection_unlocked(session_id) logger.info("Closed all WebSocket connections") # Global connection pool instance websocket_connection_pool = WebSocketConnectionPool() @auto_listen_toolkit(BaseHybridBrowserToolkit) class HybridBrowserToolkit(BaseHybridBrowserToolkit, AbstractToolkit): agent_name: str = Agents.browser_agent def __init__( self, api_task_id: str, *, headless: bool = False, user_data_dir: str | None = None, stealth: bool = True, cache_dir: Optional[str] = None, enabled_tools: List[str] | None = None, browser_log_to_file: bool = False, log_dir: Optional[str] = None, session_id: str | None = None, default_start_url: Optional[str] = None, default_timeout: int | None = None, short_timeout: int | None = None, navigation_timeout: int | None = None, network_idle_timeout: int | None = None, screenshot_timeout: int | None = None, page_stability_timeout: int | None = None, dom_content_loaded_timeout: int | None = None, viewport_limit: bool = False, connect_over_cdp: bool = True, cdp_url: str | None = "http://localhost:9222", cdp_keep_current_page: bool = False, full_visual_mode: bool = False, ) -> None: logger.info(f"[HybridBrowserToolkit] Initializing with api_task_id: {api_task_id}") self.api_task_id = api_task_id logger.debug(f"[HybridBrowserToolkit] api_task_id set to: {self.api_task_id}") # Set default user_data_dir if not provided if user_data_dir is None: # Use browser port to determine profile directory browser_port = env('browser_port', '9222') user_data_base = os.path.expanduser("~/.eigent/browser_profiles") user_data_dir = os.path.join(user_data_base, f"profile_{browser_port}") os.makedirs(user_data_dir, exist_ok=True) logger.info(f"[HybridBrowserToolkit] Using port-based user_data_dir: {user_data_dir} (port: {browser_port})") else: logger.info(f"[HybridBrowserToolkit] Using provided user_data_dir: {user_data_dir}") logger.debug(f"[HybridBrowserToolkit] Calling super().__init__ with session_id: {session_id}") super().__init__( headless=headless, user_data_dir=user_data_dir, stealth=stealth, cache_dir=cache_dir, enabled_tools=enabled_tools, browser_log_to_file=browser_log_to_file, session_id=session_id, default_start_url=default_start_url, default_timeout=default_timeout, short_timeout=short_timeout, navigation_timeout=navigation_timeout, network_idle_timeout=network_idle_timeout, screenshot_timeout=screenshot_timeout, page_stability_timeout=page_stability_timeout, dom_content_loaded_timeout=dom_content_loaded_timeout, viewport_limit=viewport_limit, connect_over_cdp=connect_over_cdp, cdp_url=cdp_url, cdp_keep_current_page=cdp_keep_current_page, full_visual_mode=full_visual_mode, ) logger.info(f"[HybridBrowserToolkit] Initialization complete for api_task_id: {self.api_task_id}") async def _ensure_ws_wrapper(self): """Ensure WebSocket wrapper is initialized using connection pool.""" logger.debug(f"[HybridBrowserToolkit] _ensure_ws_wrapper called for api_task_id: {getattr(self, 'api_task_id', 'NOT SET')}") global websocket_connection_pool # Get session ID from config or use default session_id = self._ws_config.get("session_id", "default") logger.debug(f"[HybridBrowserToolkit] Using session_id: {session_id}") # Log when connecting to browser cdp_url = self._ws_config.get("cdp_url", f"http://localhost:{env('browser_port', '9222')}") logger.info(f"[PROJECT BROWSER] Connecting to browser via CDP at {cdp_url}") # Get or create connection from pool self._ws_wrapper = await websocket_connection_pool.get_connection(session_id, self._ws_config) logger.info(f"[HybridBrowserToolkit] WebSocket wrapper initialized for session: {session_id}") # Additional health check if self._ws_wrapper.websocket is None: logger.warning(f"WebSocket connection for session {session_id} is None after pool retrieval, recreating...") await websocket_connection_pool.close_connection(session_id) self._ws_wrapper = await websocket_connection_pool.get_connection(session_id, self._ws_config) def clone_for_new_session(self, new_session_id: str | None = None) -> "HybridBrowserToolkit": import uuid if new_session_id is None: new_session_id = str(uuid.uuid4())[:8] # For cloned sessions, use the same user_data_dir to share login state # This allows multiple agents to use the same browser profile without conflicts logger.info(f"Cloning session {new_session_id} with shared user_data_dir: {self._user_data_dir}") # Use the same session_id to share the same browser instance # This ensures all clones use the same WebSocket connection and browser return HybridBrowserToolkit( self.api_task_id, headless=self._headless, user_data_dir=self._user_data_dir, # Use the same user_data_dir stealth=self._stealth, cache_dir=f"{self._cache_dir.rstrip('/')}/_clone_{new_session_id}/", enabled_tools=self.enabled_tools.copy(), browser_log_to_file=self._browser_log_to_file, log_dir=self.config_loader.get_toolkit_config().log_dir, session_id=new_session_id, default_start_url=self._default_start_url, default_timeout=self._default_timeout, short_timeout=self._short_timeout, navigation_timeout=self._navigation_timeout, network_idle_timeout=self._network_idle_timeout, screenshot_timeout=self._screenshot_timeout, page_stability_timeout=self._page_stability_timeout, dom_content_loaded_timeout=self._dom_content_loaded_timeout, viewport_limit=self._viewport_limit, connect_over_cdp=self.config_loader.get_browser_config().connect_over_cdp, cdp_url=f"http://localhost:{env('browser_port', '9222')}", cdp_keep_current_page=self.config_loader.get_browser_config().cdp_keep_current_page, full_visual_mode=self._full_visual_mode, ) @classmethod def toolkit_name(cls) -> str: return "Browser Toolkit" async def close(self): """Close the browser toolkit and release WebSocket connection.""" try: # Close browser if needed if self._ws_wrapper: await super().browser_close() except Exception as e: logger.error(f"Error closing browser: {e}") # Release connection from pool session_id = self._ws_config.get("session_id", "default") await websocket_connection_pool.close_connection(session_id) logger.info(f"Released WebSocket connection for session {session_id}") def __del__(self): """Cleanup when object is garbage collected.""" if hasattr(self, "_ws_wrapper") and self._ws_wrapper: session_id = self._ws_config.get("session_id", "default") logger.debug(f"HybridBrowserToolkit for session {session_id} is being garbage collected")