fix(frontend): hide copy for streaming assistant turn (#3176)

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,16 +186,19 @@ 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 renderAssistantCopyButton = useCallback(
-    const clipboardData = [...messages]
+    (messages: Message[], isStreaming: boolean) => {
-      .reverse()
+      const clipboardData = getAssistantTurnCopyData(messages, { isStreaming });
-      .filter((message) => message.type === "ai")
-      .map((message) => {
-        const content = extractContentFromMessage(message);
-        return content ?? extractReasoningContentFromMessage(message) ?? "";
-      })
-      .find((content) => content.length > 0);
 
       if (!clipboardData) {
         return null;
@@ -204,7 +209,9 @@ export function MessageList({
           <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") {

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 (

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);
+});