leocrm/app/plugins/builtins/ai_ui_control/websocket_manager.py

"""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 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,
)
from sqlalchemy.ext.asyncio import AsyncSession

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)