diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 30c5595cd8..2ef0dc7a1c 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -4,7 +4,7 @@ import { API } from "#/api/api"; import type * as TypesGen from "#/api/typesGenerated"; import { archiveChat, - cancelChatListQueries, + cancelChatListRefetches, chatCostSummary, chatCostSummaryKey, chatCostUsers, @@ -1080,7 +1080,7 @@ describe("sidebar title race condition", () => { expect(readTitle(queryClient, chatId)).toBe("fallback title"); }); - it("cancelChatListQueries before the update prevents the overwrite (the fix)", async () => { + it("cancelChatListRefetches before the update prevents the overwrite (the fix)", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; @@ -1104,7 +1104,7 @@ describe("sidebar title race condition", () => { }); // Cancel, then write. Matches the new WebSocket handler code. - await cancelChatListQueries(queryClient); + await cancelChatListRefetches(queryClient); updateInfiniteChatsCache(queryClient, (chats) => chats.map((c) => @@ -1118,3 +1118,195 @@ describe("sidebar title race condition", () => { expect(readTitle(queryClient, chatId)).toBe("generated title"); }); }); + +describe("cancelChatListRefetches", () => { + it("cancels a regular refetch", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + + seedInfiniteChats(queryClient, [makeChat(chatId, { title: "original" })]); + + // Start an in-flight refetch (no fetchMeta — simulates a + // regular invalidation or window-focus refetch). + const fetchDone = queryClient.prefetchQuery({ + queryKey: infiniteChatsTestKey, + queryFn: () => + new Promise((resolve) => { + setTimeout( + () => + resolve({ + pages: [[makeChat(chatId, { title: "stale" })]], + pageParams: [0], + }), + 50, + ); + }), + }); + + await cancelChatListRefetches(queryClient); + await fetchDone; + + // The refetch was cancelled and reverted, so the original + // data is preserved. + const title = readInfiniteChats(queryClient)?.find( + (c) => c.id === chatId, + )?.title; + expect(title).toBe("original"); + }); + + it("does not cancel a fetchNextPage fetch", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + + seedInfiniteChats(queryClient, [makeChat(chatId, { title: "original" })]); + + // Start an in-flight fetch. + const fetchDone = queryClient.prefetchQuery({ + queryKey: infiniteChatsTestKey, + queryFn: () => + new Promise((resolve) => { + setTimeout( + () => + resolve({ + pages: [[makeChat(chatId, { title: "page-2-data" })]], + pageParams: [0], + }), + 50, + ); + }), + }); + + // Simulate fetchNextPage via the public setState API. + // In react-query v5, fetchNextPage dispatches a fetch + // action with meta: { fetchMore: { direction: "forward" } } + // which is stored in query.state.fetchMeta. + const query = queryClient + .getQueryCache() + .find({ queryKey: infiniteChatsTestKey }); + expect(query).toBeDefined(); + query!.setState({ fetchMeta: { fetchMore: { direction: "forward" } } }); + + await cancelChatListRefetches(queryClient); + await fetchDone; + + // The fetch was NOT cancelled — the new data landed. + const title = readInfiniteChats(queryClient)?.find( + (c) => c.id === chatId, + )?.title; + expect(title).toBe("page-2-data"); + }); + + it("does not cancel a fetchPreviousPage fetch", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + + seedInfiniteChats(queryClient, [makeChat(chatId, { title: "original" })]); + + const fetchDone = queryClient.prefetchQuery({ + queryKey: infiniteChatsTestKey, + queryFn: () => + new Promise((resolve) => { + setTimeout( + () => + resolve({ + pages: [[makeChat(chatId, { title: "prev-page" })]], + pageParams: [0], + }), + 50, + ); + }), + }); + + const query = queryClient + .getQueryCache() + .find({ queryKey: infiniteChatsTestKey }); + expect(query).toBeDefined(); + query!.setState({ fetchMeta: { fetchMore: { direction: "backward" } } }); + + await cancelChatListRefetches(queryClient); + await fetchDone; + + const title = readInfiniteChats(queryClient)?.find( + (c) => c.id === chatId, + )?.title; + expect(title).toBe("prev-page"); + }); + + it("does not cancel the initial load when no data is cached yet", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + + // Do NOT seed the cache — simulate the very first fetch + // where no data exists yet. + const fetchDone = queryClient.prefetchQuery({ + queryKey: infiniteChatsTestKey, + queryFn: () => + new Promise((resolve) => { + setTimeout( + () => + resolve({ + pages: [[makeChat(chatId, { title: "first-load" })]], + pageParams: [0], + }), + 50, + ); + }), + }); + + // A WebSocket event arrives while the initial fetch is + // in-flight. Without the data guard, this would cancel + // the fetch and leave the query stuck in pending/idle. + await cancelChatListRefetches(queryClient); + await fetchDone; + + const title = readInfiniteChats(queryClient)?.find( + (c) => c.id === chatId, + )?.title; + expect(title).toBe("first-load"); + }); +}); + +describe("mutation onMutate cancels pagination fetches", () => { + it("archiveChat onMutate cancels a pagination fetch to protect optimistic updates", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + + seedInfiniteChats(queryClient, [makeChat(chatId, { archived: false })]); + + // Start a fetch and mark it as a fetchNextPage via + // fetchMeta so we can verify the broad predicate in + // mutation onMutate still cancels it (unlike the + // narrow cancelChatListRefetches used by the WS + // handler). + const fetchDone = queryClient.prefetchQuery({ + queryKey: infiniteChatsTestKey, + queryFn: () => + new Promise((resolve) => { + setTimeout( + () => + resolve({ + pages: [[makeChat(chatId, { archived: false })]], + pageParams: [0], + }), + 50, + ); + }), + }); + + const query = queryClient + .getQueryCache() + .find({ queryKey: infiniteChatsTestKey }); + expect(query).toBeDefined(); + query!.setState({ fetchMeta: { fetchMore: { direction: "forward" } } }); + + const mutation = archiveChat(queryClient); + await mutation.onMutate(chatId); + await fetchDone; + + // The optimistic archive survives because onMutate + // cancelled the pagination fetch before it could + // overwrite the cache with stale oldPages. + const chat = readInfiniteChats(queryClient)?.find((c) => c.id === chatId); + expect(chat?.archived).toBe(true); + }); +}); diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 366255e4c5..56f98d606e 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -145,16 +145,55 @@ 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. + * Predicate that matches chat-list queries performing a regular + * refetch (window-focus, invalidation, mount) but not a + * fetchNextPage or fetchPreviousPage. During pagination fetches + * react-query sets fetchMeta.fetchMore.direction to "forward" + * or "backward"; regular refetches leave fetchMeta null. + * + * Also excludes queries that have never loaded data. Cancelling + * a first-ever fetch with revert:true leaves the query stuck in + * { status: 'pending', fetchStatus: 'idle', data: undefined } + * with no automatic recovery, so the sidebar shows skeletons + * forever until the user refocuses the window. */ -export const cancelChatListQueries = (queryClient: QueryClient) => { +const isChatListRefetch = (query: { + queryKey: readonly unknown[]; + state: { data: unknown; fetchMeta: unknown }; +}): boolean => { + if (!isChatListQuery(query)) return false; + // Never cancel the initial load. Reverting a first-ever + // fetch produces a stuck pending/idle state that react-query + // does not automatically recover from. + if (query.state.data === undefined) return false; + const meta = query.state.fetchMeta as { + fetchMore?: { direction?: string }; + } | null; + if (meta?.fetchMore?.direction) return false; + return true; +}; + +/** + * Cancel in-flight background refetches for sidebar chat-list + * queries, but leave fetchNextPage / fetchPreviousPage fetches + * alone. Call this before writing WebSocket-driven cache + * updates so a concurrent refetch cannot overwrite the update + * with stale server data. + * + * Pagination fetches are intentionally excluded because + * cancelling them would prevent the sidebar from loading + * additional pages when WebSocket events arrive frequently. + * + * Mutation onMutate handlers should keep the broad + * isChatListQuery predicate instead: mutations are infrequent + * and must cancel pagination fetches to protect optimistic + * updates from being overwritten by the oldPages snapshot + * that fetchNextPage captured before the mutation. + */ +export const cancelChatListRefetches = (queryClient: QueryClient) => { return queryClient.cancelQueries({ queryKey: chatsKey, - predicate: isChatListQuery, + predicate: isChatListRefetch, }); }; diff --git a/site/src/pages/AgentsPage/AgentsPage.tsx b/site/src/pages/AgentsPage/AgentsPage.tsx index 8214b5de33..8df092d318 100644 --- a/site/src/pages/AgentsPage/AgentsPage.tsx +++ b/site/src/pages/AgentsPage/AgentsPage.tsx @@ -11,7 +11,7 @@ import { API, watchChats } from "#/api/api"; import { getErrorMessage } from "#/api/errors"; import { archiveChat, - cancelChatListQueries, + cancelChatListRefetches, chatDiffContentsKey, chatKey, chatModelConfigs, @@ -543,7 +543,7 @@ const AgentsPage: FC = () => { // the refetch may have been issued before the async // title generation finished, so its response carries // the fallback title. - void cancelChatListQueries(queryClient); + void cancelChatListRefetches(queryClient); // Only cancel a per-chat refetch when the cache // already has data. Cancelling a first-time fetch // reverts the query to pending/idle with no data