import socketio from typing import Callable, Dict, Any import logging import time from typing import Optional from ..models.websocket import WebSocketConfig, WebSocketInfo class WebSocketService: def __init__(self, client): self.client = client def set_websocket(self, instance_id: str, config: WebSocketConfig, instance_token: str): """ Configure WebSocket settings for an instance Args: instance_id (str): The instance ID config (WebSocketConfig): The WebSocket configuration instance_token (str): The instance token Returns: dict: The response from the API """ return self.client.post( f'websocket/set/{instance_id}', data=config.__dict__, instance_token=instance_token ) def find_websocket(self, instance_id: str, instance_token: str) -> WebSocketInfo: """ Get WebSocket settings for an instance Args: instance_id (str): The instance ID instance_token (str): The instance token Returns: WebSocketInfo: The WebSocket information """ response = self.client.get( f'websocket/find/{instance_id}', instance_token=instance_token ) return WebSocketInfo(**response) class WebSocketManager: def __init__(self, base_url: str, instance_id: str, api_token: str, max_retries: int = 5, retry_delay: float = 1.0): """ Initialize the WebSocket manager Args: base_url (str): Base URL of the API instance_id (str): Instance ID api_token (str): API authentication token max_retries (int): Maximum number of reconnection attempts retry_delay (float): Initial delay between attempts in seconds """ self.base_url = base_url.rstrip('/') self.instance_id = instance_id self.api_token = api_token self.max_retries = max_retries self.retry_delay = retry_delay self.retry_count = 0 self.should_reconnect = True # Socket.IO configuration self.sio = socketio.Client( ssl_verify=False, # For local development logger=False, engineio_logger=False, request_timeout=30 ) # Configure class logger to INFO self.logger = logging.getLogger(__name__) self.logger.setLevel(logging.INFO) # Dictionary to store registered handlers self._handlers = {} # Configure event handlers self.sio.on('connect', self._on_connect) self.sio.on('disconnect', self._on_disconnect) self.sio.on('error', self._on_error) # Register global handler in instance-specific namespace self.sio.on('*', self._handle_event, namespace=f'/{self.instance_id}') def _on_connect(self): """Handler for connection event""" self.logger.info("Socket.IO connected") self.retry_count = 0 # Reset retry counter after successful connection def _on_disconnect(self): """Handler for disconnection event""" self.logger.warning(f"Socket.IO disconnected. Attempt {self.retry_count + 1}/{self.max_retries}") if self.should_reconnect and self.retry_count < self.max_retries: self._attempt_reconnect() else: self.logger.error("Maximum number of reconnection attempts reached") def _on_error(self, error): """Handler for error events""" self.logger.error(f"Socket.IO error: {str(error)}", exc_info=True) def _attempt_reconnect(self): """Attempt to reconnect with exponential backoff""" try: delay = self.retry_delay * (2 ** self.retry_count) # Exponential backoff self.logger.info(f"Attempting to reconnect in {delay:.2f} seconds...") time.sleep(delay) self.connect() self.retry_count += 1 except Exception as e: self.logger.error(f"Error during reconnection attempt: {str(e)}", exc_info=True) if self.retry_count < self.max_retries: self._attempt_reconnect() else: self.logger.error("All reconnection attempts failed") def _handle_event(self, event, *args): """Global handler for all events""" # Only process registered events if event in self._handlers: self.logger.debug(f"Event received in namespace /{self.instance_id}: {event}") self.logger.debug(f"Event data: {args}") try: # Extract event data raw_data = args[0] if args else {} # Ensure we're passing the correct object to the callback if isinstance(raw_data, dict): self.logger.debug(f"Calling handler for {event} with data: {raw_data}") self._handlers[event](raw_data) else: self.logger.error(f"Invalid data received for event {event}: {raw_data}") except Exception as e: self.logger.error(f"Error processing event {event}: {str(e)}", exc_info=True) def connect(self): """Connect to Socket.IO server""" try: # Connect only to instance namespace with authentication header self.sio.connect( f"{self.base_url}?apikey={self.api_token}", transports=['websocket'], namespaces=[f'/{self.instance_id}'], wait_timeout=30 ) # Join instance-specific room self.sio.emit('subscribe', {'instance': self.instance_id}, namespace=f'/{self.instance_id}') except Exception as e: self.logger.error(f"Error connecting to Socket.IO: {str(e)}", exc_info=True) raise def disconnect(self): """Disconnect from Socket.IO server""" self.should_reconnect = False # Prevent reconnection attempts if self.sio.connected: self.sio.disconnect() def on(self, event: str, callback: Callable): """ Register a callback for a specific event Args: event (str): Event name callback (Callable): Function to be called when the event occurs """ self._handlers[event] = callback