29 lines
866 B
Python
29 lines
866 B
Python
# sanguo_api/ws.py
|
|
"""WS 连接池：task_id → 订阅者集合。简单广播，不做重连/心跳。"""
|
|
from fastapi import WebSocket
|
|
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
self._connections: dict[str, set[WebSocket]] = {}
|
|
|
|
def connect(self, task_id: str, ws: WebSocket):
|
|
self._connections.setdefault(task_id, set()).add(ws)
|
|
|
|
def disconnect(self, task_id: str, ws: WebSocket):
|
|
conns = self._connections.get(task_id)
|
|
if conns:
|
|
conns.discard(ws)
|
|
if not conns:
|
|
del self._connections[task_id]
|
|
|
|
async def broadcast(self, task_id: str, msg: dict):
|
|
for ws in list(self._connections.get(task_id, [])):
|
|
try:
|
|
await ws.send_json(msg)
|
|
except Exception:
|
|
self.disconnect(task_id, ws)
|
|
|
|
|
|
manager = ConnectionManager()
|