添加图片生成状态管理功能

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}
         />
 
         <div className="sticky bottom-0 z-1 mx-auto w-full max-w-4xl bg-background px-2 pb-3 md:px-4 md:pb-4">

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) ? <QianwenVideoLoading /> : 
-            isGeminiModel ? <GeminiImageLoading /> : 
-            isVolcanoVisionModel ? <VolcanoVisionLoading /> : 
-            <ThinkingMessage />
+          {/* 🆕 支持恢复的生成状态（用户离开再返回） */}
+          {(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 <QianwenVideoLoading />;
+              if (showGemini) return <GeminiImageLoading />;
+              if (showVolcano) return <VolcanoVisionLoading />;
+              return <ThinkingMessage />;
+            })()
           )}
 
           <div
@@ -260,6 +272,10 @@ export const Messages = memo(PureMessages, (prevProps: MessagesProps, nextProps:
   if (prevProps.selectedModelId !== nextProps.selectedModelId) {
     return false;
   }
+  // 🆕 检查持久化生成状态变化
+  if (!equal(prevProps.pendingGeneration, nextProps.pendingGeneration)) {
+    return false;
+  }
   if (prevProps.messages.length !== nextProps.messages.length) {
     return false;
   }

hooks/use-pending-generation.ts

@@ -0,0 +1,108 @@
+"use client";
+
+import { useEffect, useState, useCallback } from "react";
+
+/**
+ * 持久化的生成任务状态
+ */
+export interface PendingGeneration {
+  chatId: string;
+  modelId: string;
+  modelType: "gemini" | "volcano-vision" | "volcano-video" | "qianwen-video" | "other";
+  timestamp: number;
+  userMessageId?: string;
+}
+
+const STORAGE_KEY = "ai-chatbot-pending-generations";
+const EXPIRY_TIME = 5 * 60 * 1000; // 5 分钟超时
+
+/**
+ * 获取所有未完成的生成任务
+ */
+function getPendingGenerations(): PendingGeneration[] {
+  if (typeof window === "undefined") return [];
+  try {
+    const data = localStorage.getItem(STORAGE_KEY);
+    if (!data) return [];
+    const generations: PendingGeneration[] = JSON.parse(data);
+    // 过滤掉过期的任务
+    const now = Date.now();
+    return generations.filter((g) => 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<PendingGeneration | null>(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,
+  };
+}
+