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
This commit is contained in:
Aari 2026-08-16 23:53:22 +08:00 committed by GitHub
parent adf6c422c2
commit 16ecf7b006
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 372 additions and 27 deletions

View File

@ -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"

View File

@ -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(

View File

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

View File

@ -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);

View File

@ -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: {} },

View File

@ -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 = [
{