From e7967a7fc37547f47d305b5057ec24aae6ef1591 Mon Sep 17 00:00:00 2001
From: Admire <64821731+LittleChenLiya@users.noreply.github.com>
Date: Sat, 23 May 2026 23:29:16 +0800
Subject: [PATCH] fix(frontend): hide copy for streaming assistant turn (#3176)

---
 .../workspace/messages/message-list.tsx       |  53 ++--
 frontend/src/core/messages/utils.ts           |  80 ++++++
 .../tests/unit/core/messages/utils.test.ts    | 272 ++++++++++++++++++
 3 files changed, 385 insertions(+), 20 deletions(-)

diff --git a/frontend/src/components/workspace/messages/message-list.tsx b/frontend/src/components/workspace/messages/message-list.tsx
index 74dca2af5..ca8672a3a 100644
--- a/frontend/src/components/workspace/messages/message-list.tsx
+++ b/frontend/src/components/workspace/messages/message-list.tsx
@@ -16,13 +16,15 @@ import {
 import {
   extractContentFromMessage,
   extractPresentFilesFromMessage,
-  extractReasoningContentFromMessage,
   extractTextFromMessage,
+  getAssistantTurnCopyData,
   getAssistantTurnUsageMessages,
   getMessageGroups,
+  getStreamingMessageLookup,
   hasContent,
   hasPresentFiles,
   hasReasoning,
+  isAssistantMessageGroupStreaming,
 } from "@/core/messages/utils";
 import { useRehypeSplitWordsIntoSpans } from "@/core/rehype";
 import type { Subtask } from "@/core/tasks";
@@ -184,27 +186,32 @@ export function MessageList({
     () => buildTokenDebugSteps(messages, t),
     [messages, t],
   );
+  const streamingMessages = useMemo(
+    () =>
+      getStreamingMessageLookup(
+        messages,
+        thread.isLoading,
+        thread.getMessagesMetadata,
+      ),
+    [messages, thread.getMessagesMetadata, thread.isLoading],
+  );
 
-  const renderAssistantCopyButton = useCallback((messages: Message[]) => {
-    const clipboardData = [...messages]
-      .reverse()
-      .filter((message) => message.type === "ai")
-      .map((message) => {
-        const content = extractContentFromMessage(message);
-        return content ?? extractReasoningContentFromMessage(message) ?? "";
-      })
-      .find((content) => content.length > 0);
+  const renderAssistantCopyButton = useCallback(
+    (messages: Message[], isStreaming: boolean) => {
+      const clipboardData = getAssistantTurnCopyData(messages, { isStreaming });
 
-    if (!clipboardData) {
-      return null;
-    }
+      if (!clipboardData) {
+        return null;
+      }
 
-    return (
-      <div className="mt-2 flex justify-start opacity-0 transition-opacity delay-200 duration-300 group-hover/assistant-turn:opacity-100">
-        <CopyButton clipboardData={clipboardData} />
-      </div>
-    );
-  }, []);
+      return (
+        <div className="mt-2 flex justify-start opacity-0 transition-opacity delay-200 duration-300 group-hover/assistant-turn:opacity-100">
+          <CopyButton clipboardData={clipboardData} />
+        </div>
+      );
+    },
+    [],
+  );
 
   const renderTokenUsage = useCallback(
     ({
@@ -294,7 +301,13 @@ export function MessageList({
                   turnUsageMessages,
                 })}
                 {group.type === "assistant" &&
-                  renderAssistantCopyButton(group.messages)}
+                  renderAssistantCopyButton(
+                    group.messages,
+                    isAssistantMessageGroupStreaming(
+                      group.messages,
+                      streamingMessages,
+                    ),
+                  )}
               </div>
             );
           } else if (group.type === "assistant:clarification") {
diff --git a/frontend/src/core/messages/utils.ts b/frontend/src/core/messages/utils.ts
index 1c165fd8d..5863195b8 100644
--- a/frontend/src/core/messages/utils.ts
+++ b/frontend/src/core/messages/utils.ts
@@ -170,6 +170,86 @@ export function getAssistantTurnUsageMessages(groups: MessageGroup[]) {
   return usageMessagesByGroupIndex;
 }
 
+type MessageMetadataLookup = (
+  message: Message,
+  index: number,
+) => { streamMetadata?: Record<string, unknown> } | undefined;
+
+export type StreamingMessageLookup = {
+  ids: ReadonlySet<string>;
+  messages: ReadonlySet<Message>;
+};
+
+export function getStreamingMessageLookup(
+  messages: Message[],
+  isStreaming: boolean,
+  getMessagesMetadata?: MessageMetadataLookup,
+): StreamingMessageLookup {
+  const streamingMessageIds = new Set<string>();
+  const streamingMessages = new Set<Message>();
+
+  if (!isStreaming) {
+    return {
+      ids: streamingMessageIds,
+      messages: streamingMessages,
+    };
+  }
+
+  messages.forEach((message, index) => {
+    if (!getMessagesMetadata?.(message, index)?.streamMetadata) {
+      return;
+    }
+
+    if (typeof message.id === "string" && message.id.length > 0) {
+      streamingMessageIds.add(message.id);
+    }
+    streamingMessages.add(message);
+  });
+
+  return {
+    ids: streamingMessageIds,
+    messages: streamingMessages,
+  };
+}
+
+export function isAssistantMessageGroupStreaming(
+  groupMessages: Message[],
+  streamingMessages: StreamingMessageLookup,
+) {
+  return groupMessages.some((message) => {
+    if (message.type !== "ai") {
+      return false;
+    }
+
+    return (
+      (typeof message.id === "string" &&
+        message.id.length > 0 &&
+        streamingMessages.ids.has(message.id)) ||
+      streamingMessages.messages.has(message)
+    );
+  });
+}
+
+export function getAssistantTurnCopyData(
+  messages: Message[],
+  { isStreaming = false }: { isStreaming?: boolean } = {},
+) {
+  if (isStreaming) {
+    return null;
+  }
+
+  return (
+    [...messages]
+      .reverse()
+      .filter((message) => message.type === "ai")
+      .map((message) => {
+        const content = extractContentFromMessage(message);
+        return content ?? extractReasoningContentFromMessage(message) ?? "";
+      })
+      .find((content) => content.length > 0) ?? null
+  );
+}
+
 export function extractTextFromMessage(message: Message) {
   if (typeof message.content === "string") {
     return (
diff --git a/frontend/tests/unit/core/messages/utils.test.ts b/frontend/tests/unit/core/messages/utils.test.ts
index cbc245583..1cc456e22 100644
--- a/frontend/tests/unit/core/messages/utils.test.ts
+++ b/frontend/tests/unit/core/messages/utils.test.ts
@@ -2,8 +2,11 @@ import type { Message } from "@langchain/langgraph-sdk";
 import { expect, test } from "vitest";
 
 import {
+  getAssistantTurnCopyData,
   getAssistantTurnUsageMessages,
   getMessageGroups,
+  getStreamingMessageLookup,
+  isAssistantMessageGroupStreaming,
 } from "@/core/messages/utils";
 
 test("aggregates token usage messages once per assistant turn", () => {
@@ -97,3 +100,272 @@ test("hides internal todo reminder messages from message groups", () => {
     groups.flatMap((group) => group.messages).map((message) => message.id),
   ).toEqual(["human-1", "ai-1"]);
 });
+
+test("hides assistant copy data while that turn is streaming", () => {
+  const messages = [
+    {
+      id: "ai-1",
+      type: "ai",
+      content: "Partial answer",
+    },
+  ] as Message[];
+
+  expect(getAssistantTurnCopyData(messages)).toBe("Partial answer");
+  expect(getAssistantTurnCopyData(messages, { isStreaming: true })).toBeNull();
+});
+
+test("marks the latest assistant message as streaming", () => {
+  const messages = [
+    {
+      id: "human-1",
+      type: "human",
+      content: "Hello",
+    },
+    {
+      id: "ai-1",
+      type: "ai",
+      content: "Still generating",
+    },
+  ] as Message[];
+  const groups = getMessageGroups(messages);
+  const assistantGroupIndex = groups.findIndex(
+    (group) => group.type === "assistant",
+  );
+
+  expect(
+    isAssistantMessageGroupStreaming(
+      groups[assistantGroupIndex]?.messages ?? [],
+      getStreamingMessageLookup(messages, true, () => ({
+        streamMetadata: { langgraph_node: "agent" },
+      })),
+    ),
+  ).toBe(true);
+  expect(
+    isAssistantMessageGroupStreaming(
+      groups[assistantGroupIndex]?.messages ?? [],
+      getStreamingMessageLookup(messages, false, () => ({
+        streamMetadata: { langgraph_node: "agent" },
+      })),
+    ),
+  ).toBe(false);
+});
+
+test("keeps previous assistant copyable while waiting for a new visible answer", () => {
+  const messages = [
+    {
+      id: "human-1",
+      type: "human",
+      content: "Hello",
+    },
+    {
+      id: "ai-1",
+      type: "ai",
+      content: "Completed answer",
+    },
+    {
+      id: "opt-human-1",
+      type: "human",
+      content: "Continue",
+    },
+  ] as Message[];
+  const groups = getMessageGroups(messages);
+  const assistantGroupIndex = groups.findIndex(
+    (group) => group.type === "assistant",
+  );
+
+  expect(
+    isAssistantMessageGroupStreaming(
+      groups[assistantGroupIndex]?.messages ?? [],
+      getStreamingMessageLookup(messages, true),
+    ),
+  ).toBe(false);
+});
+
+test("keeps previous assistant copyable while a hidden send is starting", () => {
+  const messages = [
+    {
+      id: "human-1",
+      type: "human",
+      content: "Hello",
+    },
+    {
+      id: "ai-1",
+      type: "ai",
+      content: "Completed answer",
+    },
+  ] as Message[];
+  const groups = getMessageGroups(messages);
+  const assistantGroupIndex = groups.findIndex(
+    (group) => group.type === "assistant",
+  );
+
+  expect(
+    isAssistantMessageGroupStreaming(
+      groups[assistantGroupIndex]?.messages ?? [],
+      getStreamingMessageLookup(messages, true),
+    ),
+  ).toBe(false);
+});
+
+test("keeps previous assistant copyable after a hidden send is appended", () => {
+  const messages = [
+    {
+      id: "human-1",
+      type: "human",
+      content: "Hello",
+    },
+    {
+      id: "ai-1",
+      type: "ai",
+      content: "Completed answer",
+    },
+    {
+      id: "human-hidden",
+      type: "human",
+      content: "Save this agent",
+      additional_kwargs: { hide_from_ui: true },
+    },
+  ] as Message[];
+  const groups = getMessageGroups(messages);
+  const assistantGroupIndex = groups.findIndex(
+    (group) => group.type === "assistant",
+  );
+
+  expect(
+    isAssistantMessageGroupStreaming(
+      groups[assistantGroupIndex]?.messages ?? [],
+      getStreamingMessageLookup(messages, true),
+    ),
+  ).toBe(false);
+});
+
+test("uses stream metadata to identify an assistant before optimistic input", () => {
+  const messages = [
+    {
+      id: "human-1",
+      type: "human",
+      content: "Hello",
+    },
+    {
+      id: "ai-1",
+      type: "ai",
+      content: "Completed answer",
+    },
+    {
+      id: "ai-2",
+      type: "ai",
+      content: "Still generating",
+    },
+    {
+      id: "opt-human-1",
+      type: "human",
+      content: "Continue",
+    },
+  ] as Message[];
+  const assistantGroups = getMessageGroups(messages).filter(
+    (group) => group.type === "assistant",
+  );
+  const groups = getMessageGroups(messages);
+  const assistantGroupIndexes = groups
+    .map((group, index) => (group.type === "assistant" ? index : -1))
+    .filter((index) => index >= 0);
+
+  expect(
+    isAssistantMessageGroupStreaming(
+      groups[assistantGroupIndexes[0] ?? -1]?.messages ?? [],
+      getStreamingMessageLookup(messages, true, (message) =>
+        message.id === "ai-2"
+          ? { streamMetadata: { langgraph_node: "agent" } }
+          : undefined,
+      ),
+    ),
+  ).toBe(false);
+  expect(
+    isAssistantMessageGroupStreaming(
+      groups[assistantGroupIndexes[1] ?? -1]?.messages ?? [],
+      getStreamingMessageLookup(messages, true, (message) =>
+        message.id === "ai-2"
+          ? { streamMetadata: { langgraph_node: "agent" } }
+          : undefined,
+      ),
+    ),
+  ).toBe(true);
+  expect(assistantGroups.map((group) => group.id)).toEqual(["ai-1", "ai-2"]);
+});
+
+test("does not mark a completed assistant group streaming from a later processing group", () => {
+  const messages = [
+    {
+      id: "human-1",
+      type: "human",
+      content: "Hello",
+    },
+    {
+      id: "ai-1",
+      type: "ai",
+      content: "Visible answer",
+    },
+    {
+      id: "ai-2",
+      type: "ai",
+      content: "",
+      tool_calls: [{ id: "tool-1", name: "web_search", args: {} }],
+    },
+  ] as Message[];
+  const groups = getMessageGroups(messages);
+  const assistantGroupIndex = groups.findIndex(
+    (group) => group.type === "assistant",
+  );
+
+  expect(groups.map((group) => group.type)).toEqual([
+    "human",
+    "assistant",
+    "assistant:processing",
+  ]);
+  expect(
+    isAssistantMessageGroupStreaming(
+      groups[assistantGroupIndex]?.messages ?? [],
+      getStreamingMessageLookup(messages, true, (message) =>
+        message.id === "ai-2"
+          ? { streamMetadata: { langgraph_node: "agent" } }
+          : undefined,
+      ),
+    ),
+  ).toBe(false);
+});
+
+test("keeps streaming assistant hidden when a hidden control message follows it", () => {
+  const messages = [
+    {
+      id: "human-1",
+      type: "human",
+      content: "Hello",
+    },
+    {
+      id: "ai-1",
+      type: "ai",
+      content: "Still generating",
+    },
+    {
+      id: "human-hidden",
+      type: "human",
+      content: "Save this agent",
+      additional_kwargs: { hide_from_ui: true },
+    },
+  ] as Message[];
+  const groups = getMessageGroups(messages);
+  const assistantGroupIndex = groups.findIndex(
+    (group) => group.type === "assistant",
+  );
+
+  expect(
+    isAssistantMessageGroupStreaming(
+      groups[assistantGroupIndex]?.messages ?? [],
+      getStreamingMessageLookup(messages, true, (message) =>
+        message.id === "ai-1"
+          ? { streamMetadata: { langgraph_node: "agent" } }
+          : undefined,
+      ),
+    ),
+  ).toBe(true);
+});