"""WebSocket connection manager for AI UI Control. Manages WebSocket connections from frontend clients. When an AI agent sends a UI command via REST API, the command is forwarded to the frontend via these WebSocket connections. The frontend executes the command and sends feedback back through the same WebSocket. """ from __future__ import annotations import json import logging import time import uuid from typing import Any from fastapi import WebSocket from sqlalchemy.ext.asyncio import AsyncSession from app.core.ws_helpers import ( authenticate_ws, check_ws_origin, check_ws_tenant, cleanup_ws_connection, start_heartbeat, ) from app.core.ws_pubsub import ( broadcast_to_tenants, get_tenant_channel, subscribe_to_channel, ) logger = logging.getLogger(__name__) class AIUIControlWSManager: """Manages WebSocket connections for AI UI control. Connections are per-user: each authenticated user can have one or more frontend tabs connected. Commands are delivered to all tabs of the target user. Feedback from any tab is accepted and stored. """ def __init__(self) -> None: # user_id → list of WebSocket connections self._connections: dict[str, list[WebSocket]] = {} # command_id → feedback dict (stored when frontend responds) self._feedback: dict[str, dict[str, Any]] = {} # command_id → timestamp when delivered (for timeout tracking) self._delivered_at: dict[str, float] = {} # command_id → user_id (to route feedback) self._command_user: dict[str, str] = {} # user_id → list of heartbeat tasks self._heartbeat_tasks: dict[str, list] = {} # user_id → list of pubsub subscriber tasks self._pubsub_tasks: dict[str, list] = {} async def connect( self, websocket: WebSocket, db: AsyncSession, ) -> dict[str, Any] | None: """Accept, authenticate and register a new WebSocket connection. Performs origin check, session authentication and tenant validation. Returns the auth dict (user_id, tenant_id, …) on success, or ``None`` if the connection was rejected (already closed). """ # Origin / CSRF check if not await check_ws_origin(websocket): return None # Session authentication auth = await authenticate_ws(websocket, db) if auth is None: return None user_id = auth["user_id"] tenant_id_str = auth["tenant_id"] tenant_id = uuid.UUID(tenant_id_str) user_uuid = uuid.UUID(user_id) # Tenant membership check if not await check_ws_tenant(websocket, tenant_id, user_uuid, db): return None # Accept the WebSocket await websocket.accept() # Register connection if user_id not in self._connections: self._connections[user_id] = [] self._connections[user_id].append(websocket) # Start heartbeat hb_task = await start_heartbeat(websocket) self._heartbeat_tasks.setdefault(user_id, []).append(hb_task) # Start Redis Pub/Sub subscriber for tenant-wide fanout channel = get_tenant_channel(tenant_id, "ai_ui_control") pubsub_task = await subscribe_to_channel(channel, lambda msg: self._on_pubsub_message(user_id, msg)) self._pubsub_tasks.setdefault(user_id, []).append(pubsub_task) logger.debug(f"AI UI Control WS connected: user={user_id}, total={len(self._connections[user_id])}") return auth async def _on_pubsub_message(self, user_id: str, msg: dict[str, Any]) -> None: """Handle a message received via Redis Pub/Sub.""" # Forward pubsub messages to the user's connections conns = self._connections.get(user_id, []) text = json.dumps(msg, default=str) for ws in conns: try: await ws.send_text(text) except Exception: logger.warning(f"Failed to send pubsub message to user {user_id}") async def disconnect(self, websocket: WebSocket, user_id: str) -> None: """Remove a WebSocket connection and clean up resources.""" # Cancel heartbeat tasks for this user for task in self._heartbeat_tasks.pop(user_id, []): task.cancel() # Cancel pubsub tasks for this user for task in self._pubsub_tasks.pop(user_id, []): task.cancel() # Clean up connection registry await cleanup_ws_connection(websocket, user_id, self._connections) logger.debug(f"AI UI Control WS disconnected: user={user_id}") async def send_command(self, user_id: str, command: dict[str, Any]) -> str | None: """Send a UI command to all frontend connections of a user. Returns the command_id if delivered, None if user has no active connections. """ command_id = command.get("command_id") or str(uuid.uuid4()) command["command_id"] = command_id command.setdefault("timestamp", time.time()) conns = self._connections.get(user_id, []) if not conns: logger.debug(f"AI UI Control: no active connections for user {user_id}") return None self._command_user[command_id] = user_id self._delivered_at[command_id] = time.time() text = json.dumps(command, default=str) delivered = False for ws in conns: try: await ws.send_text(text) delivered = True except Exception: logger.warning(f"Failed to send command to user {user_id}, removing connection") await self.disconnect(ws, user_id) if delivered: # Initialize as pending feedback self._feedback.setdefault(command_id, { "command_id": command_id, "status": "delivered", "action": command.get("action"), }) return command_id return None def store_feedback(self, feedback: dict[str, Any]) -> None: """Store feedback from frontend after command execution. Maintains a maximum of 100 feedback entries to prevent memory exhaustion. Oldest entries are removed when the limit is reached. """ command_id = feedback.get("command_id") if command_id: self._feedback[command_id] = feedback # Enforce max feedback entries (FIFO eviction) max_feedback_entries = 100 if len(self._feedback) > max_feedback_entries: keys_to_remove = list(self._feedback.keys())[:-max_feedback_entries] for key in keys_to_remove: del self._feedback[key] logger.debug(f"AI UI Control: feedback stored for command {command_id}: {feedback.get('status')}") def get_feedback(self, command_id: str) -> dict[str, Any] | None: """Get stored feedback for a command.""" return self._feedback.get(command_id) def is_user_online(self, user_id: str) -> bool: """Check if a user has any active frontend connections.""" return user_id in self._connections and len(self._connections[user_id]) > 0 def get_online_users(self) -> list[str]: """Get list of currently connected user IDs.""" return list(self._connections.keys()) async def broadcast(self, message: dict[str, Any], tenant_id: uuid.UUID | None = None) -> None: """Broadcast a message to all connected users. If *tenant_id* is provided, publishes via Redis Pub/Sub for multi-worker fanout. Otherwise, sends directly to all locally connected users. """ if tenant_id is not None: await broadcast_to_tenants(tenant_id, "ai_ui_control", message) else: for user_id in list(self._connections.keys()): conns = self._connections.get(user_id, []) text = json.dumps(message, default=str) for ws in conns: try: await ws.send_text(text) except Exception: logger.warning(f"Failed to broadcast to user {user_id}") def cleanup_stale(self, timeout_seconds: int = 60) -> None: """Remove stale command tracking entries older than timeout.""" now = time.time() stale_ids = [ cid for cid, ts in self._delivered_at.items() if now - ts > timeout_seconds ] for cid in stale_ids: if cid not in self._feedback or self._feedback[cid].get("status") == "delivered": self._feedback[cid] = { "command_id": cid, "status": "timeout", "action": None, "message": "Command timed out waiting for frontend response", } self._delivered_at.pop(cid, None) self._command_user.pop(cid, None)