diff --git a/components/chat.tsx b/components/chat.tsx index 735e4e5..680f1d3 100644 --- a/components/chat.tsx +++ b/components/chat.tsx @@ -20,6 +20,7 @@ import { import { useArtifactSelector } from "@/hooks/use-artifact"; import { useAutoResume } from "@/hooks/use-auto-resume"; import { useChatVisibility } from "@/hooks/use-chat-visibility"; +import { usePendingGeneration, type PendingGeneration } from "@/hooks/use-pending-generation"; import type { Vote } from "@/lib/db/schema"; import type { UserType } from "@/app/(auth)/auth"; import { ChatSDKError } from "@/lib/errors"; @@ -62,6 +63,9 @@ export function Chat({ initialVisibilityType, }); + // 🆕 持久化生成任务状态（解决用户离开再返回时加载状态丢失） + const { pendingGeneration, startGeneration, completeGeneration } = usePendingGeneration(id); + const { mutate } = useSWRConfig(); // Handle browser back/forward navigation @@ -247,6 +251,9 @@ export function Chat({ console.log("[Chat] Received appendMessage:", messageData); + // 🆕 收到消息，清除持久化状态 + completeGeneration(); + setMessages((prevMessages) => { if (prevMessages.some(m => m.id === messageData.id)) { return prevMessages; @@ -272,6 +279,8 @@ export function Chat({ }, onFinish: () => { mutate(unstable_serialize(getChatHistoryPaginationKey)); + // 🆕 生成完成，清除持久化状态 + completeGeneration(); }, onError: (error) => { if (error instanceof ChatSDKError) { @@ -290,10 +299,27 @@ export function Chat({ }, }); - // 调试：监控聊天状态变化 + // 调试：监控聊天状态变化 & 持久化生成状态 useEffect(() => { console.log(`[Chat] Status: ${status}`); - }, [status]); + + // 🆕 当状态变为 submitted 时，记录生成任务 + if (status === "submitted") { + // 判断模型类型 + let modelType: PendingGeneration["modelType"] = "other"; + if (currentModel?.providerType === "gemini") { + modelType = "gemini"; + } else if (currentModel?.providerType === "volcano-vision") { + modelType = "volcano-vision"; + } else if (currentModel?.providerType === "volcano-video") { + modelType = "volcano-video"; + } else if (currentModel?.providerType === "qianwen-video") { + modelType = "qianwen-video"; + } + + startGeneration(currentModelId, modelType); + } + }, [status, currentModelId, currentModel?.providerType, startGeneration]); const searchParams = useSearchParams(); const query = searchParams.get("query"); @@ -413,12 +439,13 @@ export function Chat({ messages={messages} regenerate={regenerate} sendMessage={sendMessage} - selectedModelId={initialChatModel} + selectedModelId={currentModelId} setMessages={setMessages} status={status} votes={votes} onImageEdit={handleImageEdit} onDrawingComplete={handleDrawingComplete} + pendingGeneration={pendingGeneration} />
diff --git a/components/messages.tsx b/components/messages.tsx index bb7f82f..36b7e00 100644 --- a/components/messages.tsx +++ b/components/messages.tsx @@ -4,6 +4,7 @@ import { ArrowDownIcon } from "lucide-react"; import { memo, useEffect, useState, useRef } from "react"; import useSWR from "swr"; import { useMessages } from "@/hooks/use-messages"; +import type { PendingGeneration } from "@/hooks/use-pending-generation"; import type { Vote } from "@/lib/db/schema"; import type { ChatMessage } from "@/lib/types"; import { useDataStream } from "./data-stream-provider"; @@ -36,6 +37,7 @@ type MessagesProps = { selectedModelId: string; onImageEdit?: (imageUrl: string) => void; onDrawingComplete?: (annotatedImageUrl: string) => void; + pendingGeneration?: PendingGeneration | null; // 🆕 持久化的生成状态 }; function PureMessages({ @@ -50,6 +52,7 @@ function PureMessages({ selectedModelId, onImageEdit, onDrawingComplete, + pendingGeneration, }: MessagesProps) { const { containerRef: messagesContainerRef, @@ -219,11 +222,20 @@ function PureMessages({ ))} {/* 等待动画：根据模型类型显示不同的加载效果 */} - {showThinkingWithDelay && ( - (isQianwenVideoModel || isVideoModel) ? : - isGeminiModel ? : - isVolcanoVisionModel ? : - + {/* 🆕 支持恢复的生成状态（用户离开再返回） */} + {(showThinkingWithDelay || (pendingGeneration && status === "ready")) && ( + (() => { + // 优先使用恢复的状态，否则使用当前模型类型 + const modelType = pendingGeneration?.modelType; + const showGemini = modelType === "gemini" || (!modelType && isGeminiModel); + const showVolcano = modelType === "volcano-vision" || (!modelType && isVolcanoVisionModel); + const showVideo = modelType === "volcano-video" || modelType === "qianwen-video" || (!modelType && (isQianwenVideoModel || isVideoModel)); + + if (showVideo) return ; + if (showGemini) return ; + if (showVolcano) return ; + return ; + })() )}
now - g.timestamp < EXPIRY_TIME); + } catch { + return []; + } +} + +/** + * 保存未完成的生成任务 + */ +function savePendingGenerations(generations: PendingGeneration[]) { + if (typeof window === "undefined") return; + try { + localStorage.setItem(STORAGE_KEY, JSON.stringify(generations)); + } catch (e) { + console.error("Failed to save pending generations:", e); + } +} + +/** + * Hook: 管理持久化的生成任务状态 + * 解决用户离开页面再返回时加载状态丢失的问题 + */ +export function usePendingGeneration(chatId: string) { + const [pendingGeneration, setPendingGeneration] = useState(null); + + // 初始化：检查是否有当前 chatId 的未完成任务 + useEffect(() => { + const generations = getPendingGenerations(); + const found = generations.find((g) => g.chatId === chatId); + if (found) { + setPendingGeneration(found); + console.log("[PendingGeneration] 恢复处理中状态:", found); + } + }, [chatId]); + + // 开始生成：保存状态 + const startGeneration = useCallback( + (modelId: string, modelType: PendingGeneration["modelType"], userMessageId?: string) => { + const newGeneration: PendingGeneration = { + chatId, + modelId, + modelType, + timestamp: Date.now(), + userMessageId, + }; + + // 更新 localStorage + const generations = getPendingGenerations(); + // 移除同一 chatId 的旧记录 + const filtered = generations.filter((g) => g.chatId !== chatId); + savePendingGenerations([...filtered, newGeneration]); + + setPendingGeneration(newGeneration); + console.log("[PendingGeneration] 开始生成:", newGeneration); + }, + [chatId] + ); + + // 完成生成：清除状态 + const completeGeneration = useCallback(() => { + const generations = getPendingGenerations(); + const filtered = generations.filter((g) => g.chatId !== chatId); + savePendingGenerations(filtered); + + setPendingGeneration(null); + console.log("[PendingGeneration] 生成完成，清除状态"); + }, [chatId]); + + // 检查是否有恢复的任务（用户离开再返回） + const hasRestoredPending = pendingGeneration !== null && Date.now() - pendingGeneration.timestamp > 1000; + + return { + pendingGeneration, + hasRestoredPending, + startGeneration, + completeGeneration, + }; +} +