牧安云哨-后端

bbit_ai/app/db/postgres/ws_manager.py

@@ -0,0 +1,25 @@
+import asyncio
+from typing import List
+
+from fastapi import WebSocket
+
+
+class ConnectionManager:
+    def __init__(self):
+        self.active_connections: List[WebSocket] = []
+        self.lock = asyncio.Lock()
+
+    async def connect(self, websocket: WebSocket):
+        await websocket.accept()
+        async with self.lock:
+            self.active_connections.append(websocket)
+
+    async def disconnect(self, websocket: WebSocket):
+        async with self.lock:
+            if websocket in self.active_connections:
+                self.active_connections.remove(websocket)
+
+    async def broadcast(self, message: dict):
+        async with self.lock:
+            for ws in self.active_connections:
+                await ws.send_json(message)