From 24648194aeb1488c1351b38ca65b81c08a650050 Mon Sep 17 00:00:00 2001 From: Huixin615 Date: Tue, 14 Jul 2026 11:11:43 +0800 Subject: [PATCH] fix(frontend): show branch action only for completed turns (#4147) * fix(frontend): gate branch action by completed turn * test(frontend): clarify branch turn boundaries --- .../workspace/messages/message-list.tsx | 71 +++++++++++-------- frontend/src/core/messages/utils.ts | 37 ++++++++++ frontend/tests/e2e/branch-thread.spec.ts | 49 ++++++++++++- .../tests/unit/core/messages/utils.test.ts | 49 +++++++++++++ 4 files changed, 174 insertions(+), 32 deletions(-) diff --git a/frontend/src/components/workspace/messages/message-list.tsx b/frontend/src/components/workspace/messages/message-list.tsx index d03ef19dc..87996b042 100644 --- a/frontend/src/components/workspace/messages/message-list.tsx +++ b/frontend/src/components/workspace/messages/message-list.tsx @@ -47,6 +47,7 @@ import { extractTextFromMessage, getAssistantTurnCopyData, getAssistantTurnUsageMessages, + getBranchableAssistantGroupIds, getMessageGroups, getStreamingMessageLookup, hasContent, @@ -404,6 +405,10 @@ export function MessageList({ } return null; }, [groupedMessages, thread.isLoading]); + const branchableAssistantGroupIds = useMemo( + () => getBranchableAssistantGroupIds(groupedMessages, thread.isLoading), + [groupedMessages, thread.isLoading], + ); const clearSelectionToolbar = useCallback(() => { setSelectionToolbar(null); @@ -544,6 +549,7 @@ export function MessageList({ ( messages: Message[], isStreaming: boolean, + enableBranchForTurn: boolean, enableRegenerateForTurn: boolean, ) => { const clipboardData = getAssistantTurnCopyData(messages, { isStreaming }); @@ -561,36 +567,41 @@ export function MessageList({ return (
{clipboardData && } - {!isStreaming && actionTarget?.id && onBranchTurn && ( - - - - )} + onClick={() => { + const targetId = actionTarget.id; + if (!targetId) { + return; + } + setBranchingMessageId(targetId); + void Promise.resolve( + onBranchTurn(targetId, assistantMessageIds), + ).finally(() => { + setBranchingMessageId(null); + }); + }} + > + + + + )} {enableRegenerateForTurn && actionTarget?.id && onRegenerateMessage && ( @@ -782,6 +793,8 @@ export function MessageList({ group.messages, streamingMessages, ), + group.id !== undefined && + branchableAssistantGroupIds.has(group.id), group.id === latestAssistantGroupId, )}
diff --git a/frontend/src/core/messages/utils.ts b/frontend/src/core/messages/utils.ts index 77d7b2c1e..131290803 100644 --- a/frontend/src/core/messages/utils.ts +++ b/frontend/src/core/messages/utils.ts @@ -154,6 +154,43 @@ export function getMessageGroups(messages: Message[]): MessageGroup[] { return groups; } +export function getBranchableAssistantGroupIds( + groups: MessageGroup[], + isCurrentTurnLoading: boolean, +): Set { + // Hidden messages were already removed by getMessageGroups, matching the + // backend's branch checkpoint visibility rules. Within each visible human + // turn, branching is exposed only when the final AI-bearing group is a + // terminal assistant text group. Processing, present-files, and subagent + // groups do not render assistant actions. + const branchableGroupIds = new Set(); + let lastAIGroup: MessageGroup | null = null; + + const completeTurn = () => { + if (lastAIGroup?.type === "assistant" && lastAIGroup.id) { + branchableGroupIds.add(lastAIGroup.id); + } + lastAIGroup = null; + }; + + for (const group of groups) { + if (group.type === "human") { + completeTurn(); + continue; + } + + if (group.messages.some((message) => message.type === "ai")) { + lastAIGroup = group; + } + } + + if (!isCurrentTurnLoading) { + completeTurn(); + } + + return branchableGroupIds; +} + export function groupMessages( messages: Message[], mapper: (group: MessageGroup) => T, diff --git a/frontend/tests/e2e/branch-thread.spec.ts b/frontend/tests/e2e/branch-thread.spec.ts index f6ce07414..b879a7e77 100644 --- a/frontend/tests/e2e/branch-thread.spec.ts +++ b/frontend/tests/e2e/branch-thread.spec.ts @@ -34,7 +34,31 @@ test.describe("Branch from turn", () => { { type: "ai", id: "ai-2", - content: "Second answer", + content: "Intermediate answer", + }, + { + type: "ai", + id: "ai-3", + content: "", + tool_calls: [ + { + id: "tool-call-1", + name: "write_todos", + args: { todos: [] }, + }, + ], + }, + { + type: "tool", + id: "tool-1", + name: "write_todos", + tool_call_id: "tool-call-1", + content: "Todos updated", + }, + { + type: "ai", + id: "ai-4", + content: "Final answer", }, ], }, @@ -43,9 +67,28 @@ test.describe("Branch from turn", () => { await page.goto(`/workspace/chats/${MOCK_THREAD_ID}`); + const historicalTurn = page + .locator("[data-assistant-turn]") + .filter({ hasText: "First answer" }); + const intermediateTurn = page + .locator("[data-assistant-turn]") + .filter({ hasText: "Intermediate answer" }); const targetTurn = page .locator("[data-assistant-turn]") - .filter({ hasText: "Second answer" }); + .filter({ hasText: "Final answer" }); + + await expect(historicalTurn).toBeVisible(); + await historicalTurn.hover(); + await expect( + historicalTurn.getByRole("button", { name: /branch conversation/i }), + ).toBeVisible(); + + await expect(intermediateTurn).toBeVisible(); + await intermediateTurn.hover(); + await expect( + intermediateTurn.getByRole("button", { name: /branch conversation/i }), + ).toHaveCount(0); + await expect(targetTurn).toBeVisible(); await targetTurn.hover(); @@ -56,7 +99,7 @@ test.describe("Branch from turn", () => { await expect(page).toHaveURL( new RegExp(`/workspace/chats/${MOCK_THREAD_ID_2}$`), ); - await expect(page.getByText("Second answer")).toBeVisible(); + await expect(page.getByText("Final answer")).toBeVisible(); const branchThreadLink = page.locator( `a[href="/workspace/chats/${MOCK_THREAD_ID_2}"]`, ); diff --git a/frontend/tests/unit/core/messages/utils.test.ts b/frontend/tests/unit/core/messages/utils.test.ts index 524687bcb..750abe39e 100644 --- a/frontend/tests/unit/core/messages/utils.test.ts +++ b/frontend/tests/unit/core/messages/utils.test.ts @@ -5,6 +5,7 @@ import { extractContentFromMessage, extractTextFromMessage, extractReasoningContentFromMessage, + getBranchableAssistantGroupIds, getMessageCopyData, getAssistantTurnCopyData, getAssistantTurnUsageMessages, @@ -82,6 +83,54 @@ test("aggregates token usage messages once per assistant turn", () => { ).toEqual([null, null, ["ai-1", "ai-2"], null, ["ai-3"]]); }); +describe("branchable assistant groups", () => { + const messages = [ + { id: "human-1", type: "human", content: "First question" }, + { id: "ai-history", type: "ai", content: "Historical final answer" }, + { id: "human-2", type: "human", content: "Complex question" }, + { id: "ai-intermediate", type: "ai", content: "Intermediate answer" }, + { + id: "ai-tool", + type: "ai", + content: "", + tool_calls: [{ id: "tool-1", name: "write_todos", args: {} }], + }, + { + id: "tool-result", + type: "tool", + name: "write_todos", + tool_call_id: "tool-1", + content: "Todos updated", + }, + { id: "ai-final", type: "ai", content: "Final answer" }, + ] as Message[]; + + test("keeps historical turns branchable and selects only the final AI group in the current completed turn", () => { + const groups = getMessageGroups(messages); + + expect([...getBranchableAssistantGroupIds(groups, false)]).toEqual([ + "ai-history", + "ai-final", + ]); + }); + + test("does not expose the current turn while it is still loading", () => { + const groups = getMessageGroups(messages); + + expect([...getBranchableAssistantGroupIds(groups, true)]).toEqual([ + "ai-history", + ]); + }); + + test("does not expose a completed turn that ends in processing", () => { + const groups = getMessageGroups(messages.slice(0, -1)); + + expect([...getBranchableAssistantGroupIds(groups, false)]).toEqual([ + "ai-history", + ]); + }); +}); + test("reasoning + content (no tool calls) yields a single assistant bubble, not a duplicate processing group", () => { // Regression for #3868: in thinking/pro/ultra modes the final assistant // message carries both reasoning_content and answer text. It must surface its