86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
"""WebSocket connection manager for real-time map updates."""
|
|
from fastapi import WebSocket
|
|
from typing import Dict, List, Set
|
|
from uuid import UUID
|
|
import json
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Manages WebSocket connections for real-time map updates."""
|
|
|
|
def __init__(self):
|
|
# map_id -> Set of WebSocket connections
|
|
self.active_connections: Dict[str, Set[WebSocket]] = {}
|
|
|
|
async def connect(self, websocket: WebSocket, map_id: UUID):
|
|
"""Accept a new WebSocket connection for a map."""
|
|
await websocket.accept()
|
|
map_key = str(map_id)
|
|
|
|
if map_key not in self.active_connections:
|
|
self.active_connections[map_key] = set()
|
|
|
|
self.active_connections[map_key].add(websocket)
|
|
logger.info(f"Client connected to map {map_id}. Total connections: {len(self.active_connections[map_key])}")
|
|
|
|
def disconnect(self, websocket: WebSocket, map_id: UUID):
|
|
"""Remove a WebSocket connection."""
|
|
map_key = str(map_id)
|
|
|
|
if map_key in self.active_connections:
|
|
self.active_connections[map_key].discard(websocket)
|
|
if not self.active_connections[map_key]:
|
|
del self.active_connections[map_key]
|
|
|
|
logger.info(f"Client disconnected from map {map_id}")
|
|
|
|
async def broadcast_to_map(self, map_id: UUID, message: dict):
|
|
"""Broadcast a message to all clients connected to a specific map."""
|
|
map_key = str(map_id)
|
|
|
|
if map_key not in self.active_connections:
|
|
return
|
|
|
|
# Create a copy of the set to avoid modification during iteration
|
|
connections = self.active_connections[map_key].copy()
|
|
disconnected = []
|
|
|
|
for connection in connections:
|
|
try:
|
|
await connection.send_json(message)
|
|
except Exception as e:
|
|
logger.error(f"Error sending message to client: {e}")
|
|
disconnected.append(connection)
|
|
|
|
# Remove disconnected clients
|
|
for connection in disconnected:
|
|
self.disconnect(connection, map_id)
|
|
|
|
async def send_item_created(self, map_id: UUID, item_data: dict):
|
|
"""Notify clients that a new item was created."""
|
|
await self.broadcast_to_map(map_id, {
|
|
"type": "item_created",
|
|
"data": item_data
|
|
})
|
|
|
|
async def send_item_updated(self, map_id: UUID, item_data: dict):
|
|
"""Notify clients that an item was updated."""
|
|
await self.broadcast_to_map(map_id, {
|
|
"type": "item_updated",
|
|
"data": item_data
|
|
})
|
|
|
|
async def send_item_deleted(self, map_id: UUID, item_id: str):
|
|
"""Notify clients that an item was deleted."""
|
|
await self.broadcast_to_map(map_id, {
|
|
"type": "item_deleted",
|
|
"data": {"id": item_id}
|
|
})
|
|
|
|
|
|
# Global connection manager instance
|
|
manager = ConnectionManager()
|