diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 310d4f78d9..52cd88f8d4 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -48,6 +48,7 @@ import { unpinChat, updateChatAdvisorConfig, updateChatPlanMode, + updateChatTitle, updateChildInParentCache, updateInfiniteChatsCache, } from "./chats"; @@ -323,6 +324,65 @@ describe("updateChatPlanMode optimistic update", () => { }); }); +describe("updateChatTitle cache update", () => { + it("patches chat detail and infinite chat list caches after success", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + queryClient.setQueryData( + chatKey(chatId), + makeChat(chatId, { title: "Old" }), + ); + seedInfiniteChats(queryClient, [ + makeChat(chatId, { title: "Old" }), + makeChat("chat-2", { title: "Other" }), + ]); + seedInfiniteChats( + queryClient, + [makeChat(chatId, { archived: true, title: "Old" })], + { archived: true }, + ); + + const mutation = updateChatTitle(queryClient); + mutation.onSuccess(undefined, { chatId, title: "New" }); + + expect( + queryClient.getQueryData(chatKey(chatId))?.title, + ).toBe("New"); + expect( + readInfiniteChats(queryClient)?.find((chat) => chat.id === chatId), + ).toMatchObject({ title: "New" }); + expect( + readInfiniteChats(queryClient, { archived: true })?.find( + (chat) => chat.id === chatId, + ), + ).toMatchObject({ title: "New" }); + }); + + it("does not return pending invalidation promises from settlement", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + const invalidateSpy = vi + .spyOn(queryClient, "invalidateQueries") + .mockReturnValue(new Promise(() => {})); + + const mutation = updateChatTitle(queryClient); + const result = mutation.onSettled(undefined, undefined, { + chatId, + title: "New", + }); + + expect(result).toBeUndefined(); + expect(invalidateSpy).toHaveBeenCalledWith( + expect.objectContaining({ queryKey: chatsKey }), + ); + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: chatKey(chatId), + exact: true, + }); + invalidateSpy.mockRestore(); + }); +}); + describe("archiveChat optimistic update", () => { it("optimistically sets archived to true in the chats list", async () => { const queryClient = createTestQueryClient(); @@ -380,6 +440,52 @@ describe("archiveChat optimistic update", () => { expect(result?.[0].children?.[0].id).toBe("child-2"); }); + it("removes an archived root chat from active filtered lists after success", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + seedInfiniteChats( + queryClient, + [ + makeChat(chatId, { pin_order: 2 }), + makeChat("chat-2", { archived: false }), + ], + { archived: false }, + ); + queryClient.setQueryData( + chatKey(chatId), + makeChat(chatId, { pin_order: 2 }), + ); + + const mutation = archiveChat(queryClient); + mutation.onSuccess(undefined, chatId); + + expect( + readInfiniteChats(queryClient, { archived: false })?.map( + (chat) => chat.id, + ), + ).toEqual(["chat-2"]); + expect( + queryClient.getQueryData(chatKey(chatId)), + ).toMatchObject({ + archived: true, + pin_order: 0, + }); + }); + + it("clears pin order for archived chats that remain in unfiltered lists", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + seedInfiniteChats(queryClient, [makeChat(chatId, { pin_order: 3 })]); + + const mutation = archiveChat(queryClient); + mutation.onSuccess(undefined, chatId); + + expect(readInfiniteChats(queryClient)?.[0]).toMatchObject({ + archived: true, + pin_order: 0, + }); + }); + it("rolls back the chats list on error by invalidating", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; @@ -456,14 +562,20 @@ describe("archiveChat optimistic update", () => { expect(context?.previousChat).toBeUndefined(); }); - it("invalidates queries on settled regardless of outcome", async () => { + it("invalidates on settled without returning pending promises", () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; - const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + // Mock invalidateQueries to never resolve so a regression back to + // an awaited (async) onSettled surfaces as a pending promise return + // value, which is what keeps the mutation's loading state stuck. + const invalidateSpy = vi + .spyOn(queryClient, "invalidateQueries") + .mockReturnValue(new Promise(() => {})); const mutation = archiveChat(queryClient); - await mutation.onSettled(undefined, undefined, chatId); + const result = mutation.onSettled(undefined, undefined, chatId); + expect(result).toBeUndefined(); expect(invalidateSpy).toHaveBeenCalledWith( expect.objectContaining({ queryKey: chatsKey }), ); @@ -471,6 +583,7 @@ describe("archiveChat optimistic update", () => { queryKey: chatKey(chatId), exact: true, }); + invalidateSpy.mockRestore(); }); }); @@ -503,6 +616,37 @@ describe("unarchiveChat optimistic update", () => { ).toBe(false); }); + it("removes an unarchived root chat from archived filtered lists after success", () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + seedInfiniteChats( + queryClient, + [ + makeChat(chatId, { archived: true }), + makeChat("chat-2", { archived: true }), + ], + { archived: true }, + ); + queryClient.setQueryData( + chatKey(chatId), + makeChat(chatId, { archived: true }), + ); + + const mutation = unarchiveChat(queryClient); + mutation.onSuccess(undefined, chatId); + + expect( + readInfiniteChats(queryClient, { archived: true })?.map( + (chat) => chat.id, + ), + ).toEqual(["chat-2"]); + expect( + queryClient.getQueryData(chatKey(chatId)), + ).toMatchObject({ + archived: false, + }); + }); + it("rolls back both caches on error", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; @@ -535,14 +679,20 @@ describe("unarchiveChat optimistic update", () => { ).toBe(true); }); - it("invalidates queries on settled", async () => { + it("invalidates on settled without returning pending promises", () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; - const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + // Mock invalidateQueries to never resolve so a regression back to + // an awaited (async) onSettled surfaces as a pending promise return + // value, which is what keeps the mutation's loading state stuck. + const invalidateSpy = vi + .spyOn(queryClient, "invalidateQueries") + .mockReturnValue(new Promise(() => {})); const mutation = unarchiveChat(queryClient); - await mutation.onSettled(undefined, undefined, chatId); + const result = mutation.onSettled(undefined, undefined, chatId); + expect(result).toBeUndefined(); expect(invalidateSpy).toHaveBeenCalledWith( expect.objectContaining({ queryKey: chatsKey }), ); @@ -550,6 +700,7 @@ describe("unarchiveChat optimistic update", () => { queryKey: chatKey(chatId), exact: true, }); + invalidateSpy.mockRestore(); }); }); @@ -1483,6 +1634,18 @@ describe("mutation invalidation scope", () => { }); }); +describe("infiniteChatsKey shape", () => { + it("places the filter object one slot after the chatsKey prefix", () => { + // archivedFilterForChatListKey reads the archived filter from the + // slot immediately after the chatsKey prefix. If this layout ever + // changes, that helper silently stops removing chats from + // conflicting filtered lists, so keep the two in sync. + const key = infiniteChatsKey({ archived: true }); + expect(key.length).toBe(chatsKey.length + 1); + expect(key[chatsKey.length]).toEqual({ archived: true }); + }); +}); + describe("infiniteChats", () => { const PAGE_LIMIT = 50; diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index d3e1e94597..bd4e2ae099 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -233,6 +233,119 @@ export const removeChildFromParentInCache = ( return found; }; +// Inverse of infiniteChatsKey, which builds keys as [...chatsKey, filters?]. +// The optional filter object lives in the slot immediately after the +// chatsKey prefix, so derive both the expected length and the filter index +// from chatsKey. If infiniteChatsKey's shape changes, this must change with +// it; the "infiniteChatsKey shape" test in chats.test.ts guards that contract. +const archivedFilterForChatListKey = ( + queryKey: readonly unknown[], +): boolean | undefined => { + if (queryKey.length !== chatsKey.length + 1) { + return undefined; + } + const filters = queryKey[chatsKey.length]; + if (!filters || typeof filters !== "object") { + return undefined; + } + const archived = (filters as { archived?: unknown }).archived; + return typeof archived === "boolean" ? archived : undefined; +}; + +const isInfiniteChatsCacheData = ( + data: unknown, +): data is InfiniteChatsCacheData => { + if (!data || typeof data !== "object") { + return false; + } + const maybeData = data as { pages?: unknown; pageParams?: unknown }; + return Array.isArray(maybeData.pages) && Array.isArray(maybeData.pageParams); +}; + +const patchChatArchiveState = ( + chat: TypesGen.Chat, + archived: boolean, +): TypesGen.Chat => { + const pinOrder = archived ? 0 : chat.pin_order; + if (chat.archived === archived && chat.pin_order === pinOrder) { + return chat; + } + return { ...chat, archived, pin_order: pinOrder }; +}; + +/** + * Applies an accepted archive state to loaded sidebar and detail caches. + * Removes the chat from any filtered list whose archived filter conflicts + * with the new state, and resets pin_order to 0 when archiving. + */ +export const applyChatArchiveStateToCaches = ( + queryClient: QueryClient, + chatId: string, + archived: boolean, +) => { + queryClient.setQueryData( + chatKey(chatId), + (chat) => (chat ? patchChatArchiveState(chat, archived) : chat), + ); + + if (archived) { + removeChildFromParentInCache(queryClient, chatId); + } else { + updateChildInParentCache( + queryClient, + (child) => patchChatArchiveState(child, archived), + chatId, + ); + } + + const queries = queryClient.getQueriesData({ + queryKey: chatsKey, + predicate: isChatListQuery, + }); + + for (const [queryKey, data] of queries) { + if (!isInfiniteChatsCacheData(data)) { + continue; + } + const archivedFilter = archivedFilterForChatListKey(queryKey); + queryClient.setQueryData(queryKey, (prev) => { + if (!isInfiniteChatsCacheData(prev)) { + return prev; + } + + let changed = false; + const pages = prev.pages.map((page) => { + let pageChanged = false; + const nextPage: TypesGen.Chat[] = []; + for (const chat of page) { + if (chat.id !== chatId) { + nextPage.push(chat); + continue; + } + + if (archivedFilter !== undefined && archivedFilter !== archived) { + pageChanged = true; + continue; + } + + const updatedChat = patchChatArchiveState(chat, archived); + if (updatedChat !== chat) { + pageChanged = true; + } + nextPage.push(updatedChat); + } + if (pageChanged) { + changed = true; + return nextPage; + } + return page; + }); + + return changed ? { ...prev, pages } : prev; + }); + } +}; + const parseUpdatedAtInstant = (updatedAt: string) => { const match = updatedAt.match(/^(.*?)(?:\.(\d+))?(Z|[+-]\d\d:\d\d)$/); if (!match) { @@ -674,18 +787,20 @@ export const archiveChat = (queryClient: QueryClient) => ({ ); // Flip archived flag in the flat root list; strip the // chat from any parent's embedded children (individual - // child archive). + // child archive). Reuse patchChatArchiveState so the + // optimistic snapshot matches the confirmed onSuccess state, + // including the pin_order reset for an archived chat. updateInfiniteChatsCache(queryClient, (chats) => chats.map((chat) => - chat.id === chatId ? { ...chat, archived: true } : chat, + chat.id === chatId ? patchChatArchiveState(chat, true) : chat, ), ); removeChildFromParentInCache(queryClient, chatId); if (previousChat) { - queryClient.setQueryData(chatKey(chatId), { - ...previousChat, - archived: true, - }); + queryClient.setQueryData( + chatKey(chatId), + patchChatArchiveState(previousChat, true), + ); } return { previousChat }; }, @@ -707,13 +822,16 @@ export const archiveChat = (queryClient: QueryClient) => ({ ); } }, - onSettled: async (_data: unknown, _error: unknown, chatId: string) => { - await invalidateChatListQueries(queryClient); - await queryClient.invalidateQueries({ + onSuccess: (_data: unknown, chatId: string) => { + applyChatArchiveStateToCaches(queryClient, chatId, true); + }, + onSettled: (_data: unknown, _error: unknown, chatId: string) => { + void invalidateChatListQueries(queryClient); + void queryClient.invalidateQueries({ queryKey: chatKey(chatId), exact: true, }); - await queryClient.invalidateQueries({ + void queryClient.invalidateQueries({ queryKey: chatsByWorkspaceKeyPrefix, }); }, @@ -734,16 +852,18 @@ export const unarchiveChat = (queryClient: QueryClient) => ({ const previousChat = queryClient.getQueryData( chatKey(chatId), ); + // Reuse patchChatArchiveState so the optimistic snapshot + // matches the confirmed onSuccess state. updateInfiniteChatsCache(queryClient, (chats) => chats.map((chat) => - chat.id === chatId ? { ...chat, archived: false } : chat, + chat.id === chatId ? patchChatArchiveState(chat, false) : chat, ), ); if (previousChat) { - queryClient.setQueryData(chatKey(chatId), { - ...previousChat, - archived: false, - }); + queryClient.setQueryData( + chatKey(chatId), + patchChatArchiveState(previousChat, false), + ); } return { previousChat }; }, @@ -765,13 +885,16 @@ export const unarchiveChat = (queryClient: QueryClient) => ({ ); } }, - onSettled: async (_data: unknown, _error: unknown, chatId: string) => { - await invalidateChatListQueries(queryClient); - await queryClient.invalidateQueries({ + onSuccess: (_data: unknown, chatId: string) => { + applyChatArchiveStateToCaches(queryClient, chatId, false); + }, + onSettled: (_data: unknown, _error: unknown, chatId: string) => { + void invalidateChatListQueries(queryClient); + void queryClient.invalidateQueries({ queryKey: chatKey(chatId), exact: true, }); - await queryClient.invalidateQueries({ + void queryClient.invalidateQueries({ queryKey: chatsByWorkspaceKeyPrefix, }); }, @@ -1137,13 +1260,13 @@ export const updateChatTitle = (queryClient: QueryClient) => ({ ); }, - onSettled: async ( + onSettled: ( _data: unknown, _error: unknown, { chatId }: UpdateChatTitleVariables, ) => { - await invalidateChatListQueries(queryClient); - await queryClient.invalidateQueries({ + void invalidateChatListQueries(queryClient); + void queryClient.invalidateQueries({ queryKey: chatKey(chatId), exact: true, }); diff --git a/site/src/pages/AgentsPage/AgentsPage.tsx b/site/src/pages/AgentsPage/AgentsPage.tsx index fc0c8af888..e2b241e678 100644 --- a/site/src/pages/AgentsPage/AgentsPage.tsx +++ b/site/src/pages/AgentsPage/AgentsPage.tsx @@ -16,6 +16,7 @@ import { API, watchChats } from "#/api/api"; import { getErrorMessage } from "#/api/errors"; import { addChildToParentInCache, + applyChatArchiveStateToCaches, archiveChat, cancelChatListRefetches, chatDiffContentsKey, @@ -207,7 +208,8 @@ const AgentsPage: FC = () => { const archiveChatBase = archiveChat(queryClient); const archiveAgentMutation = useMutation({ ...archiveChatBase, - onSuccess: (_data, chatId) => { + onSuccess: (data, chatId) => { + archiveChatBase.onSuccess(data, chatId); clearChatErrorReason(chatId); clearPersistedSidebarTabId(chatId); clearPersistedRightPanelState(chatId); @@ -231,19 +233,20 @@ const AgentsPage: FC = () => { (id) => API.experimental.updateChat(id, { archived: true }), (id) => API.deleteWorkspace(id), ), - onSuccess: async ({ chatId }) => { + onSuccess: ({ chatId }) => { + applyChatArchiveStateToCaches(queryClient, chatId, true); clearChatErrorReason(chatId); clearPersistedSidebarTabId(chatId); clearPersistedRightPanelState(chatId); - await invalidateChatListQueries(queryClient); - await queryClient.invalidateQueries({ + void invalidateChatListQueries(queryClient); + void queryClient.invalidateQueries({ queryKey: chatKey(chatId), exact: true, }); - await queryClient.invalidateQueries({ + void queryClient.invalidateQueries({ queryKey: chatsByWorkspaceKeyPrefix, }); - await invalidateWorkspaceMutationQueries(queryClient, { + void invalidateWorkspaceMutationQueries(queryClient, { organizationName, username: user.username, });