core/websocket_manager.py

"""
WebSocket connection manager
"""
from fastapi import WebSocket

class WebSocketManager:
    def __init__(self):
        self.connections = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        if websocket in self.connections:
            self.connections.remove(websocket)

    async def broadcast(self, message: dict):
        for conn in list(self.connections):
            try:
                await conn.send_json(message)
            except:
                self.disconnect(conn)