fix(frontend): show branch action only for completed turns (#4147)

* fix(frontend): gate branch action by completed turn

* test(frontend): clarify branch turn boundaries

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 (
         <div className="mt-2 flex justify-start gap-1 opacity-0 transition-opacity delay-200 duration-300 group-hover/assistant-turn:opacity-100">
           {clipboardData && <CopyButton clipboardData={clipboardData} />}
-          {!isStreaming && actionTarget?.id && onBranchTurn && (
-            <Tooltip content={t.common.branch}>
-              <Button
-                aria-label={t.common.branch}
-                size="icon-sm"
-                type="button"
-                variant="ghost"
-                disabled={!canBranch || branchingMessageId === actionTarget.id}
-                onClick={() => {
-                  const targetId = actionTarget.id;
-                  if (!targetId) {
-                    return;
+          {enableBranchForTurn &&
+            !isStreaming &&
+            actionTarget?.id &&
+            onBranchTurn && (
+              <Tooltip content={t.common.branch}>
+                <Button
+                  aria-label={t.common.branch}
+                  size="icon-sm"
+                  type="button"
+                  variant="ghost"
+                  disabled={
+                    !canBranch || branchingMessageId === actionTarget.id
                   }
-                  setBranchingMessageId(targetId);
-                  void Promise.resolve(
-                    onBranchTurn(targetId, assistantMessageIds),
-                  ).finally(() => {
-                    setBranchingMessageId(null);
-                  });
-                }}
-              >
-                <GitBranchPlusIcon
-                  className={cn(
-                    "size-4",
-                    branchingMessageId === actionTarget.id && "animate-pulse",
-                  )}
-                />
-              </Button>
-            </Tooltip>
-          )}
+                  onClick={() => {
+                    const targetId = actionTarget.id;
+                    if (!targetId) {
+                      return;
+                    }
+                    setBranchingMessageId(targetId);
+                    void Promise.resolve(
+                      onBranchTurn(targetId, assistantMessageIds),
+                    ).finally(() => {
+                      setBranchingMessageId(null);
+                    });
+                  }}
+                >
+                  <GitBranchPlusIcon
+                    className={cn(
+                      "size-4",
+                      branchingMessageId === actionTarget.id && "animate-pulse",
+                    )}
+                  />
+                </Button>
+              </Tooltip>
+            )}
           {enableRegenerateForTurn &&
             actionTarget?.id &&
             onRegenerateMessage && (
@@ -782,6 +793,8 @@ export function MessageList({
                         group.messages,
                         streamingMessages,
                       ),
+                      group.id !== undefined &&
+                        branchableAssistantGroupIds.has(group.id),
                       group.id === latestAssistantGroupId,
                     )}
                 </div>

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<string> {
+  // 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<string>();
+  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<T>(
   messages: Message[],
   mapper: (group: MessageGroup) => T,

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}"]`,
     );

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