Refactor WebSocket handling with connection manager

- Introduce `WebSocketConnectionManager` for centralized connection tracking and session management.
- Add support for session-specific connection handling and message broadcasting.
- Implement `websocket_connection_context` to streamline connection lifecycle management.
- Update WebSocket endpoint to utilize the connection manager for improved reliability and efficiency.

src/voice_dialogue/api/routes/websocket_routes.py

@@ -1,6 +1,10 @@
+import asyncio
+from contextlib import asynccontextmanager
 from queue import Empty
+from typing import Set, Dict
 
 from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+from fastapi.websockets import WebSocketState
 
 from voice_dialogue.core.constants import websocket_message_queue, session_manager
 from voice_dialogue.utils.logger import logger
@@ -8,25 +12,124 @@ from voice_dialogue.utils.logger import logger
 ws = APIRouter()
 
 
+class WebSocketConnectionManager:
+    """WebSocket 连接管理器 - 管理所有活跃连接"""
+
+    def __init__(self):
+        # 使用 WeakSet 避免内存泄漏
+        self._connections: Set[WebSocket] = set()
+        # 会话ID到连接的映射
+        self._session_connections: Dict[str, Set[WebSocket]] = {}
+        self._lock = asyncio.Lock()
+
+    async def connect(self, websocket: WebSocket, session_id: str = None):
+        """建立新连接"""
+        async with self._lock:
+            await websocket.accept()
+            self._connections.add(websocket)
+
+            # 如果指定了会话ID，建立映射关系
+            if session_id:
+                if session_id not in self._session_connections:
+                    self._session_connections[session_id] = set()
+                self._session_connections[session_id].add(websocket)
+
+            logger.info(f"WebSocket连接已建立，当前活跃连接数: {len(self._connections)}")
+
+    async def disconnect(self, websocket: WebSocket):
+        """断开连接"""
+        async with self._lock:
+            self._connections.discard(websocket)
+
+            # 从会话映射中移除
+            for session_id, connections in list(self._session_connections.items()):
+                connections.discard(websocket)
+                if not connections:  # 如果该会话没有活跃连接，清理映射
+                    del self._session_connections[session_id]
+
+            logger.info(f"WebSocket连接已断开，当前活跃连接数: {len(self._connections)}")
+
+    async def close_session_connections(self, session_id: str):
+        """关闭指定会话的所有连接"""
+        async with self._lock:
+            if session_id in self._session_connections:
+                connections_to_close = list(self._session_connections[session_id])
+                for connection in connections_to_close:
+                    try:
+                        await connection.close()
+                        logger.info(f"已关闭会话 {session_id} 的一个连接")
+                    except Exception as e:
+                        logger.warning(f"关闭连接时出错: {e}")
+
+                # 清理映射
+                del self._session_connections[session_id]
+
+    async def send_to_session(self, session_id: str, message: dict):
+        """向指定会话的所有连接发送消息"""
+        async with self._lock:
+            if session_id in self._session_connections:
+                connections = list(self._session_connections[session_id])
+                disconnected_connections = []
+
+                for connection in connections:
+                    try:
+                        await connection.send_json(message)
+                    except Exception as e:
+                        logger.warning(f"发送消息失败，标记连接为断开: {e}")
+                        disconnected_connections.append(connection)
+
+                # 清理断开的连接
+                for connection in disconnected_connections:
+                    await self.disconnect(connection)
+
+    @property
+    def connection_count(self) -> int:
+        """获取当前连接数"""
+        return len(self._connections)
+
+    def get_session_connection_count(self, session_id: str) -> int:
+        """获取指定会话的连接数"""
+        return len(self._session_connections.get(session_id, set()))
+
+
+# 全局连接管理器实例
+connection_manager = WebSocketConnectionManager()
+
+
+@asynccontextmanager
+async def websocket_connection_context(websocket: WebSocket):
+    """WebSocket连接上下文管理器"""
+    current_session_id = session_manager.current_id
+
+    # 关闭同一会话的旧连接
+    if connection_manager.get_session_connection_count(current_session_id) > 0:
+        logger.info(f"检测到会话 {current_session_id} 已有连接，关闭旧连接")
+        await connection_manager.close_session_connections(current_session_id)
+
+    try:
+        # 建立新连接
+        await connection_manager.connect(websocket, current_session_id)
+        yield websocket
+    finally:
+        # 确保连接被正确清理
+        await connection_manager.disconnect(websocket)
+
+
 @ws.websocket("/api/v1/ws")
 async def websocket_endpoint(websocket: WebSocket):
     """WebSocket连接端点"""
-    try:
-        # 建立连接
-        await websocket.accept()
-        # 保持连接活跃
-        while True:
-            try:
-                message = await websocket_message_queue.get()
-            except Empty:
-                continue
-
-            if message.session_id != session_manager.current_id:
-                continue
-
-            await websocket.send_json(message.model_dump())
-
-    except WebSocketDisconnect:
-        pass
-    except Exception as e:
-        logger.error(f"WebSocket连接异常: {e}")
+    async with websocket_connection_context(websocket):
+        try:
+            # 保持连接活跃
+            while websocket.client_state == WebSocketState.CONNECTED:
+                try:
+                    message = await websocket_message_queue.get()
+                except Empty:
+                    continue
+
+                await connection_manager.send_to_session(message.session_id, message.model_dump())
+
+        except WebSocketDisconnect:
+            logger.info("WebSocket连接断开")
+        except Exception as e:
+            logger.error(f"WebSocket连接异常: {e}")