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

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.

site/src/api/queries/chatMessageEdits.test.ts

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

site/src/api/queries/chatMessageEdits.ts

@@ -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,
+	};
+};

site/src/api/queries/chats.test.ts

@@ -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 () => {

site/src/api/queries/chats.ts

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