From 16ecf7b0064280996c478ba550d736da18d33138 Mon Sep 17 00:00:00 2001
From: Aari <ansusu1025@gmail.com>
Date: Sun, 16 Aug 2026 23:53:22 +0800
Subject: [PATCH] fix(frontend): preserve completed message actions during
 streaming (#4844)

* fix(frontend): preserve completed message actions during streaming

* fix(frontend): address streaming action review feedback

* fix(frontend): reuse settled stream snapshots
---
 .../src/components/workspace/copy-button.tsx  |   3 +
 .../workspace/messages/message-list.tsx       |  40 +++++-
 frontend/src/core/messages/utils.ts           |  74 +++++++++-
 frontend/tests/e2e/chat.spec.ts               |  75 +++++++++-
 frontend/tests/e2e/utils/mock-api.ts          |  72 ++++++----
 .../tests/unit/core/messages/utils.test.ts    | 135 ++++++++++++++++++
 6 files changed, 372 insertions(+), 27 deletions(-)

diff --git a/frontend/src/components/workspace/copy-button.tsx b/frontend/src/components/workspace/copy-button.tsx
index e85313cb7..ae7aea3bc 100644
--- a/frontend/src/components/workspace/copy-button.tsx
+++ b/frontend/src/components/workspace/copy-button.tsx
@@ -33,6 +33,9 @@ export function CopyButton({
   return (
     <Tooltip content={t.clipboard.copyToClipboard}>
       <Button
+        aria-label={
+          copied ? t.clipboard.copiedToClipboard : t.clipboard.copyToClipboard
+        }
         size="icon-sm"
         type="button"
         variant="ghost"
diff --git a/frontend/src/components/workspace/messages/message-list.tsx b/frontend/src/components/workspace/messages/message-list.tsx
index f896b8120..d982ebc16 100644
--- a/frontend/src/components/workspace/messages/message-list.tsx
+++ b/frontend/src/components/workspace/messages/message-list.tsx
@@ -46,12 +46,14 @@ import {
   type TokenUsageInlineMode,
 } from "@/core/messages/usage-model";
 import {
+  areStreamMetadataSnapshotsEqual,
   extractContentFromMessage,
   extractPresentFilesFromMessage,
   extractTextFromMessage,
   getAssistantTurnCopyData,
   getBranchableAssistantGroupIds,
   getLatestEditableTurn,
+  getStreamMetadataSnapshot,
   getStreamingMessageLookup,
   hasContent,
   hasPresentFiles,
@@ -59,6 +61,7 @@ import {
   isAssistantMessageGroupStreaming,
   isHiddenFromUIMessage,
   type MessageGroup as ThreadMessageGroup,
+  type StreamMetadataSnapshot,
 } from "@/core/messages/utils";
 import { getWorkspaceChangeAnchorGroupIndices } from "@/core/messages/workspace-change-anchor";
 import {
@@ -99,6 +102,11 @@ import { VirtualMessageList } from "./virtual-message-list";
 const EMPTY_TOKEN_DEBUG_STEPS: TokenDebugStep[] = [];
 const EMPTY_ARTIFACT_PATHS: readonly string[] = [];
 
+type SettledStreamMetadataState = {
+  threadId: string;
+  snapshot: StreamMetadataSnapshot;
+};
+
 function sameStrings(previous: readonly string[], next: readonly string[]) {
   return (
     previous.length === next.length &&
@@ -506,14 +514,44 @@ export function MessageList({
     },
     [showTokenDebugSummaries, tokenDebugStepsByMessageId],
   );
+  const [settledStreamMetadataState, setSettledStreamMetadataState] =
+    useState<SettledStreamMetadataState>();
+  useEffect(() => {
+    if (thread.isLoading) {
+      return;
+    }
+    const snapshot = getStreamMetadataSnapshot(
+      messages,
+      thread.getMessagesMetadata,
+    );
+    setSettledStreamMetadataState((previous) => {
+      if (
+        previous?.threadId === threadId &&
+        areStreamMetadataSnapshotsEqual(previous.snapshot, snapshot)
+      ) {
+        return previous;
+      }
+      return { threadId, snapshot };
+    });
+  }, [messages, thread.getMessagesMetadata, thread.isLoading, threadId]);
+  const settledStreamMetadata =
+    settledStreamMetadataState?.threadId === threadId
+      ? settledStreamMetadataState.snapshot
+      : undefined;
   const streamingMessages = useMemo(
     () =>
       getStreamingMessageLookup(
         messages,
         thread.isLoading,
         thread.getMessagesMetadata,
+        settledStreamMetadata,
       ),
-    [messages, thread.getMessagesMetadata, thread.isLoading],
+    [
+      messages,
+      settledStreamMetadata,
+      thread.getMessagesMetadata,
+      thread.isLoading,
+    ],
   );
 
   const humanInputState = useMemo(
diff --git a/frontend/src/core/messages/utils.ts b/frontend/src/core/messages/utils.ts
index 2e8663b20..f10f8d821 100644
--- a/frontend/src/core/messages/utils.ts
+++ b/frontend/src/core/messages/utils.ts
@@ -363,15 +363,74 @@ type MessageMetadataLookup = (
   index: number,
 ) => { streamMetadata?: Record<string, unknown> } | undefined;
 
+export type StreamMetadataSnapshot = {
+  ids: ReadonlyMap<string, Record<string, unknown>>;
+  messages: ReadonlyMap<Message, Record<string, unknown>>;
+};
+
 export type StreamingMessageLookup = {
   ids: ReadonlySet<string>;
   messages: ReadonlySet<Message>;
 };
 
+export function areStreamMetadataSnapshotsEqual(
+  left: StreamMetadataSnapshot,
+  right: StreamMetadataSnapshot,
+) {
+  if (
+    left.ids.size !== right.ids.size ||
+    left.messages.size !== right.messages.size
+  ) {
+    return false;
+  }
+
+  for (const [id, metadata] of left.ids) {
+    if (right.ids.get(id) !== metadata) {
+      return false;
+    }
+  }
+  for (const [message, metadata] of left.messages) {
+    if (right.messages.get(message) !== metadata) {
+      return false;
+    }
+  }
+  return true;
+}
+
+export function getStreamMetadataSnapshot(
+  messages: Message[],
+  getMessagesMetadata?: MessageMetadataLookup,
+): StreamMetadataSnapshot {
+  const metadataById = new Map<string, Record<string, unknown>>();
+  const metadataByMessage = new Map<Message, Record<string, unknown>>();
+
+  messages.forEach((message, index) => {
+    const streamMetadata = getMessagesMetadata?.(
+      message,
+      index,
+    )?.streamMetadata;
+    if (!streamMetadata) {
+      return;
+    }
+
+    if (typeof message.id === "string" && message.id.length > 0) {
+      metadataById.set(message.id, streamMetadata);
+    } else {
+      metadataByMessage.set(message, streamMetadata);
+    }
+  });
+
+  return {
+    ids: metadataById,
+    messages: metadataByMessage,
+  };
+}
+
 export function getStreamingMessageLookup(
   messages: Message[],
   isStreaming: boolean,
   getMessagesMetadata?: MessageMetadataLookup,
+  settledMetadata?: StreamMetadataSnapshot,
 ): StreamingMessageLookup {
   const streamingMessageIds = new Set<string>();
   const streamingMessages = new Set<Message>();
@@ -384,12 +443,25 @@ export function getStreamingMessageLookup(
   }
 
   messages.forEach((message, index) => {
-    if (!getMessagesMetadata?.(message, index)?.streamMetadata) {
+    const streamMetadata = getMessagesMetadata?.(
+      message,
+      index,
+    )?.streamMetadata;
+    if (!streamMetadata) {
       return;
     }
 
     if (typeof message.id === "string" && message.id.length > 0) {
+      // MessageTupleManager retains metadata until the whole stream instance is
+      // cleared. A later run therefore exposes the completed turn's metadata
+      // again. Only an unchanged metadata object is stale: a new object for the
+      // same message id means that message received another stream event.
+      if (settledMetadata?.ids.get(message.id) === streamMetadata) {
+        return;
+      }
       streamingMessageIds.add(message.id);
+    } else if (settledMetadata?.messages.get(message) === streamMetadata) {
+      return;
     }
     streamingMessages.add(message);
   });
diff --git a/frontend/tests/e2e/chat.spec.ts b/frontend/tests/e2e/chat.spec.ts
index be99520d1..859e48957 100644
--- a/frontend/tests/e2e/chat.spec.ts
+++ b/frontend/tests/e2e/chat.spec.ts
@@ -1,4 +1,4 @@
-import { expect, test } from "@playwright/test";
+import { expect, test, type Route } from "@playwright/test";
 
 import { handleRunStream, mockLangGraphAPI } from "./utils/mock-api";
 
@@ -21,6 +21,79 @@ function textFromMessageContent(content: unknown) {
     .join("");
 }
 
+test.describe("Streaming message actions", () => {
+  test("keeps a completed answer copyable while the next turn starts", async ({
+    page,
+  }) => {
+    let streamCalls = 0;
+    let releaseSecondStream!: () => void;
+    const secondStreamHeld = new Promise<void>((resolve) => {
+      releaseSecondStream = resolve;
+    });
+
+    const handleCopyRegressionStream = async (route: Route) => {
+      streamCalls += 1;
+      if (streamCalls === 2) {
+        await secondStreamHeld;
+      }
+      return handleRunStream(route, {}, undefined, {
+        responseMessage: {
+          type: "ai",
+          id: `copy-regression-ai-${streamCalls}`,
+          content:
+            streamCalls === 1 ? "First completed answer" : "Second answer",
+        },
+        messageMetadata: {
+          langgraph_node: "agent",
+          langgraph_step: streamCalls,
+        },
+      });
+    };
+    mockLangGraphAPI(page, {
+      createdThreadMessages: [
+        {
+          type: "human",
+          id: "copy-regression-human-1",
+          content: "First question",
+        },
+        {
+          type: "ai",
+          id: "copy-regression-ai-1",
+          content: "First completed answer",
+        },
+      ],
+      runStreamHandler: handleCopyRegressionStream,
+    });
+
+    try {
+      await page.goto("/workspace/chats/new");
+      const textarea = page.getByPlaceholder(/how can i assist you/i);
+      await expect(textarea).toBeVisible({ timeout: 15_000 });
+
+      await textarea.fill("First question");
+      await textarea.press("Enter");
+      await expect.poll(() => streamCalls).toBe(1);
+      await expect(page.getByText("First completed answer")).toBeVisible({
+        timeout: 10_000,
+      });
+
+      await textarea.fill("Second question");
+      await textarea.press("Enter");
+      await expect.poll(() => streamCalls).toBe(2);
+
+      const completedTurn = page
+        .locator('[data-assistant-turn=""]')
+        .filter({ hasText: "First completed answer" });
+      await completedTurn.hover();
+      await expect(
+        completedTurn.getByRole("button", { name: "Copy to clipboard" }),
+      ).toBeVisible();
+    } finally {
+      releaseSecondStream();
+    }
+  });
+});
+
 test.describe("Chat workspace", () => {
   test.beforeEach(async ({ page }) => {
     mockLangGraphAPI(page);
diff --git a/frontend/tests/e2e/utils/mock-api.ts b/frontend/tests/e2e/utils/mock-api.ts
index 13daeac82..565323c53 100644
--- a/frontend/tests/e2e/utils/mock-api.ts
+++ b/frontend/tests/e2e/utils/mock-api.ts
@@ -59,6 +59,7 @@ export type MockSkill = {
 
 export type MockAPIOptions = {
   threads?: MockThread[];
+  createdThreadMessages?: unknown[];
   agents?: MockAgent[];
   skills?: MockSkill[];
   scheduledTasks?: Array<{
@@ -95,6 +96,7 @@ export type MockAPIOptions = {
     agentsApiEnabled?: boolean;
     browserControlEnabled?: boolean;
   };
+  runStreamHandler?: (route: Route) => Promise<void>;
 };
 
 const DEFAULT_SKILLS: MockSkill[] = [
@@ -174,17 +176,20 @@ function branchMessagesFromTurn(messages: unknown[], targetIds: Set<string>) {
   return targetEndIndex >= 0 ? messages.slice(0, targetEndIndex + 1) : messages;
 }
 
-function mockStreamMessages(route?: Route, inputMessages?: unknown[]) {
+function mockStreamMessages(
+  route?: Route,
+  inputMessages?: unknown[],
+  responseMessage: Record<string, unknown> = {
+    type: "ai",
+    id: "msg-ai-1",
+    content: "Hello from DeerFlow!",
+  },
+) {
   const submittedMessages = inputMessages
     ? visibleInputMessages(inputMessages)
     : route
       ? visibleRunInputMessages(route)
       : [];
-  const responseMessage = {
-    type: "ai",
-    id: "msg-ai-1",
-    content: "Hello from DeerFlow!",
-  };
   if (submittedMessages.length > 0) {
     return [...submittedMessages, responseMessage];
   }
@@ -699,7 +704,7 @@ export function mockLangGraphAPI(page: Page, options?: MockAPIOptions) {
         thread_id: MOCK_THREAD_ID,
         title: "New Chat",
         updated_at: new Date().toISOString(),
-        messages: mockStreamMessages(),
+        messages: options?.createdThreadMessages ?? mockStreamMessages(),
       });
       return route.fulfill({
         status: 200,
@@ -1115,23 +1120,25 @@ export function mockLangGraphAPI(page: Page, options?: MockAPIOptions) {
   });
 
   // Run stream — returns a minimal SSE response with an AI message
-  const handleMockRunStream = (route: Route) => {
-    const threadId = runStreamThreadId(route);
-    const existingThread = threads.find(
-      (thread) => thread.thread_id === threadId,
-    );
-    const fallbackGoal = threads.find((thread) => thread.goal)?.goal ?? null;
-    const goal = existingThread?.goal ?? fallbackGoal;
-    upsertThread({
-      thread_id: threadId,
-      title: threadId === MOCK_SIDECAR_THREAD_ID ? "Side chat" : "New Chat",
-      updated_at: new Date().toISOString(),
-      goal,
-      metadata: existingThread?.metadata,
-      messages: mockStreamMessages(route),
+  const handleMockRunStream =
+    options?.runStreamHandler ??
+    ((route: Route) => {
+      const threadId = runStreamThreadId(route);
+      const existingThread = threads.find(
+        (thread) => thread.thread_id === threadId,
+      );
+      const fallbackGoal = threads.find((thread) => thread.goal)?.goal ?? null;
+      const goal = existingThread?.goal ?? fallbackGoal;
+      upsertThread({
+        thread_id: threadId,
+        title: threadId === MOCK_SIDECAR_THREAD_ID ? "Side chat" : "New Chat",
+        updated_at: new Date().toISOString(),
+        goal,
+        metadata: existingThread?.metadata,
+        messages: mockStreamMessages(route),
+      });
+      return handleRunStream(route, { goal });
     });
-    return handleRunStream(route, { goal });
-  };
 
   void page.route("**/api/langgraph/runs/stream", handleMockRunStream);
   void page.route(
@@ -1418,18 +1425,35 @@ export function handleRunStream(
   route: Route,
   values: Record<string, unknown> = {},
   inputMessages?: unknown[],
+  options?: {
+    responseMessage?: Record<string, unknown>;
+    messageMetadata?: Record<string, unknown>;
+  },
 ) {
   const threadId = runStreamThreadId(route);
+  const responseMessage = options?.responseMessage ?? {
+    type: "ai",
+    id: "msg-ai-1",
+    content: "Hello from DeerFlow!",
+  };
   const events = [
     {
       event: "metadata",
       data: { run_id: MOCK_RUN_ID, thread_id: threadId },
     },
+    ...(options?.messageMetadata
+      ? [
+          {
+            event: "messages",
+            data: [responseMessage, options.messageMetadata],
+          },
+        ]
+      : []),
     {
       event: "values",
       data: {
         ...values,
-        messages: mockStreamMessages(route, inputMessages),
+        messages: mockStreamMessages(route, inputMessages, responseMessage),
       },
     },
     { event: "end", data: {} },
diff --git a/frontend/tests/unit/core/messages/utils.test.ts b/frontend/tests/unit/core/messages/utils.test.ts
index 2a9b2d685..0749b9c89 100644
--- a/frontend/tests/unit/core/messages/utils.test.ts
+++ b/frontend/tests/unit/core/messages/utils.test.ts
@@ -2,6 +2,7 @@ import type { Message } from "@langchain/langgraph-sdk";
 import { describe, expect, test } from "@rstest/core";
 
 import {
+  areStreamMetadataSnapshotsEqual,
   extractContentFromMessage,
   extractTextFromMessage,
   extractReasoningContentFromMessage,
@@ -11,6 +12,7 @@ import {
   getAssistantTurnCopyData,
   getAssistantTurnUsageMessages,
   getMessageGroups,
+  getStreamMetadataSnapshot,
   getStreamingMessageLookup,
   hasContent,
   hasReasoning,
@@ -789,6 +791,139 @@ test("marks the latest assistant message as streaming", () => {
   ).toBe(false);
 });
 
+test("compares stream metadata snapshots by keys and metadata identity", () => {
+  const identifiedMessage = {
+    id: "ai-1",
+    type: "ai",
+    content: "Completed answer",
+  } as Message;
+  const anonymousMessage = {
+    type: "ai",
+    content: "Anonymous answer",
+  } as Message;
+  const identifiedMetadata = { langgraph_node: "agent" };
+  const anonymousMetadata = { langgraph_node: "agent" };
+  const messages = [identifiedMessage, anonymousMessage];
+  const snapshot = getStreamMetadataSnapshot(messages, (message) => ({
+    streamMetadata:
+      message === identifiedMessage ? identifiedMetadata : anonymousMetadata,
+  }));
+  const equivalentSnapshot = getStreamMetadataSnapshot(messages, (message) => ({
+    streamMetadata:
+      message === identifiedMessage ? identifiedMetadata : anonymousMetadata,
+  }));
+  const changedSnapshot = getStreamMetadataSnapshot(messages, (message) => ({
+    streamMetadata:
+      message === identifiedMessage
+        ? { ...identifiedMetadata }
+        : anonymousMetadata,
+  }));
+  const missingSnapshot = getStreamMetadataSnapshot(
+    [identifiedMessage],
+    () => ({ streamMetadata: identifiedMetadata }),
+  );
+
+  expect(areStreamMetadataSnapshotsEqual(snapshot, equivalentSnapshot)).toBe(
+    true,
+  );
+  expect(areStreamMetadataSnapshotsEqual(snapshot, changedSnapshot)).toBe(
+    false,
+  );
+  expect(areStreamMetadataSnapshotsEqual(snapshot, missingSnapshot)).toBe(
+    false,
+  );
+});
+
+test("ignores stream metadata retained from a completed turn", () => {
+  const completedMetadata = { langgraph_node: "agent", langgraph_step: 1 };
+  const activeMetadata = { langgraph_node: "agent", langgraph_step: 2 };
+  const completedMessages = [
+    {
+      id: "human-1",
+      type: "human",
+      content: "Hello",
+    },
+    {
+      id: "ai-1",
+      type: "ai",
+      content: "Completed answer",
+    },
+  ] as Message[];
+  const settledMetadata = getStreamMetadataSnapshot(
+    completedMessages,
+    (message) =>
+      message.id === "ai-1" ? { streamMetadata: completedMetadata } : undefined,
+  );
+  const messages = [
+    ...completedMessages,
+    {
+      id: "human-2",
+      type: "human",
+      content: "Continue",
+    },
+    {
+      id: "ai-2",
+      type: "ai",
+      content: "Still generating",
+    },
+  ] as Message[];
+  const groups = getMessageGroups(messages).filter(
+    (group) => group.type === "assistant",
+  );
+  const streamingMessages = getStreamingMessageLookup(
+    messages,
+    true,
+    (message) => {
+      if (message.id === "ai-1") {
+        return { streamMetadata: completedMetadata };
+      }
+      if (message.id === "ai-2") {
+        return { streamMetadata: activeMetadata };
+      }
+      return undefined;
+    },
+    settledMetadata,
+  );
+
+  expect(
+    isAssistantMessageGroupStreaming(
+      groups[0]?.messages ?? [],
+      streamingMessages,
+    ),
+  ).toBe(false);
+  expect(
+    isAssistantMessageGroupStreaming(
+      groups[1]?.messages ?? [],
+      streamingMessages,
+    ),
+  ).toBe(true);
+});
+
+test("treats updated metadata for the same message id as active", () => {
+  const message = {
+    id: "ai-1",
+    type: "ai",
+    content: "Partial answer",
+  } as Message;
+  const completedMetadata = { langgraph_node: "agent", langgraph_step: 1 };
+  const activeMetadata = { langgraph_node: "agent", langgraph_step: 2 };
+  const settledMetadata = getStreamMetadataSnapshot([message], () => ({
+    streamMetadata: completedMetadata,
+  }));
+
+  expect(
+    isAssistantMessageGroupStreaming(
+      [message],
+      getStreamingMessageLookup(
+        [message],
+        true,
+        () => ({ streamMetadata: activeMetadata }),
+        settledMetadata,
+      ),
+    ),
+  ).toBe(true);
+});
+
 test("keeps previous assistant copyable while waiting for a new visible answer", () => {
   const messages = [
     {