mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
104 lines
3.0 KiB
Python
104 lines
3.0 KiB
Python
"""
|
|
WebSocket connection manager for the AstrBot gateway.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Any, cast
|
|
|
|
import anyio
|
|
|
|
from astrbot import logger
|
|
|
|
if TYPE_CHECKING:
|
|
from fastapi import WebSocket
|
|
else:
|
|
try:
|
|
from fastapi import WebSocket
|
|
except ImportError:
|
|
logger.warning("FastAPI not installed, WebSocketManager unavailable.")
|
|
WebSocket = cast(Any, None)
|
|
|
|
log = logger
|
|
|
|
|
|
class WebSocketManager:
|
|
"""
|
|
Manages all active WebSocket connections.
|
|
|
|
Provides connection/disconnection handling and broadcast capabilities.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._connections: set[WebSocket] = set()
|
|
self._lock = anyio.Lock()
|
|
|
|
async def connect(self, websocket: WebSocket) -> None:
|
|
"""Accept and register a new WebSocket connection."""
|
|
await websocket.accept()
|
|
async with self._lock:
|
|
self._connections.add(websocket)
|
|
log.debug(f"WebSocket connected. Total: {len(self._connections)}")
|
|
|
|
async def disconnect(self, websocket: WebSocket) -> None:
|
|
"""Remove a WebSocket connection."""
|
|
async with self._lock:
|
|
self._connections.discard(websocket)
|
|
log.debug(f"WebSocket disconnected. Total: {len(self._connections)}")
|
|
|
|
async def send_json(self, websocket: WebSocket, data: dict[str, Any]) -> None:
|
|
"""
|
|
Send JSON data to a specific WebSocket.
|
|
|
|
Args:
|
|
websocket: Target WebSocket connection
|
|
data: Data to send (must be JSON-serializable)
|
|
"""
|
|
try:
|
|
await websocket.send_json(data)
|
|
except Exception as e:
|
|
log.warning(f"Failed to send to WebSocket: {e}")
|
|
await self.disconnect(websocket)
|
|
|
|
async def broadcast(self, data: dict[str, Any]) -> None:
|
|
"""
|
|
Broadcast JSON data to all connected WebSockets.
|
|
|
|
Args:
|
|
data: Data to broadcast (must be JSON-serializable)
|
|
"""
|
|
async with self._lock:
|
|
connections = list(self._connections)
|
|
|
|
for conn in connections:
|
|
try:
|
|
await conn.send_json(data)
|
|
except Exception as e:
|
|
log.warning(f"Failed to broadcast to WebSocket: {e}")
|
|
async with self._lock:
|
|
self._connections.discard(conn)
|
|
|
|
async def send_to(
|
|
self, websocket: WebSocket, message: str | dict[str, Any]
|
|
) -> None:
|
|
"""
|
|
Send a message to a specific WebSocket.
|
|
|
|
Args:
|
|
websocket: Target WebSocket connection
|
|
message: Message to send (string or dict)
|
|
"""
|
|
try:
|
|
if isinstance(message, str):
|
|
await websocket.send_text(message)
|
|
else:
|
|
await websocket.send_json(message)
|
|
except Exception as e:
|
|
log.warning(f"Failed to send to WebSocket: {e}")
|
|
await self.disconnect(websocket)
|
|
|
|
@property
|
|
def connection_count(self) -> int:
|
|
"""Return the number of active connections."""
|
|
return len(self._connections)
|