From 7eca33c69b84017023edee71433b9fc04da370ad Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Wed, 25 Mar 2026 16:18:32 +0200 Subject: [PATCH] fix(site): cancel stale refetches before WebSocket cache writes (#23582) When a chat is created, createChat.onSuccess invalidates the sidebar list query, triggering a background refetch. The refetch can hit the server before async title generation finishes, returning the fallback (truncated) title. If the title_change WebSocket event arrives and writes the generated title into the cache, the in-flight refetch response then overwrites it with the stale fallback title. Cancel any in-flight sidebar-list and per-chat refetches before every WebSocket-driven cache write. This mirrors the existing pattern in archiveChat/unarchiveChat, which cancel queries before optimistic updates for the same reason. --- site/src/api/queries/chats.test.ts | 92 ++++++++++++++++++++++++ site/src/api/queries/chats.ts | 14 ++++ site/src/pages/AgentsPage/AgentsPage.tsx | 15 ++++ 3 files changed, 121 insertions(+) diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index abd0cab93b..8ca63da900 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -4,6 +4,7 @@ import { QueryClient } from "react-query"; import { describe, expect, it, vi } from "vitest"; import { archiveChat, + cancelChatListQueries, chatCostSummary, chatCostSummaryKey, chatCostUsers, @@ -21,6 +22,7 @@ import { invalidateChatListQueries, promoteChatQueuedMessage, unarchiveChat, + updateInfiniteChatsCache, } from "./chats"; vi.mock("api/api", () => ({ @@ -833,3 +835,93 @@ describe("diff_status_change invalidation scope", () => { ).toBe(true); }); }); + +describe("sidebar title race condition", () => { + const readTitle = ( + queryClient: QueryClient, + chatId: string, + ): string | undefined => { + const data = queryClient.getQueryData(infiniteChatsTestKey); + return data?.pages.flat().find((c) => c.id === chatId)?.title; + }; + + it("in-flight refetch overwrites a WebSocket title update (the bug)", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + + seedInfiniteChats(queryClient, [ + makeChat(chatId, { title: "fallback title" }), + ]); + + // Simulate invalidateChatListQueries triggering a refetch that + // returns stale data (the server hadn't generated the title yet + // when it processed this request). + const fetchDone = queryClient.prefetchQuery({ + queryKey: infiniteChatsTestKey, + queryFn: () => + new Promise((resolve) => { + setTimeout( + () => + resolve({ + pages: [[makeChat(chatId, { title: "fallback title" })]], + pageParams: [0], + }), + 50, + ); + }), + }); + + // Simulate the title_change WebSocket event arriving while the + // refetch is in flight. This mirrors what AgentsPage does. + updateInfiniteChatsCache(queryClient, (chats) => + chats.map((c) => + c.id === chatId ? { ...c, title: "generated title" } : c, + ), + ); + + // The cache shows the generated title immediately. + expect(readTitle(queryClient, chatId)).toBe("generated title"); + + // After the refetch settles, it overwrites with stale data. + await fetchDone; + expect(readTitle(queryClient, chatId)).toBe("fallback title"); + }); + + it("cancelChatListQueries before the update prevents the overwrite (the fix)", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + + seedInfiniteChats(queryClient, [ + makeChat(chatId, { title: "fallback title" }), + ]); + + const fetchDone = queryClient.prefetchQuery({ + queryKey: infiniteChatsTestKey, + queryFn: () => + new Promise((resolve) => { + setTimeout( + () => + resolve({ + pages: [[makeChat(chatId, { title: "fallback title" })]], + pageParams: [0], + }), + 50, + ); + }), + }); + + // Cancel, then write. Matches the new WebSocket handler code. + await cancelChatListQueries(queryClient); + + updateInfiniteChatsCache(queryClient, (chats) => + chats.map((c) => + c.id === chatId ? { ...c, title: "generated title" } : c, + ), + ); + + expect(readTitle(queryClient, chatId)).toBe("generated title"); + + await fetchDone; + expect(readTitle(queryClient, chatId)).toBe("generated title"); + }); +}); diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 6b6e2ae591..991d75e5aa 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -104,6 +104,20 @@ export const invalidateChatListQueries = (queryClient: QueryClient) => { }); }; +/** + * Cancel in-flight refetches for sidebar chat-list queries. + * Call this before writing WebSocket-driven cache updates so a + * concurrent refetch (e.g. from createChat.onSuccess or the + * watchChats onOpen handler) cannot overwrite the update with + * stale server data that predates async title generation. + */ +export const cancelChatListQueries = (queryClient: QueryClient) => { + return queryClient.cancelQueries({ + queryKey: chatsKey, + predicate: isChatListQuery, + }); +}; + const DEFAULT_CHAT_PAGE_LIMIT = 50; export const infiniteChats = (opts?: { q?: string; archived?: boolean }) => { diff --git a/site/src/pages/AgentsPage/AgentsPage.tsx b/site/src/pages/AgentsPage/AgentsPage.tsx index fd7afb1d0e..f33768eb31 100644 --- a/site/src/pages/AgentsPage/AgentsPage.tsx +++ b/site/src/pages/AgentsPage/AgentsPage.tsx @@ -2,6 +2,7 @@ import { API, watchChats } from "api/api"; import { getErrorMessage } from "api/errors"; import { archiveChat, + cancelChatListQueries, chatDiffContentsKey, chatKey, chatModelConfigs, @@ -411,6 +412,20 @@ const AgentsPage: FC = () => { const isStatusEvent = chatEvent.kind === "status_change"; const isDiffStatusEvent = chatEvent.kind === "diff_status_change"; + // Cancel in-flight list and per-chat refetches so + // they cannot overwrite the cache update below with + // stale server data. This matters when a title_change + // event races with a refetch triggered by + // createChat.onSuccess or the onOpen invalidation: + // the refetch may have been issued before the async + // title generation finished, so its response carries + // the fallback title. + void cancelChatListQueries(queryClient); + void queryClient.cancelQueries({ + queryKey: chatKey(updatedChat.id), + exact: true, + }); + // For "created" events, use a cross-page existence // check and prepend only to the first page. // updateInfiniteChatsCache runs the updater per