46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
"""WebSocket 连接管理器：仅负责连接注册/注销和消息收发"""
|
|
from app.core.logging import get_logger
|
|
from typing import Dict
|
|
|
|
from fastapi import HTTPException, WebSocket
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class ConnectionManager:
|
|
"""WebSocket 连接管理器"""
|
|
|
|
def __init__(self):
|
|
self.active_connections: Dict[str, WebSocket] = {}
|
|
|
|
async def connect(self, websocket: WebSocket, conversation_id: str):
|
|
"""建立连接"""
|
|
await websocket.accept()
|
|
self.active_connections[conversation_id] = websocket
|
|
|
|
async def disconnect(self, conversation_id: str):
|
|
"""断开连接"""
|
|
if conversation_id in self.active_connections:
|
|
del self.active_connections[conversation_id]
|
|
|
|
async def send_message(self, conversation_id: str, message: dict):
|
|
"""发送消息"""
|
|
if conversation_id in self.active_connections:
|
|
websocket = self.active_connections[conversation_id]
|
|
try:
|
|
await websocket.send_json(message)
|
|
except (RuntimeError, Exception) as e:
|
|
logger.warning(f"发送消息失败 (conversation_id={conversation_id}): {e}")
|
|
if conversation_id in self.active_connections:
|
|
del self.active_connections[conversation_id]
|
|
|
|
async def receive_message(self, conversation_id: str) -> dict:
|
|
"""接收消息"""
|
|
if conversation_id in self.active_connections:
|
|
websocket = self.active_connections[conversation_id]
|
|
return await websocket.receive_json()
|
|
raise HTTPException(status_code=404, detail="Connection not found")
|
|
|
|
|
|
manager = ConnectionManager()
|