增加流式信息中断处理机制

src/components/chat/hooks/useAgentChatSession.ts

@@ -2,7 +2,12 @@
 
 import { useCallback, useEffect, useRef, useState } from "react";
 
-import { abortAgentChat, forkAgentChat, streamAgentChat } from "@/lib/chatStream";
+import {
+  abortAgentChat,
+  forkAgentChat,
+  resumeAgentChatStream,
+  streamAgentChat,
+} from "@/lib/chatStream";
 import type { AgentModel, StreamEvent } from "@/lib/chatStream";
 import type {
   AgentArtifact,
@@ -164,6 +169,8 @@ export const useAgentChatSession = ({
   const [isHydrating, setIsHydrating] = useState(true);
   const abortRef = useRef<AbortController | null>(null);
   const sessionIdRef = useRef<string | undefined>(undefined);
+  const messagesRef = useRef<Message[]>([]);
+  const resumeStreamingSessionRef = useRef<((sessionId: string) => void) | null>(null);
   const isSessionTitleManuallyEditedRef = useRef(false);
   const cancelPromiseRef = useRef<Promise<void> | null>(null);
   const titleUpdateNonceRef = useRef(0);
@@ -181,6 +188,10 @@ export const useAgentChatSession = ({
     sessionIdRef.current = sessionId;
   }, [sessionId]);
 
+  useEffect(() => {
+    messagesRef.current = messages;
+  }, [messages]);
+
   useEffect(() => {
     isSessionTitleManuallyEditedRef.current = isSessionTitleManuallyEdited;
   }, [isSessionTitleManuallyEdited]);
@@ -216,10 +227,11 @@ export const useAgentChatSession = ({
       }
 
       try {
-        const [loadedState, sessions] = await Promise.all([
-          Promise.resolve(createEmptyChatState()),
-          listChatSessions(),
-        ]);
+        const sessions = await listChatSessions();
+        const streamingSession = sessions.find((session) => session.isStreaming);
+        const loadedState = streamingSession
+          ? await loadChatSessionById(streamingSession.id)
+          : createEmptyChatState();
         if (cancelled) return;
 
         sessionIdRef.current = loadedState.sessionId;
@@ -234,6 +246,12 @@ export const useAgentChatSession = ({
         setSessionId(loadedState.sessionId);
         setBranchGroups(loadedState.branchGroups);
         setChatSessions(sessions);
+        if (
+          loadedState.sessionId &&
+          (loadedState.isStreaming || streamingSession?.isStreaming)
+        ) {
+          resumeStreamingSessionRef.current?.(loadedState.sessionId);
+        }
       } catch (error) {
         console.error("[GlobalChatbox] Failed to hydrate chat state:", error);
       } finally {
@@ -255,6 +273,10 @@ export const useAgentChatSession = ({
 
     const currentHydrationNonce = hydrationNonceRef.current;
     const persistTimer = window.setTimeout(() => {
+      if (isStreaming) {
+        return;
+      }
+
       const state: LoadedChatState = {
         title: sessionTitle,
         isTitleManuallyEdited: isSessionTitleManuallyEdited,
@@ -262,13 +284,6 @@ export const useAgentChatSession = ({
         sessionId,
         branchGroups,
       };
-      if (
-        isStreaming &&
-        !state.sessionId &&
-        state.messages.length > 0
-      ) {
-        return;
-      }
 
       const currentStateKey = createPersistedStateKey(state);
       if (currentStateKey === lastPersistedStateKeyRef.current) {
@@ -351,6 +366,150 @@ export const useAgentChatSession = ({
     );
   }, []);
 
+  const getLastAssistantMessageId = useCallback((fallback?: string) => {
+    const assistant = [...messagesRef.current]
+      .reverse()
+      .find((message) => message.role === "assistant");
+    return assistant?.id ?? fallback;
+  }, []);
+
+  const applyStreamEvent = useCallback(
+    (
+      event: StreamEvent,
+      options?: {
+        assistantMessageId?: string;
+      },
+    ) => {
+      if ("sessionId" in event && event.sessionId && event.sessionId !== sessionIdRef.current) {
+        sessionIdRef.current = event.sessionId;
+        setSessionId(event.sessionId);
+      }
+
+      if (event.type === "state") {
+        const nextMessages = cloneMessages(event.messages as Message[]);
+        messagesRef.current = nextMessages;
+        setMessages(nextMessages);
+        setIsStreaming(event.isStreaming);
+        return;
+      }
+
+      const assistantMessageId = getLastAssistantMessageId(options?.assistantMessageId);
+      if (!assistantMessageId) {
+        return;
+      }
+
+      if (event.type === "token") {
+        setMessages((prev) =>
+          prev.map((message) =>
+            message.id === assistantMessageId
+              ? {
+                  ...message,
+                  content: message.content + event.content,
+                  isError: false,
+                }
+              : message,
+          ),
+        );
+      } else if (event.type === "progress") {
+        setMessages((prev) =>
+          prev.map((message) =>
+            message.id === assistantMessageId
+              ? { ...message, progress: upsertProgress(message.progress, event) }
+              : message,
+          ),
+        );
+      } else if (event.type === "tool_call") {
+        onToolCall(event, {
+          assistantMessageId,
+          appendArtifact,
+        });
+      } else if (event.type === "session_title") {
+        const nextTitle = event.title.trim();
+        if (nextTitle && !isSessionTitleManuallyEditedRef.current) {
+          setSessionTitle(nextTitle);
+          const currentSessionId = sessionIdRef.current;
+          if (currentSessionId) {
+            const currentNonce = ++titleUpdateNonceRef.current;
+            void updateChatSessionTitle(currentSessionId, nextTitle, {
+              isTitleManuallyEdited: false,
+            })
+              .then(() => listChatSessions())
+              .then((sessions) => {
+                if (titleUpdateNonceRef.current !== currentNonce) return;
+                setChatSessions(sessions);
+              })
+              .catch((error) => {
+                console.error("[GlobalChatbox] Failed to persist session title:", error);
+              });
+          }
+        }
+      } else if (event.type === "done") {
+        setMessages((prev) =>
+          prev.map((message) => {
+            if (message.id !== assistantMessageId) return message;
+            const completedProgress = completeRunningProgress(message.progress);
+            if (
+              message.content.trim().length === 0 &&
+              !(message.artifacts?.length)
+            ) {
+              return {
+                ...message,
+                content:
+                  "Agent 已完成处理，但没有生成文本回答。请查看过程记录，或换个更具体的问题重试。",
+                progress: completedProgress,
+              };
+            }
+            return { ...message, progress: completedProgress };
+          }),
+        );
+        setIsStreaming(false);
+      } else if (event.type === "error") {
+        setMessages((prev) =>
+          prev.map((message) =>
+            message.id === assistantMessageId
+              ? {
+                  ...message,
+                  content: message.content || `⚠️ **错误：** ${event.message}`,
+                  isError: true,
+                  progress: completeRunningProgress(message.progress),
+                }
+              : message,
+          ),
+        );
+        setIsStreaming(false);
+      }
+    },
+    [appendArtifact, getLastAssistantMessageId, onToolCall],
+  );
+
+  const resumeStreamingSession = useCallback(
+    (nextSessionId: string) => {
+      const controller = new AbortController();
+      abortRef.current?.abort();
+      abortRef.current = controller;
+      setIsStreaming(true);
+
+      void resumeAgentChatStream({
+        sessionId: nextSessionId,
+        signal: controller.signal,
+        onEvent: (event) => applyStreamEvent(event),
+      })
+        .catch((error) => {
+          if (!controller.signal.aborted) {
+            console.error("[GlobalChatbox] Failed to resume chat stream:", error);
+            setIsStreaming(false);
+          }
+        })
+        .finally(() => {
+          if (abortRef.current === controller) {
+            abortRef.current = null;
+          }
+        });
+    },
+    [applyStreamEvent],
+  );
+  resumeStreamingSessionRef.current = resumeStreamingSession;
+
   const runPrompt = useCallback(
     async ({
       prompt: rawPrompt,
@@ -372,8 +531,10 @@ export const useAgentChatSession = ({
         preparedMessages ??
         [...messages, nextUserMessage, nextAssistantMessage];
 
+      const clonedNextMessages = cloneMessages(nextMessages);
       setIsStreaming(true);
-      setMessages(cloneMessages(nextMessages));
+      messagesRef.current = clonedNextMessages;
+      setMessages(clonedNextMessages);
       if (sessionIdOverride !== undefined) {
         sessionIdRef.current = sessionIdOverride;
         setSessionId(sessionIdOverride);
@@ -388,93 +549,10 @@ export const useAgentChatSession = ({
           sessionId: sessionIdOverride ?? sessionIdRef.current,
           model: getModel?.(),
           signal: controller.signal,
-          onEvent: (event) => {
-            if ("sessionId" in event && event.sessionId && event.sessionId !== sessionIdRef.current) {
-              sessionIdRef.current = event.sessionId;
-              setSessionId(event.sessionId);
-            }
-
-            if (event.type === "token") {
-              setMessages((prev) =>
-                prev.map((message) =>
-                  message.id === nextAssistantMessage.id
-                    ? {
-                        ...message,
-                        content: message.content + event.content,
-                        isError: false,
-                      }
-                    : message,
-                ),
-              );
-            } else if (event.type === "progress") {
-              setMessages((prev) =>
-                prev.map((message) =>
-                  message.id === nextAssistantMessage.id
-                    ? { ...message, progress: upsertProgress(message.progress, event) }
-                    : message,
-                ),
-              );
-            } else if (event.type === "tool_call") {
-              onToolCall(event, {
-                assistantMessageId: nextAssistantMessage.id,
-                appendArtifact,
-              });
-            } else if (event.type === "session_title") {
-              const nextTitle = event.title.trim();
-              if (nextTitle && !isSessionTitleManuallyEditedRef.current) {
-                setSessionTitle(nextTitle);
-                const currentSessionId = sessionIdRef.current;
-                if (currentSessionId) {
-                  const currentNonce = ++titleUpdateNonceRef.current;
-                  void updateChatSessionTitle(currentSessionId, nextTitle, {
-                    isTitleManuallyEdited: false,
-                  })
-                    .then(() => listChatSessions())
-                    .then((sessions) => {
-                      if (titleUpdateNonceRef.current !== currentNonce) return;
-                      setChatSessions(sessions);
-                    })
-                    .catch((error) => {
-                      console.error("[GlobalChatbox] Failed to persist session title:", error);
-                    });
-                }
-              }
-            } else if (event.type === "done") {
-              setMessages((prev) =>
-                prev.map((message) => {
-                  if (message.id !== nextAssistantMessage.id) return message;
-                  const completedProgress = completeRunningProgress(message.progress);
-                  if (
-                    message.content.trim().length === 0 &&
-                    !(message.artifacts?.length)
-                  ) {
-                    return {
-                      ...message,
-                      content:
-                        "Agent 已完成处理，但没有生成文本回答。请查看过程记录，或换个更具体的问题重试。",
-                      progress: completedProgress,
-                    };
-                  }
-                  return { ...message, progress: completedProgress };
-                }),
-              );
-              setIsStreaming(false);
-            } else if (event.type === "error") {
-              setMessages((prev) =>
-                prev.map((message) =>
-                  message.id === nextAssistantMessage.id
-                    ? {
-                        ...message,
-                        content: message.content || `⚠️ **错误：** ${event.message}`,
-                        isError: true,
-                        progress: completeRunningProgress(message.progress),
-                      }
-                    : message,
-                ),
-              );
-              setIsStreaming(false);
-            }
-          },
+          onEvent: (event) =>
+            applyStreamEvent(event, {
+              assistantMessageId: nextAssistantMessage.id,
+            }),
         });
       } catch (error) {
         if (controller.signal.aborted) {
@@ -520,7 +598,7 @@ export const useAgentChatSession = ({
         setIsStreaming(false);
       }
     },
-    [appendArtifact, getModel, isHydrating, isStreaming, messages, onBeforeSend, onToolCall],
+    [applyStreamEvent, getModel, isHydrating, isStreaming, messages, onBeforeSend],
   );
 
   const abort = useCallback(() => {
@@ -587,13 +665,18 @@ export const useAgentChatSession = ({
         setSessionId(nextState.sessionId);
         setBranchGroups(nextState.branchGroups);
         setChatSessions(sessions);
+        if (nextState.sessionId && nextState.isStreaming) {
+          resumeStreamingSession(nextState.sessionId);
+        } else {
+          setIsStreaming(false);
+        }
       } catch (error) {
         console.error("[GlobalChatbox] Failed to switch chat session:", error);
       } finally {
         setIsHydrating(false);
       }
     },
-    [isHydrating, isStreaming],
+    [isHydrating, isStreaming, resumeStreamingSession],
   );
 
   const removeSession = useCallback(
@@ -683,7 +766,6 @@ export const useAgentChatSession = ({
             title: normalizedTitle,
             isTitleManuallyEdited: true,
             messages,
-            sessionId: sessionIdRef.current,
             branchGroups,
           });
         }