Enhance WebSocket handling for connection management and reliability

- Add message broadcasting support to `WebSocketConnectionManager`.
- Improve connection cleanup for disconnected clients.
- Refine WebSocket endpoint to handle client disconnects and queue tasks more gracefully.

src/voice_dialogue/api/routes/websocket_routes.py

@@ -1,10 +1,8 @@
 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
@@ -82,6 +80,26 @@ class WebSocketConnectionManager:
                 for connection in disconnected_connections:
                     await self.disconnect(connection)
 
+    async def broadcast(self, message: dict):
+        """向所有活跃连接广播消息"""
+        async with self._lock:
+            if not self._connections:
+                return
+
+            connections = list(self._connections)
+            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:
         """获取当前连接数"""
@@ -118,18 +136,43 @@ async def websocket_connection_context(websocket: WebSocket):
 @ws.websocket("/api/v1/ws")
 async def websocket_endpoint(websocket: WebSocket):
     """WebSocket连接端点"""
-    async with websocket_connection_context(websocket):
+    async with websocket_connection_context(websocket) as websocket_connection:
+        disconnect_task = asyncio.create_task(websocket_connection.receive_text())
+        get_message_task = None
+
         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())
+            while True:
+                get_message_task = asyncio.create_task(websocket_message_queue.get())
+
+                done, pending = await asyncio.wait(
+                    {disconnect_task, get_message_task},
+                    return_when=asyncio.FIRST_COMPLETED
+                )
+
+                # 如果是 disconnect_task 完成，说明客户端已断开连接。
+                if disconnect_task in done:
+                    # 这将重新引发 WebSocketDisconnect 异常并跳出循环。
+                    if get_message_task in done:
+                        # 如果有，将消息放回队列，然后让连接断开
+                        message = get_message_task.result()
+                        websocket_message_queue.put_nowait(message)
+                        logger.info("连接已关闭，将消息重新入队。")
+
+                    disconnect_task.result()
+
+                # 如果是 queue_task 完成，说明有可用的消息。
+                if get_message_task in done:
+                    message = get_message_task.result()
+                    await connection_manager.send_to_session(message.session_id, message.model_dump())
+                    # queue_task 现已完成，我们将在循环的下一次迭代中创建一个新的。
 
         except WebSocketDisconnect:
-            logger.info("WebSocket连接断开")
+            logger.info("WebSocket连接已断开")
         except Exception as e:
             logger.error(f"WebSocket连接异常: {e}")
+        finally:
+            # 确保如果循环因其他原因退出，disconnect_task 会被取消。
+            disconnect_task.cancel()
+            if get_message_task and not get_message_task.done():
+                get_message_task.cancel()