""" WebSocket server implementation for Extendium plugin. """ import threading import time import uuid from typing import Optional from logger.logger import logger # pylint: disable=import-error from websockets.exceptions import ConnectionClosed from websockets.sync.server import Server, ServerConnection, serve from .client_manager import ClientManager from .message_handler import MessageHandler from .request_handler import RequestHandler # Global instances _server: Optional[Server] = None _client_manager: Optional[ClientManager] = None _request_handler: Optional[RequestHandler] = None _message_handler: Optional[MessageHandler] = None _cleanup_thread: Optional[threading.Thread] = None _server_thread: Optional[threading.Thread] = None _running = False class WebSocketClientAdapter: """ Adapter class to maintain compatibility with the existing code that expects the client structure from the websocket-server package. """ def __init__(self, websocket: ServerConnection): self.websocket = websocket self.id = str(uuid.uuid4()) self.data = {'id': self.id} def get(self, key, default=None): return self.data.get(key, default) def __getitem__(self, key): return self.data[key] def __setitem__(self, key, value): self.data[key] = value class MessageAdapter: """ Adapter to maintain compatibility with the existing MessageHandler that expects a server with a send_message method. """ def send_message(self, client: WebSocketClientAdapter, message: str) -> None: try: client.websocket.send(message) except ConnectionClosed: # logger.error("Failed to send message to client: connection closed") pass except Exception as e: logger.error(f"Error sending message: {e}") def handle_client(websocket: ServerConnection) -> None: """ Main handler for each client connection. Args: websocket: The WebSocket connection object """ # Create client adapter client = WebSocketClientAdapter(websocket) try: # Process messages in a loop for message in websocket: if _message_handler: _message_handler.process_message(client, message) except ConnectionClosed: pass except Exception as e: logger.error(f"Error handling client: {e}") finally: # Handle client disconnection (equivalent to _client_left) if _client_manager: _client_manager.remove_client(client.id) def _cleanup_routine() -> None: """Periodically clean up old pending requests.""" global _request_handler, _running while _running and _request_handler: # Clean up requests older than 30 seconds _request_handler.cleanup_old_requests(30) # split wait into intervals to allow faster shutdown detection for _ in range(20): if not _running: return time.sleep(0.25) def initialize_server() -> None: """ Initialize the WebSocket server. Args: port: The port to listen on host: The host address to bind to Returns: The initialized WebSocket server """ global _server, _client_manager, _request_handler, _message_handler if _server: logger.warn("WebSocket server already initialized") return _server # Create instances _client_manager = ClientManager() _request_handler = RequestHandler() # Create message adapter for compatibility message_adapter = MessageAdapter() # Create message handler with the adapter _message_handler = MessageHandler(message_adapter, _client_manager, _request_handler) def run_server(port: int, host: str = "localhost") -> None: """Run the WebSocket server in a separate thread.""" global _server, _cleanup_thread, _server_thread, _running if _running: logger.log("WebSocket server already running") return _running = True # Start cleanup thread as non-daemon for proper shutdown _cleanup_thread = threading.Thread(target=_cleanup_routine) _cleanup_thread.daemon = False # non-daemon so we can actually clean it up. _cleanup_thread.start() # Start server in a separate thread def start_server(): global _server try: _server = serve(handle_client, host, port) _server.serve_forever() except Exception as e: if _running: # Only log if we're not shutting down logger.error(f"Error in WebSocket server: {e}") finally: logger.log("WebSocket server thread ending") _server_thread = threading.Thread(target=start_server) _server_thread.daemon = False # Non-daemon so we can wait for proper termination _server_thread.start() logger.log(f"WebSocket server is running on {host}:{port}") def shutdown_server() -> None: """Shutdown the WebSocket server.""" global _server, _running, _cleanup_thread, _server_thread, _client_manager, _request_handler, _message_handler if not _running: logger.log("WebSocket server not running") return logger.log("Shutting down WebSocket server...") _running = False # disconnect all connected/dangling clients if _client_manager: try: _client_manager.disconnect_all_clients() except Exception as e: logger.error(f"Error disconnecting clients: {e}") # shutdown the server if _server: try: _server.shutdown() except Exception as e: logger.error(f"Error shutting down WebSocket server: {e}") try: if hasattr(_server, 'close'): _server.close() except Exception as e: logger.error(f"Error closing WebSocket server: {e}") # give some time for the server to shutdown, this is somewhat just a safeguard. time.sleep(0.2) # wait for cleanup thread to finish if _cleanup_thread and _cleanup_thread.is_alive(): _cleanup_thread.join(timeout=5) if _cleanup_thread.is_alive(): logger.error("Cleanup thread did not finish within timeout") # wait for server thread to finish if _server_thread and _server_thread.is_alive(): _server_thread.join(timeout=5) if _server_thread.is_alive(): logger.error("Server thread did not finish within timeout") # force inline gc cleanup from the gil. _server = None _client_manager = None _request_handler = None _message_handler = None _cleanup_thread = None _server_thread = None logger.log("WebSocket server shutdown complete")