Bernt
bae705aa97
- Add NFC ePassport roadmap (ICAO 9303, eIDAS) - Add TensorFlow.js edge face detection (BlazeFace) - Add structured audit logger (GDPR-compliant) - Risk scoring support Part of KYC Apple Native UX v1.1.0
128 lines
4.3 KiB
Python
128 lines
4.3 KiB
Python
"""
|
|
WebSocket support for real-time updates
|
|
"""
|
|
|
|
from typing import Dict, List, Set
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
import json
|
|
import asyncio
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Manage WebSocket connections"""
|
|
|
|
def __init__(self):
|
|
self.active_connections: Dict[str, List[WebSocket]] = {}
|
|
self.user_connections: Dict[str, WebSocket] = {}
|
|
|
|
async def connect(self, websocket: WebSocket, client_id: str, room: str = "default"):
|
|
"""Connect a client"""
|
|
await websocket.accept()
|
|
|
|
if room not in self.active_connections:
|
|
self.active_connections[room] = []
|
|
|
|
self.active_connections[room].append(websocket)
|
|
self.user_connections[client_id] = websocket
|
|
|
|
def disconnect(self, websocket: WebSocket, client_id: str, room: str = "default"):
|
|
"""Disconnect a client"""
|
|
if room in self.active_connections:
|
|
if websocket in self.active_connections[room]:
|
|
self.active_connections[room].remove(websocket)
|
|
|
|
if client_id in self.user_connections:
|
|
del self.user_connections[client_id]
|
|
|
|
async def broadcast(self, message: dict, room: str = "default"):
|
|
"""Broadcast message to all clients in a room"""
|
|
if room not in self.active_connections:
|
|
return
|
|
|
|
disconnected = []
|
|
for connection in self.active_connections[room]:
|
|
try:
|
|
await connection.send_json(message)
|
|
except:
|
|
disconnected.append(connection)
|
|
|
|
# Clean up disconnected clients
|
|
for conn in disconnected:
|
|
if conn in self.active_connections[room]:
|
|
self.active_connections[room].remove(conn)
|
|
|
|
async def send_to_client(self, client_id: str, message: dict):
|
|
"""Send message to specific client"""
|
|
if client_id in self.user_connections:
|
|
await self.user_connections[client_id].send_json(message)
|
|
|
|
async def broadcast_observation(self, observation: dict):
|
|
"""Broadcast new observation"""
|
|
await self.broadcast({
|
|
"type": "new_observation",
|
|
"data": observation
|
|
}, room="observations")
|
|
|
|
async def broadcast_signal(self, signal: dict):
|
|
"""Broadcast new signal"""
|
|
await self.broadcast({
|
|
"type": "new_signal",
|
|
"data": signal
|
|
}, room="signals")
|
|
|
|
async def broadcast_decision(self, decision: dict):
|
|
"""Broadcast new decision"""
|
|
await self.broadcast({
|
|
"type": "new_decision",
|
|
"data": decision
|
|
}, room="decisions")
|
|
|
|
async def broadcast_outcome(self, outcome: dict):
|
|
"""Broadcast new outcome"""
|
|
await self.broadcast({
|
|
"type": "new_outcome",
|
|
"data": outcome
|
|
}, room="outcomes")
|
|
|
|
def get_stats(self) -> dict:
|
|
"""Get connection statistics"""
|
|
return {
|
|
"total_connections": sum(len(conns) for conns in self.active_connections.values()),
|
|
"rooms": {
|
|
room: len(conns)
|
|
for room, conns in self.active_connections.items()
|
|
},
|
|
"unique_clients": len(self.user_connections)
|
|
}
|
|
|
|
|
|
# Global connection manager
|
|
manager = ConnectionManager()
|
|
|
|
|
|
async def websocket_endpoint(websocket: WebSocket, client_id: str, room: str = "default"):
|
|
"""WebSocket endpoint"""
|
|
await manager.connect(websocket, client_id, room)
|
|
|
|
try:
|
|
while True:
|
|
# Receive message from client
|
|
data = await websocket.receive_text()
|
|
message = json.loads(data)
|
|
|
|
# Handle different message types
|
|
if message.get("type") == "subscribe":
|
|
# Client wants to subscribe to a room
|
|
new_room = message.get("room", "default")
|
|
# In production, validate subscription permissions
|
|
pass
|
|
|
|
elif message.get("type") == "ping":
|
|
await websocket.send_json({"type": "pong"})
|
|
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(websocket, client_id, room)
|
|
except Exception as e:
|
|
manager.disconnect(websocket, client_id, room)
|
|
print(f"WebSocket error: {e}")
|