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 && (
-
-
+
+ )}
{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