145 lines
4.5 KiB
Python
145 lines
4.5 KiB
Python
"""
|
|
ConnectionManager — unified WebSocket manager for all real-time features.
|
|
|
|
Currently handles two types of connectivity:
|
|
1. Room-based (Chat): file-based collaboration where multiple users share a room.
|
|
2. Direct/Task-based: 1:1 connections for tracking specific backend task progress.
|
|
|
|
This replaces the old infrastructure/websocket/ws_manager.py for a
|
|
unified, production-ready implementation.
|
|
"""
|
|
import logging
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from typing import Any, Optional, Set, Dict
|
|
|
|
from fastapi import WebSocket
|
|
from starlette.websockets import WebSocketState
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(eq=False)
|
|
class Connection:
|
|
"""Metadata associated with a single WebSocket connection."""
|
|
websocket: WebSocket
|
|
room_id: str
|
|
user_id: Optional[int] = None
|
|
conn_type: str = 'chat'
|
|
|
|
|
|
class ConnectionManager:
|
|
"""
|
|
Manages WebSocket connections partitioned by room_id.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._rooms: Dict[str, Set[Connection]] = defaultdict(set)
|
|
self._ws_index: Dict[int, Connection] = {}
|
|
|
|
|
|
async def connect(
|
|
self,
|
|
room_id: str,
|
|
websocket: WebSocket,
|
|
user_id: Optional[int] = None,
|
|
conn_type: str = 'chat'
|
|
) -> Connection:
|
|
"""Accept the WebSocket and register it in the specified room."""
|
|
await websocket.accept()
|
|
|
|
conn = Connection(
|
|
websocket=websocket,
|
|
room_id=room_id,
|
|
user_id=user_id,
|
|
conn_type=conn_type
|
|
)
|
|
self._rooms[room_id].add(conn)
|
|
self._ws_index[id(websocket)] = conn
|
|
|
|
logger.info(
|
|
f"WS connected: type={conn_type} room={room_id} user={user_id} "
|
|
f"(room size: {len(self._rooms[room_id])})"
|
|
)
|
|
return conn
|
|
|
|
async def connect_chat(self, user_id: int, file_id: int, websocket: WebSocket) -> Connection:
|
|
"""Helper to connect to a chat room."""
|
|
return await self.connect(
|
|
room_id=f"chat:{file_id}",
|
|
websocket=websocket,
|
|
user_id=user_id,
|
|
conn_type='chat'
|
|
)
|
|
|
|
async def connect_task(self, task_id: str, websocket: WebSocket) -> Connection:
|
|
"""Helper to connect to a task status channel."""
|
|
return await self.connect(
|
|
room_id=f"task:{task_id}",
|
|
websocket=websocket,
|
|
conn_type='task'
|
|
)
|
|
|
|
def disconnect(self, websocket: WebSocket) -> None:
|
|
"""Remove a WebSocket from its room."""
|
|
conn = self._ws_index.pop(id(websocket), None)
|
|
if conn is None:
|
|
return
|
|
|
|
room = self._rooms.get(conn.room_id)
|
|
if room:
|
|
room.discard(conn)
|
|
if not room:
|
|
del self._rooms[conn.room_id]
|
|
|
|
logger.info(
|
|
f"WS disconnected: room={conn.room_id} user={conn.user_id} "
|
|
f"(room size: {len(self._rooms.get(conn.room_id, []))})"
|
|
)
|
|
|
|
async def broadcast(self, room_id: str, message: Dict[str, Any]) -> None:
|
|
"""Send a JSON payload to every connection in the room."""
|
|
room = self._rooms.get(room_id)
|
|
if not room:
|
|
return
|
|
|
|
stale: list[Connection] = []
|
|
|
|
for conn in room:
|
|
try:
|
|
if conn.websocket.client_state == WebSocketState.CONNECTED:
|
|
await conn.websocket.send_json(message)
|
|
else:
|
|
stale.append(conn)
|
|
except Exception:
|
|
stale.append(conn)
|
|
|
|
for conn in stale:
|
|
self.disconnect(conn.websocket)
|
|
|
|
async def broadcast_chat(self, file_id: int, data: Dict[str, Any]) -> None:
|
|
"""Legacy helper for broadcast on chat:file:{file_id}."""
|
|
await self.broadcast(f"chat:{file_id}", data)
|
|
|
|
async def send_task(self, task_id: str, data: Dict[str, Any]) -> None:
|
|
"""Send message to a specific task status channel."""
|
|
await self.broadcast(f"task:{task_id}", data)
|
|
|
|
|
|
def get_room_users(self, file_id: int) -> list[int]:
|
|
"""Return a list of user IDs currently in a file's chat room."""
|
|
room_id = f"chat:{file_id}"
|
|
return list({conn.user_id for conn in self._rooms.get(room_id, set()) if conn.user_id})
|
|
|
|
def get_room_size(self, room_id: str) -> int:
|
|
"""Number of active connections in a specific room."""
|
|
return len(self._rooms.get(room_id, set()))
|
|
|
|
@property
|
|
def total_connections(self) -> int:
|
|
"""Total number of active WebSocket connections across all rooms."""
|
|
return len(self._ws_index)
|
|
|
|
|
|
manager = ConnectionManager()
|