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

2026-06-04 16:27:15 +08:00
parent e60e1f6453
commit 7764e25398
6 changed files with 559 additions and 197 deletions
@@ -71,6 +71,8 @@ export type ChatSessionSummary = {
  title: string;
  createdAt: number;
  updatedAt: number;
+  isStreaming?: boolean;
+  runStatus?: string;
 };

 export type LoadedChatState = {
@@ -79,4 +81,6 @@ export type LoadedChatState = {
  isTitleManuallyEdited?: boolean;
  messages: Message[];
  branchGroups: BranchGroup[];
+  isStreaming?: boolean;
+  runStatus?: string;
 };
@@ -14,6 +14,8 @@ type BackendSessionPayload = {
  title?: string;
  created_at?: string | number;
  updated_at?: string | number;
+  is_streaming?: boolean;
+  run_status?: string;
 };

 export const createEmptyChatState = (): LoadedChatState => ({
@@ -76,6 +78,8 @@ const fetchBackendChatSessions = async (): Promise<ChatSessionSummary[]> => {
      title: normalizeTitle(session.title),
      createdAt: toMillis(session.created_at),
      updatedAt: toMillis(session.updated_at),
+      isStreaming: session.is_streaming,
+      runStatus: session.run_status,
    }))
    .filter((session) => Boolean(session.id))
    .sort(compareSessionsByAnchorTime);
@@ -104,6 +108,8 @@ const fetchBackendChatSession = async (sessionId: string): Promise<LoadedChatSta
    session_id?: string;
    messages?: Message[];
    branch_groups?: BranchGroup[];
+    is_streaming?: boolean;
+    run_status?: string;
  };
  return {
    title: normalizeTitle(payload.title),
@@ -111,6 +117,8 @@ const fetchBackendChatSession = async (sessionId: string): Promise<LoadedChatSta
    messages: sanitizeMessages(payload.messages),
    sessionId: payload.session_id ?? payload.id,
    branchGroups: sanitizeBranchGroups(payload.branch_groups),
+    isStreaming: payload.is_streaming ?? false,
+    runStatus: payload.run_status,
  };
 };

@@ -3,12 +3,13 @@
 import { act, renderHook, waitFor } from "@testing-library/react";

 import { useAgentChatSession } from "./useAgentChatSession";
-import { streamAgentChat } from "@/lib/chatStream";
+import { abortAgentChat, resumeAgentChatStream, streamAgentChat } from "@/lib/chatStream";
 import type { StreamEvent } from "@/lib/chatStream";

 jest.mock("@/lib/chatStream", () => ({
  abortAgentChat: jest.fn(async () => undefined),
  forkAgentChat: jest.fn(async () => "forked-session"),
+  resumeAgentChatStream: jest.fn(async () => undefined),
  streamAgentChat: jest.fn(async () => undefined),
 }));

@@ -42,7 +43,12 @@ describe("useAgentChatSession", () => {
    listChatSessions.mockReset();
    saveActiveChatState.mockReset();
    updateChatSessionTitle.mockReset();
+    jest.mocked(abortAgentChat).mockReset();
+    jest.mocked(resumeAgentChatStream).mockReset();
    jest.mocked(streamAgentChat).mockReset();
+    jest.mocked(abortAgentChat).mockImplementation(async () => undefined);
+    jest.mocked(resumeAgentChatStream).mockImplementation(async () => undefined);
+    jest.mocked(streamAgentChat).mockImplementation(async () => undefined);
    saveActiveChatState.mockImplementation(async (state) => state.sessionId);
  });

@@ -103,7 +109,7 @@ describe("useAgentChatSession", () => {
    ]);
  });

-  it("waits for the stream session id before persisting a new streaming conversation", async () => {
+  it("persists a new conversation only after the stream is done", async () => {
    listChatSessions.mockResolvedValue([]);
    let emitStreamEvent: ((event: StreamEvent) => void) | undefined;
    jest.mocked(streamAgentChat).mockImplementationOnce(async ({ onEvent }) => {
@@ -147,15 +153,140 @@ describe("useAgentChatSession", () => {
        jest.advanceTimersByTime(200);
      });

-      expect(saveActiveChatState).toHaveBeenCalledTimes(1);
+      expect(saveActiveChatState).not.toHaveBeenCalled();
+
+      act(() => {
+        emitStreamEvent?.({
+          type: "done",
+          sessionId: "chat-stream-1",
+        });
+      });
+
+      await act(async () => {
+        jest.advanceTimersByTime(200);
+      });
+
+      await waitFor(() => expect(saveActiveChatState).toHaveBeenCalledTimes(1));
      expect(saveActiveChatState.mock.calls[0][0]).toMatchObject({
        sessionId: "chat-stream-1",
+        messages: [
+          expect.objectContaining({ role: "user", content: "第一条消息" }),
+          expect.objectContaining({ role: "assistant", content: "收到" }),
+        ],
      });
    } finally {
      jest.useRealTimers();
    }
  });

+  it("hydrates a backend streaming session and resumes its stream", async () => {
+    listChatSessions.mockResolvedValue([
+      {
+        id: "session-streaming",
+        title: "运行中",
+        createdAt: 1,
+        updatedAt: 2,
+        isStreaming: true,
+        runStatus: "running",
+      },
+    ]);
+
+    const { result } = renderHook(() =>
+      useAgentChatSession({
+        projectId: "project-1",
+        onToolCall: jest.fn(),
+      }),
+    );
+
+    await waitFor(() => expect(result.current.isHydrating).toBe(false));
+
+    expect(result.current.isStreaming).toBe(true);
+    expect(result.current.activeSessionId).toBe("session-loaded");
+    expect(resumeAgentChatStream).toHaveBeenCalledWith(
+      expect.objectContaining({
+        sessionId: "session-loaded",
+      }),
+    );
+  });
+
+  it("updates resumed messages from state, token, and done events", async () => {
+    listChatSessions.mockResolvedValue([
+      {
+        id: "session-streaming",
+        title: "运行中",
+        createdAt: 1,
+        updatedAt: 2,
+        isStreaming: true,
+      },
+    ]);
+    jest.mocked(resumeAgentChatStream).mockImplementationOnce(async ({ onEvent }) => {
+      onEvent({
+        type: "state",
+        sessionId: "session-loaded",
+        messages: [
+          { id: "u1", role: "user", content: "继续分析" },
+          { id: "a1", role: "assistant", content: "已有" },
+        ],
+        isStreaming: true,
+        runStatus: "running",
+      });
+      onEvent({
+        type: "token",
+        sessionId: "session-loaded",
+        content: "输出",
+      });
+      onEvent({
+        type: "done",
+        sessionId: "session-loaded",
+      });
+    });
+
+    const { result } = renderHook(() =>
+      useAgentChatSession({
+        projectId: "project-1",
+        onToolCall: jest.fn(),
+      }),
+    );
+
+    await waitFor(() => expect(result.current.isHydrating).toBe(false));
+    await waitFor(() => expect(result.current.isStreaming).toBe(false));
+
+    expect(result.current.messages).toEqual([
+      expect.objectContaining({ id: "u1", role: "user", content: "继续分析" }),
+      expect.objectContaining({ id: "a1", role: "assistant", content: "已有输出" }),
+    ]);
+  });
+
+  it("aborts a resumed streaming session through the backend abort endpoint", async () => {
+    listChatSessions.mockResolvedValue([
+      {
+        id: "session-streaming",
+        title: "运行中",
+        createdAt: 1,
+        updatedAt: 2,
+        isStreaming: true,
+      },
+    ]);
+    jest.mocked(resumeAgentChatStream).mockImplementationOnce(async () => {
+      await new Promise<void>(() => undefined);
+    });
+
+    const { result } = renderHook(() =>
+      useAgentChatSession({
+        projectId: "project-1",
+        onToolCall: jest.fn(),
+      }),
+    );
+
+    await waitFor(() => expect(result.current.isStreaming).toBe(true));
+
+    act(() => {
+      result.current.abort();
+    });
+
+    expect(abortAgentChat).toHaveBeenCalledWith("session-loaded");
+  });
+
  it("ignores generated session titles after the title was edited manually", async () => {
    listChatSessions.mockResolvedValue([]);
    jest.mocked(streamAgentChat).mockImplementationOnce(async ({ onEvent }) => {
@@ -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,
          });
        }