diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 5e6b173381..8959617eb7 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -11,6 +11,34 @@ vi.mock("api/api", () => ({ }, })); +// The infinite query key used by useInfiniteQuery(infiniteChats()) +// is [...chatsKey, undefined] = ["chats", undefined]. +const infiniteChatsTestKey = [...chatsKey, undefined]; + +type InfiniteData = { + pages: TypesGen.Chat[][]; + pageParams: unknown[]; +}; + +/** Seed the infinite chats cache in the format TanStack Query expects. */ +const seedInfiniteChats = ( + queryClient: QueryClient, + chats: TypesGen.Chat[], +) => { + queryClient.setQueryData(infiniteChatsTestKey, { + pages: [chats], + pageParams: [0], + }); +}; + +/** Read chats back from the infinite query cache. */ +const readInfiniteChats = ( + queryClient: QueryClient, +): TypesGen.Chat[] | undefined => { + const data = queryClient.getQueryData(infiniteChatsTestKey); + return data?.pages.flat(); +}; + const makeChat = ( id: string, overrides?: Partial, @@ -53,14 +81,14 @@ describe("archiveChat optimistic update", () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; const initialChats = [makeChat(chatId), makeChat("chat-2")]; - queryClient.setQueryData(chatsKey, initialChats); + seedInfiniteChats(queryClient, initialChats); vi.mocked(API.archiveChat).mockResolvedValue(); const mutation = archiveChat(queryClient); await mutation.onMutate(chatId); - const updatedChats = queryClient.getQueryData(chatsKey); + const updatedChats = readInfiniteChats(queryClient); expect(updatedChats).toHaveLength(2); expect(updatedChats?.find((c) => c.id === chatId)?.archived).toBe(true); // Other chats are unchanged. @@ -70,7 +98,7 @@ describe("archiveChat optimistic update", () => { it("optimistically sets archived to true in the individual chat cache", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; - queryClient.setQueryData(chatsKey, [makeChat(chatId)]); + seedInfiniteChats(queryClient, [makeChat(chatId)]); queryClient.setQueryData(chatKey(chatId), makeChatWithMessages(chatId)); vi.mocked(API.archiveChat).mockResolvedValue(); @@ -84,33 +112,33 @@ describe("archiveChat optimistic update", () => { expect(cachedChat?.chat.archived).toBe(true); }); - it("rolls back the chats list on error", async () => { + it("rolls back the chats list on error by invalidating", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; const initialChats = [makeChat(chatId)]; - queryClient.setQueryData(chatsKey, initialChats); + seedInfiniteChats(queryClient, initialChats); queryClient.setQueryData(chatKey(chatId), makeChatWithMessages(chatId)); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); const mutation = archiveChat(queryClient); const context = await mutation.onMutate(chatId); // Verify the optimistic update took effect. - expect( - queryClient.getQueryData(chatsKey)?.[0].archived, - ).toBe(true); + expect(readInfiniteChats(queryClient)?.[0].archived).toBe(true); - // Simulate an error — the onError handler should restore original - // data. + // Simulate an error — the onError handler invalidates the + // cache so a re-fetch restores the correct state. mutation.onError(new Error("server error"), chatId, context); - const rolledBack = queryClient.getQueryData(chatsKey); - expect(rolledBack?.[0].archived).toBe(false); + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: chatsKey, + }); }); it("rolls back the individual chat cache on error", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; - queryClient.setQueryData(chatsKey, [makeChat(chatId)]); + seedInfiniteChats(queryClient, [makeChat(chatId)]); queryClient.setQueryData(chatKey(chatId), makeChatWithMessages(chatId)); const mutation = archiveChat(queryClient); @@ -132,7 +160,8 @@ describe("archiveChat optimistic update", () => { it("handles error rollback gracefully when context is undefined", () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; - queryClient.setQueryData(chatsKey, [makeChat(chatId, { archived: true })]); + seedInfiniteChats(queryClient, [makeChat(chatId, { archived: true })]); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); const mutation = archiveChat(queryClient); @@ -141,26 +170,23 @@ describe("archiveChat optimistic update", () => { mutation.onError(new Error("fail"), chatId, undefined); }).not.toThrow(); - // Data should remain unchanged since there was nothing to roll - // back to. - expect( - queryClient.getQueryData(chatsKey)?.[0].archived, - ).toBe(true); + // The handler should still invalidate to trigger a refetch. + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: chatsKey, + }); }); it("handles onMutate when no individual chat cache exists", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; - queryClient.setQueryData(chatsKey, [makeChat(chatId)]); + seedInfiniteChats(queryClient, [makeChat(chatId)]); // Deliberately do NOT set chatKey(chatId) data. const mutation = archiveChat(queryClient); const context = await mutation.onMutate(chatId); // The list should still be optimistically updated. - expect( - queryClient.getQueryData(chatsKey)?.[0].archived, - ).toBe(true); + expect(readInfiniteChats(queryClient)?.[0].archived).toBe(true); // previousChat should be undefined. expect(context?.previousChat).toBeUndefined(); }); @@ -186,20 +212,18 @@ describe("unarchiveChat optimistic update", () => { it("optimistically sets archived to false in the chats list", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; - queryClient.setQueryData(chatsKey, [makeChat(chatId, { archived: true })]); + seedInfiniteChats(queryClient, [makeChat(chatId, { archived: true })]); const mutation = unarchiveChat(queryClient); await mutation.onMutate(chatId); - expect( - queryClient.getQueryData(chatsKey)?.[0].archived, - ).toBe(false); + expect(readInfiniteChats(queryClient)?.[0].archived).toBe(false); }); it("optimistically sets archived to false in the individual chat cache", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; - queryClient.setQueryData(chatsKey, [makeChat(chatId, { archived: true })]); + seedInfiniteChats(queryClient, [makeChat(chatId, { archived: true })]); queryClient.setQueryData( chatKey(chatId), makeChatWithMessages(chatId, { archived: true }), @@ -217,19 +241,18 @@ describe("unarchiveChat optimistic update", () => { it("rolls back both caches on error", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; - queryClient.setQueryData(chatsKey, [makeChat(chatId, { archived: true })]); + seedInfiniteChats(queryClient, [makeChat(chatId, { archived: true })]); queryClient.setQueryData( chatKey(chatId), makeChatWithMessages(chatId, { archived: true }), ); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); const mutation = unarchiveChat(queryClient); const context = await mutation.onMutate(chatId); // Verify optimistic update. - expect( - queryClient.getQueryData(chatsKey)?.[0].archived, - ).toBe(false); + expect(readInfiniteChats(queryClient)?.[0].archived).toBe(false); expect( queryClient.getQueryData(chatKey(chatId))?.chat .archived, @@ -238,9 +261,11 @@ describe("unarchiveChat optimistic update", () => { // Roll back. mutation.onError(new Error("server error"), chatId, context); - expect( - queryClient.getQueryData(chatsKey)?.[0].archived, - ).toBe(true); + // The chats list is rolled back via invalidation. + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: chatsKey, + }); + // The individual chat cache is restored directly. expect( queryClient.getQueryData(chatKey(chatId))?.chat .archived, diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 5266375d2a..a68b087b3c 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -5,6 +5,48 @@ import type { QueryClient, UseInfiniteQueryOptions } from "react-query"; export const chatsKey = ["chats"] as const; export const chatKey = (chatId: string) => ["chats", chatId] as const; +/** + * Updates a single chat inside every page of the infinite chats query + * cache. Use this instead of setQueryData(chatsKey, ...) which writes + * to the wrong key (the flat list key, not the infinite query key). + */ +export const updateInfiniteChatsCache = ( + queryClient: QueryClient, + updater: (chats: TypesGen.Chat[]) => TypesGen.Chat[], +) => { + // Update ALL infinite chat queries regardless of their filter opts. + queryClient.setQueriesData<{ + pages: TypesGen.Chat[][]; + pageParams: unknown[]; + }>({ queryKey: chatsKey }, (prev) => { + if (!prev) return prev; + if (!prev.pages) return prev; + const nextPages = prev.pages.map((page) => updater(page)); + // Only return a new reference if something actually changed. + const changed = nextPages.some((page, i) => page !== prev.pages[i]); + return changed ? { ...prev, pages: nextPages } : prev; + }); +}; + +/** + * Reads the flat list of chats from the first matching infinite query + * in the cache. Returns undefined when no data is cached yet. + */ +export const readInfiniteChatsCache = ( + queryClient: QueryClient, +): TypesGen.Chat[] | undefined => { + const queries = queryClient.getQueriesData<{ + pages: TypesGen.Chat[][]; + pageParams: unknown[]; + }>({ queryKey: chatsKey }); + for (const [, data] of queries) { + if (data?.pages) { + return data.pages.flat(); + } + } + return undefined; +}; + const DEFAULT_CHAT_PAGE_LIMIT = 50; export const infiniteChats = (opts?: { archived?: boolean }) => { @@ -49,12 +91,11 @@ export const archiveChat = (queryClient: QueryClient) => ({ onMutate: async (chatId: string) => { await queryClient.cancelQueries({ queryKey: chatsKey }); await queryClient.cancelQueries({ queryKey: chatKey(chatId) }); - const previousChats = queryClient.getQueryData(chatsKey); const previousChat = queryClient.getQueryData( chatKey(chatId), ); - queryClient.setQueryData(chatsKey, (old) => - old?.map((chat) => + updateInfiniteChatsCache(queryClient, (chats) => + chats.map((chat) => chat.id === chatId ? { ...chat, archived: true } : chat, ), ); @@ -64,24 +105,19 @@ export const archiveChat = (queryClient: QueryClient) => ({ chat: { ...previousChat.chat, archived: true }, }); } - return { previousChats, previousChat }; + return { previousChat }; }, onError: ( _error: unknown, chatId: string, context: | { - previousChats?: TypesGen.Chat[]; previousChat?: TypesGen.ChatWithMessages; } | undefined, ) => { - if (context?.previousChats) { - queryClient.setQueryData( - chatsKey, - context.previousChats, - ); - } + // Rollback: invalidate to re-fetch the correct state. + void queryClient.invalidateQueries({ queryKey: chatsKey }); if (context?.previousChat) { queryClient.setQueryData( chatKey(chatId), @@ -100,12 +136,11 @@ export const unarchiveChat = (queryClient: QueryClient) => ({ onMutate: async (chatId: string) => { await queryClient.cancelQueries({ queryKey: chatsKey }); await queryClient.cancelQueries({ queryKey: chatKey(chatId) }); - const previousChats = queryClient.getQueryData(chatsKey); const previousChat = queryClient.getQueryData( chatKey(chatId), ); - queryClient.setQueryData(chatsKey, (old) => - old?.map((chat) => + updateInfiniteChatsCache(queryClient, (chats) => + chats.map((chat) => chat.id === chatId ? { ...chat, archived: false } : chat, ), ); @@ -115,24 +150,19 @@ export const unarchiveChat = (queryClient: QueryClient) => ({ chat: { ...previousChat.chat, archived: false }, }); } - return { previousChats, previousChat }; + return { previousChat }; }, onError: ( _error: unknown, chatId: string, context: | { - previousChats?: TypesGen.Chat[]; previousChat?: TypesGen.ChatWithMessages; } | undefined, ) => { - if (context?.previousChats) { - queryClient.setQueryData( - chatsKey, - context.previousChats, - ); - } + // Rollback: invalidate to re-fetch the correct state. + void queryClient.invalidateQueries({ queryKey: chatsKey }); if (context?.previousChat) { queryClient.setQueryData( chatKey(chatId), diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx index 80b76bf3e5..b4244bb406 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx @@ -1,6 +1,35 @@ import { act, render, renderHook, waitFor } from "@testing-library/react"; import { watchChat } from "api/api"; import { chatKey, chatsKey } from "api/queries/chats"; + +// The infinite query key used by useInfiniteQuery(infiniteChats()) +// is [...chatsKey, undefined] = ["chats", undefined]. +const infiniteChatsTestKey = [...chatsKey, undefined]; + +type InfiniteData = { + pages: TypesGen.Chat[][]; + pageParams: unknown[]; +}; + +/** Seed the infinite chats cache in the format TanStack Query expects. */ +const seedInfiniteChats = ( + queryClient: QueryClient, + chats: TypesGen.Chat[], +) => { + queryClient.setQueryData(infiniteChatsTestKey, { + pages: [chats], + pageParams: [0], + }); +}; + +/** Read chats back from the infinite query cache. */ +const readInfiniteChats = ( + queryClient: QueryClient, +): TypesGen.Chat[] | undefined => { + const data = queryClient.getQueryData(infiniteChatsTestKey); + return data?.pages.flat(); +}; + import type * as TypesGen from "api/typesGenerated"; import type { FC, PropsWithChildren } from "react"; import { QueryClient, QueryClientProvider } from "react-query"; @@ -2188,7 +2217,7 @@ describe("updateSidebarChat via stream events", () => { }); const initialChat = makeChat(chatID); // Seed the chats list so updateSidebarChat can find it. - queryClient.setQueryData(chatsKey, [initialChat]); + seedInfiniteChats(queryClient, [initialChat]); const wrapper = ({ children }: PropsWithChildren) => ( {children} @@ -2229,7 +2258,7 @@ describe("updateSidebarChat via stream events", () => { }); await waitFor(() => { - const sidebarChats = queryClient.getQueryData(chatsKey); + const sidebarChats = readInfiniteChats(queryClient); expect(sidebarChats?.[0].status).toBe("completed"); }); }); @@ -2252,7 +2281,7 @@ describe("updateSidebarChat via stream events", () => { }, }); const initialChat = makeChat(chatID); - queryClient.setQueryData(chatsKey, [initialChat]); + seedInfiniteChats(queryClient, [initialChat]); const wrapper = ({ children }: PropsWithChildren) => ( {children} @@ -2302,7 +2331,7 @@ describe("updateSidebarChat via stream events", () => { // global chat-list WebSocket delivers the authoritative server // timestamp. Verify it stays at the original value. await waitFor(() => { - const sidebarChats = queryClient.getQueryData(chatsKey); + const sidebarChats = readInfiniteChats(queryClient); expect(sidebarChats?.[0].updated_at).toBe(initialChat.updated_at); }); }); @@ -2325,7 +2354,7 @@ describe("updateSidebarChat via stream events", () => { }, }); const initialChat = makeChat(chatID); - queryClient.setQueryData(chatsKey, [initialChat]); + seedInfiniteChats(queryClient, [initialChat]); const wrapper = ({ children }: PropsWithChildren) => ( {children} @@ -2366,7 +2395,7 @@ describe("updateSidebarChat via stream events", () => { }); await waitFor(() => { - const sidebarChats = queryClient.getQueryData(chatsKey); + const sidebarChats = readInfiniteChats(queryClient); expect(sidebarChats?.[0].status).toBe("error"); }); }); @@ -2391,7 +2420,7 @@ describe("updateSidebarChat via stream events", () => { }); const activeChat = makeChat(chatID); const otherChat = makeChat(otherChatID); - queryClient.setQueryData(chatsKey, [activeChat, otherChat]); + seedInfiniteChats(queryClient, [activeChat, otherChat]); const wrapper = ({ children }: PropsWithChildren) => ( {children} @@ -2432,14 +2461,14 @@ describe("updateSidebarChat via stream events", () => { }); await waitFor(() => { - const sidebarChats = queryClient.getQueryData(chatsKey); + const sidebarChats = readInfiniteChats(queryClient); expect(sidebarChats?.find((c) => c.id === chatID)?.status).toBe( "completed", ); }); // The other chat should remain unchanged. - const sidebarChats = queryClient.getQueryData(chatsKey); + const sidebarChats = readInfiniteChats(queryClient); expect(sidebarChats?.find((c) => c.id === otherChatID)?.status).toBe( "running", ); @@ -2464,7 +2493,7 @@ describe("updateSidebarChat via stream events", () => { }); const futureTimestamp = "2099-01-01T00:00:00.000Z"; const initialChat = { ...makeChat(chatID), updated_at: futureTimestamp }; - queryClient.setQueryData(chatsKey, [initialChat]); + seedInfiniteChats(queryClient, [initialChat]); const wrapper = ({ children }: PropsWithChildren) => ( {children} @@ -2512,7 +2541,7 @@ describe("updateSidebarChat via stream events", () => { }); await waitFor(() => { - const sidebarChats = queryClient.getQueryData(chatsKey); + const sidebarChats = readInfiniteChats(queryClient); expect(sidebarChats?.[0].updated_at).toBe(futureTimestamp); }); }); @@ -2535,7 +2564,7 @@ describe("updateSidebarChat via stream events", () => { }, }); const initialChat = makeChat(chatID); - queryClient.setQueryData(chatsKey, [initialChat]); + seedInfiniteChats(queryClient, [initialChat]); const wrapper = ({ children }: PropsWithChildren) => ( {children} @@ -2576,7 +2605,7 @@ describe("updateSidebarChat via stream events", () => { }); await waitFor(() => { - const sidebarChats = queryClient.getQueryData(chatsKey); + const sidebarChats = readInfiniteChats(queryClient); // Status should update, but updated_at must stay untouched. expect(sidebarChats?.[0].status).toBe("completed"); expect(sidebarChats?.[0].updated_at).toBe(initialChat.updated_at); @@ -2601,7 +2630,7 @@ describe("updateSidebarChat via stream events", () => { }, }); const initialChat = makeChat(chatID); - queryClient.setQueryData(chatsKey, [initialChat]); + seedInfiniteChats(queryClient, [initialChat]); const wrapper = ({ children }: PropsWithChildren) => ( {children} @@ -2642,7 +2671,7 @@ describe("updateSidebarChat via stream events", () => { }); await waitFor(() => { - const sidebarChats = queryClient.getQueryData(chatsKey); + const sidebarChats = readInfiniteChats(queryClient); expect(sidebarChats?.[0].status).toBe("error"); expect(sidebarChats?.[0].updated_at).toBe(initialChat.updated_at); }); diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts index 4df139f967..6c571d9b73 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts @@ -1,5 +1,5 @@ import { watchChat } from "api/api"; -import { chatKey, chatsKey } from "api/queries/chats"; +import { chatKey, updateInfiniteChatsCache } from "api/queries/chats"; import type * as TypesGen from "api/typesGenerated"; import { asRecord, asString } from "components/ai-elements/runtimeTypeUtils"; import { @@ -475,23 +475,17 @@ export const useChatStore = ( if (!chatID) { return; } - queryClient.setQueryData( - chatsKey, - (currentChats) => { - if (!currentChats) { - return currentChats; + updateInfiniteChatsCache(queryClient, (chats) => { + let didUpdate = false; + const nextChats = chats.map((chat) => { + if (chat.id !== chatID) { + return chat; } - let didUpdate = false; - const nextChats = currentChats.map((chat) => { - if (chat.id !== chatID) { - return chat; - } - didUpdate = true; - return updater(chat); - }); - return didUpdate ? nextChats : currentChats; - }, - ); + didUpdate = true; + return updater(chat); + }); + return didUpdate ? nextChats : chats; + }); }, [chatID, queryClient], ); diff --git a/site/src/pages/AgentsPage/AgentsPage.tsx b/site/src/pages/AgentsPage/AgentsPage.tsx index 06af742c6e..a8cdd0e499 100644 --- a/site/src/pages/AgentsPage/AgentsPage.tsx +++ b/site/src/pages/AgentsPage/AgentsPage.tsx @@ -11,8 +11,10 @@ import { chatsKey, createChat, infiniteChats, + readInfiniteChatsCache, unarchiveChat, updateChatSystemPrompt, + updateInfiniteChatsCache, } from "api/queries/chats"; import { workspaces } from "api/queries/workspaces"; import type * as TypesGen from "api/typesGenerated"; @@ -380,12 +382,10 @@ const AgentsPage: FC = () => { // is synchronously updated by both the per-chat WebSocket // (via updateSidebarChat) and this handler. This avoids // the async-lag of a useEffect-based status map. - const currentChats = - queryClient.getQueryData(chatsKey); + const currentChats = readInfiniteChatsCache(queryClient); const prevStatus = currentChats?.find( (c) => c.id === updatedChat.id, - )?.status; - // Only play the chime for top-level chats, not sub-agents. + )?.status; // Only play the chime for top-level chats, not sub-agents. if (!updatedChat.parent_chat_id) { maybePlayChime( prevStatus, @@ -396,14 +396,11 @@ const AgentsPage: FC = () => { } if (chatEvent.kind === "deleted") { - queryClient.setQueryData( - chatsKey, - (prev: TypesGen.Chat[] | undefined) => - prev?.filter( - (c) => - c.id !== updatedChat.id && - c.root_chat_id !== updatedChat.id, - ), + updateInfiniteChatsCache(queryClient, (chats) => + chats.filter( + (c) => + c.id !== updatedChat.id && c.root_chat_id !== updatedChat.id, + ), ); queryClient.removeQueries({ queryKey: chatKey(updatedChat.id), @@ -435,31 +432,27 @@ const AgentsPage: FC = () => { const isTitleEvent = chatEvent.kind === "title_change"; const isStatusEvent = chatEvent.kind === "status_change"; - queryClient.setQueryData( - chatsKey, - (prev: TypesGen.Chat[] | undefined) => { - if (!prev) return prev; - const exists = prev.some((c) => c.id === updatedChat.id); - if (exists) { - return prev.map((c) => { - if (c.id !== updatedChat.id) return c; - return { - ...c, - ...(isStatusEvent && { status: updatedChat.status }), - ...(isTitleEvent && { title: updatedChat.title }), - updated_at: - c.updated_at > updatedChat.updated_at - ? c.updated_at - : updatedChat.updated_at, - }; - }); - } - if (chatEvent.kind === "created") { - return [updatedChat, ...prev]; - } - return prev; - }, - ); + updateInfiniteChatsCache(queryClient, (chats) => { + const exists = chats.some((c) => c.id === updatedChat.id); + if (exists) { + return chats.map((c) => { + if (c.id !== updatedChat.id) return c; + return { + ...c, + ...(isStatusEvent && { status: updatedChat.status }), + ...(isTitleEvent && { title: updatedChat.title }), + updated_at: + c.updated_at > updatedChat.updated_at + ? c.updated_at + : updatedChat.updated_at, + }; + }); + } + if (chatEvent.kind === "created") { + return [updatedChat, ...chats]; + } + return chats; + }); queryClient.setQueryData( chatKey(updatedChat.id), (previousChat) => {