From 34fd5bfb1aa8ed1b72121c47489e1eea9d02a0dd Mon Sep 17 00:00:00 2001
From: Huarch <huarch97@outlook.com>
Date: Mon, 8 Jun 2026 15:13:21 +0800
Subject: [PATCH] fix(chat): guard generated title events

---
 .../chat/hooks/useAgentChatSession.test.tsx   | 59 +++++++++++++++++
 .../chat/hooks/useAgentChatSession.ts         | 65 +++++++++++++------
 2 files changed, 103 insertions(+), 21 deletions(-)

diff --git a/src/components/chat/hooks/useAgentChatSession.test.tsx b/src/components/chat/hooks/useAgentChatSession.test.tsx
index d7322b5..f7eb1b7 100644
--- a/src/components/chat/hooks/useAgentChatSession.test.tsx
+++ b/src/components/chat/hooks/useAgentChatSession.test.tsx
@@ -61,6 +61,7 @@ describe("useAgentChatSession", () => {
     jest.mocked(streamAgentChat).mockImplementation(async () => undefined);
     deleteChatSession.mockImplementation(async () => undefined);
     saveActiveChatState.mockImplementation(async (state) => state.sessionId);
+    updateChatSessionTitle.mockImplementation(async () => undefined);
   });
 
   it("does not add a new empty session to history until there is actual chat content", async () => {
@@ -522,6 +523,64 @@ describe("useAgentChatSession", () => {
     );
   });
 
+  it("does not apply a late generated title to a newly created session", async () => {
+    listChatSessions.mockResolvedValue([]);
+    let emitStreamEvent: ((event: StreamEvent) => void) | undefined;
+    let resolveStream: (() => void) | undefined;
+    jest.mocked(streamAgentChat).mockImplementationOnce(async ({ onEvent }) => {
+      emitStreamEvent = onEvent;
+      await new Promise<void>((resolve) => {
+        resolveStream = resolve;
+      });
+    });
+
+    const { result } = renderHook(() =>
+      useAgentChatSession({
+        projectId: "project-1",
+        onToolCall: jest.fn(),
+      }),
+    );
+
+    await waitFor(() => expect(result.current.isHydrating).toBe(false));
+
+    await act(async () => {
+      void result.current.sendPrompt("帮我分析一下");
+      await Promise.resolve();
+    });
+
+    act(() => {
+      emitStreamEvent?.({
+        type: "done",
+        sessionId: "old-session",
+      });
+    });
+
+    await waitFor(() => expect(result.current.isStreaming).toBe(false));
+
+    act(() => {
+      result.current.createSession();
+    });
+
+    expect(result.current.sessionTitle).toBe("新对话");
+
+    await act(async () => {
+      emitStreamEvent?.({
+        type: "session_title",
+        sessionId: "old-session",
+        title: "旧请求标题",
+      });
+      resolveStream?.();
+      await Promise.resolve();
+    });
+
+    expect(result.current.sessionTitle).toBe("新对话");
+    expect(updateChatSessionTitle).toHaveBeenCalledWith(
+      "old-session",
+      "旧请求标题",
+      { isTitleManuallyEdited: false },
+    );
+  });
+
   it("asks the backend to undo the previous user turn before regenerating", async () => {
     listChatSessions.mockResolvedValue([]);
 
diff --git a/src/components/chat/hooks/useAgentChatSession.ts b/src/components/chat/hooks/useAgentChatSession.ts
index f9145d2..f046270 100644
--- a/src/components/chat/hooks/useAgentChatSession.ts
+++ b/src/components/chat/hooks/useAgentChatSession.ts
@@ -234,6 +234,7 @@ export const useAgentChatSession = ({
   const abortRef = useRef<AbortController | null>(null);
   const sessionIdRef = useRef<string | undefined>(undefined);
   const messagesRef = useRef<Message[]>([]);
+  const branchGroupsRef = useRef<BranchGroup[]>([]);
   const resumeStreamingSessionRef = useRef<((sessionId: string) => void) | null>(null);
   const isSessionTitleManuallyEditedRef = useRef(false);
   const cancelPromiseRef = useRef<Promise<void> | null>(null);
@@ -256,6 +257,10 @@ export const useAgentChatSession = ({
     messagesRef.current = messages;
   }, [messages]);
 
+  useEffect(() => {
+    branchGroupsRef.current = branchGroups;
+  }, [branchGroups]);
+
   useEffect(() => {
     isSessionTitleManuallyEditedRef.current = isSessionTitleManuallyEdited;
   }, [isSessionTitleManuallyEdited]);
@@ -444,7 +449,12 @@ export const useAgentChatSession = ({
         assistantMessageId?: string;
       },
     ) => {
-      if ("sessionId" in event && event.sessionId && event.sessionId !== sessionIdRef.current) {
+      if (
+        event.type !== "session_title" &&
+        "sessionId" in event &&
+        event.sessionId &&
+        event.sessionId !== sessionIdRef.current
+      ) {
         sessionIdRef.current = event.sessionId;
         setSessionId(event.sessionId);
       }
@@ -457,6 +467,39 @@ export const useAgentChatSession = ({
         return;
       }
 
+      if (event.type === "session_title") {
+        const nextTitle = event.title.trim();
+        if (nextTitle && !isSessionTitleManuallyEditedRef.current) {
+          const currentSessionId = sessionIdRef.current;
+          const targetSessionId = event.sessionId || currentSessionId;
+          if (targetSessionId === currentSessionId) {
+            setSessionTitle(nextTitle);
+            lastPersistedStateKeyRef.current = createPersistedStateKey({
+              sessionId: targetSessionId,
+              title: nextTitle,
+              isTitleManuallyEdited: false,
+              messages: messagesRef.current,
+              branchGroups: branchGroupsRef.current,
+            });
+          }
+          if (targetSessionId) {
+            const currentNonce = ++titleUpdateNonceRef.current;
+            void updateChatSessionTitle(targetSessionId, 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);
+              });
+          }
+        }
+        return;
+      }
+
       const assistantMessageId = getLastAssistantMessageId(options?.assistantMessageId);
       if (!assistantMessageId) {
         return;
@@ -519,26 +562,6 @@ export const useAgentChatSession = ({
             };
           }),
         );
-      } 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) => {