petru/coder-server
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

perf(site/src): optimistically edit chat messages (#23976)

Browse Source 
Previously, editing a past user message in Agents chat waited for the
PATCH round-trip and cache reconciliation before the conversation
visibly settled. The edited bubble and truncated tail could briefly fall
back to older fetched state, and a failed edit did not restore the full
local editing context cleanly.

Keep history editing optimistic end-to-end: update the edited user
bubble and truncate the tail immediately, preserve that visible
conversation until the authoritative replacement message and cache catch
up, and restore the draft/editor/attachment state on failure. The route
already scopes each `agentId` to a keyed `AgentChatPage` instance with
its own store/cache-writing closures, so navigating between chats does
not need an extra post-await active-chat guard to keep one chat's edit
response out of another chat.
This commit is contained in:
Ethan
2026-04-10 23:40:49 +10:00
committed by GitHub
parent 0a14bb529e
commit a0ea71b74c
15 changed files with 793 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
site/src/api/queries/chatMessageEdits.test.ts Normal file
View File

@@ -0,0 +1,44 @@
import { describe, expect, it } from "vitest";
import type * as TypesGen from "#/api/typesGenerated";
import { buildOptimisticEditedMessage } from "./chatMessageEdits";
const makeUserMessage = (
content: readonly TypesGen.ChatMessagePart[] = [
{ type: "text", text: "original" },
],
): TypesGen.ChatMessage => ({
id: 1,
chat_id: "chat-1",
created_at: "2025-01-01T00:00:00.000Z",
role: "user",
content,
});
describe("buildOptimisticEditedMessage", () => {
it("preserves image MIME types for newly attached files", () => {
const message = buildOptimisticEditedMessage({
requestContent: [{ type: "file", file_id: "image-1" }],
originalMessage: makeUserMessage(),
attachmentMediaTypes: new Map([["image-1", "image/png"]]),
});
expect(message.content).toEqual([
{ type: "file", file_id: "image-1", media_type: "image/png" },
]);
});
it("reuses existing file parts before local attachment metadata", () => {
const existingFilePart: TypesGen.ChatFilePart = {
type: "file",
file_id: "existing-1",
media_type: "image/jpeg",
};
const message = buildOptimisticEditedMessage({
requestContent: [{ type: "file", file_id: "existing-1" }],
originalMessage: makeUserMessage([existingFilePart]),
attachmentMediaTypes: new Map([["existing-1", "text/plain"]]),
});
expect(message.content).toEqual([existingFilePart]);
});
});

148
site/src/api/queries/chatMessageEdits.ts Normal file
View File

@@ -0,0 +1,148 @@
import type { InfiniteData } from "react-query";
import type * as TypesGen from "#/api/typesGenerated";
const buildOptimisticEditedContent = ({
requestContent,
originalMessage,
attachmentMediaTypes,
}: {
requestContent: readonly TypesGen.ChatInputPart[];
originalMessage: TypesGen.ChatMessage;
attachmentMediaTypes?: ReadonlyMap<string, string>;
}): readonly TypesGen.ChatMessagePart[] => {
const existingFilePartsByID = new Map<string, TypesGen.ChatFilePart>();
for (const part of originalMessage.content ?? []) {
if (part.type === "file" && part.file_id) {
existingFilePartsByID.set(part.file_id, part);
}
}
return requestContent.map((part): TypesGen.ChatMessagePart => {
if (part.type === "text") {
return { type: "text", text: part.text ?? "" };
}
if (part.type === "file-reference") {
return {
type: "file-reference",
file_name: part.file_name ?? "",
start_line: part.start_line ?? 1,
end_line: part.end_line ?? 1,
content: part.content ?? "",
};
}
const fileId = part.file_id ?? "";
return (
existingFilePartsByID.get(fileId) ?? {
type: "file",
file_id: part.file_id,
media_type:
attachmentMediaTypes?.get(fileId) ?? "application/octet-stream",
}
);
});
};
export const buildOptimisticEditedMessage = ({
requestContent,
originalMessage,
attachmentMediaTypes,
}: {
requestContent: readonly TypesGen.ChatInputPart[];
originalMessage: TypesGen.ChatMessage;
attachmentMediaTypes?: ReadonlyMap<string, string>;
}): TypesGen.ChatMessage => ({
...originalMessage,
content: buildOptimisticEditedContent({
requestContent,
originalMessage,
attachmentMediaTypes,
}),
});
const sortMessagesDescending = (
messages: readonly TypesGen.ChatMessage[],
): TypesGen.ChatMessage[] => [...messages].sort((a, b) => b.id - a.id);
const upsertFirstPageMessage = (
messages: readonly TypesGen.ChatMessage[],
message: TypesGen.ChatMessage,
): TypesGen.ChatMessage[] => {
const byID = new Map(
messages.map((existingMessage) => [existingMessage.id, existingMessage]),
);
byID.set(message.id, message);
return sortMessagesDescending(Array.from(byID.values()));
};
export const projectEditedConversationIntoCache = ({
currentData,
editedMessageId,
replacementMessage,
queuedMessages,
}: {
currentData: InfiniteData<TypesGen.ChatMessagesResponse> | undefined;
editedMessageId: number;
replacementMessage?: TypesGen.ChatMessage;
queuedMessages?: readonly TypesGen.ChatQueuedMessage[];
}): InfiniteData<TypesGen.ChatMessagesResponse> | undefined => {
if (!currentData?.pages?.length) {
return currentData;
}
const truncatedPages = currentData.pages.map((page, pageIndex) => {
const truncatedMessages = page.messages.filter(
(message) => message.id < editedMessageId,
);
const nextPage = {
...page,
...(pageIndex === 0 && queuedMessages !== undefined
? { queued_messages: queuedMessages }
: {}),
};
if (pageIndex !== 0 || !replacementMessage) {
return { ...nextPage, messages: truncatedMessages };
}
return {
...nextPage,
messages: upsertFirstPageMessage(truncatedMessages, replacementMessage),
};
});
return {
...currentData,
pages: truncatedPages,
};
};
export const reconcileEditedMessageInCache = ({
currentData,
optimisticMessageId,
responseMessage,
}: {
currentData: InfiniteData<TypesGen.ChatMessagesResponse> | undefined;
optimisticMessageId: number;
responseMessage: TypesGen.ChatMessage;
}): InfiniteData<TypesGen.ChatMessagesResponse> | undefined => {
if (!currentData?.pages?.length) {
return currentData;
}
const replacedPages = currentData.pages.map((page, pageIndex) => {
const preservedMessages = page.messages.filter(
(message) =>
message.id !== optimisticMessageId && message.id !== responseMessage.id,
);
if (pageIndex !== 0) {
return { ...page, messages: preservedMessages };
}
return {
...page,
messages: upsertFirstPageMessage(preservedMessages, responseMessage),
};
});
return {
...currentData,
pages: replacedPages,
};
};

196
site/src/api/queries/chats.test.ts
View File

@@ -2,6 +2,7 @@ import { QueryClient } from "react-query";
import { describe, expect, it, vi } from "vitest";
import { API } from "#/api/api";
import type * as TypesGen from "#/api/typesGenerated";
import { buildOptimisticEditedMessage } from "./chatMessageEdits";
import {
archiveChat,
cancelChatListRefetches,
@@ -795,14 +796,44 @@ describe("mutation invalidation scope", () => {
content: [{ type: "text" as const, text: `msg ${id}` }],
});
const makeQueuedMessage = (
chatId: string,
id: number,
): TypesGen.ChatQueuedMessage => ({
id,
chat_id: chatId,
created_at: `2025-01-01T00:10:${String(id).padStart(2, "0")}Z`,
content: [{ type: "text" as const, text: `queued ${id}` }],
});
const editReq = {
content: [{ type: "text" as const, text: "edited" }],
};
it("editChatMessage optimistically removes truncated messages from cache", async () => {
const requireMessage = (
messages: readonly TypesGen.ChatMessage[],
messageId: number,
): TypesGen.ChatMessage => {
const message = messages.find((candidate) => candidate.id === messageId);
if (!message) {
throw new Error(`missing message ${messageId}`);
}
return message;
};
const buildOptimisticMessage = (message: TypesGen.ChatMessage) =>
buildOptimisticEditedMessage({
originalMessage: message,
requestContent: editReq.content,
});
it("editChatMessage writes the optimistic replacement into cache", async () => {
const queryClient = createTestQueryClient();
const chatId = "chat-1";
const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id));
const optimisticMessage = buildOptimisticMessage(
requireMessage(messages, 3),
);
queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
pages: [{ messages, queued_messages: [], has_more: false }],
@@ -812,18 +843,58 @@ describe("mutation invalidation scope", () => {
const mutation = editChatMessage(queryClient, chatId);
const context = await mutation.onMutate({
messageId: 3,
optimisticMessage,
req: editReq,
});
const data = queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId));
expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([2, 1]);
expect(data?.pages[0]?.messages.map((message) => message.id)).toEqual([
3, 2, 1,
]);
expect(data?.pages[0]?.messages[0]?.content).toEqual(
optimisticMessage.content,
);
expect(context?.previousData?.pages[0]?.messages).toHaveLength(5);
});
it("editChatMessage clears queued messages in cache during optimistic history edit", async () => {
const queryClient = createTestQueryClient();
const chatId = "chat-1";
const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id));
const optimisticMessage = buildOptimisticMessage(
requireMessage(messages, 3),
);
const queuedMessages = [makeQueuedMessage(chatId, 11)];
queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
pages: [
{
messages,
queued_messages: queuedMessages,
has_more: false,
},
],
pageParams: [undefined],
});
const mutation = editChatMessage(queryClient, chatId);
await mutation.onMutate({
messageId: 3,
optimisticMessage,
req: editReq,
});
const data = queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId));
expect(data?.pages[0]?.queued_messages).toEqual([]);
});
it("editChatMessage restores cache on error", async () => {
const queryClient = createTestQueryClient();
const chatId = "chat-1";
const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id));
const optimisticMessage = buildOptimisticMessage(
requireMessage(messages, 3),
);
queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
pages: [{ messages, queued_messages: [], has_more: false }],
@@ -833,22 +904,85 @@ describe("mutation invalidation scope", () => {
const mutation = editChatMessage(queryClient, chatId);
const context = await mutation.onMutate({
messageId: 3,
optimisticMessage,
req: editReq,
});
expect(
queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId))?.pages[0]
?.messages,
).toHaveLength(2);
).toHaveLength(3);
mutation.onError(
new Error("network failure"),
{ messageId: 3, req: editReq },
{ messageId: 3, optimisticMessage, req: editReq },
context,
);
const data = queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId));
expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([5, 4, 3, 2, 1]);
expect(data?.pages[0]?.messages.map((message) => message.id)).toEqual([
5, 4, 3, 2, 1,
]);
});
it("editChatMessage preserves websocket-upserted newer messages on success", async () => {
const queryClient = createTestQueryClient();
const chatId = "chat-1";
const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id));
const optimisticMessage = buildOptimisticMessage(
requireMessage(messages, 3),
);
const responseMessage = {
...makeMsg(chatId, 9),
content: [{ type: "text" as const, text: "edited authoritative" }],
};
const websocketMessage = {
...makeMsg(chatId, 10),
content: [{ type: "text" as const, text: "assistant follow-up" }],
role: "assistant" as const,
};
queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
pages: [{ messages, queued_messages: [], has_more: false }],
pageParams: [undefined],
});
const mutation = editChatMessage(queryClient, chatId);
await mutation.onMutate({
messageId: 3,
optimisticMessage,
req: editReq,
});
queryClient.setQueryData<InfMessages | undefined>(
chatMessagesKey(chatId),
(current) => {
if (!current) {
return current;
}
return {
...current,
pages: [
{
...current.pages[0],
messages: [websocketMessage, ...current.pages[0].messages],
},
...current.pages.slice(1),
],
};
},
);
mutation.onSuccess(
{ message: responseMessage },
{ messageId: 3, optimisticMessage, req: editReq },
);
const data = queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId));
expect(data?.pages[0]?.messages.map((message) => message.id)).toEqual([
10, 9, 2, 1,
]);
expect(data?.pages[0]?.messages[1]?.content).toEqual(
responseMessage.content,
);
});
it("editChatMessage onMutate is a no-op when cache is empty", async () => {
@@ -890,13 +1024,14 @@ describe("mutation invalidation scope", () => {
expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([3, 2, 1]);
});
it("editChatMessage onMutate filters across multiple pages", async () => {
it("editChatMessage onMutate updates the first page and preserves older pages", async () => {
const queryClient = createTestQueryClient();
const chatId = "chat-1";
// Page 0 (newest): IDs 10–6. Page 1 (older): IDs 5–1.
const page0 = [10, 9, 8, 7, 6].map((id) => makeMsg(chatId, id));
const page1 = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id));
const optimisticMessage = buildOptimisticMessage(requireMessage(page0, 7));
queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
pages: [
@@ -907,19 +1042,28 @@ describe("mutation invalidation scope", () => {
});
const mutation = editChatMessage(queryClient, chatId);
await mutation.onMutate({ messageId: 7, req: editReq });
await mutation.onMutate({
messageId: 7,
optimisticMessage,
req: editReq,
});
const data = queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId));
// Page 0: only ID 6 survives (< 7).
expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([6]);
// Page 1: all survive (all < 7).
expect(data?.pages[1]?.messages.map((m) => m.id)).toEqual([5, 4, 3, 2, 1]);
expect(data?.pages[0]?.messages.map((message) => message.id)).toEqual([
7, 6,
]);
expect(data?.pages[1]?.messages.map((message) => message.id)).toEqual([
5, 4, 3, 2, 1,
]);
});
it("editChatMessage onMutate editing the first message empties all pages", async () => {
it("editChatMessage onMutate keeps the optimistic replacement when editing the first message", async () => {
const queryClient = createTestQueryClient();
const chatId = "chat-1";
const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id));
const optimisticMessage = buildOptimisticMessage(
requireMessage(messages, 1),
);
queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
pages: [{ messages, queued_messages: [], has_more: false }],
@@ -927,20 +1071,25 @@ describe("mutation invalidation scope", () => {
});
const mutation = editChatMessage(queryClient, chatId);
await mutation.onMutate({ messageId: 1, req: editReq });
await mutation.onMutate({
messageId: 1,
optimisticMessage,
req: editReq,
});
const data = queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId));
// All messages have id >= 1, so the page is empty.
expect(data?.pages[0]?.messages).toHaveLength(0);
// Sibling fields survive the spread.
expect(data?.pages[0]?.messages.map((message) => message.id)).toEqual([1]);
expect(data?.pages[0]?.queued_messages).toEqual([]);
expect(data?.pages[0]?.has_more).toBe(false);
});
it("editChatMessage onMutate editing the latest message keeps earlier ones", async () => {
it("editChatMessage onMutate keeps earlier messages when editing the latest message", async () => {
const queryClient = createTestQueryClient();
const chatId = "chat-1";
const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id));
const optimisticMessage = buildOptimisticMessage(
requireMessage(messages, 5),
);
queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
pages: [{ messages, queued_messages: [], has_more: false }],
@@ -948,10 +1097,19 @@ describe("mutation invalidation scope", () => {
});
const mutation = editChatMessage(queryClient, chatId);
await mutation.onMutate({ messageId: 5, req: editReq });
await mutation.onMutate({
messageId: 5,
optimisticMessage,
req: editReq,
});
const data = queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId));
expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([4, 3, 2, 1]);
expect(data?.pages[0]?.messages.map((message) => message.id)).toEqual([
5, 4, 3, 2, 1,
]);
expect(data?.pages[0]?.messages[0]?.content).toEqual(
optimisticMessage.content,
);
});
it("interruptChat does not invalidate unrelated queries", async () => {

63
site/src/api/queries/chats.ts
View File

@@ -6,6 +6,10 @@ import type {
import { API } from "#/api/api";
import type * as TypesGen from "#/api/typesGenerated";
import type { UsePaginatedQueryOptions } from "#/hooks/usePaginatedQuery";
import {
projectEditedConversationIntoCache,
reconcileEditedMessageInCache,
} from "./chatMessageEdits";
export const chatsKey = ["chats"] as const;
export const chatKey = (chatId: string) => ["chats", chatId] as const;
@@ -601,13 +605,21 @@ export const createChatMessage = (
type EditChatMessageMutationArgs = {
messageId: number;
optimisticMessage?: TypesGen.ChatMessage;
req: TypesGen.EditChatMessageRequest;
};
type EditChatMessageMutationContext = {
previousData?: InfiniteData<TypesGen.ChatMessagesResponse> | undefined;
};
export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({
mutationFn: ({ messageId, req }: EditChatMessageMutationArgs) =>
API.experimental.editChatMessage(chatId, messageId, req),
onMutate: async ({ messageId }: EditChatMessageMutationArgs) => {
onMutate: async ({
messageId,
optimisticMessage,
}: EditChatMessageMutationArgs): Promise<EditChatMessageMutationContext> => {
// Cancel in-flight refetches so they don't overwrite the
// optimistic update before the mutation completes.
await queryClient.cancelQueries({
@@ -619,40 +631,23 @@ export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({
InfiniteData<TypesGen.ChatMessagesResponse>
>(chatMessagesKey(chatId));
// Optimistically remove the edited message and everything
// after it. The server soft-deletes these and inserts a
// replacement with a new ID. Without this, the WebSocket
// handler's upsertCacheMessages adds new messages to the
// React Query cache without removing the soft-deleted ones,
// causing deleted messages to flash back into view until
// the full REST refetch resolves.
queryClient.setQueryData<
InfiniteData<TypesGen.ChatMessagesResponse> | undefined
>(chatMessagesKey(chatId), (current) => {
if (!current?.pages?.length) {
return current;
}
return {
...current,
pages: current.pages.map((page) => ({
...page,
messages: page.messages.filter((m) => m.id < messageId),
})),
};
});
>(chatMessagesKey(chatId), (current) =>
projectEditedConversationIntoCache({
currentData: current,
editedMessageId: messageId,
replacementMessage: optimisticMessage,
queuedMessages: [],
}),
);
return { previousData };
},
onError: (
_error: unknown,
_variables: EditChatMessageMutationArgs,
context:
| {
previousData?:
| InfiniteData<TypesGen.ChatMessagesResponse>
| undefined;
}
| undefined,
context: EditChatMessageMutationContext | undefined,
) => {
// Restore the cache on failure so the user sees the
// original messages again.
@@ -660,6 +655,20 @@ export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({
queryClient.setQueryData(chatMessagesKey(chatId), context.previousData);
}
},
onSuccess: (
response: TypesGen.EditChatMessageResponse,
variables: EditChatMessageMutationArgs,
) => {
queryClient.setQueryData<
InfiniteData<TypesGen.ChatMessagesResponse> | undefined
>(chatMessagesKey(chatId), (current) =>
reconcileEditedMessageInCache({
currentData: current,
optimisticMessageId: variables.messageId,
responseMessage: response.message,
}),
);
},
onSettled: () => {
// Always reconcile with the server regardless of whether
// the mutation succeeded or failed. On success this picks

96
site/src/pages/AgentsPage/AgentChatPage.test.ts
View File

@@ -4,9 +4,12 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
import {
draftInputStorageKeyPrefix,
getPersistedDraftInputValue,
restoreOptimisticRequestSnapshot,
useConversationEditingState,
} from "./AgentChatPage";
import type { ChatMessageInputRef } from "./components/AgentChatInput";
import { createChatStore } from "./components/ChatConversation/chatStore";
import type { PendingAttachment } from "./components/ChatPageContent";
type MockChatInputHandle = {
handle: ChatMessageInputRef;
@@ -84,6 +87,41 @@ describe("getPersistedDraftInputValue", () => {
});
});
describe("restoreOptimisticRequestSnapshot", () => {
it("restores queued messages, stream output, status, and stream error", () => {
const store = createChatStore();
store.setQueuedMessages([
{
id: 9,
chat_id: "chat-abc-123",
created_at: "2025-01-01T00:00:00.000Z",
content: [{ type: "text" as const, text: "queued" }],
},
]);
store.setChatStatus("running");
store.applyMessagePart({ type: "text", text: "partial response" });
store.setStreamError({ kind: "generic", message: "old error" });
const previousSnapshot = store.getSnapshot();
store.batch(() => {
store.setQueuedMessages([]);
store.setChatStatus("pending");
store.clearStreamState();
store.clearStreamError();
});
restoreOptimisticRequestSnapshot(store, previousSnapshot);
const restoredSnapshot = store.getSnapshot();
expect(restoredSnapshot.queuedMessages).toEqual(
previousSnapshot.queuedMessages,
);
expect(restoredSnapshot.chatStatus).toBe(previousSnapshot.chatStatus);
expect(restoredSnapshot.streamState).toBe(previousSnapshot.streamState);
expect(restoredSnapshot.streamError).toEqual(previousSnapshot.streamError);
});
});
describe("useConversationEditingState", () => {
const chatID = "chat-abc-123";
const expectedKey = `${draftInputStorageKeyPrefix}${chatID}`;
@@ -327,6 +365,64 @@ describe("useConversationEditingState", () => {
unmount();
});
it("forwards pending attachments through history-edit send", async () => {
const { result, onSend, unmount } = renderEditing();
const attachments: PendingAttachment[] = [
{ fileId: "file-1", mediaType: "image/png" },
];
act(() => {
result.current.handleEditUserMessage(7, "hello");
});
await act(async () => {
await result.current.handleSendFromInput("hello", attachments);
});
expect(onSend).toHaveBeenCalledWith("hello", attachments, 7);
unmount();
});
it("restores the edit draft and file-block seed when an edit submission fails", async () => {
const { result, onSend, unmount } = renderEditing();
const mockInput = createMockChatInputHandle("edited message");
const fileBlocks = [
{ type: "file", file_id: "file-1", media_type: "image/png" },
] as const;
result.current.chatInputRef.current = mockInput.handle;
onSend.mockRejectedValueOnce(new Error("boom"));
const editorState = JSON.stringify({
root: {
children: [
{
children: [{ text: "edited message" }],
type: "paragraph",
},
],
type: "root",
},
});
act(() => {
result.current.handleEditUserMessage(7, "edited message", fileBlocks);
result.current.handleContentChange("edited message", editorState, false);
});
await act(async () => {
await expect(
result.current.handleSendFromInput("edited message"),
).rejects.toThrow("boom");
});
expect(mockInput.clear).toHaveBeenCalled();
expect(result.current.inputValueRef.current).toBe("edited message");
expect(result.current.editingMessageId).toBe(7);
expect(result.current.editingFileBlocks).toEqual(fileBlocks);
expect(result.current.editorInitialValue).toBe("edited message");
expect(result.current.initialEditorState).toBe(editorState);
unmount();
});
it("clears the composer and persisted draft after a successful send", async () => {
localStorage.setItem(expectedKey, "draft to clear");
const { result, onSend, unmount } = renderEditing();

166
site/src/pages/AgentsPage/AgentChatPage.tsx
View File

@@ -11,6 +11,7 @@ import { toast } from "sonner";
import type { UrlTransform } from "streamdown";
import { API, watchWorkspace } from "#/api/api";
import { isApiError } from "#/api/errors";
import { buildOptimisticEditedMessage } from "#/api/queries/chatMessageEdits";
import {
chat,
chatDesktopEnabled,
@@ -51,11 +52,14 @@ import {
getWorkspaceAgent,
} from "./components/ChatConversation/chatHelpers";
import {
type ChatStore,
type ChatStoreState,
selectChatStatus,
useChatSelector,
useChatStore,
} from "./components/ChatConversation/chatStore";
import { useWorkspaceCreationWatcher } from "./components/ChatConversation/useWorkspaceCreationWatcher";
import type { PendingAttachment } from "./components/ChatPageContent";
import {
getDefaultMCPSelection,
getSavedMCPSelection,
@@ -101,12 +105,47 @@ export function getPersistedDraftInputValue(
).text;
}
/** @internal Exported for testing. */
export const restoreOptimisticRequestSnapshot = (
store: Pick<
ChatStore,
| "batch"
| "setChatStatus"
| "setQueuedMessages"
| "setStreamError"
| "setStreamState"
>,
snapshot: Pick<
ChatStoreState,
"chatStatus" | "queuedMessages" | "streamError" | "streamState"
>,
): void => {
store.batch(() => {
store.setQueuedMessages(snapshot.queuedMessages);
store.setChatStatus(snapshot.chatStatus);
store.setStreamState(snapshot.streamState);
store.setStreamError(snapshot.streamError);
});
};
const buildAttachmentMediaTypes = (
attachments?: readonly PendingAttachment[],
): ReadonlyMap<string, string> | undefined => {
if (!attachments?.length) {
return undefined;
}
return new Map(
attachments.map(({ fileId, mediaType }) => [fileId, mediaType]),
);
};
/** @internal Exported for testing. */
export function useConversationEditingState(deps: {
chatID: string | undefined;
onSend: (
message: string,
fileIds?: string[],
attachments?: readonly PendingAttachment[],
editedMessageID?: number,
) => Promise<void>;
onDeleteQueuedMessage: (id: number) => Promise<void>;
@@ -130,6 +169,9 @@ export function useConversationEditingState(deps: {
};
},
);
const serializedEditorStateRef = useRef<string | undefined>(
initialEditorState,
);
// Monotonic counter to force LexicalComposer remount.
const [remountKey, setRemountKey] = useState(0);
@@ -176,6 +218,7 @@ export function useConversationEditingState(deps: {
editorInitialValue: text,
initialEditorState: undefined,
});
serializedEditorStateRef.current = undefined;
setRemountKey((k) => k + 1);
inputValueRef.current = text;
setEditingFileBlocks(fileBlocks ?? []);
@@ -188,6 +231,7 @@ export function useConversationEditingState(deps: {
editorInitialValue: savedText,
initialEditorState: savedState,
});
serializedEditorStateRef.current = savedState;
setRemountKey((k) => k + 1);
inputValueRef.current = savedText;
setEditingMessageId(null);
@@ -221,6 +265,7 @@ export function useConversationEditingState(deps: {
editorInitialValue: text,
initialEditorState: undefined,
});
serializedEditorStateRef.current = undefined;
setRemountKey((k) => k + 1);
inputValueRef.current = text;
setEditingFileBlocks(fileBlocks);
@@ -233,6 +278,7 @@ export function useConversationEditingState(deps: {
editorInitialValue: savedText,
initialEditorState: savedState,
});
serializedEditorStateRef.current = savedState;
setRemountKey((k) => k + 1);
inputValueRef.current = savedText;
setEditingQueuedMessageID(null);
@@ -240,25 +286,48 @@ export function useConversationEditingState(deps: {
setEditingFileBlocks([]);
};
// Wraps the parent onSend to clear local input/editing state
// and handle queue-edit deletion.
const handleSendFromInput = async (message: string, fileIds?: string[]) => {
const editedMessageID =
editingMessageId !== null ? editingMessageId : undefined;
const queueEditID = editingQueuedMessageID;
// Clears the composer for an in-flight history edit and
// returns a rollback function that restores the editing draft
// if the send fails.
const clearInputForHistoryEdit = (message: string) => {
const snapshot = {
editorState: serializedEditorStateRef.current,
fileBlocks: editingFileBlocks,
messageId: editingMessageId,
};
await onSend(message, fileIds, editedMessageID);
// Clear input and editing state on success.
chatInputRef.current?.clear();
inputValueRef.current = "";
setEditingMessageId(null);
return () => {
setDraftState({
editorInitialValue: message,
initialEditorState: snapshot.editorState,
});
serializedEditorStateRef.current = snapshot.editorState;
setRemountKey((k) => k + 1);
inputValueRef.current = message;
setEditingMessageId(snapshot.messageId);
setEditingFileBlocks(snapshot.fileBlocks);
};
};
// Clears all input and editing state after a successful send.
const finalizeSuccessfulSend = (
editedMessageID: number | undefined,
queueEditID: number | null,
) => {
chatInputRef.current?.clear();
if (!isMobileViewport()) {
chatInputRef.current?.focus();
}
inputValueRef.current = "";
serializedEditorStateRef.current = undefined;
if (draftStorageKey) {
localStorage.removeItem(draftStorageKey);
}
if (editingMessageId !== null) {
setEditingMessageId(null);
if (editedMessageID !== undefined) {
setDraftBeforeHistoryEdit(null);
setEditingFileBlocks([]);
}
@@ -270,12 +339,41 @@ export function useConversationEditingState(deps: {
}
};
// Wraps the parent onSend to clear local input/editing state
// and handle queue-edit deletion.
const handleSendFromInput = async (
message: string,
attachments?: readonly PendingAttachment[],
) => {
const editedMessageID =
editingMessageId !== null ? editingMessageId : undefined;
const queueEditID = editingQueuedMessageID;
const sendPromise = onSend(message, attachments, editedMessageID);
// For history edits, clear input immediately and prepare
// a rollback in case the send fails.
const rollback =
editedMessageID !== undefined
? clearInputForHistoryEdit(message)
: undefined;
try {
await sendPromise;
} catch (error) {
rollback?.();
throw error;
}
finalizeSuccessfulSend(editedMessageID, queueEditID);
};
const handleContentChange = (
content: string,
serializedEditorState: string,
hasFileReferences: boolean,
) => {
inputValueRef.current = content;
serializedEditorStateRef.current = serializedEditorState;
// Don't overwrite the persisted draft while editing a
// history or queued message — the original draft (possibly
@@ -430,9 +528,6 @@ const AgentChatPage: FC = () => {
} = useOutletContext<AgentsOutletContext>();
const queryClient = useQueryClient();
const [selectedModel, setSelectedModel] = useState("");
const [pendingEditMessageId, setPendingEditMessageId] = useState<
number | null
>(null);
const scrollToBottomRef = useRef<(() => void) | null>(null);
const chatInputRef = useRef<ChatMessageInputRef | null>(null);
const inputValueRef = useRef(
@@ -775,7 +870,7 @@ const AgentChatPage: FC = () => {
const handleSend = async (
message: string,
fileIds?: string[],
attachments?: readonly PendingAttachment[],
editedMessageID?: number,
) => {
const chatInputHandle = (
@@ -790,7 +885,9 @@ const AgentChatPage: FC = () => {
(p) => p.type === "file-reference",
);
const hasContent =
message.trim() || (fileIds && fileIds.length > 0) || hasFileReferences;
message.trim() ||
(attachments && attachments.length > 0) ||
hasFileReferences;
if (!hasContent || isSubmissionPending || !agentId || !hasModelOptions) {
return;
}
@@ -818,28 +915,41 @@ const AgentChatPage: FC = () => {
}
}
// Add pre-uploaded file references.
if (fileIds && fileIds.length > 0) {
for (const fileId of fileIds) {
// Add pre-uploaded file attachments.
if (attachments && attachments.length > 0) {
for (const { fileId } of attachments) {
content.push({ type: "file", file_id: fileId });
}
}
if (editedMessageID !== undefined) {
const request: TypesGen.EditChatMessageRequest = { content };
const originalEditedMessage = chatMessagesList?.find(
(existingMessage) => existingMessage.id === editedMessageID,
);
const optimisticMessage = originalEditedMessage
? buildOptimisticEditedMessage({
requestContent: request.content,
originalMessage: originalEditedMessage,
attachmentMediaTypes: buildAttachmentMediaTypes(attachments),
})
: undefined;
const previousSnapshot = store.getSnapshot();
clearChatErrorReason(agentId);
clearStreamError();
setPendingEditMessageId(editedMessageID);
store.batch(() => {
store.setQueuedMessages([]);
store.setChatStatus("running");
store.clearStreamState();
});
scrollToBottomRef.current?.();
try {
await editMessage({
messageId: editedMessageID,
optimisticMessage,
req: request,
});
store.clearStreamState();
store.setChatStatus("running");
setPendingEditMessageId(null);
} catch (error) {
setPendingEditMessageId(null);
restoreOptimisticRequestSnapshot(store, previousSnapshot);
handleUsageLimitError(error);
throw error;
}
@@ -918,10 +1028,8 @@ const AgentChatPage: FC = () => {
const handlePromoteQueuedMessage = async (id: number) => {
const previousSnapshot = store.getSnapshot();
const previousQueuedMessages = previousSnapshot.queuedMessages;
const previousChatStatus = previousSnapshot.chatStatus;
store.setQueuedMessages(
previousQueuedMessages.filter((message) => message.id !== id),
previousSnapshot.queuedMessages.filter((message) => message.id !== id),
);
store.clearStreamState();
if (agentId) {
@@ -937,8 +1045,7 @@ const AgentChatPage: FC = () => {
store.upsertDurableMessage(promotedMessage);
upsertCacheMessages([promotedMessage]);
} catch (error) {
store.setQueuedMessages(previousQueuedMessages);
store.setChatStatus(previousChatStatus);
restoreOptimisticRequestSnapshot(store, previousSnapshot);
handleUsageLimitError(error);
throw error;
}
@@ -1133,7 +1240,6 @@ const AgentChatPage: FC = () => {
workspaceAgent={workspaceAgent}
store={store}
editing={editing}
pendingEditMessageId={pendingEditMessageId}
effectiveSelectedModel={effectiveSelectedModel}
setSelectedModel={setSelectedModel}
modelOptions={modelOptions}

47
site/src/pages/AgentsPage/AgentChatPageView.stories.tsx
View File

@@ -113,7 +113,6 @@ const StoryAgentChatPageView: FC<StoryProps> = ({ editing, ...overrides }) => {
parentChat: undefined as TypesGen.Chat | undefined,
isArchived: false,
store: createChatStore(),
pendingEditMessageId: null as number | null,
effectiveSelectedModel: defaultModelConfigID,
setSelectedModel: fn(),
modelOptions: defaultModelOptions,
@@ -474,16 +473,6 @@ const buildStoreWithMessages = (
return store;
};
const gapTestStore = createChatStore();
gapTestStore.replaceMessages([
buildMessage(1, "user", "Explain the layout."),
buildMessage(2, "assistant", "Here is the explanation."),
buildMessage(3, "user", "Can you elaborate?"),
]);
gapTestStore.applyMessageParts([
{ type: "text", text: "Certainly, here are more details..." },
]);
// ---------------------------------------------------------------------------
// Editing flow stories
// ---------------------------------------------------------------------------
@@ -515,42 +504,6 @@ export const EditingMessage: Story = {
),
};
/** The saving state while an edit is in progress — shows the pending
* indicator on the message being saved. */
export const EditingSaving: Story = {
render: () => (
<StoryAgentChatPageView
store={buildStoreWithMessages(editingMessages)}
editing={{
editingMessageId: 3,
editorInitialValue: "Now tell me a better joke",
}}
pendingEditMessageId={3}
isSubmissionPending
/>
),
};
export const ConsistentGapBetweenTimelineAndStream: Story = {
render: () => <StoryAgentChatPageView store={gapTestStore} />,
play: async ({ canvasElement }) => {
const wrapper = canvasElement.querySelector(
'[data-testid="chat-timeline-wrapper"]',
);
expect(wrapper).not.toBeNull();
const outerGap = window.getComputedStyle(wrapper!).rowGap;
expect(outerGap).toBe("8px");
const timeline = wrapper!.querySelector(
'[data-testid="conversation-timeline"]',
);
expect(timeline).not.toBeNull();
const innerGap = window.getComputedStyle(timeline!).rowGap;
expect(innerGap).toBe("8px");
},
};
// ---------------------------------------------------------------------------
// AgentChatPageNotFoundView stories
// ---------------------------------------------------------------------------

9
site/src/pages/AgentsPage/AgentChatPageView.tsx
View File

@@ -36,6 +36,7 @@ import {
import type { useChatStore } from "./components/ChatConversation/chatStore";
import type { ModelSelectorOption } from "./components/ChatElements";
import { DesktopPanelContext } from "./components/ChatElements/tools/DesktopPanelContext";
import type { PendingAttachment } from "./components/ChatPageContent";
import { ChatPageInput, ChatPageTimeline } from "./components/ChatPageContent";
import { ChatScrollContainer } from "./components/ChatScrollContainer";
import { ChatTopBar } from "./components/ChatTopBar";
@@ -69,7 +70,10 @@ interface EditingState {
fileBlocks: readonly ChatMessagePart[],
) => void;
handleCancelQueueEdit: () => void;
handleSendFromInput: (message: string, fileIds?: string[]) => void;
handleSendFromInput: (
message: string,
attachments?: readonly PendingAttachment[],
) => void;
handleContentChange: (
content: string,
serializedEditorState: string,
@@ -92,7 +96,6 @@ interface AgentChatPageViewProps {
// Editing state.
editing: EditingState;
pendingEditMessageId: number | null;
// Model/input configuration.
effectiveSelectedModel: string;
@@ -179,7 +182,6 @@ export const AgentChatPageView: FC<AgentChatPageViewProps> = ({
workspace,
store,
editing,
pendingEditMessageId,
effectiveSelectedModel,
setSelectedModel,
modelOptions,
@@ -387,7 +389,6 @@ export const AgentChatPageView: FC<AgentChatPageViewProps> = ({
persistedError={persistedError}
onEditUserMessage={editing.handleEditUserMessage}
editingMessageId={editing.editingMessageId}
savingMessageId={pendingEditMessageId}
urlTransform={urlTransform}
mcpServers={mcpServers}
/>

5
site/src/pages/AgentsPage/components/AgentChatInput.tsx
View File

@@ -668,9 +668,8 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
<div className="flex items-center justify-between border-b border-border-warning/50 px-3 py-1.5">
<span className="flex items-center gap-1.5 text-xs font-medium text-content-warning">
<PencilIcon className="h-3.5 w-3.5" />
{isLoading
? "Saving edit..."
: "Editing will delete all subsequent messages and restart the conversation here."}
Editing will delete all subsequent messages and restart the
conversation here.
</span>
<Button
type="button"

20
site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx
View File

@@ -12,7 +12,6 @@ import type { UrlTransform } from "streamdown";
import type * as TypesGen from "#/api/typesGenerated";
import { Button } from "#/components/Button/Button";
import { CopyButton } from "#/components/CopyButton/CopyButton";
import { Spinner } from "#/components/Spinner/Spinner";
import {
Tooltip,
TooltipContent,
@@ -427,7 +426,6 @@ const ChatMessageItem = memo<{
fileBlocks?: readonly TypesGen.ChatMessagePart[],
) => void;
editingMessageId?: number | null;
savingMessageId?: number | null;
isAfterEditingMessage?: boolean;
hideActions?: boolean;
@@ -446,7 +444,6 @@ const ChatMessageItem = memo<{
parsed,
onEditUserMessage,
editingMessageId,
savingMessageId,
isAfterEditingMessage = false,
hideActions = false,
fadeFromBottom = false,
@@ -458,7 +455,6 @@ const ChatMessageItem = memo<{
showDesktopPreviews,
}) => {
const isUser = message.role === "user";
const isSavingMessage = savingMessageId === message.id;
const [previewImage, setPreviewImage] = useState<string | null>(null);
const [previewText, setPreviewText] = useState<string | null>(null);
if (
@@ -541,7 +537,6 @@ const ChatMessageItem = memo<{
"rounded-lg border border-solid border-border-default bg-surface-secondary px-3 py-2 font-sans shadow-sm transition-shadow",
editingMessageId === message.id &&
"border-surface-secondary shadow-[0_0_0_2px_hsla(var(--border-warning),0.6)]",
isSavingMessage && "ring-2 ring-content-secondary/40",
fadeFromBottom && "relative overflow-hidden",
)}
style={
@@ -572,13 +567,6 @@ const ChatMessageItem = memo<{
: parsed.markdown || ""}
</span>
)}
{isSavingMessage && (
<Spinner
className="mt-0.5 h-3.5 w-3.5 shrink-0 text-content-secondary"
aria-label="Saving message edit"
loading
/>
)}
</div>
)}
{hasFileBlocks && (
@@ -704,7 +692,6 @@ const StickyUserMessage = memo<{
fileBlocks?: readonly TypesGen.ChatMessagePart[],
) => void;
editingMessageId?: number | null;
savingMessageId?: number | null;
isAfterEditingMessage?: boolean;
}>(
({
@@ -712,7 +699,6 @@ const StickyUserMessage = memo<{
parsed,
onEditUserMessage,
editingMessageId,
savingMessageId,
isAfterEditingMessage = false,
}) => {
const [isStuck, setIsStuck] = useState(false);
@@ -937,7 +923,6 @@ const StickyUserMessage = memo<{
parsed={parsed}
onEditUserMessage={handleEditUserMessage}
editingMessageId={editingMessageId}
savingMessageId={savingMessageId}
isAfterEditingMessage={isAfterEditingMessage}
/>
</div>
@@ -980,7 +965,6 @@ const StickyUserMessage = memo<{
parsed={parsed}
onEditUserMessage={handleEditUserMessage}
editingMessageId={editingMessageId}
savingMessageId={savingMessageId}
isAfterEditingMessage={isAfterEditingMessage}
fadeFromBottom
/>
@@ -1002,7 +986,6 @@ interface ConversationTimelineProps {
fileBlocks?: readonly TypesGen.ChatMessagePart[],
) => void;
editingMessageId?: number | null;
savingMessageId?: number | null;
urlTransform?: UrlTransform;
mcpServers?: readonly TypesGen.MCPServerConfig[];
computerUseSubagentIds?: Set<string>;
@@ -1016,7 +999,6 @@ export const ConversationTimeline = memo<ConversationTimelineProps>(
subagentTitles,
onEditUserMessage,
editingMessageId,
savingMessageId,
urlTransform,
mcpServers,
computerUseSubagentIds,
@@ -1053,7 +1035,6 @@ export const ConversationTimeline = memo<ConversationTimelineProps>(
parsed={parsed}
onEditUserMessage={onEditUserMessage}
editingMessageId={editingMessageId}
savingMessageId={savingMessageId}
isAfterEditingMessage={afterEditingMessageIds.has(message.id)}
/>
);
@@ -1067,7 +1048,6 @@ export const ConversationTimeline = memo<ConversationTimelineProps>(
key={message.id}
message={message}
parsed={parsed}
savingMessageId={savingMessageId}
urlTransform={urlTransform}
isAfterEditingMessage={afterEditingMessageIds.has(message.id)}
hideActions={!isLastInChain}

21
site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts
View File

@@ -190,6 +190,27 @@ describe("setChatStatus", () => {
});
});
// ---------------------------------------------------------------------------
// setStreamState
// ---------------------------------------------------------------------------
describe("setStreamState", () => {
it("does not notify when setting the same stream state reference", () => {
const store = createChatStore();
store.applyMessagePart({ type: "text", text: "hello" });
const streamState = store.getSnapshot().streamState;
expect(streamState).not.toBeNull();
let notified = false;
store.subscribe(() => {
notified = true;
});
store.setStreamState(streamState);
expect(notified).toBe(false);
});
});
// ---------------------------------------------------------------------------
// setStreamError / clearStreamError
// ---------------------------------------------------------------------------

90
site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx
View File

@@ -4213,6 +4213,96 @@ describe("store/cache desync protection", () => {
expect(result.current.orderedMessageIDs).toEqual([1]);
});
});
it("reflects optimistic and authoritative history-edit cache updates through the normal sync effect", async () => {
immediateAnimationFrame();
const chatID = "chat-local-edit-sync";
const msg1 = makeMessage(chatID, 1, "user", "first");
const msg2 = makeMessage(chatID, 2, "assistant", "second");
const msg3 = makeMessage(chatID, 3, "user", "third");
const optimisticReplacement = {
...msg3,
content: [{ type: "text" as const, text: "edited draft" }],
};
const authoritativeReplacement = makeMessage(chatID, 9, "user", "edited");
const mockSocket = createMockSocket();
mockWatchChatReturn(mockSocket);
const queryClient = createTestQueryClient();
const wrapper: FC<PropsWithChildren> = ({ children }) => (
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
);
const initialOptions = {
chatID,
chatMessages: [msg1, msg2, msg3],
chatRecord: makeChat(chatID),
chatMessagesData: {
messages: [msg1, msg2, msg3],
queued_messages: [],
has_more: false,
},
chatQueuedMessages: [] as TypesGen.ChatQueuedMessage[],
setChatErrorReason: vi.fn(),
clearChatErrorReason: vi.fn(),
};
const { result, rerender } = renderHook(
(options: Parameters<typeof useChatStore>[0]) => {
const { store } = useChatStore(options);
return {
store,
messagesByID: useChatSelector(store, selectMessagesByID),
orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs),
};
},
{ initialProps: initialOptions, wrapper },
);
await waitFor(() => {
expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]);
});
act(() => {
mockSocket.emitOpen();
});
rerender({
...initialOptions,
chatMessages: [msg1, msg2, optimisticReplacement],
chatMessagesData: {
messages: [msg1, msg2, optimisticReplacement],
queued_messages: [],
has_more: false,
},
});
await waitFor(() => {
expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]);
expect(result.current.messagesByID.get(3)?.content).toEqual(
optimisticReplacement.content,
);
});
rerender({
...initialOptions,
chatMessages: [msg1, msg2, authoritativeReplacement],
chatMessagesData: {
messages: [msg1, msg2, authoritativeReplacement],
queued_messages: [],
has_more: false,
},
});
await waitFor(() => {
expect(result.current.orderedMessageIDs).toEqual([1, 2, 9]);
expect(result.current.messagesByID.has(3)).toBe(false);
expect(result.current.messagesByID.get(9)?.content).toEqual(
authoritativeReplacement.content,
);
});
});
});
describe("parse errors", () => {

15
site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts
View File

@@ -174,6 +174,7 @@ export type ChatStore = {
queuedMessages: readonly TypesGen.ChatQueuedMessage[] | undefined,
) => void;
setChatStatus: (status: TypesGen.ChatStatus | null) => void;
setStreamState: (streamState: StreamState | null) => void;
setStreamError: (reason: ChatDetailError | null) => void;
clearStreamError: () => void;
setRetryState: (state: RetryState | null) => void;
@@ -412,6 +413,20 @@ export const createChatStore = (): ChatStore => {
chatStatus: status,
}));
},
setStreamState: (streamState) => {
if (state.streamState === streamState) {
return;
}
setState((current) => {
if (current.streamState === streamState) {
return current;
}
return {
...current,
streamState,
};
});
},
setStreamError: (reason) => {
setState((current) => {
if (chatDetailErrorsEqual(current.streamError, reason)) {

7
site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts
View File

@@ -206,10 +206,9 @@ export const useChatStore = (
const fetchedIDs = new Set(chatMessages.map((m) => m.id));
// Only classify a store-held ID as stale if it was
// present in the PREVIOUS sync's fetched data. IDs
// added to the store after the last sync (by the WS
// handler or handleSend) are new, not stale, and
// must not trigger the destructive replaceMessages
// path.
// added to the store after the last sync (for example
// by the WS handler) are new, not stale, and must not
// trigger the destructive replaceMessages path.
const prevIDs = new Set(prev.map((m) => m.id));
const hasStaleEntries =
contentChanged &&

30
site/src/pages/AgentsPage/components/ChatPageContent.tsx
View File

@@ -48,7 +48,6 @@ interface ChatPageTimelineProps {
fileBlocks?: readonly TypesGen.ChatMessagePart[],
) => void;
editingMessageId?: number | null;
savingMessageId?: number | null;
urlTransform?: UrlTransform;
mcpServers?: readonly TypesGen.MCPServerConfig[];
}
@@ -59,7 +58,6 @@ export const ChatPageTimeline: FC<ChatPageTimelineProps> = ({
persistedError,
onEditUserMessage,
editingMessageId,
savingMessageId,
urlTransform,
mcpServers,
}) => {
@@ -100,7 +98,6 @@ export const ChatPageTimeline: FC<ChatPageTimelineProps> = ({
subagentTitles={subagentTitles}
onEditUserMessage={onEditUserMessage}
editingMessageId={editingMessageId}
savingMessageId={savingMessageId}
urlTransform={urlTransform}
mcpServers={mcpServers}
computerUseSubagentIds={computerUseSubagentIds}
@@ -121,10 +118,18 @@ export const ChatPageTimeline: FC<ChatPageTimelineProps> = ({
);
};
export type PendingAttachment = {
fileId: string;
mediaType: string;
};
interface ChatPageInputProps {
store: ChatStoreHandle;
compressionThreshold: number | undefined;
onSend: (message: string, fileIds?: string[]) => void;
onSend: (
message: string,
attachments?: readonly PendingAttachment[],
) => Promise<void> | void;
onDeleteQueuedMessage: (id: number) => Promise<void>;
onPromoteQueuedMessage: (id: number) => Promise<void>;
onInterrupt: () => void;
@@ -315,9 +320,10 @@ export const ChatPageInput: FC<ChatPageInputProps> = ({
<AgentChatInput
onSend={(message) => {
void (async () => {
// Collect file IDs from already-uploaded attachments.
// Skip files in error state (e.g. too large).
const fileIds: string[] = [];
// Collect uploaded attachment metadata for the optimistic
// transcript builder while keeping the server payload
// shape unchanged downstream.
const pendingAttachments: PendingAttachment[] = [];
let skippedErrors = 0;
for (const file of attachments) {
const state = uploadStates.get(file);
@@ -326,7 +332,10 @@ export const ChatPageInput: FC<ChatPageInputProps> = ({
continue;
}
if (state?.status === "uploaded" && state.fileId) {
fileIds.push(state.fileId);
pendingAttachments.push({
fileId: state.fileId,
mediaType: file.type || "application/octet-stream",
});
}
}
if (skippedErrors > 0) {
@@ -334,9 +343,10 @@ export const ChatPageInput: FC<ChatPageInputProps> = ({
`${skippedErrors} attachment${skippedErrors > 1 ? "s" : ""} could not be sent (upload failed)`,
);
}
const fileArg = fileIds.length > 0 ? fileIds : undefined;
const attachmentArg =
pendingAttachments.length > 0 ? pendingAttachments : undefined;
try {
await onSend(message, fileArg);
await onSend(message, attachmentArg);
} catch {
// Attachments preserved for retry on failure.
return;