From 1221622bf0d50cccedd24f880a0b0a15cf93e909 Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Thu, 2 Apr 2026 12:41:33 +0100 Subject: [PATCH] fix(site/src/api/queries): optimistically truncate cache on chat message edit (#23864) --- site/src/api/queries/chats.test.ts | 178 ++++++++++++++++++++++++++++- site/src/api/queries/chats.ts | 70 +++++++++++- 2 files changed, 240 insertions(+), 8 deletions(-) diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 2ef0dc7a1c..0620b176a9 100644 --- a/site/src/api/queries/chats.test.ts +++ b/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(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(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(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(chatMessagesKey(chatId))?.pages[0] + ?.messages, + ).toHaveLength(2); + + mutation.onError( + new Error("network failure"), + { messageId: 3, req: editReq }, + context, + ); + + const data = queryClient.getQueryData(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(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(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(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(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(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(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(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(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"; diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 56f98d606e..77cc3d2130 100644 --- a/site/src/api/queries/chats.ts +++ b/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 + >(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 | 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 + | 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),