Merge branch 'main' of https://github.com/jihe520/MathModelAgent

backend/app/routers/common_router.py

@@ -1,6 +1,9 @@
-from fastapi import APIRouter
+import json
+from pathlib import Path
+
+from fastapi import APIRouter, HTTPException
 from app.config.setting import settings
-from app.utils.common_utils import get_config_template
+from app.utils.common_utils import ensure_safe_task_id, get_config_template
 from app.schemas.enums import CompTemplate
 from app.services.redis_manager import redis_manager
 from app.utils.log_util import logger
@@ -8,6 +11,28 @@ from app.utils.log_util import logger
 router = APIRouter()
 
 
+def _require_safe_task_id(task_id: str) -> str:
+    try:
+        return ensure_safe_task_id(task_id)
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail="非法任务ID") from exc
+
+
+def _load_task_messages_from_file(task_id: str) -> list[dict]:
+    safe_task_id = _require_safe_task_id(task_id)
+    message_file = Path("logs/messages") / f"{safe_task_id}.json"
+    if not message_file.exists():
+        return []
+
+    try:
+        with open(message_file, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        return data if isinstance(data, list) else []
+    except Exception as e:
+        logger.error(f"读取任务消息文件失败: {str(e)}")
+        return []
+
+
 @router.get("/")
 async def root():
     return {"message": "Hello World"}
@@ -32,6 +57,11 @@ async def get_writer_seque():
     return list(config_template.keys())
 
 
+@router.get("/messages")
+async def get_task_messages(task_id: str):
+    return _load_task_messages_from_file(task_id)
+
+
 @router.get("/track")
 async def track(task_id: str):
     # 获取任务的token使用情况

backend/app/routers/ws_router.py

@@ -1,66 +1,119 @@
-from fastapi import WebSocket, WebSocketDisconnect, APIRouter
-from app.services.redis_manager import redis_manager
-from app.schemas.response import SystemMessage
 import asyncio
-from app.services.ws_manager import ws_manager
 import json
 
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+from starlette.websockets import WebSocketState
+
+from app.schemas.response import SystemMessage
+from app.services.redis_manager import redis_manager
+from app.services.ws_manager import ws_manager
+from app.utils.common_utils import ensure_safe_task_id
+from app.utils.log_util import logger
+
 router = APIRouter()
 
 
+def _is_websocket_closed(websocket: WebSocket) -> bool:
+    return (
+        websocket.client_state == WebSocketState.DISCONNECTED
+        or websocket.application_state == WebSocketState.DISCONNECTED
+    )
+
+
+def _is_closed_send_error(error: Exception) -> bool:
+    text = str(error)
+    return (
+        "Cannot call \"send\" once a close message has been sent" in text
+        or "Unexpected ASGI message 'websocket.send'" in text
+    )
+
+
 @router.websocket("/task/{task_id}")
 async def websocket_endpoint(websocket: WebSocket, task_id: str):
-    print(f"WebSocket 尝试连接 task_id: {task_id}")
+    try:
+        safe_task_id = ensure_safe_task_id(task_id)
+    except ValueError:
+        logger.warning(f"WebSocket task_id 非法: {task_id}")
+        await websocket.close(code=1008, reason="Invalid task id")
+        return
+
+    logger.info(f"WebSocket 尝试连接 task_id: {safe_task_id}")
 
     redis_async_client = await redis_manager.get_client()
-    if not await redis_async_client.exists(f"task_id:{task_id}"):
-        print(f"Task not found: {task_id}")
+    if not await redis_async_client.exists(f"task_id:{safe_task_id}"):
+        logger.warning(f"Task not found: {safe_task_id}")
         await websocket.close(code=1008, reason="Task not found")
         return
-    print(f"WebSocket connected for task: {task_id}")
+    logger.info(f"WebSocket connected for task: {safe_task_id}")
 
     # 建立 WebSocket 连接
     await ws_manager.connect(websocket)
     websocket.timeout = 500
-    print(f"WebSocket connection status: {websocket.client}")
+    logger.debug(f"WebSocket connection status: {websocket.client}")
 
     # 订阅 Redis 频道
-    pubsub = await redis_manager.subscribe_to_task(task_id)
-    print(f"Subscribed to Redis channel: task:{task_id}:messages")
-
-    await redis_manager.publish_message(
-        task_id,
-        SystemMessage(content="任务开始处理"),
-    )
+    pubsub = await redis_manager.subscribe_to_task(safe_task_id)
+    logger.debug(f"Subscribed to Redis channel: task:{safe_task_id}:messages")
 
     try:
         while True:
+            if _is_websocket_closed(websocket):
+                logger.info(f"WebSocket 已关闭，停止转发 task_id: {safe_task_id}")
+                break
             try:
                 msg = await pubsub.get_message(ignore_subscribe_messages=True)
                 if msg:
-                    print(f"Received message: {msg}")
                     try:
                         msg_dict = json.loads(msg["data"])
-                        await ws_manager.send_personal_message_json(msg_dict, websocket)
-                        print(f"Sent message to WebSocket: {msg_dict}")
                     except Exception as e:
-                        print(f"Error parsing message: {e}")
-                        await ws_manager.send_personal_message_json(
-                            {"error": str(e)}, websocket
-                        )
+                        logger.error(f"Error parsing websocket payload: {e}")
+                        if _is_websocket_closed(websocket):
+                            break
+                        try:
+                            await ws_manager.send_personal_message_json(
+                                SystemMessage(
+                                    content="实时消息解析失败，已忽略异常数据。",
+                                    type="error",
+                                ).model_dump(),
+                                websocket,
+                            )
+                        except WebSocketDisconnect:
+                            logger.info("WebSocket disconnected while sending parse error notice")
+                            break
+                        except RuntimeError as send_error:
+                            if _is_closed_send_error(send_error):
+                                logger.info("WebSocket 已关闭，跳过解析失败提示发送")
+                                break
+                            raise
+                    else:
+                        try:
+                            await ws_manager.send_personal_message_json(msg_dict, websocket)
+                        except WebSocketDisconnect:
+                            logger.info("WebSocket disconnected while sending message")
+                            break
+                        except RuntimeError as send_error:
+                            if _is_closed_send_error(send_error):
+                                logger.info(
+                                    f"WebSocket 已关闭，停止发送后续消息 task_id: {safe_task_id}"
+                                )
+                                break
+                            raise
                 await asyncio.sleep(0.1)
 
             except WebSocketDisconnect:
-                print("WebSocket disconnected")
+                logger.info("WebSocket disconnected")
                 break
             except Exception as e:
-                print(f"Error in websocket loop: {e}")
+                if _is_closed_send_error(e) or _is_websocket_closed(websocket):
+                    logger.info(f"WebSocket 发送通道已关闭，结束循环 task_id: {safe_task_id}")
+                    break
+                logger.error(f"Error in websocket loop: {e}")
                 await asyncio.sleep(1)
                 continue
 
     except Exception as e:
-        print(f"WebSocket error: {e}")
+        logger.error(f"WebSocket error: {e}")
     finally:
-        await pubsub.unsubscribe(f"task:{task_id}:messages")
+        await pubsub.unsubscribe(f"task:{safe_task_id}:messages")
         ws_manager.disconnect(websocket)
-        print(f"WebSocket connection closed for task: {task_id}")
+        logger.info(f"WebSocket connection closed for task: {safe_task_id}")

backend/app/services/ws_manager.py

@@ -10,7 +10,8 @@ class WebSocketManager:
         self.active_connections.append(websocket)
 
     def disconnect(self, websocket: WebSocket):
-        self.active_connections.remove(websocket)
+        if websocket in self.active_connections:
+            self.active_connections.remove(websocket)
 
     async def send_personal_message(self, message: str, websocket: WebSocket):
         await websocket.send_text(message)
@@ -19,7 +20,7 @@ class WebSocketManager:
         await websocket.send_json(message)
 
     async def broadcast(self, message: str):
-        for connection in self.active_connections:
+        for connection in list(self.active_connections):
             await connection.send_text(message)
 
 

backend/app/utils/common_utils.py

@@ -9,6 +9,8 @@ import pypandoc
 from app.config.setting import settings
 from icecream import ic
 
+TASK_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$")
+
 
 def create_task_id() -> str:
     """生成任务ID"""
@@ -18,6 +20,13 @@ def create_task_id() -> str:
     return f"{timestamp}-{random_hash}"
 
 
+def ensure_safe_task_id(task_id: str) -> str:
+    normalized = (task_id or "").strip()
+    if not normalized or not TASK_ID_PATTERN.fullmatch(normalized):
+        raise ValueError("非法 task_id")
+    return normalized
+
+
 def create_work_dir(task_id: str) -> str:
     # 设置主工作目录和子目录
     work_dir = os.path.join("project", "work_dir", task_id)