Initial commit - Phase 3/4
🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
71
app/backend/websocket/manager.py
Normal file
71
app/backend/websocket/manager.py
Normal file
@@ -0,0 +1,71 @@
|
||||
"""WebSocket connection manager for real-time notifications."""
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import Set
|
||||
from fastapi import WebSocket
|
||||
|
||||
|
||||
class ConnectionManager:
|
||||
"""Manages WebSocket connections and broadcasts."""
|
||||
|
||||
def __init__(self):
|
||||
self._connections: Set[WebSocket] = set()
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def connect(self, websocket: WebSocket) -> None:
|
||||
"""Accept and track a new WebSocket connection."""
|
||||
await websocket.accept()
|
||||
async with self._lock:
|
||||
self._connections.add(websocket)
|
||||
|
||||
# Send initial connected event
|
||||
await self._send_to_client(websocket, {
|
||||
"type": "event",
|
||||
"event": "connected",
|
||||
"data": {"message": "Connected to Dangerous Pi event stream"},
|
||||
"timestamp": datetime.now(timezone.utc).isoformat()
|
||||
})
|
||||
|
||||
async def disconnect(self, websocket: WebSocket) -> None:
|
||||
"""Remove a disconnected WebSocket."""
|
||||
async with self._lock:
|
||||
self._connections.discard(websocket)
|
||||
|
||||
async def broadcast(self, event_type: str, data: dict) -> None:
|
||||
"""Broadcast an event to all connected clients."""
|
||||
message = {
|
||||
"type": "event",
|
||||
"event": event_type,
|
||||
"data": data,
|
||||
"timestamp": datetime.now(timezone.utc).isoformat()
|
||||
}
|
||||
|
||||
async with self._lock:
|
||||
dead_connections: Set[WebSocket] = set()
|
||||
for connection in self._connections:
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._send_to_client(connection, message),
|
||||
timeout=5.0
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
dead_connections.add(connection)
|
||||
except Exception:
|
||||
dead_connections.add(connection)
|
||||
|
||||
# Remove dead connections
|
||||
self._connections -= dead_connections
|
||||
|
||||
async def _send_to_client(self, websocket: WebSocket, message: dict) -> None:
|
||||
"""Send a message to a specific client."""
|
||||
await websocket.send_json(message)
|
||||
|
||||
@property
|
||||
def connection_count(self) -> int:
|
||||
"""Return number of active connections."""
|
||||
return len(self._connections)
|
||||
|
||||
|
||||
# Global singleton instance
|
||||
ws_manager = ConnectionManager()
|
||||
Reference in New Issue
Block a user