backend/app/core/websocket.py

"""
WebSocket broadcasting — single source of truth.
Extracted from _legacy_main.py to break circular imports.

Usage:
    from app.core.websocket import broadcast_alert, broadcast_scan
    await broadcast_alert({"severity": "critical", "title": "..."})
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any

logger = logging.getLogger(__name__)

# In-memory set of active WebSocket connections
_connections: set[Any] = set()


async def broadcast_alert(alert_data: dict[str, Any]) -> None:
    """Broadcast an alert to all connected WebSocket clients."""
    payload = json.dumps({"event": "alert", "data": alert_data})
    dead: list[Any] = []
    for ws in _connections:
        try:
            await ws.send_text(payload)
        except Exception:
            dead.append(ws)
    for ws in dead:
        _connections.discard(ws)
    logger.debug("alert_broadcast", connections=len(_connections), alert=alert_data.get("title", "")[:50])


async def broadcast_scan(scan_data: dict[str, Any]) -> None:
    """Broadcast a scan result to all connected WebSocket clients."""
    payload = json.dumps({"event": "scan", "data": scan_data})
    dead: list[Any] = []
    for ws in _connections:
        try:
            await ws.send_text(payload)
        except Exception:
            dead.append(ws)
    for ws in dead:
        _connections.discard(ws)


async def register_connection(ws: Any) -> None:
    """Register a new WebSocket connection."""
    _connections.add(ws)


async def unregister_connection(ws: Any) -> None:
    """Remove a WebSocket connection."""
    _connections.discard(ws)


def active_connections() -> int:
    """Return count of active WebSocket connections."""
    return len(_connections)