from dataclasses import dataclass from typing import Any from fastapi import WebSocket @dataclass class ClientConnection: """A connected browser inside one chat room.""" websocket: WebSocket client_id: str class ConnectionManager: """Keeps all active WebSocket connections grouped by room name.""" def __init__(self): # Example: # {"عمومی": [ClientConnection(...), ClientConnection(...)]} self.active_connections: dict[str, list[ClientConnection]] = {} async def connect(self, websocket: WebSocket, room_id: str, client_id: str): """Accept the WebSocket handshake and remember this user in the room.""" await websocket.accept() self.active_connections.setdefault(room_id, []) self.active_connections[room_id].append( ClientConnection(websocket=websocket, client_id=client_id) ) def disconnect(self, websocket: WebSocket, room_id: str): """Remove a socket from a room and delete the room when it becomes empty.""" if room_id not in self.active_connections: return self.active_connections[room_id] = [ connection for connection in self.active_connections[room_id] if connection.websocket is not websocket ] if not self.active_connections[room_id]: del self.active_connections[room_id] async def send_json(self, websocket: WebSocket, payload: dict[str, Any]): await websocket.send_json(payload) async def broadcast_json(self, payload: dict[str, Any], room_id: str): """Send one JSON message to every connected client in the selected room.""" if room_id not in self.active_connections: return dead_connections: list[WebSocket] = [] # Iterate over a copy so disconnected clients can be removed safely later. for connection in list(self.active_connections[room_id]): try: await connection.websocket.send_json(payload) except Exception: dead_connections.append(connection.websocket) for websocket in dead_connections: self.disconnect(websocket, room_id) def room_users(self, room_id: str) -> list[str]: """Return display names of users currently connected to a room.""" return [ connection.client_id for connection in self.active_connections.get(room_id, []) ] def rooms_overview(self) -> dict[str, Any]: """Build a small debug-friendly snapshot of all active rooms.""" return { "rooms": [ { "room_id": room_id, "online_count": len(connections), "users": [connection.client_id for connection in connections], } for room_id, connections in self.active_connections.items() ] } manager = ConnectionManager()