life-echo/api/app/features/conversation/ws/connection_manager.py

"""WebSocket 连接管理器：仅负责连接注册/注销和消息收发"""

from typing import Dict

from fastapi import HTTPException, WebSocket

from app.core.logging import get_logger

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(
                    "发送消息失败 (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()