From f96671d588c7d7b17e43df77398712be7a53f21f Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Thu, 6 Aug 2026 08:29:10 +0100 Subject: [PATCH] refactor(site/src): centralize chat cache operations in semantic cache functions (#27869) --- site/src/api/queries/chats.test.ts | 319 +++++++++++++++ site/src/api/queries/chats.ts | 373 +++++++++--------- site/src/pages/AgentsPage/AgentChatPage.tsx | 19 +- .../pages/AgentsPage/AgentChatPageView.tsx | 7 +- .../src/pages/AgentsPage/AgentsPageLayout.tsx | 50 +-- .../ChatConversation/useChatStore.ts | 20 +- .../useChatToolInvalidations.test.tsx | 39 +- .../useChatToolInvalidations.ts | 6 +- 8 files changed, 567 insertions(+), 266 deletions(-) diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 64ec5979b2..5f4afc546d 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -11,19 +11,25 @@ import { addChildToParentInCache, archiveChat, type ChatListInput, + cancelChatEntity, + cancelChatListQueries, cancelChatListRefetches, + cancelChatMessages, + cancelLoadedChatEntityRefetch, chatACL, chatACLKey, chatAdvisorConfig, chatAdvisorConfigKey, chatCost, chatCostTreeKey, + chatDebugRunKey, chatDebugRunsKey, chatDiffContentsKey, chatEntityKey, chatListFamilyKey, chatListKey, chatMessagesKey, + chatPromptsKey, chatSearch, chatsByWorkspace, createChat, @@ -32,13 +38,24 @@ import { editChatMessage, infiniteChats, interruptChat, + invalidateChatACL, + invalidateChatCostTree, + invalidateChatDebugRuns, + invalidateChatDiffContents, + invalidateChatEntity, invalidateChatListQueries, + invalidateChatMessages, + invalidateChatPrompts, + invalidateChatsByWorkspace, mergeWatchedChatIntoCaches, mergeWatchedChatSummary, + patchChatEntity, + patchChatMessages, pinChat, prependToInfiniteChatsCache, promoteChatQueuedMessage, proposeChatTitle, + removeChatEntity, removeChildFromParentInCache, reorderPinnedChat, setChatGroupRole, @@ -2859,3 +2876,305 @@ describe("chat ACL query factories", () => { ); }); }); + +describe("semantic cache operations: exact invalidations", () => { + it("invalidateChatEntity touches only the detail entry", async () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + queryClient.setQueryData(chatMessagesKey("chat-1"), []); + queryClient.setQueryData(chatEntityKey("chat-2"), makeChat("chat-2")); + + await invalidateChatEntity(queryClient, "chat-1"); + + expect( + queryClient.getQueryState(chatEntityKey("chat-1"))?.isInvalidated, + ).toBe(true); + expect( + queryClient.getQueryState(chatMessagesKey("chat-1"))?.isInvalidated, + "messages entry should NOT be invalidated", + ).not.toBe(true); + expect( + queryClient.getQueryState(chatEntityKey("chat-2"))?.isInvalidated, + "other chat's detail entry should NOT be invalidated", + ).not.toBe(true); + }); + + it("invalidateChatDiffContents touches only the diff-contents entry", async () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatDiffContentsKey("chat-1"), { files: [] }); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + + await invalidateChatDiffContents(queryClient, "chat-1"); + + expect( + queryClient.getQueryState(chatDiffContentsKey("chat-1"))?.isInvalidated, + ).toBe(true); + expect( + queryClient.getQueryState(chatEntityKey("chat-1"))?.isInvalidated, + "detail entry should NOT be invalidated", + ).not.toBe(true); + }); + + it("invalidateChatPrompts touches only the prompts entry", async () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatPromptsKey("chat-1"), { prompts: [] }); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + + await invalidateChatPrompts(queryClient, "chat-1"); + + expect( + queryClient.getQueryState(chatPromptsKey("chat-1"))?.isInvalidated, + ).toBe(true); + expect( + queryClient.getQueryState(chatEntityKey("chat-1"))?.isInvalidated, + "detail entry should NOT be invalidated", + ).not.toBe(true); + }); + + it("invalidateChatMessages touches only the messages entry", async () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatMessagesKey("chat-1"), []); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + + await invalidateChatMessages(queryClient, "chat-1"); + + expect( + queryClient.getQueryState(chatMessagesKey("chat-1"))?.isInvalidated, + ).toBe(true); + expect( + queryClient.getQueryState(chatEntityKey("chat-1"))?.isInvalidated, + "detail entry should NOT be invalidated", + ).not.toBe(true); + }); + + it("invalidateChatACL touches only the ACL entry", async () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatACLKey("chat-1"), {}); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + + await invalidateChatACL(queryClient, "chat-1"); + + expect(queryClient.getQueryState(chatACLKey("chat-1"))?.isInvalidated).toBe( + true, + ); + expect( + queryClient.getQueryState(chatEntityKey("chat-1"))?.isInvalidated, + "detail entry should NOT be invalidated", + ).not.toBe(true); + }); + + it("invalidateChatCostTree touches only the matching cost tree entry", async () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatCostTreeKey("chat-1"), {}); + queryClient.setQueryData(chatCostTreeKey("chat-2"), {}); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + + await invalidateChatCostTree(queryClient, "chat-1"); + + expect( + queryClient.getQueryState(chatCostTreeKey("chat-1"))?.isInvalidated, + ).toBe(true); + expect( + queryClient.getQueryState(chatCostTreeKey("chat-2"))?.isInvalidated, + ).not.toBe(true); + expect( + queryClient.getQueryState(chatEntityKey("chat-1"))?.isInvalidated, + "detail entry should NOT be invalidated", + ).not.toBe(true); + }); +}); + +describe("semantic cache operations: prefix invalidations", () => { + it("invalidateChatListQueries touches every list entry and nothing outside the family", async () => { + const queryClient = createTestQueryClient(); + seedInfiniteChats(queryClient, [makeChat("chat-1")]); + seedInfiniteChats(queryClient, [makeChat("chat-1")], { archived: true }); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + queryClient.setQueryData(chatsByWorkspace(["ws-1"]).queryKey, {}); + + await invalidateChatListQueries(queryClient); + + expect(queryClient.getQueryState(infiniteChatsTestKey)?.isInvalidated).toBe( + true, + ); + expect( + queryClient.getQueryState( + chatListKey(toChatListParams({ archived: true })), + )?.isInvalidated, + ).toBe(true); + expect( + queryClient.getQueryState(chatEntityKey("chat-1"))?.isInvalidated, + ).not.toBe(true); + expect( + queryClient.getQueryState(chatsByWorkspace(["ws-1"]).queryKey) + ?.isInvalidated, + ).not.toBe(true); + }); + + it("invalidateChatsByWorkspace touches by-workspace entries only", async () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatsByWorkspace(["ws-1"]).queryKey, {}); + seedInfiniteChats(queryClient, [makeChat("chat-1")]); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + + await invalidateChatsByWorkspace(queryClient); + + expect( + queryClient.getQueryState(chatsByWorkspace(["ws-1"]).queryKey) + ?.isInvalidated, + ).toBe(true); + expect( + queryClient.getQueryState(infiniteChatsTestKey)?.isInvalidated, + ).not.toBe(true); + expect( + queryClient.getQueryState(chatEntityKey("chat-1"))?.isInvalidated, + ).not.toBe(true); + }); + + it("invalidateChatDebugRuns touches the runs list and run details only", async () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatDebugRunsKey("chat-1"), []); + queryClient.setQueryData(chatDebugRunKey("chat-1", "run-1"), {}); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + queryClient.setQueryData(chatMessagesKey("chat-1"), []); + + await invalidateChatDebugRuns(queryClient, "chat-1"); + + expect( + queryClient.getQueryState(chatDebugRunsKey("chat-1"))?.isInvalidated, + ).toBe(true); + expect( + queryClient.getQueryState(chatDebugRunKey("chat-1", "run-1")) + ?.isInvalidated, + "run detail entry should be invalidated by the family prefix", + ).toBe(true); + expect( + queryClient.getQueryState(chatEntityKey("chat-1"))?.isInvalidated, + "detail entry should NOT be invalidated", + ).not.toBe(true); + expect( + queryClient.getQueryState(chatMessagesKey("chat-1"))?.isInvalidated, + "messages entry should NOT be invalidated", + ).not.toBe(true); + }); +}); + +describe("semantic cache operations: cancellation", () => { + it("cancelChatListQueries cancels unconditionally across the list family", async () => { + const queryClient = createTestQueryClient(); + const cancelSpy = vi.spyOn(queryClient, "cancelQueries"); + + await cancelChatListQueries(queryClient); + + expect(cancelSpy).toHaveBeenCalledWith({ + queryKey: chatListFamilyKey, + }); + }); + + it("cancelChatEntity cancels the exact detail entry unconditionally", async () => { + const queryClient = createTestQueryClient(); + const cancelSpy = vi.spyOn(queryClient, "cancelQueries"); + + await cancelChatEntity(queryClient, "chat-1"); + + expect(cancelSpy).toHaveBeenCalledWith({ + queryKey: chatEntityKey("chat-1"), + exact: true, + }); + }); + + it("cancelLoadedChatEntityRefetch is a no-op when detail data is absent", async () => { + const queryClient = createTestQueryClient(); + const cancelSpy = vi.spyOn(queryClient, "cancelQueries"); + + await cancelLoadedChatEntityRefetch(queryClient, "chat-1"); + + expect(cancelSpy).not.toHaveBeenCalled(); + }); + + it("cancelLoadedChatEntityRefetch cancels exactly when detail data exists", async () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + const cancelSpy = vi.spyOn(queryClient, "cancelQueries"); + + await cancelLoadedChatEntityRefetch(queryClient, "chat-1"); + + expect(cancelSpy).toHaveBeenCalledWith({ + queryKey: chatEntityKey("chat-1"), + exact: true, + }); + }); + + it("cancelChatMessages cancels the exact messages entry", async () => { + const queryClient = createTestQueryClient(); + const cancelSpy = vi.spyOn(queryClient, "cancelQueries"); + + await cancelChatMessages(queryClient, "chat-1"); + + expect(cancelSpy).toHaveBeenCalledWith({ + queryKey: chatMessagesKey("chat-1"), + exact: true, + }); + }); +}); + +describe("semantic cache operations: removal and patching", () => { + it("removeChatEntity removes only the exact detail entry", () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + queryClient.setQueryData(chatMessagesKey("chat-1"), []); + queryClient.setQueryData(chatPromptsKey("chat-1"), { prompts: [] }); + queryClient.setQueryData(chatACLKey("chat-1"), {}); + queryClient.setQueryData(chatDiffContentsKey("chat-1"), { files: [] }); + queryClient.setQueryData(chatDebugRunsKey("chat-1"), []); + queryClient.setQueryData(chatDebugRunKey("chat-1", "run-1"), {}); + seedInfiniteChats(queryClient, [makeChat("chat-1")]); + + removeChatEntity(queryClient, "chat-1"); + + expect(queryClient.getQueryData(chatEntityKey("chat-1"))).toBeUndefined(); + for (const [label, key] of [ + ["messages", chatMessagesKey("chat-1")], + ["prompts", chatPromptsKey("chat-1")], + ["acl", chatACLKey("chat-1")], + ["diff-contents", chatDiffContentsKey("chat-1")], + ["debug-runs", chatDebugRunsKey("chat-1")], + ["debug-run detail", chatDebugRunKey("chat-1", "run-1")], + ] as const) { + expect( + queryClient.getQueryData(key), + `${label} entry should survive removeChatEntity`, + ).toBeDefined(); + } + expect( + queryClient.getQueryData(infiniteChatsTestKey), + "list entry should survive removeChatEntity", + ).toBeDefined(); + }); + + it("patchChatEntity applies the updater to the exact detail entry", () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + + patchChatEntity(queryClient, "chat-1", (chat) => + chat ? { ...chat, title: "Patched" } : chat, + ); + + expect( + queryClient.getQueryData(chatEntityKey("chat-1"))?.title, + ).toBe("Patched"); + }); + + it("patchChatMessages preserves the previous reference when the updater is a no-op", () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatMessagesKey("chat-1"), { + pages: [{ messages: [], queued_messages: [], has_more: false }], + pageParams: [undefined], + }); + const before = queryClient.getQueryData(chatMessagesKey("chat-1")); + + patchChatMessages(queryClient, "chat-1", (data) => data); + + expect(queryClient.getQueryData(chatMessagesKey("chat-1"))).toBe(before); + }); +}); diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 34d2e7bf0b..49a4896c6a 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -22,7 +22,7 @@ export const chatListFamilyKey = [...chatCollectionsKey, "list"] as const; const chatSearchFamilyKey = [...chatCollectionsKey, "search"] as const; -export const chatsByWorkspaceFamilyKey = [ +const chatsByWorkspaceFamilyKey = [ ...chatCollectionsKey, "by-workspace", ] as const; @@ -601,12 +601,6 @@ const getNextOptimisticPinOrder = (queryClient: QueryClient): number => { return maxPinOrder + 1; }; -export const invalidateChatListQueries = (queryClient: QueryClient) => { - return queryClient.invalidateQueries({ - queryKey: chatListFamilyKey, - }); -}; - /** * Predicate that matches chat-list queries performing a regular * refetch (window-focus, invalidation, mount) but not a @@ -635,24 +629,145 @@ const isChatListRefetch = (query: { return true; }; +export const invalidateChatEntity = ( + queryClient: QueryClient, + chatId: string, +) => + queryClient.invalidateQueries({ + queryKey: chatEntityKey(chatId), + exact: true, + }); + +export const invalidateChatListQueries = (queryClient: QueryClient) => + queryClient.invalidateQueries({ + queryKey: chatListFamilyKey, + }); + +export const invalidateChatsByWorkspace = (queryClient: QueryClient) => + queryClient.invalidateQueries({ + queryKey: chatsByWorkspaceFamilyKey, + }); + +export const invalidateChatDebugRuns = ( + queryClient: QueryClient, + chatId: string, +) => + queryClient.invalidateQueries({ + queryKey: chatDebugRunsKey(chatId), + }); + +export const invalidateChatDiffContents = ( + queryClient: QueryClient, + chatId: string, +) => + queryClient.invalidateQueries({ + queryKey: chatDiffContentsKey(chatId), + exact: true, + }); + +export const invalidateChatPrompts = ( + queryClient: QueryClient, + chatId: string, +) => + queryClient.invalidateQueries({ + queryKey: chatPromptsKey(chatId), + exact: true, + }); + +export const invalidateChatMessages = ( + queryClient: QueryClient, + chatId: string, +) => + queryClient.invalidateQueries({ + queryKey: chatMessagesKey(chatId), + exact: true, + }); + +export const invalidateChatACL = (queryClient: QueryClient, chatId: string) => + queryClient.invalidateQueries({ + queryKey: chatACLKey(chatId), + exact: true, + }); + +export const invalidateChatCostTree = ( + queryClient: QueryClient, + rootChatId: string, +) => + queryClient.invalidateQueries({ + queryKey: chatCostTreeKey(rootChatId), + exact: true, + }); + +export const cancelChatListQueries = (queryClient: QueryClient) => + queryClient.cancelQueries({ + queryKey: chatListFamilyKey, + }); + /** - * 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. + * Cancel background chat-list refetches, leaving pagination fetches alone. + * Call before applying WebSocket-driven cache updates, or a concurrent + * refetch may overwrite them with stale data. */ -export const cancelChatListRefetches = (queryClient: QueryClient) => { - return queryClient.cancelQueries({ +export const cancelChatListRefetches = (queryClient: QueryClient) => + queryClient.cancelQueries({ queryKey: chatListFamilyKey, predicate: isChatListRefetch, }); + +export const cancelChatEntity = (queryClient: QueryClient, chatId: string) => + queryClient.cancelQueries({ + queryKey: chatEntityKey(chatId), + exact: true, + }); + +// Cancelling a first-time fetch leaves the query pending with no retry, +// which the page shows as "Chat not found". +export const cancelLoadedChatEntityRefetch = ( + queryClient: QueryClient, + chatId: string, +) => { + if (queryClient.getQueryData(chatEntityKey(chatId)) === undefined) { + return; + } + return queryClient.cancelQueries({ + queryKey: chatEntityKey(chatId), + exact: true, + }); }; +export const cancelChatMessages = (queryClient: QueryClient, chatId: string) => + queryClient.cancelQueries({ + queryKey: chatMessagesKey(chatId), + exact: true, + }); + +export const removeChatEntity = (queryClient: QueryClient, chatId: string) => + queryClient.removeQueries({ + queryKey: chatEntityKey(chatId), + exact: true, + }); + +export const patchChatEntity = ( + queryClient: QueryClient, + chatId: string, + updater: (chat: TypesGen.Chat | undefined) => TypesGen.Chat | undefined, +) => + queryClient.setQueryData( + chatEntityKey(chatId), + updater, + ); + +export const patchChatMessages = ( + queryClient: QueryClient, + chatId: string, + updater: ( + data: InfiniteData | undefined, + ) => InfiniteData | undefined, +) => + queryClient.setQueryData< + InfiniteData | undefined + >(chatMessagesKey(chatId), updater); + const DEFAULT_CHAT_PAGE_LIMIT = 50; export const CHAT_SEARCH_LIMIT = 50; @@ -832,13 +947,8 @@ export const archiveChat = (queryClient: QueryClient) => ({ mutationFn: (chatId: string) => API.experimental.updateChat(chatId, { archived: true }), onMutate: async (chatId: string) => { - await queryClient.cancelQueries({ - queryKey: chatListFamilyKey, - }); - await queryClient.cancelQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); + await cancelChatListQueries(queryClient); + await cancelChatEntity(queryClient, chatId); const previousChat = queryClient.getQueryData( chatEntityKey(chatId), ); @@ -873,10 +983,7 @@ export const archiveChat = (queryClient: QueryClient) => ({ // Rollback: invalidate to re-fetch the correct state. void invalidateChatListQueries(queryClient); if (context?.previousChat) { - queryClient.setQueryData( - chatEntityKey(chatId), - context.previousChat, - ); + patchChatEntity(queryClient, chatId, () => context.previousChat); } }, onSuccess: (_data: unknown, chatId: string) => { @@ -884,13 +991,8 @@ export const archiveChat = (queryClient: QueryClient) => ({ }, onSettled: (_data: unknown, _error: unknown, chatId: string) => { void invalidateChatListQueries(queryClient); - void queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); - void queryClient.invalidateQueries({ - queryKey: chatsByWorkspaceFamilyKey, - }); + void invalidateChatEntity(queryClient, chatId); + void invalidateChatsByWorkspace(queryClient); }, }); @@ -898,13 +1000,8 @@ export const unarchiveChat = (queryClient: QueryClient) => ({ mutationFn: (chatId: string) => API.experimental.updateChat(chatId, { archived: false }), onMutate: async (chatId: string) => { - await queryClient.cancelQueries({ - queryKey: chatListFamilyKey, - }); - await queryClient.cancelQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); + await cancelChatListQueries(queryClient); + await cancelChatEntity(queryClient, chatId); const previousChat = queryClient.getQueryData( chatEntityKey(chatId), ); @@ -935,10 +1032,7 @@ export const unarchiveChat = (queryClient: QueryClient) => ({ // Rollback: invalidate to re-fetch the correct state. void invalidateChatListQueries(queryClient); if (context?.previousChat) { - queryClient.setQueryData( - chatEntityKey(chatId), - context.previousChat, - ); + patchChatEntity(queryClient, chatId, () => context.previousChat); } }, onSuccess: (_data: unknown, chatId: string) => { @@ -946,13 +1040,8 @@ export const unarchiveChat = (queryClient: QueryClient) => ({ }, onSettled: (_data: unknown, _error: unknown, chatId: string) => { void invalidateChatListQueries(queryClient); - void queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); - void queryClient.invalidateQueries({ - queryKey: chatsByWorkspaceFamilyKey, - }); + void invalidateChatEntity(queryClient, chatId); + void invalidateChatsByWorkspace(queryClient); }, }); @@ -962,13 +1051,8 @@ export const updateChatPlanMode = (queryClient: QueryClient) => ({ plan_mode: toChatPlanModePayload(planMode), }), onMutate: async ({ chatId, planMode }: UpdateChatPlanModeVariables) => { - await queryClient.cancelQueries({ - queryKey: chatListFamilyKey, - }); - await queryClient.cancelQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); + await cancelChatListQueries(queryClient); + await cancelChatEntity(queryClient, chatId); const previousChat = queryClient.getQueryData( chatEntityKey(chatId), ); @@ -1009,10 +1093,7 @@ export const updateChatPlanMode = (queryClient: QueryClient) => ({ : chat, ), ); - queryClient.setQueryData( - chatEntityKey(chatId), - previousChat, - ); + patchChatEntity(queryClient, chatId, () => previousChat); }, }); @@ -1025,13 +1106,8 @@ export const updateChatWorkspace = (queryClient: QueryClient) => ({ "00000000-0000-0000-0000-000000000000", }), onMutate: async ({ chatId, workspaceId }: UpdateChatWorkspaceVariables) => { - await queryClient.cancelQueries({ - queryKey: chatListFamilyKey, - }); - await queryClient.cancelQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); + await cancelChatListQueries(queryClient); + await cancelChatEntity(queryClient, chatId); const previousChat = queryClient.getQueryData( chatEntityKey(chatId), ); @@ -1072,10 +1148,7 @@ export const updateChatWorkspace = (queryClient: QueryClient) => ({ : chat, ), ); - queryClient.setQueryData( - chatEntityKey(chatId), - previousChat, - ); + patchChatEntity(queryClient, chatId, () => previousChat); } }, onSettled: async ( @@ -1084,13 +1157,8 @@ export const updateChatWorkspace = (queryClient: QueryClient) => ({ { chatId }: UpdateChatWorkspaceVariables, ) => { await invalidateChatListQueries(queryClient); - await queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); - await queryClient.invalidateQueries({ - queryKey: chatsByWorkspaceFamilyKey, - }); + await invalidateChatEntity(queryClient, chatId); + await invalidateChatsByWorkspace(queryClient); }, }); @@ -1098,13 +1166,8 @@ export const pinChat = (queryClient: QueryClient) => ({ mutationFn: (chatId: string) => API.experimental.updateChat(chatId, { pin_order: 1 }), onMutate: async (chatId: string) => { - await queryClient.cancelQueries({ - queryKey: chatListFamilyKey, - }); - await queryClient.cancelQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); + await cancelChatListQueries(queryClient); + await cancelChatEntity(queryClient, chatId); const previousChat = queryClient.getQueryData( chatEntityKey(chatId), ); @@ -1134,18 +1197,12 @@ export const pinChat = (queryClient: QueryClient) => ({ // Rollback: invalidate to re-fetch the correct state. void invalidateChatListQueries(queryClient); if (context?.previousChat) { - queryClient.setQueryData( - chatEntityKey(chatId), - context.previousChat, - ); + patchChatEntity(queryClient, chatId, () => context.previousChat); } }, onSettled: async (_data: unknown, _error: unknown, chatId: string) => { await invalidateChatListQueries(queryClient); - await queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); + await invalidateChatEntity(queryClient, chatId); }, }); @@ -1153,13 +1210,8 @@ export const unpinChat = (queryClient: QueryClient) => ({ mutationFn: (chatId: string) => API.experimental.updateChat(chatId, { pin_order: 0 }), onMutate: async (chatId: string) => { - await queryClient.cancelQueries({ - queryKey: chatListFamilyKey, - }); - await queryClient.cancelQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); + await cancelChatListQueries(queryClient); + await cancelChatEntity(queryClient, chatId); const previousChat = queryClient.getQueryData( chatEntityKey(chatId), ); @@ -1188,18 +1240,12 @@ export const unpinChat = (queryClient: QueryClient) => ({ // Rollback: invalidate to re-fetch the correct state. void invalidateChatListQueries(queryClient); if (context?.previousChat) { - queryClient.setQueryData( - chatEntityKey(chatId), - context.previousChat, - ); + patchChatEntity(queryClient, chatId, () => context.previousChat); } }, onSettled: async (_data: unknown, _error: unknown, chatId: string) => { await invalidateChatListQueries(queryClient); - await queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); + await invalidateChatEntity(queryClient, chatId); }, }); @@ -1213,13 +1259,8 @@ export const reorderPinnedChat = (queryClient: QueryClient) => ({ chatId: string; pinOrder: number; }) => { - await queryClient.cancelQueries({ - queryKey: chatListFamilyKey, - }); - await queryClient.cancelQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); + await cancelChatListQueries(queryClient); + await cancelChatEntity(queryClient, chatId); // Optimistically reorder pinned chats in the cache so the // sidebar reflects the new order immediately without waiting @@ -1247,10 +1288,7 @@ export const reorderPinnedChat = (queryClient: QueryClient) => ({ { chatId }: { chatId: string; pinOrder: number }, ) => { await invalidateChatListQueries(queryClient); - await queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); + await invalidateChatEntity(queryClient, chatId); }, }); @@ -1276,9 +1314,8 @@ export const updateChatTitle = (queryClient: QueryClient) => ({ API.experimental.updateChat(chatId, { title }), onSuccess: (_data: unknown, { chatId, title }: UpdateChatTitleVariables) => { - queryClient.setQueryData( - chatEntityKey(chatId), - (chat) => (chat ? { ...chat, title } : chat), + patchChatEntity(queryClient, chatId, (chat) => + chat ? { ...chat, title } : chat, ); updateInfiniteChatsCache(queryClient, (chats) => chats.map((chat) => (chat.id === chatId ? { ...chat, title } : chat)), @@ -1291,10 +1328,7 @@ export const updateChatTitle = (queryClient: QueryClient) => ({ { chatId }: UpdateChatTitleVariables, ) => { void invalidateChatListQueries(queryClient); - void queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); + void invalidateChatEntity(queryClient, chatId); }, }); @@ -1369,20 +1403,12 @@ export const chatDebugRun = (chatId: string, runId: string) => refetchIntervalInBackground: false, }); -const invalidateChatDebugRuns = (queryClient: QueryClient, chatId: string) => { - return queryClient.invalidateQueries({ - queryKey: chatDebugRunsKey(chatId), - }); -}; - export const createChat = (queryClient: QueryClient) => ({ mutationFn: (req: TypesGen.CreateChatRequest) => API.experimental.createChat(req), onSuccess: () => { void invalidateChatListQueries(queryClient); - void queryClient.invalidateQueries({ - queryKey: chatsByWorkspaceFamilyKey, - }); + void invalidateChatsByWorkspace(queryClient); }, }); @@ -1394,14 +1420,8 @@ export const createChatMessage = ( API.experimental.createChatMessage(chatId, req), onSuccess: () => { void invalidateChatDebugRuns(queryClient, chatId); - void queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); - void queryClient.invalidateQueries({ - queryKey: chatPromptsKey(chatId), - exact: true, - }); + void invalidateChatEntity(queryClient, chatId); + void invalidateChatPrompts(queryClient, chatId); }, }); @@ -1424,18 +1444,13 @@ export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({ }: EditChatMessageMutationArgs): Promise => { // Cancel in-flight refetches so they don't overwrite the // optimistic update before the mutation completes. - await queryClient.cancelQueries({ - queryKey: chatMessagesKey(chatId), - exact: true, - }); + await cancelChatMessages(queryClient, chatId); const previousData = queryClient.getQueryData< InfiniteData >(chatMessagesKey(chatId)); - queryClient.setQueryData< - InfiniteData | undefined - >(chatMessagesKey(chatId), (current) => + patchChatMessages(queryClient, chatId, (current) => projectEditedConversationIntoCache({ currentData: current, editedMessageId: messageId, @@ -1454,23 +1469,18 @@ export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({ // Restore the cache on failure so the user sees the // original messages again. if (context?.previousData) { - queryClient.setQueryData(chatMessagesKey(chatId), context.previousData); + patchChatMessages(queryClient, chatId, () => context.previousData); } // Invalidate messages as a safety net: the restored snapshot // may be missing WebSocket-delivered messages that arrived // during the mutation's flight time. - void queryClient.invalidateQueries({ - queryKey: chatMessagesKey(chatId), - exact: true, - }); + void invalidateChatMessages(queryClient, chatId); }, onSuccess: ( response: TypesGen.EditChatMessageResponse, variables: EditChatMessageMutationArgs, ) => { - queryClient.setQueryData< - InfiniteData | undefined - >(chatMessagesKey(chatId), (current) => + patchChatMessages(queryClient, chatId, (current) => reconcileEditedMessageInCache({ currentData: current, optimisticMessageId: variables.messageId, @@ -1488,14 +1498,8 @@ export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({ // refetch that causes extra store mutations while the // sticky user message is settling after the optimistic // truncation. - void queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); - void queryClient.invalidateQueries({ - queryKey: chatPromptsKey(chatId), - exact: true, - }); + void invalidateChatEntity(queryClient, chatId); + void invalidateChatPrompts(queryClient, chatId); void invalidateChatDebugRuns(queryClient, chatId); }, }); @@ -1512,10 +1516,7 @@ export const compactChat = (queryClient: QueryClient, chatId: string) => ({ onSuccess: () => { // The compaction transitions the chat to running; the summary // rows stream in over the websocket like any other turn. - void queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); + void invalidateChatEntity(queryClient, chatId); void invalidateChatDebugRuns(queryClient, chatId); }, }); @@ -1533,7 +1534,7 @@ export const refreshChatContext = ( ) => ({ mutationFn: () => API.experimental.refreshChatContext(chatId), onSuccess: (updatedChat: TypesGen.Chat) => { - queryClient.setQueryData(chatEntityKey(chatId), (cached) => + patchChatEntity(queryClient, chatId, (cached) => cached ? { ...cached, context: updatedChat.context } : updatedChat, ); const applyContext = (chat: TypesGen.Chat): TypesGen.Chat => @@ -1560,14 +1561,8 @@ export const deleteChatQueuedMessage = ( mutationFn: (queuedMessageId: number) => API.experimental.deleteChatQueuedMessage(chatId, queuedMessageId), onSuccess: async () => { - await queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); - await queryClient.invalidateQueries({ - queryKey: chatMessagesKey(chatId), - exact: true, - }); + await invalidateChatEntity(queryClient, chatId); + await invalidateChatMessages(queryClient, chatId); }, }); @@ -2106,10 +2101,7 @@ export const setChatUserRole = (queryClient: QueryClient) => ({ user_roles: { [userId]: role }, }), onSuccess: async (_data: unknown, { chatId }: SetChatUserRoleVariables) => { - await queryClient.invalidateQueries({ - queryKey: chatACLKey(chatId), - exact: true, - }); + await invalidateChatACL(queryClient, chatId); }, }); @@ -2119,9 +2111,6 @@ export const setChatGroupRole = (queryClient: QueryClient) => ({ group_roles: { [groupId]: role }, }), onSuccess: async (_data: unknown, { chatId }: SetChatGroupRoleVariables) => { - await queryClient.invalidateQueries({ - queryKey: chatACLKey(chatId), - exact: true, - }); + await invalidateChatACL(queryClient, chatId); }, }); diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index c9d00b7ac8..27d2ff021d 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -28,7 +28,6 @@ import { checkAuthorization } from "#/api/queries/authCheck"; import { buildOptimisticEditedMessage } from "#/api/queries/chatMessageEdits"; import { chat, - chatEntityKey, chatMessagesForInfiniteScroll, chatModelConfigs, chatModels, @@ -38,7 +37,9 @@ import { deleteChatQueuedMessage, editChatMessage, interruptChat, + invalidateChatEntity, mcpServerConfigs, + patchChatEntity, promoteChatQueuedMessage, updateChatPlanMode, updateChatWorkspace, @@ -1182,10 +1183,8 @@ const AgentChatPage: FC = () => { chat.id === chatId ? { ...chat, plan_mode: planMode } : chat, ), ); - queryClient.setQueryData( - chatEntityKey(chatId), - (previousChat) => - previousChat ? { ...previousChat, plan_mode: planMode } : previousChat, + patchChatEntity(queryClient, chatId, (previousChat) => + previousChat ? { ...previousChat, plan_mode: planMode } : previousChat, ); }; @@ -1694,10 +1693,7 @@ const AgentChatPage: FC = () => { handleRequestError(error); // Hook dispatch failures can park an idle chat in error before returning the request error. acceptServerChatStatus(); - void queryClient.invalidateQueries({ - queryKey: chatEntityKey(agentId), - exact: true, - }); + void invalidateChatEntity(queryClient, agentId); }, }); if (editSelectedModelConfigID) { @@ -1745,10 +1741,7 @@ const AgentChatPage: FC = () => { handleRequestError(error); // Hook dispatch failures can park an idle chat in error before returning the request error. acceptServerChatStatus(); - void queryClient.invalidateQueries({ - queryKey: chatEntityKey(agentId), - exact: true, - }); + void invalidateChatEntity(queryClient, agentId); throw error; } const isActiveChat = store.getActiveChatID() === agentId; diff --git a/site/src/pages/AgentsPage/AgentChatPageView.tsx b/site/src/pages/AgentsPage/AgentChatPageView.tsx index bb772d3dba..d6fc345683 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.tsx @@ -11,7 +11,7 @@ import { import { useQueryClient } from "react-query"; import type { UrlTransform } from "streamdown"; import { v4 as uuidv4 } from "uuid"; -import { chatDiffContentsKey } from "#/api/queries/chats"; +import { invalidateChatDiffContents } from "#/api/queries/chats"; import type * as TypesGen from "#/api/typesGenerated"; import type { AgentChatSendShortcut, @@ -407,10 +407,7 @@ export const AgentChatPageView: FC = ({ const handleRefresh = () => { const sent = gitWatcher.refresh(); if (sent && agentId) { - void queryClient.invalidateQueries({ - queryKey: chatDiffContentsKey(agentId), - exact: true, - }); + void invalidateChatDiffContents(queryClient, agentId); } return sent; }; diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.tsx b/site/src/pages/AgentsPage/AgentsPageLayout.tsx index bacd512843..bfaf3eacc9 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.tsx +++ b/site/src/pages/AgentsPage/AgentsPageLayout.tsx @@ -20,19 +20,22 @@ import { applyChatArchiveStateToCaches, archiveChat, cancelChatListRefetches, - chatCostTreeKey, - chatDiffContentsKey, + cancelLoadedChatEntityRefetch, chatEntityKey, chatModelConfigs, chatModels, - chatsByWorkspaceFamilyKey, infiniteChats, + invalidateChatCostTree, + invalidateChatDiffContents, + invalidateChatEntity, invalidateChatListQueries, + invalidateChatsByWorkspace, mergeWatchedChatIntoCaches, pinChat, prependToInfiniteChatsCache, proposeChatTitle, readInfiniteChatsCache, + removeChatEntity, removeChildFromParentInCache, reorderPinnedChat, unarchiveChat, @@ -303,13 +306,8 @@ const AgentsPageLayout: FC = () => { clearPersistedSidebarTabId(chatId); clearPersistedRightPanelState(chatId); void invalidateChatListQueries(queryClient); - void queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatId), - exact: true, - }); - void queryClient.invalidateQueries({ - queryKey: chatsByWorkspaceFamilyKey, - }); + void invalidateChatEntity(queryClient, chatId); + void invalidateChatsByWorkspace(queryClient); void invalidateWorkspaceMutationQueries(queryClient, { organizationName, username: user.username, @@ -616,20 +614,14 @@ const AgentsPageLayout: FC = () => { ), ); removeChildFromParentInCache(queryClient, updatedChat.id); - queryClient.removeQueries({ - queryKey: chatEntityKey(updatedChat.id), - exact: true, - }); + removeChatEntity(queryClient, updatedChat.id); return; } if (chatEvent.kind === "diff_status_change") { // Only refetch the diff file contents. The chat's // diff_status field is already written into the // chatKey and infinite-list caches below. - void queryClient.invalidateQueries({ - queryKey: chatDiffContentsKey(updatedChat.id), - exact: true, - }); + void invalidateChatDiffContents(queryClient, updatedChat.id); } // Merge watch payloads by event kind so stale field // snapshots do not clobber fresher cached metadata. @@ -643,17 +635,7 @@ const AgentsPageLayout: FC = () => { // title generation finished, so its response carries // the fallback title. 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 - // and no retry, which AgentChatPage shows as - // "Chat not found". - if (queryClient.getQueryData(chatEntityKey(updatedChat.id))) { - void queryClient.cancelQueries({ - queryKey: chatEntityKey(updatedChat.id), - exact: true, - }); - } + void cancelLoadedChatEntityRefetch(queryClient, updatedChat.id); if (chatEvent.kind === "created") { if (updatedChat.parent_chat_id) { @@ -682,10 +664,7 @@ const AgentsPageLayout: FC = () => { chatEvent.kind, ); if (costChatId) { - void queryClient.invalidateQueries({ - queryKey: chatCostTreeKey(costChatId), - exact: true, - }); + void invalidateChatCostTree(queryClient, costChatId); } if (chatEvent.kind === "context_dirty") { // The watch payload carries only the lightweight @@ -694,10 +673,7 @@ const AgentsPageLayout: FC = () => { // resources the single-chat GET computes. Only the // active chat has an observer, so other chats are // merely marked stale. - void queryClient.invalidateQueries({ - queryKey: chatEntityKey(updatedChat.id), - exact: true, - }); + void invalidateChatEntity(queryClient, updatedChat.id); } } }); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index 367430a5c8..d56fc8b5fa 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts @@ -13,7 +13,8 @@ import { import { watchChat } from "#/api/api"; import { chatMessagesKey, - chatPromptsKey, + invalidateChatPrompts, + patchChatMessages, updateInfiniteChatsCache, } from "#/api/queries/chats"; import type * as TypesGen from "#/api/typesGenerated"; @@ -40,9 +41,7 @@ const writeQueuedMessagesToCache = ( return; } const nextQueuedMessages = queuedMessages ?? []; - queryClient.setQueryData< - InfiniteData | undefined - >(chatMessagesKey(chatID), (currentData) => { + patchChatMessages(queryClient, chatID, (currentData) => { if (!currentData?.pages?.length) { return currentData; } @@ -197,9 +196,7 @@ export const useChatStore = ( if (!chatID || messages.length === 0) { return; } - queryClient.setQueryData< - InfiniteData | undefined - >(chatMessagesKey(chatID), (currentData) => { + patchChatMessages(queryClient, chatID, (currentData) => { if (!currentData?.pages?.length) { return currentData; } @@ -235,10 +232,7 @@ export const useChatStore = ( // Refresh the dedicated prompt-history cache when a user message arrives. const hasNewUserPrompt = messages.some((msg) => msg.role === "user"); if (hasNewUserPrompt) { - void queryClient.invalidateQueries({ - queryKey: chatPromptsKey(chatID), - exact: true, - }); + void invalidateChatPrompts(queryClient, chatID); } }, [chatID, queryClient], @@ -249,9 +243,7 @@ export const useChatStore = ( if (!chatID) { return; } - queryClient.setQueryData< - InfiniteData | undefined - >(chatMessagesKey(chatID), (currentData) => { + patchChatMessages(queryClient, chatID, (currentData) => { if (!currentData?.pages?.length) { return currentData; } diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.test.tsx index 562f7ba1c4..4d01aa6d95 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.test.tsx @@ -3,7 +3,11 @@ import type { FC, PropsWithChildren } from "react"; import { act } from "react"; import { QueryClient, QueryClientProvider } from "react-query"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { chatEntityKey } from "#/api/queries/chats"; +import { + chatEntityKey, + chatMessagesKey, + chatPromptsKey, +} from "#/api/queries/chats"; import { getWorkspaceQuotaQueryKey } from "#/api/queries/workspaceQuota"; import { workspacesQueryKeyPrefix } from "#/api/queries/workspaces"; import { createChatStore } from "./chatStore"; @@ -108,6 +112,7 @@ describe("useChatToolInvalidations", () => { await waitFor(() => { expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: chatEntityKey("chat-1"), + exact: true, }); expect(invalidateSpy).toHaveBeenCalledWith( expect.objectContaining({ @@ -123,6 +128,35 @@ describe("useChatToolInvalidations", () => { }); }); + it("does not invalidate chat sub-resources on create_workspace completion", async () => { + queryClient.setQueryData(chatEntityKey("chat-1"), { id: "chat-1" }); + queryClient.setQueryData(chatMessagesKey("chat-1"), { + pages: [], + pageParams: [], + }); + queryClient.setQueryData(chatPromptsKey("chat-1"), { prompts: [] }); + const { setStreamState } = renderInvalidations(); + + await act(async () => { + setStreamState(createStreamState("create_workspace")); + }); + + await waitFor(() => { + expect( + queryClient.getQueryState(chatEntityKey("chat-1"))?.isInvalidated, + "detail entry should be invalidated", + ).toBe(true); + }); + expect( + queryClient.getQueryState(chatMessagesKey("chat-1"))?.isInvalidated, + "messages entry should NOT be invalidated", + ).not.toBe(true); + expect( + queryClient.getQueryState(chatPromptsKey("chat-1"))?.isInvalidated, + "prompts entry should NOT be invalidated", + ).not.toBe(true); + }); + it("dispatches workspace mutation invalidations on start_workspace completion", async () => { const { invalidateSpy, setStreamState } = renderInvalidations(); @@ -146,6 +180,7 @@ describe("useChatToolInvalidations", () => { expect(invalidateSpy).not.toHaveBeenCalledWith({ queryKey: chatEntityKey("chat-1"), + exact: true, }); }); @@ -185,6 +220,7 @@ describe("useChatToolInvalidations", () => { await waitFor(() => { expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: chatEntityKey("chat-1"), + exact: true, }); expect(invalidateSpy).toHaveBeenCalledWith( expect.objectContaining({ @@ -252,6 +288,7 @@ describe("useChatToolInvalidations", () => { expect(invalidateSpy).toHaveBeenCalledTimes(6); expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: chatEntityKey("chat-2"), + exact: true, }); }); }); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.ts index a6b57fca2c..d4f2dd0bcc 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.ts @@ -1,6 +1,6 @@ import { useEffect, useRef } from "react"; import { useQueryClient } from "react-query"; -import { chatEntityKey } from "#/api/queries/chats"; +import { invalidateChatEntity } from "#/api/queries/chats"; import { invalidateWorkspaceMutationQueries } from "#/api/queries/workspaces"; import { type ChatStore, useChatSelector } from "./chatStore"; import type { StreamState } from "./types"; @@ -87,9 +87,7 @@ export function useChatToolInvalidations({ } if (shouldInvalidateChat) { - void queryClient.invalidateQueries({ - queryKey: chatEntityKey(chatID), - }); + void invalidateChatEntity(queryClient, chatID); } if (shouldInvalidateWorkspace) {