From 949cfad54821719fbf5849657a72db69d72f0bb5 Mon Sep 17 00:00:00 2001
From: chengyongru <chengyongru.ai@gmail.com>
Date: Mon, 20 Jul 2026 13:15:01 +0800
Subject: [PATCH] fix(webui): show copy action on every assistant message

---
 .../src/components/thread/ThreadMessages.tsx  | 11 +++------
 webui/src/tests/thread-messages.test.tsx      | 23 ++++++++++++-------
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/webui/src/components/thread/ThreadMessages.tsx b/webui/src/components/thread/ThreadMessages.tsx
index 2aabb25f..9271fd74 100644
--- a/webui/src/components/thread/ThreadMessages.tsx
+++ b/webui/src/components/thread/ThreadMessages.tsx
@@ -29,7 +29,7 @@ export function buildDisplayUnits(
   });
 }
 
-export function assistantCopyFlags(units: DisplayUnit[]): boolean[] {
+export function assistantForkFlags(units: DisplayUnit[]): boolean[] {
   const flags = new Array<boolean>(units.length).fill(true);
   let hasLaterUnitBeforeUser = false;
   for (let i = units.length - 1; i >= 0; i -= 1) {
@@ -63,7 +63,7 @@ export function ThreadMessages({
     () => unitIndexAfterMessageCount(units, forkBoundaryMessageCount),
     [forkBoundaryMessageCount, units],
   );
-  const copyFlags = useMemo(() => assistantCopyFlags(units), [units]);
+  const forkFlags = useMemo(() => assistantForkFlags(units), [units]);
   const liveActivityClusterIndices = useMemo(
     () => isStreaming ? currentActivityClusterIndices(units) : new Set<number>(),
     [isStreaming, units],
@@ -90,7 +90,7 @@ export function ThreadMessages({
             ? unit.message.id
             : undefined;
         const forkIndex =
-          unit.type === "message" && unit.message.role === "assistant" && copyFlags[index]
+          unit.type === "message" && unit.message.role === "assistant" && forkFlags[index]
             ? nextUserIndex
             : undefined;
         if (unit.type === "message" && unit.message.role === "user") nextUserIndex += 1;
@@ -112,11 +112,6 @@ export function ThreadMessages({
               ) : (
                 <MessageBubble
                   message={unit.message}
-                  showCopyAction={
-                    unit.message.role === "assistant"
-                      ? copyFlags[index]
-                      : true
-                  }
                   cliApps={cliApps}
                   mcpPresets={mcpPresets}
                   slashCommands={slashCommands}
diff --git a/webui/src/tests/thread-messages.test.tsx b/webui/src/tests/thread-messages.test.tsx
index 76152941..73b8c34d 100644
--- a/webui/src/tests/thread-messages.test.tsx
+++ b/webui/src/tests/thread-messages.test.tsx
@@ -2,7 +2,7 @@ import { render, screen } from "@testing-library/react";
 import { afterEach, describe, expect, it, vi } from "vitest";
 
 import {
-  assistantCopyFlags,
+  assistantForkFlags,
   buildDisplayUnits,
   ThreadMessages,
   unitKeysForDisplay,
@@ -747,7 +747,7 @@ describe("ThreadMessages", () => {
     expect(screen.queryByText("Worked for 0s")).not.toBeInTheDocument();
   });
 
-  it("shows copy only on the last assistant slice before the next user turn", () => {
+  it("shows copy on every assistant slice while keeping fork on the last slice", () => {
     const messages: UIMessage[] = [
       {
         id: "early",
@@ -771,19 +771,26 @@ describe("ThreadMessages", () => {
       },
     ];
 
-    render(<ThreadMessages messages={messages} isStreaming={false} />);
+    render(
+      <ThreadMessages
+        messages={messages}
+        isStreaming={false}
+        onForkFromMessage={vi.fn()}
+      />,
+    );
 
-    expect(screen.getAllByRole("button", { name: "Copy" })).toHaveLength(1);
+    expect(screen.getAllByRole("button", { name: "Copy" })).toHaveLength(2);
+    expect(screen.getAllByRole("button", { name: "Fork" })).toHaveLength(1);
     expect(screen.getByText("final reply")).toBeInTheDocument();
   });
 
-  it("shows copy only on the second assistant when two text slices appear before user", () => {
+  it("shows copy on adjacent assistant text slices", () => {
     const messages: UIMessage[] = [
       { id: "a1", role: "assistant", content: "part one", createdAt: 1 },
       { id: "a2", role: "assistant", content: "part two", createdAt: 2 },
     ];
     render(<ThreadMessages messages={messages} isStreaming={false} />);
-    expect(screen.getAllByRole("button", { name: "Copy" })).toHaveLength(1);
+    expect(screen.getAllByRole("button", { name: "Copy" })).toHaveLength(2);
   });
 
   it("uses turn ids as activity grouping boundaries when available", () => {
@@ -810,7 +817,7 @@ describe("ThreadMessages", () => {
     ]);
   });
 
-  it("computes final assistant copy flags with user-boundary semantics", () => {
+  it("computes final assistant fork flags with user-boundary semantics", () => {
     const units = buildDisplayUnits([
       { id: "u1", role: "user", content: "one", createdAt: 1 },
       { id: "a1", role: "assistant", content: "draft", createdAt: 2 },
@@ -827,7 +834,7 @@ describe("ThreadMessages", () => {
       { id: "a3", role: "assistant", content: "next", createdAt: 6 },
     ]);
 
-    const flags = assistantCopyFlags(units);
+    const flags = assistantForkFlags(units);
     const assistantFlags = units
       .map((unit, index) =>
         unit.type === "message" && unit.message.role === "assistant"