fix(site/src/api/queries): optimistically truncate cache on chat message edit (#23864)

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

@@ -741,7 +741,7 @@ describe("mutation invalidation scope", () => {
 		seedAllActiveQueries(queryClient, chatId);
 
 		const mutation = editChatMessage(queryClient, chatId);
-		mutation.onSuccess();
+		mutation.onSettled();
 
 		await new Promise((r) => setTimeout(r, 0));
 
@@ -760,7 +760,7 @@ describe("mutation invalidation scope", () => {
 		seedAllActiveQueries(queryClient, chatId);
 
 		const mutation = editChatMessage(queryClient, chatId);
-		mutation.onSuccess();
+		mutation.onSettled();
 
 		await new Promise((r) => setTimeout(r, 0));
 
@@ -778,6 +778,180 @@ describe("mutation invalidation scope", () => {
 		).toBe(true);
 	});
 
+	// Shared type for the infinite messages cache shape used by
+	// editChatMessage tests below.
+	type InfMessages = {
+		pages: TypesGen.ChatMessagesResponse[];
+		pageParams: (number | undefined)[];
+	};
+
+	const makeMsg = (chatId: string, id: number): TypesGen.ChatMessage => ({
+		id,
+		chat_id: chatId,
+		created_at: `2025-01-01T00:00:${String(id).padStart(2, "0")}Z`,
+		role: "user" as const,
+		content: [{ type: "text" as const, text: `msg ${id}` }],
+	});
+
+	const editReq = {
+		content: [{ type: "text" as const, text: "edited" }],
+	};
+
+	it("editChatMessage optimistically removes truncated messages from cache", async () => {
+		const queryClient = createTestQueryClient();
+		const chatId = "chat-1";
+		const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id));
+
+		queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
+			pages: [{ messages, queued_messages: [], has_more: false }],
+			pageParams: [undefined],
+		});
+
+		const mutation = editChatMessage(queryClient, chatId);
+		const context = await mutation.onMutate({
+			messageId: 3,
+			req: editReq,
+		});
+
+		const data = queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId));
+		expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([2, 1]);
+		expect(context?.previousData?.pages[0]?.messages).toHaveLength(5);
+	});
+
+	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));
+
+		queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
+			pages: [{ messages, queued_messages: [], has_more: false }],
+			pageParams: [undefined],
+		});
+
+		const mutation = editChatMessage(queryClient, chatId);
+		const context = await mutation.onMutate({
+			messageId: 3,
+			req: editReq,
+		});
+
+		expect(
+			queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId))?.pages[0]
+				?.messages,
+		).toHaveLength(2);
+
+		mutation.onError(
+			new Error("network failure"),
+			{ messageId: 3, 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]);
+	});
+
+	it("editChatMessage onMutate is a no-op when cache is empty", async () => {
+		const queryClient = createTestQueryClient();
+		const chatId = "chat-1";
+
+		const mutation = editChatMessage(queryClient, chatId);
+		const context = await mutation.onMutate({
+			messageId: 3,
+			req: editReq,
+		});
+
+		expect(context.previousData).toBeUndefined();
+		expect(queryClient.getQueryData(chatMessagesKey(chatId))).toBeUndefined();
+	});
+
+	it("editChatMessage onError handles undefined context gracefully", async () => {
+		const queryClient = createTestQueryClient();
+		const chatId = "chat-1";
+		const messages = [3, 2, 1].map((id) => makeMsg(chatId, id));
+
+		queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
+			pages: [{ messages, queued_messages: [], has_more: false }],
+			pageParams: [undefined],
+		});
+
+		const mutation = editChatMessage(queryClient, chatId);
+
+		// Pass undefined context — simulates onMutate throwing before
+		// it could return a snapshot.
+		mutation.onError(
+			new Error("fail"),
+			{ messageId: 2, req: editReq },
+			undefined,
+		);
+
+		// Cache should be untouched — no crash, no corruption.
+		const data = queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId));
+		expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([3, 2, 1]);
+	});
+
+	it("editChatMessage onMutate filters across multiple 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));
+
+		queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
+			pages: [
+				{ messages: page0, queued_messages: [], has_more: true },
+				{ messages: page1, queued_messages: [], has_more: false },
+			],
+			pageParams: [undefined, 6],
+		});
+
+		const mutation = editChatMessage(queryClient, chatId);
+		await mutation.onMutate({ messageId: 7, 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]);
+	});
+
+	it("editChatMessage onMutate editing the first message empties all pages", async () => {
+		const queryClient = createTestQueryClient();
+		const chatId = "chat-1";
+		const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id));
+
+		queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
+			pages: [{ messages, queued_messages: [], has_more: false }],
+			pageParams: [undefined],
+		});
+
+		const mutation = editChatMessage(queryClient, chatId);
+		await mutation.onMutate({ messageId: 1, 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]?.queued_messages).toEqual([]);
+		expect(data?.pages[0]?.has_more).toBe(false);
+	});
+
+	it("editChatMessage onMutate editing the latest message keeps earlier ones", async () => {
+		const queryClient = createTestQueryClient();
+		const chatId = "chat-1";
+		const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id));
+
+		queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
+			pages: [{ messages, queued_messages: [], has_more: false }],
+			pageParams: [undefined],
+		});
+
+		const mutation = editChatMessage(queryClient, chatId);
+		await mutation.onMutate({ messageId: 5, req: editReq });
+
+		const data = queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId));
+		expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([4, 3, 2, 1]);
+	});
+
 	it("interruptChat does not invalidate unrelated queries", async () => {
 		const queryClient = createTestQueryClient();
 		const chatId = "chat-1";

site/src/api/queries/chats.ts

@@ -1,4 +1,8 @@
-import type { QueryClient, UseInfiniteQueryOptions } from "react-query";
+import type {
+	InfiniteData,
+	QueryClient,
+	UseInfiniteQueryOptions,
+} from "react-query";
 import { API } from "#/api/api";
 import type * as TypesGen from "#/api/typesGenerated";
 
@@ -602,11 +606,65 @@ type EditChatMessageMutationArgs = {
 export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({
 	mutationFn: ({ messageId, req }: EditChatMessageMutationArgs) =>
 		API.experimental.editChatMessage(chatId, messageId, req),
-	onSuccess: () => {
-		// Editing truncates all messages after the edited one on the
-		// server. The WebSocket can insert/update messages but cannot
-		// remove stale ones, so a full messages refetch is required.
-		// Use exact matching to avoid cascading to unrelated queries
+	onMutate: async ({ messageId }: EditChatMessageMutationArgs) => {
+		// Cancel in-flight refetches so they don't overwrite the
+		// optimistic update before the mutation completes.
+		await queryClient.cancelQueries({
+			queryKey: chatMessagesKey(chatId),
+			exact: true,
+		});
+
+		const previousData = queryClient.getQueryData<
+			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),
+				})),
+			};
+		});
+
+		return { previousData };
+	},
+	onError: (
+		_error: unknown,
+		_variables: EditChatMessageMutationArgs,
+		context:
+			| {
+					previousData?:
+						| InfiniteData<TypesGen.ChatMessagesResponse>
+						| undefined;
+			  }
+			| undefined,
+	) => {
+		// Restore the cache on failure so the user sees the
+		// original messages again.
+		if (context?.previousData) {
+			queryClient.setQueryData(chatMessagesKey(chatId), context.previousData);
+		}
+	},
+	onSettled: () => {
+		// Always reconcile with the server regardless of whether
+		// the mutation succeeded or failed. On success this picks
+		// up the replacement message; on failure it confirms the
+		// restore from onError matches the server state. Use exact
+		// matching to avoid cascading to unrelated queries
 		// (diff-status, diff-contents, cost summaries, etc.).
 		void queryClient.invalidateQueries({
 			queryKey: chatKey(chatId),