90 lines
3.6 KiB
Python
90 lines
3.6 KiB
Python
|
"""WebSocket connection manager for the kommunikation plugin."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import logging
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from fastapi import WebSocket
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class WebSocketManager:
|
||
|
|
"""Manages WebSocket connections per user for real-time messaging."""
|
||
|
|
|
||
|
|
def __init__(self) -> None:
|
||
|
|
# user_id (str) → list of WebSocket connections
|
||
|
|
self._connections: dict[str, list[WebSocket]] = {}
|
||
|
|
# conversation_id (str) → set of user_ids subscribed
|
||
|
|
self._subscriptions: dict[str, set[str]] = {}
|
||
|
|
|
||
|
|
async def connect(self, websocket: WebSocket, user_id: str) -> None:
|
||
|
|
"""Accept and register a new WebSocket connection."""
|
||
|
|
await websocket.accept()
|
||
|
|
if user_id not in self._connections:
|
||
|
|
self._connections[user_id] = []
|
||
|
|
self._connections[user_id].append(websocket)
|
||
|
|
logger.debug(f"WebSocket connected: user={user_id}, total={len(self._connections[user_id])}")
|
||
|
|
|
||
|
|
async def disconnect(self, websocket: WebSocket, user_id: str) -> None:
|
||
|
|
"""Remove a WebSocket connection."""
|
||
|
|
conns = self._connections.get(user_id, [])
|
||
|
|
if websocket in conns:
|
||
|
|
conns.remove(websocket)
|
||
|
|
if not conns:
|
||
|
|
self._connections.pop(user_id, None)
|
||
|
|
# Remove from all subscriptions
|
||
|
|
for conv_id, users in self._subscriptions.items():
|
||
|
|
users.discard(user_id)
|
||
|
|
logger.debug(f"WebSocket disconnected: user={user_id}, remaining={len(conns)}")
|
||
|
|
|
||
|
|
def subscribe(self, conversation_id: str, user_id: str) -> None:
|
||
|
|
"""Subscribe a user to a conversation's updates."""
|
||
|
|
if conversation_id not in self._subscriptions:
|
||
|
|
self._subscriptions[conversation_id] = set()
|
||
|
|
self._subscriptions[conversation_id].add(user_id)
|
||
|
|
|
||
|
|
def unsubscribe(self, conversation_id: str, user_id: str) -> None:
|
||
|
|
"""Unsubscribe a user from a conversation."""
|
||
|
|
if conversation_id in self._subscriptions:
|
||
|
|
self._subscriptions[conversation_id].discard(user_id)
|
||
|
|
|
||
|
|
async def send_to_user(self, user_id: str, message: dict[str, Any]) -> None:
|
||
|
|
"""Send a message to all connections of a specific user."""
|
||
|
|
conns = self._connections.get(user_id, [])
|
||
|
|
text = json.dumps(message, default=str)
|
||
|
|
for ws in conns:
|
||
|
|
try:
|
||
|
|
await ws.send_text(text)
|
||
|
|
except Exception:
|
||
|
|
logger.warning(f"Failed to send to user {user_id}, removing connection")
|
||
|
|
await self.disconnect(ws, user_id)
|
||
|
|
|
||
|
|
async def send_to_conversation(
|
||
|
|
self,
|
||
|
|
conversation_id: str,
|
||
|
|
message: dict[str, Any],
|
||
|
|
exclude_user: str | None = None,
|
||
|
|
) -> None:
|
||
|
|
"""Send a message to all users subscribed to a conversation."""
|
||
|
|
user_ids = self._subscriptions.get(conversation_id, set())
|
||
|
|
for user_id in list(user_ids):
|
||
|
|
if exclude_user and user_id == exclude_user:
|
||
|
|
continue
|
||
|
|
await self.send_to_user(user_id, message)
|
||
|
|
|
||
|
|
async def broadcast(self, message: dict[str, Any]) -> None:
|
||
|
|
"""Broadcast a message to all connected users."""
|
||
|
|
for user_id in list(self._connections.keys()):
|
||
|
|
await self.send_to_user(user_id, message)
|
||
|
|
|
||
|
|
def get_online_users(self) -> list[str]:
|
||
|
|
"""Get list of currently connected user IDs."""
|
||
|
|
return list(self._connections.keys())
|
||
|
|
|
||
|
|
def is_user_online(self, user_id: str) -> bool:
|
||
|
|
"""Check if a user has any active connections."""
|
||
|
|
return user_id in self._connections and len(self._connections[user_id]) > 0
|