#!/usr/bin/env python3 """ Setup WebSocket Server for Real-Time Features This script implements: - WebSocket server for real-time communication - Client-side WebSocket management - Real-time event handling and broadcasting - Connection management and reconnection logic - Live status updates for workflows """ import asyncio from collections import defaultdict from dataclasses import dataclass, field from datetime import datetime, timedelta from enum import Enum import json import logging import os import sys import threading import time from typing import Any, Callable, Dict, List, Optional, Set import uuid import websockets from websockets.server import WebSocketServerProtocol # Add backend directory to Python path sys.path.append(os.path.dirname(os.path.abspath(__file__))) # Import our working systems from working_enhanced_workflow_engine import working_enhanced_workflow_engine logger = logging.getLogger(__name__) class WebSocketEventType(Enum): """WebSocket event types""" WORKFLOW_UPDATE = "workflow_update" EXECUTION_STATUS = "execution_status" SERVICE_STATUS = "service_status" NOTIFICATION = "notification" COLLABORATION = "collaboration" SYSTEM_UPDATE = "system_update" USER_ACTIVITY = "user_activity" ERROR = "error" HEARTBEAT = "heartbeat" class ConnectionState(Enum): """WebSocket connection state""" CONNECTING = "connecting" CONNECTED = "connected" DISCONNECTING = "disconnecting" DISCONNECTED = "disconnected" RECONNECTING = "reconnecting" ERROR = "error" @dataclass class WebSocketConnection: """WebSocket connection information""" connection_id: str websocket: WebSocketServerProtocol user_id: str session_id: Optional[str] = None connected_at: datetime = field(default_factory=datetime.now) last_activity: datetime = field(default_factory=datetime.now) subscriptions: Set[str] = field(default_factory=set) state: ConnectionState = ConnectionState.CONNECTED metadata: Dict[str, Any] = field(default_factory=dict) @dataclass class WebSocketEvent: """WebSocket event message""" event_id: str event_type: WebSocketEventType payload: Dict[str, Any] timestamp: datetime = field(default_factory=datetime.now) user_id: Optional[str] = None session_id: Optional[str] = None target_connections: List[str] = field(default_factory=list) broadcast: bool = False class RealTimeWebSocketServer: """Real-time WebSocket server for workflow updates""" def __init__(self, host: str = "localhost", port: int = 8765): self.host = host self.port = port self.connections: Dict[str, WebSocketConnection] = {} self.user_connections: Dict[str, Set[str]] = defaultdict(set) # user_id -> connection_ids self.session_connections: Dict[str, Set[str]] = defaultdict(set) # session_id -> connection_ids self.subscriptions: Dict[str, Set[str]] = defaultdict(set) # subscription -> connection_ids self.server = None self.running = False self.event_handlers = {} self.connection_handlers = {} # Performance metrics self.metrics = { "total_connections": 0, "active_connections": 0, "events_sent": 0, "events_received": 0, "errors": 0, "start_time": None } # Initialize event handlers self._initialize_event_handlers() self._initialize_connection_handlers() def _initialize_event_handlers(self): """Initialize WebSocket event handlers""" self.event_handlers = { WebSocketEventType.WORKFLOW_UPDATE: self._handle_workflow_update, WebSocketEventType.EXECUTION_STATUS: self._handle_execution_status, WebSocketEventType.SERVICE_STATUS: self._handle_service_status, WebSocketEventType.NOTIFICATION: self._handle_notification, WebSocketEventType.COLLABORATION: self._handle_collaboration, WebSocketEventType.SYSTEM_UPDATE: self._handle_system_update, WebSocketEventType.USER_ACTIVITY: self._handle_user_activity, WebSocketEventType.ERROR: self._handle_error, WebSocketEventType.HEARTBEAT: self._handle_heartbeat } logger.info(f"Initialized {len(self.event_handlers)} event handlers") def _initialize_connection_handlers(self): """Initialize connection lifecycle handlers""" self.connection_handlers = { "on_connect": self._on_connect, "on_disconnect": self._on_disconnect, "on_message": self._on_message, "on_error": self._on_error } logger.info("Initialized connection handlers") async def start_server(self): """Start the WebSocket server""" try: logger.info(f"Starting WebSocket server on {self.host}:{self.port}") self.running = True self.metrics["start_time"] = datetime.now() # Start WebSocket server self.server = await websockets.serve( self._handle_connection, self.host, self.port, ping_interval=20, ping_timeout=10, close_timeout=10, max_size=10_000_000, # 10MB max message size max_queue=1000 # Max 1000 queued messages ) logger.info(f"WebSocket server started successfully on ws://{self.host}:{self.port}") # Start background tasks asyncio.create_task(self._heartbeat_monitor()) asyncio.create_task(self._cleanup_connections()) asyncio.create_task(self._performance_monitor()) except Exception as e: logger.error(f"Failed to start WebSocket server: {str(e)}") self.running = False raise async def stop_server(self): """Stop the WebSocket server""" try: logger.info("Stopping WebSocket server...") self.running = False # Close all connections for connection_id, connection in list(self.connections.items()): try: await connection.websocket.close() except: pass # Stop the server if self.server: self.server.close() await self.server.wait_closed() logger.info("WebSocket server stopped successfully") except Exception as e: logger.error(f"Error stopping WebSocket server: {str(e)}") async def _handle_connection(self, websocket: WebSocketServerProtocol, path: str): """Handle new WebSocket connection""" connection_id = str(uuid.uuid4()) connection = WebSocketConnection( connection_id=connection_id, websocket=websocket, user_id="", # Will be set during authentication session_id=None, state=ConnectionState.CONNECTING ) try: # Add to connections self.connections[connection_id] = connection self.metrics["total_connections"] += 1 self.metrics["active_connections"] = len(self.connections) logger.info(f"New WebSocket connection: {connection_id} from {path}") # Wait for authentication message auth_message = await websocket.recv() auth_data = json.loads(auth_message) # Validate authentication if auth_data.get("type") == "auth" and auth_data.get("user_id"): connection.user_id = auth_data["user_id"] connection.session_id = auth_data.get("session_id") connection.metadata.update(auth_data.get("metadata", {})) connection.state = ConnectionState.CONNECTED # Add to user and session mappings self.user_connections[connection.user_id].add(connection_id) if connection.session_id: self.session_connections[connection.session_id].add(connection_id) # Send authentication success await self._send_to_connection(connection_id, { "type": "auth_success", "connection_id": connection_id, "timestamp": datetime.now().isoformat() }) # Call connection handler await self.connection_handlers["on_connect"](connection) logger.info(f"WebSocket {connection_id} authenticated for user {connection.user_id}") else: # Authentication failed await websocket.close(1008, "Authentication failed") logger.warning(f"WebSocket {connection_id} authentication failed") return # Main message loop connection.state = ConnectionState.CONNECTED while connection.state == ConnectionState.CONNECTED and not websocket.closed: try: # Set timeout for receiving messages message = await asyncio.wait_for(websocket.recv(), timeout=30.0) # Update activity timestamp connection.last_activity = datetime.now() # Parse and handle message try: data = json.loads(message) await self.connection_handlers["on_message"](connection, data) self.metrics["events_received"] += 1 except json.JSONDecodeError: logger.error(f"Invalid JSON from connection {connection_id}") await self._send_error(connection_id, "Invalid message format") except asyncio.TimeoutError: # Check if connection is still alive with ping try: await websocket.ping() connection.last_activity = datetime.now() except: break # Connection is dead except websockets.exceptions.ConnectionClosed: break except Exception as e: logger.error(f"Error handling message from {connection_id}: {str(e)}") await self._send_error(connection_id, f"Message handling error: {str(e)}") except Exception as e: logger.error(f"Error in WebSocket connection {connection_id}: {str(e)}") await self.connection_handlers["on_error"](connection, e) finally: # Cleanup connection await self._cleanup_connection(connection_id) logger.info(f"WebSocket connection {connection_id} closed") async def _cleanup_connection(self, connection_id: str): """Clean up closed connection""" try: if connection_id not in self.connections: return connection = self.connections[connection_id] connection.state = ConnectionState.DISCONNECTED # Remove from user and session mappings if connection.user_id: self.user_connections[connection.user_id].discard(connection_id) if not self.user_connections[connection.user_id]: del self.user_connections[connection.user_id] if connection.session_id: self.session_connections[connection.session_id].discard(connection_id) if not self.session_connections[connection.session_id]: del self.session_connections[connection.session_id] # Remove from subscriptions for subscription in connection.subscriptions: self.subscriptions[subscription].discard(connection_id) if not self.subscriptions[subscription]: del self.subscriptions[subscription] # Remove from connections del self.connections[connection_id] self.metrics["active_connections"] = len(self.connections) # Call disconnect handler await self.connection_handlers["on_disconnect"](connection) except Exception as e: logger.error(f"Error cleaning up connection {connection_id}: {str(e)}") async def _send_to_connection(self, connection_id: str, data: Dict[str, Any]): """Send message to specific connection""" try: if connection_id not in self.connections: logger.warning(f"Attempted to send to unknown connection: {connection_id}") return False connection = self.connections[connection_id] if connection.websocket.closed: return False message = json.dumps(data) await connection.websocket.send(message) self.metrics["events_sent"] += 1 return True except Exception as e: logger.error(f"Error sending to connection {connection_id}: {str(e)}") return False async def _send_error(self, connection_id: str, error_message: str): """Send error message to connection""" error_data = { "type": "error", "error": error_message, "timestamp": datetime.now().isoformat() } await self._send_to_connection(connection_id, error_data) # Connection Handlers async def _on_connect(self, connection: WebSocketConnection): """Handle new connection""" try: logger.info(f"Connection established: {connection.connection_id} for user {connection.user_id}") # Send initial status await self._send_to_connection(connection.connection_id, { "type": "connection_established", "connection_id": connection.connection_id, "timestamp": datetime.now().isoformat(), "features": [ "workflow_updates", "execution_status", "collaboration", "notifications" ] }) # Subscribe user to default channels await self._subscribe_to_channel(connection.connection_id, f"user:{connection.user_id}") if connection.session_id: await self._subscribe_to_channel(connection.connection_id, f"session:{connection.session_id}") except Exception as e: logger.error(f"Error in connect handler: {str(e)}") async def _on_disconnect(self, connection: WebSocketConnection): """Handle connection disconnection""" try: logger.info(f"Connection disconnected: {connection.connection_id}") # Broadcast user activity await self._broadcast_event(WebSocketEvent( event_id=str(uuid.uuid4()), event_type=WebSocketEventType.USER_ACTIVITY, payload={ "user_id": connection.user_id, "activity": "disconnected", "timestamp": datetime.now().isoformat() }, user_id=connection.user_id, target_connections=list(self.user_connections.get(connection.user_id, [])), broadcast=False )) except Exception as e: logger.error(f"Error in disconnect handler: {str(e)}") async def _on_message(self, connection: WebSocketConnection, data: Dict[str, Any]): """Handle incoming message""" try: message_type = data.get("type") if message_type == "subscribe": # Subscribe to channel channel = data.get("channel") if channel: await self._subscribe_to_channel(connection.connection_id, channel) elif message_type == "unsubscribe": # Unsubscribe from channel channel = data.get("channel") if channel: await self._unsubscribe_from_channel(connection.connection_id, channel) elif message_type == "workflow_command": # Handle workflow command await self._handle_workflow_command(connection, data) elif message_type == "collaboration": # Handle collaboration event await self._handle_collaboration_message(connection, data) elif message_type == "ping": # Respond to ping with pong await self._send_to_connection(connection.connection_id, { "type": "pong", "timestamp": datetime.now().isoformat() }) else: logger.warning(f"Unknown message type: {message_type}") except Exception as e: logger.error(f"Error handling message: {str(e)}") async def _on_error(self, connection: WebSocketConnection, error: Exception): """Handle connection error""" try: logger.error(f"Connection error for {connection.connection_id}: {str(error)}") self.metrics["errors"] += 1 except Exception as e: logger.error(f"Error in error handler: {str(e)}") # Event Handlers async def _handle_workflow_update(self, event: WebSocketEvent): """Handle workflow update event""" try: # Broadcast to relevant users if event.user_id: target_connections = list(self.user_connections.get(event.user_id, [])) else: target_connections = list(self.connections.keys()) for connection_id in target_connections: await self._send_to_connection(connection_id, { "type": "workflow_update", "event_id": event.event_id, "payload": event.payload, "timestamp": event.timestamp.isoformat() }) except Exception as e: logger.error(f"Error handling workflow update: {str(e)}") async def _handle_execution_status(self, event: WebSocketEvent): """Handle execution status event""" try: # Send to user who started the execution user_id = event.payload.get("user_id") if user_id: target_connections = list(self.user_connections.get(user_id, [])) for connection_id in target_connections: await self._send_to_connection(connection_id, { "type": "execution_status", "event_id": event.event_id, "payload": event.payload, "timestamp": event.timestamp.isoformat() }) except Exception as e: logger.error(f"Error handling execution status: {str(e)}") async def _handle_service_status(self, event: WebSocketEvent): """Handle service status event""" try: # Broadcast to all connected users for connection_id in self.connections.keys(): await self._send_to_connection(connection_id, { "type": "service_status", "event_id": event.event_id, "payload": event.payload, "timestamp": event.timestamp.isoformat() }) except Exception as e: logger.error(f"Error handling service status: {str(e)}") async def _handle_notification(self, event: WebSocketEvent): """Handle notification event""" try: # Send to specific user or session if event.user_id: target_connections = list(self.user_connections.get(event.user_id, [])) elif event.session_id: target_connections = list(self.session_connections.get(event.session_id, [])) else: target_connections = list(self.connections.keys()) for connection_id in target_connections: await self._send_to_connection(connection_id, { "type": "notification", "event_id": event.event_id, "payload": event.payload, "timestamp": event.timestamp.isoformat() }) except Exception as e: logger.error(f"Error handling notification: {str(e)}") async def _handle_collaboration(self, event: WebSocketEvent): """Handle collaboration event""" try: # Send to all users in the session if event.session_id: target_connections = list(self.session_connections.get(event.session_id, [])) for connection_id in target_connections: await self._send_to_connection(connection_id, { "type": "collaboration", "event_id": event.event_id, "payload": event.payload, "timestamp": event.timestamp.isoformat() }) except Exception as e: logger.error(f"Error handling collaboration: {str(e)}") async def _handle_system_update(self, event: WebSocketEvent): """Handle system update event""" try: # Broadcast to all connections for connection_id in self.connections.keys(): await self._send_to_connection(connection_id, { "type": "system_update", "event_id": event.event_id, "payload": event.payload, "timestamp": event.timestamp.isoformat() }) except Exception as e: logger.error(f"Error handling system update: {str(e)}") async def _handle_user_activity(self, event: WebSocketEvent): """Handle user activity event""" try: # Send to other users in the same session if event.session_id: target_connections = list(self.session_connections.get(event.session_id, [])) for connection_id in target_connections: # Don't send back to the same user connection = self.connections.get(connection_id) if connection and connection.user_id != event.user_id: await self._send_to_connection(connection_id, { "type": "user_activity", "event_id": event.event_id, "payload": event.payload, "timestamp": event.timestamp.isoformat() }) except Exception as e: logger.error(f"Error handling user activity: {str(e)}") async def _handle_error(self, event: WebSocketEvent): """Handle error event""" try: # Send error to specific user if available if event.user_id: target_connections = list(self.user_connections.get(event.user_id, [])) for connection_id in target_connections: await self._send_to_connection(connection_id, { "type": "error", "event_id": event.event_id, "payload": event.payload, "timestamp": event.timestamp.isoformat() }) except Exception as e: logger.error(f"Error handling error event: {str(e)}") async def _handle_heartbeat(self, event: WebSocketEvent): """Handle heartbeat event""" try: # Respond with heartbeat to keep connection alive if event.user_id: target_connections = list(self.user_connections.get(event.user_id, [])) for connection_id in target_connections: await self._send_to_connection(connection_id, { "type": "heartbeat_response", "event_id": event.event_id, "timestamp": event.timestamp.isoformat() }) except Exception as e: logger.error(f"Error handling heartbeat: {str(e)}") # Helper Methods async def _subscribe_to_channel(self, connection_id: str, channel: str): """Subscribe connection to channel""" try: if connection_id in self.connections: connection = self.connections[connection_id] connection.subscriptions.add(channel) self.subscriptions[channel].add(connection_id) await self._send_to_connection(connection_id, { "type": "subscription_confirmed", "channel": channel, "timestamp": datetime.now().isoformat() }) logger.info(f"Connection {connection_id} subscribed to {channel}") except Exception as e: logger.error(f"Error subscribing to channel: {str(e)}") async def _unsubscribe_from_channel(self, connection_id: str, channel: str): """Unsubscribe connection from channel""" try: if connection_id in self.connections: connection = self.connections[connection_id] connection.subscriptions.discard(channel) self.subscriptions[channel].discard(connection_id) if not self.subscriptions[channel]: del self.subscriptions[channel] await self._send_to_connection(connection_id, { "type": "unsubscription_confirmed", "channel": channel, "timestamp": datetime.now().isoformat() }) logger.info(f"Connection {connection_id} unsubscribed from {channel}") except Exception as e: logger.error(f"Error unsubscribing from channel: {str(e)}") async def _handle_workflow_command(self, connection: WebSocketConnection, data: Dict[str, Any]): """Handle workflow command from client""" try: command = data.get("command") workflow_data = data.get("data", {}) if command == "create_workflow": # Create new workflow result = working_enhanced_workflow_engine.create_workflow_from_template( template_id=workflow_data.get("template_id"), parameters=workflow_data.get("parameters", {}), user_id=connection.user_id ) await self._send_to_connection(connection.connection_id, { "type": "workflow_command_response", "command": command, "success": result.get("success", False), "data": result, "timestamp": datetime.now().isoformat() }) elif command == "execute_workflow": # Execute workflow result = working_enhanced_workflow_engine.execute_workflow( workflow_id=workflow_data.get("workflow_id"), input_data=workflow_data.get("input_data", {}) ) if result.get("success"): # Start monitoring execution status execution_id = result.get("execution_id") # Broadcast execution start await self._broadcast_event(WebSocketEvent( event_id=str(uuid.uuid4()), event_type=WebSocketEventType.EXECUTION_STATUS, payload={ "execution_id": execution_id, "status": "started", "workflow_id": workflow_data.get("workflow_id"), "user_id": connection.user_id }, user_id=connection.user_id )) await self._send_to_connection(connection.connection_id, { "type": "workflow_command_response", "command": command, "success": result.get("success", False), "data": result, "timestamp": datetime.now().isoformat() }) elif command == "get_execution_status": # Get execution status execution_id = workflow_data.get("execution_id") result = working_enhanced_workflow_engine.get_execution_status(execution_id) await self._send_to_connection(connection.connection_id, { "type": "workflow_command_response", "command": command, "success": result.get("success", False), "data": result, "timestamp": datetime.now().isoformat() }) else: await self._send_error(connection.connection_id, f"Unknown command: {command}") except Exception as e: logger.error(f"Error handling workflow command: {str(e)}") await self._send_error(connection.connection_id, f"Command error: {str(e)}") async def _handle_collaboration_message(self, connection: WebSocketConnection, data: Dict[str, Any]): """Handle collaboration message""" try: collaboration_type = data.get("collaboration_type") message_data = data.get("data", {}) # Create collaboration event event = WebSocketEvent( event_id=str(uuid.uuid4()), event_type=WebSocketEventType.COLLABORATION, payload={ "collaboration_type": collaboration_type, "data": message_data, "user_id": connection.user_id, "session_id": connection.session_id }, user_id=connection.user_id, session_id=connection.session_id ) await self.event_handlers[WebSocketEventType.COLLABORATION](event) except Exception as e: logger.error(f"Error handling collaboration message: {str(e)}") async def _broadcast_event(self, event: WebSocketEvent): """Broadcast event to appropriate connections""" try: handler = self.event_handlers.get(event.event_type) if handler: await handler(event) else: logger.warning(f"No handler for event type: {event.event_type}") except Exception as e: logger.error(f"Error broadcasting event: {str(e)}") # Background Tasks async def _heartbeat_monitor(self): """Monitor connection heartbeats""" while self.running: try: await asyncio.sleep(60) # Check every minute current_time = datetime.now() dead_connections = [] for connection_id, connection in self.connections.items(): # Check if connection is dead (no activity for 5 minutes) if (current_time - connection.last_activity).total_seconds() > 300: dead_connections.append(connection_id) # Clean up dead connections for connection_id in dead_connections: logger.warning(f"Cleaning up dead connection: {connection_id}") await self._cleanup_connection(connection_id) except Exception as e: logger.error(f"Error in heartbeat monitor: {str(e)}") async def _cleanup_connections(self): """Periodic cleanup of connections""" while self.running: try: await asyncio.sleep(300) # Every 5 minutes # Clean up empty subscription mappings empty_subscriptions = [ sub for sub, conns in self.subscriptions.items() if not conns ] for sub in empty_subscriptions: del self.subscriptions[sub] logger.debug(f"Cleaned up {len(empty_subscriptions)} empty subscriptions") except Exception as e: logger.error(f"Error in cleanup task: {str(e)}") async def _performance_monitor(self): """Monitor server performance""" while self.running: try: await asyncio.sleep(300) # Every 5 minutes if self.metrics["start_time"]: uptime = datetime.now() - self.metrics["start_time"] performance_report = { "uptime_seconds": uptime.total_seconds(), "total_connections": self.metrics["total_connections"], "active_connections": len(self.connections), "events_sent": self.metrics["events_sent"], "events_received": self.metrics["events_received"], "errors": self.metrics["errors"], "avg_events_per_minute": ( self.metrics["events_sent"] / max(uptime.total_seconds() / 60, 1) ) } logger.info(f"WebSocket Server Performance: {performance_report}") # Broadcast performance to admin users await self._broadcast_event(WebSocketEvent( event_id=str(uuid.uuid4()), event_type=WebSocketEventType.SYSTEM_UPDATE, payload={ "type": "performance_report", "data": performance_report }, broadcast=True )) except Exception as e: logger.error(f"Error in performance monitor: {str(e)}") # Public API Methods async def send_notification(self, user_id: str, message: str, notification_type: str = "info"): """Send notification to specific user""" try: event = WebSocketEvent( event_id=str(uuid.uuid4()), event_type=WebSocketEventType.NOTIFICATION, payload={ "message": message, "type": notification_type, "user_id": user_id }, user_id=user_id ) await self._broadcast_event(event) except Exception as e: logger.error(f"Error sending notification: {str(e)}") async def broadcast_workflow_update(self, user_id: str, workflow_id: str, update_data: Dict[str, Any]): """Broadcast workflow update to user""" try: event = WebSocketEvent( event_id=str(uuid.uuid4()), event_type=WebSocketEventType.WORKFLOW_UPDATE, payload={ "workflow_id": workflow_id, "user_id": user_id, "update": update_data }, user_id=user_id ) await self._broadcast_event(event) except Exception as e: logger.error(f"Error broadcasting workflow update: {str(e)}") async def broadcast_execution_status(self, user_id: str, execution_id: str, status: str, details: Dict[str, Any] = None): """Broadcast execution status update""" try: event = WebSocketEvent( event_id=str(uuid.uuid4()), event_type=WebSocketEventType.EXECUTION_STATUS, payload={ "execution_id": execution_id, "status": status, "user_id": user_id, "details": details or {} }, user_id=user_id ) await self._broadcast_event(event) except Exception as e: logger.error(f"Error broadcasting execution status: {str(e)}") def get_metrics(self) -> Dict[str, Any]: """Get server performance metrics""" uptime = None if self.metrics["start_time"]: uptime = datetime.now() - self.metrics["start_time"] return { "server_running": self.running, "uptime_seconds": uptime.total_seconds() if uptime else 0, "total_connections": self.metrics["total_connections"], "active_connections": len(self.connections), "events_sent": self.metrics["events_sent"], "events_received": self.metrics["events_received"], "errors": self.metrics["errors"], "subscriptions": len(self.subscriptions), "avg_events_per_minute": ( self.metrics["events_sent"] / max(uptime.total_seconds() / 60, 1) if uptime else 0 ), "memory_usage": len(self.connections) * 1024, # Rough estimate "host": self.host, "port": self.port } # Global WebSocket server instance websocket_server = RealTimeWebSocketServer() logger.info("Real-Time WebSocket Server initialized") def start_websocket_server(): """Start WebSocket server in background""" try: # Start server loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) # Run server loop.run_until_complete(websocket_server.start_server()) loop.run_forever() except Exception as e: logger.error(f"Error starting WebSocket server: {str(e)}") # Start WebSocket server in background thread websocket_thread = threading.Thread(target=start_websocket_server, daemon=True) websocket_thread.start() logger.info("WebSocket server background thread started")