74 lines
3.3 KiB
Python
74 lines
3.3 KiB
Python
"""
|
|
WebSocket Connection Manager (Section 18, 19, 28).
|
|
Manages active WebSocket instances and per-session processing locks.
|
|
Ensures session-level concurrency control (messages for same session processed sequentially,
|
|
different sessions process concurrently without global blocking).
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
from typing import Dict, Set, Optional
|
|
from fastapi import WebSocket, status
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
# Maps session_id to set of active WebSocket instances
|
|
self.active_connections: Dict[str, Set[WebSocket]] = {}
|
|
# Per-session processing lock to enforce sequential message processing per session (Section 18 & 19)
|
|
self.session_locks: Dict[str, asyncio.Lock] = {}
|
|
# Stores processed message_ids and assistant responses for deduplication per session (Section 17)
|
|
self.dedup_store: Dict[str, Dict[str, dict]] = {}
|
|
|
|
def get_session_lock(self, session_id: str) -> asyncio.Lock:
|
|
"""Returns or creates a per-session lock."""
|
|
if session_id not in self.session_locks:
|
|
self.session_locks[session_id] = asyncio.Lock()
|
|
return self.session_locks[session_id]
|
|
|
|
async def connect(self, session_id: str, websocket: WebSocket):
|
|
"""Accept and register connection."""
|
|
if session_id not in self.active_connections:
|
|
self.active_connections[session_id] = set()
|
|
self.active_connections[session_id].add(websocket)
|
|
logger.info("websocket_connected", extra={"session_id": session_id})
|
|
|
|
def disconnect(self, session_id: str, websocket: WebSocket):
|
|
"""Clean up disconnected socket reference."""
|
|
if session_id in self.active_connections:
|
|
self.active_connections[session_id].discard(websocket)
|
|
if not self.active_connections[session_id]:
|
|
del self.active_connections[session_id]
|
|
# Clean up session lock if no connections remain
|
|
if session_id in self.session_locks:
|
|
del self.session_locks[session_id]
|
|
logger.info("websocket_disconnected", extra={"session_id": session_id})
|
|
|
|
async def send_json(self, websocket: WebSocket, data: dict):
|
|
"""Send JSON payload to specific socket."""
|
|
try:
|
|
await websocket.send_json(data)
|
|
except Exception as e:
|
|
logger.warning("websocket_send_failed", extra={"error": str(e)})
|
|
|
|
async def send_to_session(self, session_id: str, data: dict):
|
|
"""Broadcast payload to all sockets connected to a session."""
|
|
if session_id in self.active_connections:
|
|
for ws in list(self.active_connections[session_id]):
|
|
await self.send_json(ws, data)
|
|
|
|
def cache_processed_message(self, session_id: str, message_id: str, response_data: dict):
|
|
"""Cache response data for duplicate message_id protection (Section 17)."""
|
|
if session_id not in self.dedup_store:
|
|
self.dedup_store[session_id] = {}
|
|
self.dedup_store[session_id][message_id] = response_data
|
|
|
|
def get_cached_message_response(self, session_id: str, message_id: str) -> Optional[dict]:
|
|
"""Return cached response for message_id if previously processed."""
|
|
return self.dedup_store.get(session_id, {}).get(message_id)
|
|
|
|
|
|
manager = ConnectionManager()
|