Nexus/server/api/websocket.py

"""Nexus — WebSocket Alert Push (Two-Layer Architecture)

ADR-010: Memory ConnectionManager + Redis Pub/Sub two-layer architecture.
  - Layer 1: In-memory ConnectionManager tracks all local WebSocket connections
  - Layer 2: Redis Pub/Sub relays messages across multiple uvicorn workers
  - Alert/recovery/system events published to Redis channel `nexus:alerts`
  - Each worker subscribes to `nexus:alerts` and pushes to local clients

ADR-011: WebSocket JWT authentication via query parameter `?token=xxx`

R4: Ping/pong heartbeat every 30s + zombie connection cleanup (60s idle → disconnect)
"""

import asyncio
import json
import logging
import time
from typing import Dict, Set, Optional

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query

from server.config import settings

logger = logging.getLogger("nexus.websocket")

router = APIRouter()

# Redis Pub/Sub channel name
REDIS_CHANNEL = "nexus:alerts"


class ConnectionManager:
    """Layer 1: In-memory WebSocket connection manager.

    Tracks all connected clients for this worker process.
    Handles ping/pong heartbeat and zombie connection cleanup.
    """

    def __init__(self):
        # client_id → {websocket, last_pong_time}
        self._connections: Dict[str, dict] = {}
        self._heartbeat_task: Optional[asyncio.Task] = None

    async def connect(self, client_id: str, websocket: WebSocket):
        """Accept and register a new WebSocket connection"""
        await websocket.accept()
        self._connections[client_id] = {
            "websocket": websocket,
            "last_pong": time.monotonic(),
        }
        logger.info(f"WebSocket connected: {client_id}, total: {len(self._connections)}")

        # Start heartbeat task if first connection
        if len(self._connections) == 1 and self._heartbeat_task is None:
            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop(), name="ws_heartbeat")

    def disconnect(self, client_id: str):
        """Remove a WebSocket connection"""
        if client_id in self._connections:
            del self._connections[client_id]
            logger.info(f"WebSocket disconnected: {client_id}, total: {len(self._connections)}")

        # Stop heartbeat task if no connections
        if not self._connections and self._heartbeat_task:
            self._heartbeat_task.cancel()
            self._heartbeat_task = None

    async def send_to_client(self, client_id: str, message: dict) -> bool:
        """Send JSON message to a specific client. Returns False if send failed."""
        conn = self._connections.get(client_id)
        if not conn:
            return False
        try:
            await conn["websocket"].send_json(message)
            return True
        except Exception:
            self.disconnect(client_id)
            return False

    async def broadcast_local(self, message: dict):
        """Broadcast message to all locally connected clients (Layer 1 only)"""
        stale = []
        for client_id, conn in self._connections.items():
            try:
                await conn["websocket"].send_json(message)
            except Exception:
                stale.append(client_id)
        for client_id in stale:
            self.disconnect(client_id)

    def update_pong(self, client_id: str):
        """Update last pong time for a client (received pong response)"""
        conn = self._connections.get(client_id)
        if conn:
            conn["last_pong"] = time.monotonic()

    @property
    def client_count(self) -> int:
        return len(self._connections)

    async def _heartbeat_loop(self):
        """R4: Send ping every 30s, cleanup connections idle > 60s"""
        while True:
            await asyncio.sleep(30)
            now = time.monotonic()
            stale = []

            for client_id, conn in self._connections.items():
                # Check for zombie connections (no pong for 60s)
                if now - conn["last_pong"] > 60:
                    stale.append(client_id)
                    continue

                # Send ping
                try:
                    await conn["websocket"].send_json({"type": "ping"})
                except Exception:
                    stale.append(client_id)

            # Cleanup stale connections
            for client_id in stale:
                self.disconnect(client_id)
                logger.info(f"WebSocket zombie cleanup: {client_id}")

            if stale:
                logger.info(f"Heartbeat cleanup: removed {len(stale)} stale connections")


# Global connection manager for this worker
manager = ConnectionManager()


# ── Redis Pub/Sub Layer 2 ──

_pubsub_task: Optional[asyncio.Task] = None
_redis_subscriber = None


async def start_redis_subscriber():
    """Start Redis Pub/Sub subscriber (Layer 2) — called during app startup"""
    global _pubsub_task, _redis_subscriber

    try:
        from server.infrastructure.redis.client import get_redis
        redis = get_redis()
        _redis_subscriber = redis.pubsub()
        await _redis_subscriber.subscribe(REDIS_CHANNEL)

        _pubsub_task = asyncio.create_task(_redis_listen_loop(), name="ws_redis_sub")
        logger.info(f"Redis Pub/Sub subscriber started on channel: {REDIS_CHANNEL}")
    except Exception as e:
        logger.error(f"Redis Pub/Sub subscriber failed to start: {e}")
        # Non-fatal — WebSocket still works locally


async def stop_redis_subscriber():
    """Stop Redis Pub/Sub subscriber — called during app shutdown"""
    global _pubsub_task, _redis_subscriber

    if _pubsub_task:
        _pubsub_task.cancel()
        try:
            await _pubsub_task
        except asyncio.CancelledError:
            pass
        _pubsub_task = None

    if _redis_subscriber:
        try:
            await _redis_subscriber.unsubscribe(REDIS_CHANNEL)
            await _redis_subscriber.aclose()
        except Exception:
            pass
        _redis_subscriber = None

    logger.info("Redis Pub/Sub subscriber stopped")


async def _redis_listen_loop():
    """Listen for messages on Redis Pub/Sub channel and broadcast to local clients"""
    while True:
        try:
            message = await _redis_subscriber.get_message(
                ignore_subscribe_messages=True, timeout=1.0
            )
            if message and message.get("type") == "message":
                data = message.get("data")
                if isinstance(data, str):
                    msg = json.loads(data)
                    await manager.broadcast_local(msg)
        except asyncio.CancelledError:
            break
        except Exception as e:
            logger.error(f"Redis Pub/Sub listen error: {e}")
            await asyncio.sleep(1)  # Back off on error


async def _publish_to_redis(message: dict):
    """Publish message to Redis channel (Layer 2) — other workers will receive it"""
    try:
        from server.infrastructure.redis.client import get_redis
        redis = get_redis()
        await redis.publish(REDIS_CHANNEL, json.dumps(message))
    except Exception as e:
        logger.error(f"Redis Pub/Sub publish failed: {e}")


# ── WebSocket Endpoint ──

@router.websocket("/ws/alerts")
async def alert_ws(
    websocket: WebSocket,
    token: Optional[str] = Query(None, description="JWT access token"),
):
    """WebSocket endpoint for real-time alert push to frontend

    ADR-011: JWT authentication via query parameter `?token=xxx`
    R4: Ping/pong heartbeat + zombie connection cleanup
    """
    # ── JWT Authentication (ADR-011) ──
    if not token:
        await websocket.close(code=4001, reason="Missing JWT token")
        return

    admin = await _verify_ws_token(token)
    if not admin:
        await websocket.close(code=4001, reason="Invalid or expired JWT token")
        return

    # ── Accept and register ──
    client_id = f"{admin.id}:{id(websocket)}"
    await manager.connect(client_id, websocket)

    try:
        while True:
            data = await websocket.receive_text()
            if data == "pong":
                manager.update_pong(client_id)
            elif data == "ping":
                await websocket.send_json({"type": "pong"})
    except WebSocketDisconnect:
        manager.disconnect(client_id)
    except Exception as e:
        logger.warning(f"WebSocket error for {client_id}: {e}")
        manager.disconnect(client_id)


async def _verify_ws_token(token: str):
    """Verify JWT token for WebSocket connection (same logic as get_current_admin)"""
    try:
        import jwt as pyjwt
        from server.infrastructure.database.session import AsyncSessionLocal
        from server.infrastructure.database.admin_repo import AdminRepositoryImpl

        payload = pyjwt.decode(token, settings.SECRET_KEY, algorithms=["HS256"],
                                options={"require": ["exp", "sub"]})
        admin_id = payload.get("sub")
        if not admin_id:
            return None

        try:
            admin_id = int(admin_id)
        except (ValueError, TypeError):
            return None

        async with AsyncSessionLocal() as session:
            admin_repo = AdminRepositoryImpl(session)
            admin = await admin_repo.get_by_id(admin_id)
            if admin and admin.is_active:
                return admin
        return None
    except Exception:
        return None


# ── Alert Aggregation / Dedup ──
# Track last alert time per (server_id, alert_type) to avoid Telegram alert storms.
# WebSocket broadcasts are always sent (real-time), but Telegram pushes are deduplicated.
_ALERT_COOLDOWN_SECONDS = 300  # 5 minutes — same server+metric won't re-push Telegram
_last_alert_time: Dict[str, float] = {}  # key="alert:{server_id}:{type}" → monotonic timestamp


def _should_push_telegram(alert_key: str) -> bool:
    """Check if enough time has passed since last Telegram push for this alert key.

    Returns True if we should send the Telegram notification, False if suppressed (cooldown).
    """
    now = time.monotonic()
    last = _last_alert_time.get(alert_key, 0)
    if now - last < _ALERT_COOLDOWN_SECONDS:
        return False  # Suppressed — too soon
    _last_alert_time[alert_key] = now
    return True


# ── Broadcast Functions (called from other modules) ──

async def broadcast_alert(server_id: int, alert_type: str, alert_value: float, server_name: str = ""):
    """Broadcast alert to all WebSocket clients + Telegram

    Called when Agent heartbeat detects CPU/memory/disk > threshold.
    Publishes to Redis Pub/Sub (multi-worker) + Telegram.

    Alert aggregation: WebSocket always broadcasts (real-time),
    but Telegram push is deduplicated (5min cooldown per server+metric).
    """
    msg = {
        "type": "alert",
        "server_id": server_id,
        "server_name": server_name,
        "alert_type": alert_type,
        "alert_value": alert_value,
    }

    # Broadcast to local clients + publish to Redis for other workers
    await manager.broadcast_local(msg)
    await _publish_to_redis(msg)

    # Telegram push — deduplicated per (server_id, alert_type) with cooldown
    alert_key = f"alert:{server_id}:{alert_type}"
    if _should_push_telegram(alert_key):
        from server.infrastructure.telegram import send_telegram_alert
        await send_telegram_alert(server_name or str(server_id), alert_type, alert_value)
    else:
        logger.debug(f"Alert suppressed (cooldown): {alert_key}")


async def broadcast_recovery(server_id: int, metric: str, value: float, server_name: str = ""):
    """Broadcast recovery notification when alert metric returns to normal

    Recovery alerts are always pushed (clears the alert state).
    Also clears the alert cooldown so next alert will be sent immediately.
    """
    msg = {
        "type": "recovery",
        "server_id": server_id,
        "server_name": server_name,
        "metric": metric,
        "value": value,
    }

    await manager.broadcast_local(msg)
    await _publish_to_redis(msg)

    # Recovery always sends Telegram + clears cooldown for next alert cycle
    alert_key = f"alert:{server_id}:{metric}"
    _last_alert_time.pop(alert_key, None)  # Clear cooldown

    from server.infrastructure.telegram import send_telegram_recovery
    await send_telegram_recovery(server_name or str(server_id), metric, value)


async def broadcast_system_event(event_type: str, message: str):
    """Broadcast system-level events (Python restart, Redis reconnect, etc.)"""
    msg = {"type": "system", "event_type": event_type, "message": message}

    await manager.broadcast_local(msg)
    await _publish_to_redis(msg)


# Backward compatibility alias for health.py
connected_clients = []  # Deprecated — use manager.client_count instead