| 1 | import type { TaskDetail, TaskMessage } from '@/api/ai/ai.types' |
| 2 | import type { IDisplayMessage } from '@/store/agent' |
| 3 | /** |
| 4 | * 任务轮询 Hook |
| 5 | * 在页面刷新后任务未完成时，通过轮询获取最新状态 |
| 6 | */ |
| 7 | import { useCallback, useEffect, useRef, useState } from 'react' |
| 8 | import { agentApi } from '@/api/ai/ai.api' |
| 9 | |
| 10 | import { AgentTaskStatus } from '@/api/ai/ai.constants' |
| 11 | import { useUserStore } from '@/store/user' |
| 12 | import { convertMessages, isTaskCompleted } from '../utils' |
| 13 | |
| 14 | export interface ITaskPollingOptions { |
| 15 | /** 任务 ID */ |
| 16 | taskId: string |
| 17 | /** 是否为活跃任务（Store 中有对应的实时消息） */ |
| 18 | isActiveTask: boolean |
| 19 | /** 获取当前原始消息列表（用于增量轮询） */ |
| 20 | getCurrentRawMessages?: () => TaskMessage[] |
| 21 | /** 轮询间隔（ms） */ |
| 22 | pollingInterval?: number |
| 23 | /** 消息更新回调 */ |
| 24 | onMessagesUpdate: (messages: IDisplayMessage[], rawMessages: TaskMessage[]) => void |
| 25 | /** 任务更新回调 */ |
| 26 | onTaskUpdate?: (task: TaskDetail) => void |
| 27 | /** 任务状态变化回调 */ |
| 28 | onTaskStatusChange?: (status: string) => void |
| 29 | } |
| 30 | |
| 31 | export interface ITaskPollingReturn { |
| 32 | /** 是否正在轮询 */ |
| 33 | isPolling: boolean |
| 34 | /** 开始轮询 */ |
| 35 | startPolling: () => void |
| 36 | /** 停止轮询 */ |
| 37 | stopPolling: () => void |
| 38 | } |
| 39 | |
| 40 | /** |
| 41 | * 任务轮询 Hook |
| 42 | */ |
| 43 | export function useTaskPolling(options: ITaskPollingOptions): ITaskPollingReturn { |
| 44 | const { |
| 45 | taskId, |
| 46 | isActiveTask, |
| 47 | getCurrentRawMessages, |
| 48 | pollingInterval = 3000, |
| 49 | onMessagesUpdate, |
| 50 | onTaskUpdate, |
| 51 | onTaskStatusChange, |
| 52 | } = options |
| 53 | |
| 54 | const [isPolling, setIsPolling] = useState(false) |
| 55 | const pollingTimerRef = useRef<NodeJS.Timeout | null>(null) |
| 56 | |
| 57 | // 获取 Credits 余额 |
| 58 | const fetchCreditsBalance = useUserStore(state => state.fetchCreditsBalance) |
| 59 | |
| 60 | /** 开始轮询 */ |
| 61 | const startPolling = useCallback(() => { |
| 62 | setIsPolling(true) |
| 63 | }, []) |
| 64 | |
| 65 | /** 停止轮询 */ |
| 66 | const stopPolling = useCallback(() => { |
| 67 | setIsPolling(false) |
| 68 | }, []) |
| 69 | |
| 70 | /** 轮询逻辑 */ |
| 71 | useEffect(() => { |
| 72 | if (!isPolling || !taskId || isActiveTask) { |
| 73 | return |
| 74 | } |
| 75 | |
| 76 | // [TaskPolling] Starting polling for task:', taskId) |
| 77 | |
| 78 | // 用于跟踪当前是否已有最后一条消息（决定使用增量拉取或全量拉取） |
| 79 | const hasLastMessageRef = { current: false as boolean } |
| 80 | // 防止并发的轮询请求：若上一次请求尚未完成，则跳过本次 tick |
| 81 | let isRequestRunning = false |
| 82 | |
| 83 | // 读取当前已有的原始消息，决定初始轮询间隔（无 lastMessageId 则使用 5s 全量拉取） |
| 84 | const initialRawMessages = getCurrentRawMessages ? getCurrentRawMessages() : [] |
| 85 | let initialLastMessageId: string | undefined |
| 86 | for (let i = initialRawMessages.length - 1; i >= 0; i--) { |
| 87 | const uuid = initialRawMessages[i]?.uuid |
| 88 | if (uuid) { |
| 89 | initialLastMessageId = uuid |
| 90 | break |
| 91 | } |
| 92 | } |
| 93 | hasLastMessageRef.current = Boolean(initialLastMessageId) |
| 94 | |
| 95 | const pollTask = async () => { |
| 96 | // 如果上一次请求还在进行中，直接跳过本次轮询，避免并发请求 |
| 97 | if (isRequestRunning) { |
| 98 | // [TaskPolling] previous poll still running, skipping this tick |
| 99 | return |
| 100 | } |
| 101 | |
| 102 | isRequestRunning = true |
| 103 | try { |
| 104 | // [TaskPolling] pollTask tick, taskId:', taskId) |
| 105 | // 获取当前已有的原始消息列表 |
| 106 | const currentRawMessages = getCurrentRawMessages ? getCurrentRawMessages() : [] |
| 107 | // [TaskPolling] currentRawMessages length:', currentRawMessages.length) |
| 108 | // 计算最后一条消息的 UUID（用于增量拉取） |
| 109 | let lastMessageId: string | undefined |
| 110 | for (let i = currentRawMessages.length - 1; i >= 0; i--) { |
| 111 | const uuid = currentRawMessages[i]?.uuid |
| 112 | if (uuid) { |
| 113 | lastMessageId = uuid |
| 114 | break |
| 115 | } |
| 116 | } |
| 117 | |
| 118 | // 如果没有 lastMessageId，执行每 5 秒一次的全量拉取逻辑（接口支持不传 lastMessageId 获取全部消息） |
| 119 | if (!lastMessageId) { |
| 120 | const result = await agentApi.getTaskMessages(taskId) |
| 121 | if (result?.code === 0 && result.data) { |
| 122 | // 检查任务状态 |
| 123 | if (result.data.status === AgentTaskStatus.Aborted) { |
| 124 | // [TaskPolling] Task aborted, stopping polling') |
| 125 | setIsPolling(false) |
| 126 | onTaskStatusChange?.('aborted') |
| 127 | return |
| 128 | } |
| 129 | |
| 130 | const newMessages = result.data.messages |
| 131 | // [TaskPolling] fetched fullMessages length:', newMessages.length) |
| 132 | if (!newMessages.length) { |
| 133 | return |
| 134 | } |
| 135 | |
| 136 | const mergedMessages = [...newMessages] |
| 137 | |
| 138 | // 如果新消息与当前原始消息完全相同（数量和内容），跳过更新避免不必要的重渲染 |
| 139 | // 这种情况常见于任务初始化阶段，接口持续返回相同的用户消息 |
| 140 | if ( |
| 141 | currentRawMessages.length > 0 |
| 142 | && currentRawMessages.length === mergedMessages.length |
| 143 | && mergedMessages.every((msg, i) => { |
| 144 | const curr = currentRawMessages[i] |
| 145 | return curr && curr.type === msg.type && curr.uuid === msg.uuid |
| 146 | }) |
| 147 | ) { |
| 148 | // [TaskPolling] Messages unchanged, skipping update') |
| 149 | return |
| 150 | } |
| 151 | |
| 152 | // 更新消息 |
| 153 | const converted = convertMessages(mergedMessages) |
| 154 | onMessagesUpdate(converted, mergedMessages) |
| 155 | |
| 156 | // 如果本次拉取获得了最后一条消息（即 messages 中存在 uuid），则切换为增量拉取间隔 |
| 157 | let nowHasLast = false |
| 158 | for (let i = mergedMessages.length - 1; i >= 0; i--) { |
| 159 | if (mergedMessages[i]?.uuid) { |
| 160 | nowHasLast = true |
| 161 | break |
| 162 | } |
| 163 | } |
| 164 | if (!hasLastMessageRef.current && nowHasLast) { |
| 165 | // 切换为增量拉取间隔（使用传入的 pollingInterval） |
| 166 | hasLastMessageRef.current = true |
| 167 | if (pollingTimerRef.current) { |
| 168 | clearInterval(pollingTimerRef.current) |
| 169 | } |
| 170 | pollingTimerRef.current = setInterval(pollTask, pollingInterval) |
| 171 | } |
| 172 | |
| 173 | // 检测任务是否完成（此处没有最新的 TaskDetail，只能基于消息做兜底判断） |
| 174 | if (isTaskCompleted(mergedMessages)) { |
| 175 | // [TaskPolling] Task completed, stopping polling') |
| 176 | setIsPolling(false) |
| 177 | // 任务完成时刷新 Credits 余额 |
| 178 | fetchCreditsBalance() |
| 179 | } |
| 180 | } |
| 181 | |
| 182 | return |
| 183 | } |
| 184 | |
| 185 | // 如果存在 lastMessageId，调用增量消息接口，仅获取 lastMessageId 之后的新消息 |
| 186 | const result = await agentApi.getTaskMessages(taskId, lastMessageId) |
| 187 | if (result?.code === 0 && result.data) { |
| 188 | // 检查任务状态 |
| 189 | if ( |
| 190 | result.data.status === AgentTaskStatus.Aborted |
| 191 | || result.data.status === AgentTaskStatus.Completed |
| 192 | || result.data.status === AgentTaskStatus.Error |
| 193 | ) { |
| 194 | // [TaskPolling] Task aborted, stopping polling') |
| 195 | setIsPolling(false) |
| 196 | onTaskStatusChange?.('aborted') |
| 197 | return |
| 198 | } |
| 199 | |
| 200 | const newMessages = result.data.messages |
| 201 | // [TaskPolling] fetched newMessages |
| 202 | if (!newMessages.length) { |
| 203 | return |
| 204 | } |
| 205 | |
| 206 | const mergedMessages = [...currentRawMessages, ...newMessages] |
| 207 | |
| 208 | // 更新消息 |
| 209 | const converted = convertMessages(mergedMessages) |
| 210 | onMessagesUpdate(converted, mergedMessages) |
| 211 | |
| 212 | // 注意：轮询接口 getTaskMessages 只返回消息列表（TaskMessagesVo），不返回完整任务详情（TaskDetail） |
| 213 | // 如果需要更新任务详情，应该调用 getTaskDetail 接口 |
| 214 | // 这里不调用 onTaskUpdate，因为类型不匹配 |
| 215 | |
| 216 | // 检测任务是否完成（此处没有最新的 TaskDetail，只能基于消息做兜底判断） |
| 217 | if (isTaskCompleted(mergedMessages)) { |
| 218 | // [TaskPolling] Task completed, stopping polling') |
| 219 | setIsPolling(false) |
| 220 | // 任务完成时刷新 Credits 余额 |
| 221 | fetchCreditsBalance() |
| 222 | } |
| 223 | } |
| 224 | } |
| 225 | catch (error) { |
| 226 | console.error('[TaskPolling] Polling failed:', error) |
| 227 | // 轮询失败不停止，继续尝试 |
| 228 | } |
| 229 | finally { |
| 230 | // 本次请求结束，允许下一个轮询执行 |
| 231 | isRequestRunning = false |
| 232 | } |
| 233 | } |
| 234 | |
| 235 | // 根据初始是否有 lastMessageId 决定初始轮询间隔（无 lastMessageId 使用 5000ms） |
| 236 | const initialInterval = hasLastMessageRef.current ? pollingInterval : 5000 |
| 237 | pollingTimerRef.current = setInterval(pollTask, initialInterval) |
| 238 | |
| 239 | return () => { |
| 240 | if (pollingTimerRef.current) { |
| 241 | clearInterval(pollingTimerRef.current) |
| 242 | pollingTimerRef.current = null |
| 243 | } |
| 244 | } |
| 245 | }, [ |
| 246 | isPolling, |
| 247 | taskId, |
| 248 | isActiveTask, |
| 249 | pollingInterval, |
| 250 | onMessagesUpdate, |
| 251 | fetchCreditsBalance, |
| 252 | getCurrentRawMessages, |
| 253 | ]) |
| 254 | |
| 255 | /** 清理定时器 */ |
| 256 | useEffect(() => { |
| 257 | return () => { |
| 258 | if (pollingTimerRef.current) { |
| 259 | clearInterval(pollingTimerRef.current) |
| 260 | pollingTimerRef.current = null |
| 261 | } |
| 262 | } |
| 263 | }, []) |
| 264 | |
| 265 | return { |
| 266 | isPolling, |
| 267 | startPolling, |
| 268 | stopPolling, |
| 269 | } |
| 270 | } |
| 271 |