refactor: simplify chat fork flow

src/components/chat/hooks/useAgentChatSession.test.tsx

@@ -5,6 +5,7 @@ import { act, renderHook, waitFor } from "@testing-library/react";
 import { useAgentChatSession } from "./useAgentChatSession";
 import {
   abortAgentChat,
+  forkAgentChat,
   replyAgentPermission,
   resumeAgentChatStream,
   streamAgentChat,
@@ -30,7 +31,6 @@ jest.mock("../chatStorage", () => ({
     isTitleManuallyEdited: false,
     messages: [],
     sessionId: undefined,
-    branchGroups: [],
   })),
   deleteChatSession: (...args: unknown[]) => deleteChatSession(...args),
   listChatSessions: (...args: unknown[]) => listChatSessions(...args),
@@ -39,7 +39,6 @@ jest.mock("../chatStorage", () => ({
     isTitleManuallyEdited: false,
     messages: [],
     sessionId: "session-loaded",
-    branchGroups: [],
   })),
   saveActiveChatState: (...args: unknown[]) => saveActiveChatState(...args),
   updateChatSessionTitle: (...args: unknown[]) => updateChatSessionTitle(...args),
@@ -52,10 +51,12 @@ describe("useAgentChatSession", () => {
     saveActiveChatState.mockReset();
     updateChatSessionTitle.mockReset();
     jest.mocked(abortAgentChat).mockReset();
+    jest.mocked(forkAgentChat).mockReset();
     jest.mocked(replyAgentPermission).mockReset();
     jest.mocked(resumeAgentChatStream).mockReset();
     jest.mocked(streamAgentChat).mockReset();
     jest.mocked(abortAgentChat).mockImplementation(async () => undefined);
+    jest.mocked(forkAgentChat).mockImplementation(async () => "forked-session");
     jest.mocked(replyAgentPermission).mockImplementation(async () => undefined);
     jest.mocked(resumeAgentChatStream).mockImplementation(async () => undefined);
     jest.mocked(streamAgentChat).mockImplementation(async () => undefined);
@@ -596,9 +597,10 @@ describe("useAgentChatSession", () => {
     await act(async () => {
       await result.current.sendPrompt("重新分析压力异常");
     });
+    const assistantMessageId = result.current.messages[1]?.id ?? "";
 
     await act(async () => {
-      await result.current.regenerate();
+      await result.current.regenerate(assistantMessageId);
     });
 
     expect(streamAgentChat).toHaveBeenNthCalledWith(
@@ -609,4 +611,91 @@ describe("useAgentChatSession", () => {
       }),
     );
   });
+
+  it("replaces the current chain when regenerating a middle assistant message", async () => {
+    listChatSessions.mockResolvedValue([]);
+
+    const { result } = renderHook(() =>
+      useAgentChatSession({
+        projectId: "project-1",
+        onToolCall: jest.fn(),
+      }),
+    );
+
+    await waitFor(() => expect(result.current.isHydrating).toBe(false));
+
+    await act(async () => {
+      await result.current.sendPrompt("第一轮");
+    });
+
+    await act(async () => {
+      await result.current.sendPrompt("第二轮");
+    });
+
+    const firstAssistantMessageId = result.current.messages[1]?.id ?? "";
+
+    await act(async () => {
+      await result.current.regenerate(firstAssistantMessageId);
+    });
+
+    expect(result.current.messages).toHaveLength(2);
+    expect(result.current.messages[0]).toEqual(
+      expect.objectContaining({
+        role: "user",
+        content: "第一轮",
+      }),
+    );
+    expect(result.current.messages[1]).toEqual(
+      expect.objectContaining({
+        role: "assistant",
+        content: "",
+      }),
+    );
+    expect(streamAgentChat).toHaveBeenNthCalledWith(
+      3,
+      expect.objectContaining({
+        message: "第一轮",
+        regenerateFromMessageIndex: 0,
+      }),
+    );
+  });
+
+  it("forks a copied conversation from an assistant message", async () => {
+    listChatSessions.mockResolvedValue([]);
+
+    const { result } = renderHook(() =>
+      useAgentChatSession({
+        projectId: "project-1",
+        onToolCall: jest.fn(),
+      }),
+    );
+
+    await waitFor(() => expect(result.current.isHydrating).toBe(false));
+
+    await act(async () => {
+      await result.current.sendPrompt("第一轮");
+    });
+
+    const firstAssistantMessageId = result.current.messages[1]?.id ?? "";
+
+    await act(async () => {
+      await result.current.createBranch(firstAssistantMessageId);
+    });
+
+    expect(forkAgentChat).toHaveBeenCalledWith(undefined, 2);
+    expect(result.current.activeSessionId).toBe("forked-session");
+    expect(result.current.messages).toHaveLength(2);
+    expect(result.current.messages[0]).toEqual(
+      expect.objectContaining({
+        role: "user",
+        content: "第一轮",
+      }),
+    );
+    expect(result.current.messages[1]).toEqual(
+      expect.objectContaining({
+        role: "assistant",
+      }),
+    );
+    expect(streamAgentChat).toHaveBeenCalledTimes(1);
+  });
 });