fix: 刷新页面后恢复任务状态 (#58)

Co-authored-by: grey <grey@greydeMacBook-Air.local>
2026-09-19 08:09:48 +08:00 · 2026-04-30 23:18:41 +08:00
parent 3e8a0d6cac
commit c65a7eb92d
9 changed files with 275 additions and 43 deletions
@@ -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使用情况
@@ -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}")
@@ -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)


@@ -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)
@@ -0,0 +1,9 @@
+services:
+  redis:
+    container_name: mathmodelagent_jihe520_redis
+
+  backend:
+    container_name: mathmodelagent_jihe520_backend
+
+  frontend:
+    container_name: mathmodelagent_jihe520_frontend
@@ -1,4 +1,5 @@
 import request from "@/utils/request";
+import type { Message } from "@/utils/response";

 export function getHelloWorld() {
 	return request.get<{ message: string }>("/");
@@ -9,6 +10,14 @@ export function getWriterSeque() {
 	return request.get<{ writer_seque: string[] }>("/writer_seque");
 }

+export function getTaskMessages(task_id: string) {
+	return request.get<Message[]>("/messages", {
+		params: {
+			task_id,
+		},
+	});
+}
+

 export function openFolderAPI(task_id: string) {
 	return request.get<{ message: string }>("/open_folder", {
@@ -35,4 +44,4 @@ export function getServiceStatus() {
 		backend: { status: string; message: string };
 		redis: { status: string; message: string };
 	}>("/status");
-}
+}
@@ -18,9 +18,7 @@ import { onMounted, onBeforeUnmount, ref } from 'vue'
 import { useTaskStore } from '@/stores/task'
 import { getWriterSeque } from '@/apis/commonApi';
 import { Button } from '@/components/ui/button';
-import { useToast } from '@/components/ui/toast/use-toast'
 import FilesSheet from '@/pages/task/components/FileSheet.vue'
-const { toast } = useToast()


 const props = defineProps<{ task_id: string }>()
@@ -59,6 +57,7 @@ const updateDuration = () => {
 console.log('Task ID:', props.task_id)

 onMounted(async () => {
+  await taskStore.loadTaskMessages(props.task_id)
  taskStore.connectWebSocket(props.task_id)
  const res = await getWriterSeque();
  writerSequence.value = Array.isArray(res.data) ? res.data : [];
@@ -140,4 +139,4 @@ onBeforeUnmount(() => {
  </div>
 </template>

-<style scoped></style>
+<style scoped></style>
@@ -4,34 +4,154 @@ import { TaskWebSocket } from '@/utils/websocket'
 import type { Message, CoderMessage, WriterMessage, UserMessage, ModelerMessage, CoordinatorMessage, InterpreterMessage } from '@/utils/response'
 // import messageData from '@/test/20250524-115938-d4c84576.json'
 import { AgentType } from '@/utils/enum'
+import { getTaskMessages } from '@/apis/commonApi'

 export const useTaskStore = defineStore('task', () => {
  // 初始化时直接加载测试数据，确保页面首次渲染时有数据
  // const messages = ref<Message[]>(messageData as Message[])
-  const messages = ref<Message[]>([])
+  const messagesByTask = ref<Record<string, Message[]>>({})
+  const currentTaskId = ref<string | null>(null)
+  const messages = computed<Message[]>(() => {
+    if (!currentTaskId.value) {
+      return []
+    }
+    return messagesByTask.value[currentTaskId.value] ?? []
+  })
+  const seenMessageIdsByTask = new Map<string, Set<string>>()
  let ws: TaskWebSocket | null = null

+  function getMessageTimestamp(message: Message): number | null {
+    if (!message.created_at) {
+      return null
+    }
+    const timestamp = Date.parse(message.created_at)
+    return Number.isNaN(timestamp) ? null : timestamp
+  }
+
+  function sortMessages(items: Message[]) {
+    return [...items].sort((left, right) => {
+      const leftTs = getMessageTimestamp(left)
+      const rightTs = getMessageTimestamp(right)
+      if (leftTs == null || rightTs == null || leftTs === rightTs) {
+        return 0
+      }
+      return leftTs - rightTs
+    })
+  }
+
+  function isMessagePayload(payload: unknown): payload is Message {
+    if (!payload || typeof payload !== 'object') {
+      return false
+    }
+    const msgType = Reflect.get(payload, 'msg_type')
+    return (
+      typeof Reflect.get(payload, 'id') === 'string' &&
+      typeof msgType === 'string' &&
+      ['system', 'agent', 'user', 'tool'].includes(msgType)
+    )
+  }
+
+  function setCurrentTask(taskId: string) {
+    currentTaskId.value = taskId
+    if (typeof window !== 'undefined') {
+      window.localStorage.setItem('currentTaskId', taskId)
+    }
+  }
+
+  function ensureTaskBucket(taskId: string) {
+    if (!messagesByTask.value[taskId]) {
+      messagesByTask.value[taskId] = []
+    }
+    if (!seenMessageIdsByTask.has(taskId)) {
+      seenMessageIdsByTask.set(taskId, new Set())
+    }
+  }
+
+  function appendMessage(taskId: string, message: Message) {
+    ensureTaskBucket(taskId)
+    const seenIds = seenMessageIdsByTask.get(taskId)
+    if (message.id && seenIds?.has(message.id)) {
+      messagesByTask.value[taskId] = sortMessages(
+        messagesByTask.value[taskId].map((existing) =>
+          existing.id === message.id ? message : existing,
+        ),
+      )
+      return
+    }
+    if (message.id) {
+      seenIds?.add(message.id)
+    }
+    messagesByTask.value[taskId] = sortMessages([
+      ...messagesByTask.value[taskId],
+      message,
+    ])
+  }
+
+  function mergeMessages(taskId: string, incomingMessages: Message[]) {
+    ensureTaskBucket(taskId)
+    const existingMessages = messagesByTask.value[taskId]
+    const mergedById = new Map<string, Message>()
+
+    for (const message of [...existingMessages, ...incomingMessages]) {
+      if (!message.id) {
+        continue
+      }
+      mergedById.set(message.id, message)
+    }
+
+    const mergedMessages = Array.from(mergedById.values())
+    messagesByTask.value[taskId] = sortMessages(mergedMessages)
+    seenMessageIdsByTask.set(
+      taskId,
+      new Set(mergedMessages.map((message) => message.id)),
+    )
+  }
+
  // 连接 WebSocket
  function connectWebSocket(taskId: string) {
+    if (ws) {
+      ws.close()
+      ws = null
+    }
+    setCurrentTask(taskId)
+    ensureTaskBucket(taskId)
+
    const baseUrl = import.meta.env.VITE_WS_URL
    const wsUrl = `${baseUrl}/task/${taskId}`

    ws = new TaskWebSocket(wsUrl, (data) => {
-      console.log(data)
-      messages.value.push(data)
+      if (!isMessagePayload(data)) {
+        console.warn('忽略非标准任务消息:', data)
+        return
+      }
+      appendMessage(taskId, data)
    })
    // 初始化测试数据（已在上面初始化，这里可以注释掉）
    // messages.value = messageData as Message[]
    ws.connect()
  }

+  async function loadTaskMessages(taskId: string) {
+    setCurrentTask(taskId)
+    ensureTaskBucket(taskId)
+    try {
+      const response = await getTaskMessages(taskId)
+      const validMessages = (response.data ?? []).filter(isMessagePayload)
+      mergeMessages(taskId, validMessages)
+    } catch (error) {
+      console.error('加载任务历史消息失败:', error)
+    }
+  }
+
  // 关闭 WebSocket
  function closeWebSocket() {
    ws?.close()
+    ws = null
  }

  function addUserMessage(content: string) {
-    messages.value.push({
+    const taskId = currentTaskId.value ?? 'local'
+    appendMessage(taskId, {
      id: Date.now().toString(),
      msg_type: 'user',
      content: content,
@@ -43,7 +163,7 @@ export const useTaskStore = defineStore('task', () => {
    const dataStr = "data:text/json;charset=utf-8," + encodeURIComponent(JSON.stringify(messages.value, null, 2))
    const downloadAnchorNode = document.createElement('a')
    downloadAnchorNode.setAttribute("href", dataStr)
-    downloadAnchorNode.setAttribute("download", "message.json")
+    downloadAnchorNode.setAttribute("download", `${currentTaskId.value ?? 'task'}-messages.json`)
    document.body.appendChild(downloadAnchorNode)
    downloadAnchorNode.click()
    downloadAnchorNode.remove()
@@ -143,9 +263,10 @@ export const useTaskStore = defineStore('task', () => {
    writerMessages,
    interpreterMessage,
    files,
+    loadTaskMessages,
    connectWebSocket,
    closeWebSocket,
    downloadMessages,
    addUserMessage
  }
-}) 
+}) 
@@ -7,6 +7,7 @@ import { AgentType } from './enum';

 export interface BaseMessage {
  id: string;
+  created_at?: string;
  msg_type: 'system' | 'agent' | 'user' | 'tool';
  content?: string | null;
 }
@@ -105,4 +106,4 @@ export interface WriterMessage extends AgentMessage {
  sub_title?: string;
 }

-export type Message = SystemMessage | UserMessage | CoderMessage | WriterMessage | ModelerMessage | CoordinatorMessage | ToolMessage;
+export type Message = SystemMessage | UserMessage | CoderMessage | WriterMessage | ModelerMessage | CoordinatorMessage | ToolMessage;