| 1 | import type { TaskDetail, TaskMessage } from '@/api/ai/ai.types' |
| 2 | import type { IDisplayMessage, IWorkflowStep } from '@/store/agent' |
| 3 | import { useCallback, useEffect, useRef, useState } from 'react' |
| 4 | import { useShallow } from 'zustand/react/shallow' |
| 5 | import { agentApi } from '@/api/ai/ai.api' |
| 6 | import { useAgentStore } from '@/store/agent' |
| 7 | import { getDefaultTaskData } from '@/store/agent/agent.state' |
| 8 | import { useUserStore } from '@/store/user' |
| 9 | import { toast } from '@/utils/ui/toast' |
| 10 | import { convertMessages, isTaskCompleted } from '../utils' |
| 11 | import { useTaskPolling } from './useTaskPolling' |
| 12 | |
| 13 | export interface IChatStateOptions { |
| 14 | /** 任务 ID */ |
| 15 | taskId: string |
| 16 | /** 翻译函数 */ |
| 17 | t: (key: string) => string |
| 18 | } |
| 19 | |
| 20 | export interface IChatStateReturn { |
| 21 | /** 任务详情 */ |
| 22 | task: TaskDetail | null |
| 23 | /** 当前显示的消息列表 */ |
| 24 | displayMessages: IDisplayMessage[] |
| 25 | /** 工作流步骤（仅实时生成时有效） */ |
| 26 | workflowSteps: IWorkflowStep[] |
| 27 | /** 是否正在加载 */ |
| 28 | isLoading: boolean |
| 29 | /** 是否正在生成 */ |
| 30 | isGenerating: boolean |
| 31 | /** 进度百分比 */ |
| 32 | progress: number |
| 33 | /** 是否为活跃任务 */ |
| 34 | isActiveTask: boolean |
| 35 | /** 更新本地消息（供子组件使用） */ |
| 36 | setLocalMessages: React.Dispatch<React.SetStateAction<IDisplayMessage[]>> |
| 37 | /** 设置本地生成状态 */ |
| 38 | setLocalIsGenerating: React.Dispatch<React.SetStateAction<boolean>> |
| 39 | /** 更新任务标题 */ |
| 40 | updateTaskTitle: (newTitle: string) => void |
| 41 | } |
| 42 | |
| 43 | /** |
| 44 | * 聊天状态管理 Hook |
| 45 | */ |
| 46 | export function useChatState(options: IChatStateOptions): IChatStateReturn { |
| 47 | const { taskId, t } = options |
| 48 | |
| 49 | // 全局 Store 状态 - 获取任务级数据 |
| 50 | const { currentTaskId, taskMessages, setMessages, debugFiles } = useAgentStore( |
| 51 | useShallow(state => ({ |
| 52 | currentTaskId: state.currentTaskId, |
| 53 | taskMessages: state.taskMessages, |
| 54 | setMessages: state.setMessages, |
| 55 | debugFiles: state.debugFiles, |
| 56 | })), |
| 57 | ) |
| 58 | |
| 59 | // 判断是否处于 debug 模式 |
| 60 | const isDebugMode = debugFiles.length > 0 |
| 61 | |
| 62 | // 获取当前任务的数据（按 taskId 隔离） |
| 63 | const currentTaskData = taskMessages[taskId] || getDefaultTaskData() |
| 64 | const storeMessages = currentTaskData.messages |
| 65 | const storeWorkflowSteps = currentTaskData.workflowSteps |
| 66 | const storeIsGenerating = currentTaskData.isGenerating |
| 67 | const storeProgress = currentTaskData.progress |
| 68 | |
| 69 | // 获取 Credits 余额 |
| 70 | const fetchCreditsBalance = useUserStore(state => state.fetchCreditsBalance) |
| 71 | |
| 72 | // 判断是否为活跃任务 |
| 73 | const isActiveTask = currentTaskId === taskId |
| 74 | const isRealtimeGenerating = isActiveTask && storeIsGenerating |
| 75 | |
| 76 | // 本地状态 |
| 77 | const [task, setTask] = useState<TaskDetail | null>(null) |
| 78 | const [localMessages, setLocalMessages] = useState<IDisplayMessage[]>([]) |
| 79 | const [isLoading, setIsLoading] = useState(true) |
| 80 | const [localIsGenerating, setLocalIsGenerating] = useState(false) |
| 81 | |
| 82 | // Refs |
| 83 | const hasLoadedRef = useRef(false) |
| 84 | const rawMessagesRef = useRef<TaskMessage[]>([]) |
| 85 | |
| 86 | // 当 taskId 变化时，重置与任务相关的本地状态，确保不会错误地使用上一次任务的缓存 |
| 87 | useEffect(() => { |
| 88 | // 重置已加载标记，强制重新从 API 拉取数据 |
| 89 | hasLoadedRef.current = false |
| 90 | // 清空上一次的原始消息，避免后续基于旧数据的短路逻辑 |
| 91 | rawMessagesRef.current = [] |
| 92 | // 清空本地 task 与消息状态，显示 loading，等待新的加载逻辑触发 |
| 93 | setTask(null) |
| 94 | setLocalMessages([]) |
| 95 | setIsLoading(true) |
| 96 | // 注意：不主动调用 startPolling，这里只做重置，后续 loadTask 会根据新任务情况自行决定是否启动轮询 |
| 97 | }, [taskId]) |
| 98 | |
| 99 | // 注意：由于现在消息是按 taskId 隔离存储的，不再需要复杂的护栏清理逻辑 |
| 100 | // 每个任务的消息独立存储在 taskMessages[taskId] 中，不会互相干扰 |
| 101 | |
| 102 | // 轮询 Hook |
| 103 | const { isPolling, startPolling } = useTaskPolling({ |
| 104 | taskId, |
| 105 | isActiveTask, |
| 106 | // 轮询间隔（ms） |
| 107 | pollingInterval: 1500, |
| 108 | getCurrentRawMessages: useCallback(() => rawMessagesRef.current, []), |
| 109 | onMessagesUpdate: useCallback( |
| 110 | (messages, rawMessages) => { |
| 111 | rawMessagesRef.current = rawMessages |
| 112 | setLocalMessages(messages) |
| 113 | |
| 114 | // 防御性检查：只有当新消息数量 >= 当前消息数量时才更新 store |
| 115 | // 防止轮询返回的不完整数据覆盖 SSE 实时追加的消息 |
| 116 | const currentMessages = taskMessages[taskId]?.messages || [] |
| 117 | if (messages.length >= currentMessages.length) { |
| 118 | setMessages(messages, taskId) |
| 119 | } |
| 120 | else { |
| 121 | console.warn( |
| 122 | '[ChatState] Skipping setMessages: new messages count is less than current', |
| 123 | { |
| 124 | current: currentMessages.length, |
| 125 | new: messages.length, |
| 126 | }, |
| 127 | ) |
| 128 | } |
| 129 | }, |
| 130 | [setMessages, taskId, taskMessages], |
| 131 | ), |
| 132 | onTaskUpdate: useCallback((taskData: TaskDetail) => { |
| 133 | setTask(taskData) |
| 134 | }, []), |
| 135 | onTaskStatusChange: useCallback((status: string) => { |
| 136 | if (status === 'aborted') { |
| 137 | // Task aborted, clearing generating state |
| 138 | // 任务被中止时，停止本地生成状态 |
| 139 | setLocalIsGenerating(false) |
| 140 | } |
| 141 | }, []), |
| 142 | }) |
| 143 | |
| 144 | /** |
| 145 | * 加载任务详情 |
| 146 | */ |
| 147 | useEffect(() => { |
| 148 | // 如果是 "new" 任务，不加载历史数据，等待创建 |
| 149 | if (taskId === 'new') { |
| 150 | setIsLoading(false) |
| 151 | return |
| 152 | } |
| 153 | |
| 154 | // 如果已经加载过，不再重复加载 |
| 155 | if (hasLoadedRef.current) { |
| 156 | setIsLoading(false) |
| 157 | return |
| 158 | } |
| 159 | |
| 160 | // 如果 Store 中已有该任务的消息，优先使用（支持任务缓存） |
| 161 | if (storeMessages.length > 0) { |
| 162 | setIsLoading(false) |
| 163 | hasLoadedRef.current = true |
| 164 | return |
| 165 | } |
| 166 | |
| 167 | const loadTask = async () => { |
| 168 | if (!taskId) |
| 169 | return |
| 170 | |
| 171 | setIsLoading(true) |
| 172 | try { |
| 173 | const result = await agentApi.getTaskDetail(taskId) |
| 174 | if (!result) { |
| 175 | toast.error(t('message.error')) |
| 176 | return |
| 177 | } |
| 178 | if (result.code === 0 && result.data) { |
| 179 | setTask(result.data) |
| 180 | |
| 181 | if (result.data.messages) { |
| 182 | rawMessagesRef.current = result.data.messages |
| 183 | const converted = convertMessages(result.data.messages) |
| 184 | setLocalMessages(converted) |
| 185 | setMessages(converted, taskId) |
| 186 | |
| 187 | // 检测任务是否完成，如果未完成则启动轮询 |
| 188 | if (!isTaskCompleted(result.data.messages, result.data)) { |
| 189 | // Task not completed, starting polling |
| 190 | startPolling() |
| 191 | } |
| 192 | } |
| 193 | |
| 194 | // 获取到 result 后，刷新 Credits 余额 |
| 195 | fetchCreditsBalance() |
| 196 | |
| 197 | hasLoadedRef.current = true |
| 198 | } |
| 199 | else { |
| 200 | toast.error(result.message || t('message.error')) |
| 201 | } |
| 202 | } |
| 203 | catch (error) { |
| 204 | console.error('Load task detail failed:', error) |
| 205 | toast.error(t('message.error')) |
| 206 | } |
| 207 | finally { |
| 208 | setIsLoading(false) |
| 209 | } |
| 210 | } |
| 211 | |
| 212 | loadTask() |
| 213 | }, [taskId, storeMessages.length, t, setMessages, startPolling, fetchCreditsBalance]) |
| 214 | |
| 215 | // 计算最终显示的消息和状态 |
| 216 | // 优先使用 store 中的消息（支持任务缓存和实时更新） |
| 217 | // debug 模式下强制使用 store 的消息，用于调试回放 |
| 218 | // taskId='new' 时，使用 currentTaskId 对应的消息（临时任务的消息） |
| 219 | // 这样可以在创建任务后立即显示用户消息和 AI 思考状态 |
| 220 | const displayMessages = (() => { |
| 221 | if (taskId === 'new') { |
| 222 | // 获取当前活跃任务的消息（临时任务 temp-xxx） |
| 223 | const currentData = taskMessages[currentTaskId] |
| 224 | return currentData?.messages || [] |
| 225 | } |
| 226 | return storeMessages.length > 0 || isActiveTask || isDebugMode ? storeMessages : localMessages |
| 227 | })() |
| 228 | const isGenerating = isRealtimeGenerating || localIsGenerating || isPolling |
| 229 | // taskId='new' 时，工作流步骤存储在 currentTaskId（临时任务）中 |
| 230 | const workflowSteps: IWorkflowStep[] = (() => { |
| 231 | if (taskId === 'new') { |
| 232 | const currentData = taskMessages[currentTaskId] |
| 233 | return currentData?.workflowSteps || [] |
| 234 | } |
| 235 | return isActiveTask || isDebugMode ? storeWorkflowSteps : [] |
| 236 | })() |
| 237 | |
| 238 | /** |
| 239 | * 更新任务标题（用于编辑标题成功后更新本地状态） |
| 240 | */ |
| 241 | const updateTaskTitle = useCallback((newTitle: string) => { |
| 242 | setTask(prev => (prev ? { ...prev, title: newTitle } : prev)) |
| 243 | }, []) |
| 244 | |
| 245 | return { |
| 246 | task, |
| 247 | displayMessages, |
| 248 | workflowSteps, |
| 249 | isLoading, |
| 250 | isGenerating, |
| 251 | progress: storeProgress, |
| 252 | isActiveTask, |
| 253 | setLocalMessages, |
| 254 | setLocalIsGenerating, |
| 255 | updateTaskTitle, |
| 256 | } |
| 257 | } |
| 258 |