diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 5f6e68ce70b..8fe6e17d094 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -58,11 +58,13 @@ import { promoteChatQueuedMessage, proposeChatTitle, removeChatEntity, + removeChatFromChatsByWorkspace, removeChildFromParentInCache, reorderPinnedChat, setChatGroupRole, setChatUserRole, shouldInvalidateChatSearches, + shouldInvalidateChatsByWorkspace, TERMINAL_RUN_STATUSES, toChatListParams, unarchiveChat, @@ -70,6 +72,7 @@ import { updateChatAdvisorConfig, updateChatPlanMode, updateChatTitle, + updateChatWorkspace, updateChildInParentCache, updateInfiniteChatsCache, } from "./chats"; @@ -1607,6 +1610,97 @@ describe("mutation invalidation scope", () => { "chat search entry should be invalidated", ).toBe(true); }); + + it.each<{ + name: string; + settle: (queryClient: QueryClient) => unknown; + }>([ + { + name: "archiveChat onSettled", + settle: (queryClient) => + archiveChat(queryClient).onSettled(undefined, undefined, "chat-1"), + }, + { + name: "unarchiveChat onSettled", + settle: (queryClient) => + unarchiveChat(queryClient).onSettled(undefined, undefined, "chat-1"), + }, + { + name: "updateChatWorkspace onSettled", + settle: (queryClient) => + updateChatWorkspace(queryClient).onSettled(undefined, undefined, { + chatId: "chat-1", + workspaceId: "ws-1", + }), + }, + { + name: "createChat onSuccess", + settle: (queryClient) => createChat(queryClient).onSuccess(), + }, + ])("$name invalidates chats by workspace", async ({ settle }) => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatsByWorkspace(["ws-1"]).queryKey, { + "ws-1": "chat-1", + }); + + settle(queryClient); + await new Promise((r) => setTimeout(r, 0)); + + expect( + queryClient.getQueryState(chatsByWorkspace(["ws-1"]).queryKey) + ?.isInvalidated, + "by-workspace entry should be invalidated", + ).toBe(true); + }); + + it("archiveChat onSuccess synchronously removes the chat's by-workspace mappings", () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatsByWorkspace(["ws-1"]).queryKey, { + "ws-1": "chat-1", + "ws-2": "chat-2", + }); + + archiveChat(queryClient).onSuccess(undefined, "chat-1"); + + // Assert before any timer flush: the archived chat must be gone + // from the mapping without waiting for the onSettled refetch. + expect( + queryClient.getQueryData(chatsByWorkspace(["ws-1"]).queryKey), + ).toEqual({ "ws-2": "chat-2" }); + }); + + it.each<{ + name: string; + settle: (queryClient: QueryClient) => unknown; + }>([ + { + name: "updateChatTitle onSettled", + settle: (queryClient) => + updateChatTitle(queryClient).onSettled(undefined, undefined, { + chatId: "chat-1", + title: "New", + }), + }, + { + name: "createChatMessage onSuccess", + settle: (queryClient) => + createChatMessage(queryClient, "chat-1").onSuccess?.(), + }, + ])("$name does not invalidate chats by workspace", async ({ settle }) => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatsByWorkspace(["ws-1"]).queryKey, { + "ws-1": "chat-1", + }); + + settle(queryClient); + await new Promise((r) => setTimeout(r, 0)); + + expect( + queryClient.getQueryState(chatsByWorkspace(["ws-1"]).queryKey) + ?.isInvalidated, + "by-workspace entry should NOT be invalidated", + ).not.toBe(true); + }); }); describe("chatListKey shape", () => { @@ -3088,6 +3182,26 @@ describe("semantic cache operations: prefix invalidations", () => { ).not.toBe(true); }); + describe(shouldInvalidateChatsByWorkspace.name, () => { + // created/deleted have their own watch branches; title, summary, + // diff, and context events do not move updated_at ordering. + const expectedByKind: Record = { + action_required: true, + chat_summary_change: false, + context_dirty: false, + created: false, + deleted: false, + diff_status_change: false, + status_change: true, + summary_change: false, + title_change: false, + }; + + it.each(ChatWatchEventKinds)("%s", (kind) => { + expect(shouldInvalidateChatsByWorkspace(kind)).toBe(expectedByKind[kind]); + }); + }); + it("invalidateChatDebugRuns touches the runs list and run details only", async () => { const queryClient = createTestQueryClient(); queryClient.setQueryData(chatDebugRunsKey("chat-1"), []); @@ -3291,4 +3405,57 @@ describe("semantic cache operations: removal and patching", () => { expect(queryClient.getQueryData(chatMessagesKey("chat-1"))).toBe(before); }); + + it("removeChatFromChatsByWorkspace removes only mappings pointing at the chat", () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatsByWorkspace(["ws-1"]).queryKey, { + "ws-1": "chat-1", + "ws-2": "chat-2", + }); + queryClient.setQueryData(chatsByWorkspace(["ws-3"]).queryKey, { + "ws-3": "chat-1", + }); + seedInfiniteChats(queryClient, [makeChat("chat-1")]); + queryClient.setQueryData(chatEntityKey("chat-1"), makeChat("chat-1")); + queryClient.setQueryData(chatSearch({ q: "alpha" }).queryKey, []); + + removeChatFromChatsByWorkspace(queryClient, "chat-1"); + + expect( + queryClient.getQueryData(chatsByWorkspace(["ws-1"]).queryKey), + ).toEqual({ "ws-2": "chat-2" }); + expect( + queryClient.getQueryData(chatsByWorkspace(["ws-3"]).queryKey), + ).toEqual({}); + for (const [label, key] of [ + ["chat list", infiniteChatsTestKey], + ["chat detail", chatEntityKey("chat-1")], + ["chat search", chatSearch({ q: "alpha" }).queryKey], + ] as const) { + expect( + queryClient.getQueryData(key), + `${label} entry should survive removeChatFromChatsByWorkspace`, + ).toBeDefined(); + expect( + queryClient.getQueryState(key)?.isInvalidated, + `${label} entry should NOT be invalidated`, + ).not.toBe(true); + } + }); + + it("removeChatFromChatsByWorkspace preserves the previous reference when the chat is absent", () => { + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatsByWorkspace(["ws-1"]).queryKey, { + "ws-1": "chat-2", + }); + const before = queryClient.getQueryData( + chatsByWorkspace(["ws-1"]).queryKey, + ); + + removeChatFromChatsByWorkspace(queryClient, "chat-1"); + + expect(queryClient.getQueryData(chatsByWorkspace(["ws-1"]).queryKey)).toBe( + before, + ); + }); }); diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 5a0d4b364e1..f81d109c709 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -643,6 +643,15 @@ export const invalidateChatListQueries = (queryClient: QueryClient) => queryKey: chatListFamilyKey, }); +// Event kinds that can change which chat is newest for a workspace. +const BY_WORKSPACE_AFFECTING_EVENT_KINDS = new Set( + ["status_change", "action_required"], +); + +export const shouldInvalidateChatsByWorkspace = ( + eventKind: TypesGen.ChatWatchEventKind, +): boolean => BY_WORKSPACE_AFFECTING_EVENT_KINDS.has(eventKind); + export const invalidateChatsByWorkspace = (queryClient: QueryClient) => queryClient.invalidateQueries({ queryKey: chatsByWorkspaceFamilyKey, @@ -767,6 +776,25 @@ export const removeChatEntity = (queryClient: QueryClient, chatId: string) => exact: true, }); +export const removeChatFromChatsByWorkspace = ( + queryClient: QueryClient, + chatId: string, +) => + queryClient.setQueriesData>( + { queryKey: chatsByWorkspaceFamilyKey }, + (prev) => { + if (!prev) { + return prev; + } + const next = Object.fromEntries( + Object.entries(prev).filter(([, id]) => id !== chatId), + ); + return Object.keys(next).length === Object.keys(prev).length + ? prev + : next; + }, + ); + export const patchChatEntity = ( queryClient: QueryClient, chatId: string, @@ -1008,6 +1036,7 @@ export const archiveChat = (queryClient: QueryClient) => ({ }, onSuccess: (_data: unknown, chatId: string) => { applyChatArchiveStateToCaches(queryClient, chatId, true); + removeChatFromChatsByWorkspace(queryClient, chatId); }, onSettled: (_data: unknown, _error: unknown, chatId: string) => { void invalidateChatListQueries(queryClient); diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.tsx b/site/src/pages/AgentsPage/AgentsPageLayout.tsx index fa4abedeba3..16f66ca937a 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.tsx +++ b/site/src/pages/AgentsPage/AgentsPageLayout.tsx @@ -37,9 +37,11 @@ import { proposeChatTitle, readInfiniteChatsCache, removeChatEntity, + removeChatFromChatsByWorkspace, removeChildFromParentInCache, reorderPinnedChat, shouldInvalidateChatSearches, + shouldInvalidateChatsByWorkspace, unarchiveChat, unpinChat, updateChatTitle, @@ -304,6 +306,7 @@ const AgentsPageLayout: FC = () => { ), onSuccess: ({ chatId, workspaceId, deleteBuild }) => { applyChatArchiveStateToCaches(queryClient, chatId, true); + removeChatFromChatsByWorkspace(queryClient, chatId); clearChatErrorReason(chatId); clearPersistedSidebarTabId(chatId); clearPersistedRightPanelState(chatId); @@ -619,6 +622,8 @@ const AgentsPageLayout: FC = () => { ); removeChildFromParentInCache(queryClient, updatedChat.id); removeChatEntity(queryClient, updatedChat.id); + removeChatFromChatsByWorkspace(queryClient, updatedChat.id); + void invalidateChatsByWorkspace(queryClient); void invalidateChatSearches(queryClient); return; } @@ -655,6 +660,7 @@ const AgentsPageLayout: FC = () => { } else { prependToInfiniteChatsCache(queryClient, updatedChat); void invalidateChatListQueries(queryClient); + void invalidateChatsByWorkspace(queryClient); void invalidateChatSearches(queryClient); } } else { @@ -668,6 +674,9 @@ const AgentsPageLayout: FC = () => { if (shouldInvalidateChatSearches(chatEvent.kind)) { void invalidateChatSearches(queryClient); } + if (shouldInvalidateChatsByWorkspace(chatEvent.kind)) { + void invalidateChatsByWorkspace(queryClient); + } const costChatId = chatCostIdToInvalidate( updatedChat, chatEvent.kind, @@ -690,6 +699,7 @@ const AgentsPageLayout: FC = () => { }, onOpen() { void invalidateChatListQueries(queryClient); + void invalidateChatsByWorkspace(queryClient); void invalidateChatSearches(queryClient); }, }); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.test.tsx index 4d01aa6d95a..92524039f62 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.test.tsx @@ -7,6 +7,7 @@ import { chatEntityKey, chatMessagesKey, chatPromptsKey, + chatsByWorkspace, } from "#/api/queries/chats"; import { getWorkspaceQuotaQueryKey } from "#/api/queries/workspaceQuota"; import { workspacesQueryKeyPrefix } from "#/api/queries/workspaces"; @@ -124,7 +125,7 @@ describe("useChatToolInvalidations", () => { queryKey: getWorkspaceQuotaQueryKey(ORGANIZATION_NAME, USERNAME), exact: true, }); - expect(invalidateSpy).toHaveBeenCalledTimes(3); + expect(invalidateSpy).toHaveBeenCalledTimes(4); }); }); @@ -135,6 +136,9 @@ describe("useChatToolInvalidations", () => { pageParams: [], }); queryClient.setQueryData(chatPromptsKey("chat-1"), { prompts: [] }); + queryClient.setQueryData(chatsByWorkspace(["ws-1"]).queryKey, { + "ws-1": "chat-1", + }); const { setStreamState } = renderInvalidations(); await act(async () => { @@ -147,6 +151,11 @@ describe("useChatToolInvalidations", () => { "detail entry should be invalidated", ).toBe(true); }); + expect( + queryClient.getQueryState(chatsByWorkspace(["ws-1"]).queryKey) + ?.isInvalidated, + "by-workspace entry should be invalidated", + ).toBe(true); expect( queryClient.getQueryState(chatMessagesKey("chat-1"))?.isInvalidated, "messages entry should NOT be invalidated", @@ -228,7 +237,7 @@ describe("useChatToolInvalidations", () => { predicate: expect.any(Function), }), ); - expect(invalidateSpy).toHaveBeenCalledTimes(2); + expect(invalidateSpy).toHaveBeenCalledTimes(3); }); }); @@ -251,7 +260,7 @@ describe("useChatToolInvalidations", () => { }); await waitFor(() => { - expect(invalidateSpy).toHaveBeenCalledTimes(3); + expect(invalidateSpy).toHaveBeenCalledTimes(4); }); await act(async () => { @@ -259,7 +268,7 @@ describe("useChatToolInvalidations", () => { }); await waitFor(() => { - expect(invalidateSpy).toHaveBeenCalledTimes(3); + expect(invalidateSpy).toHaveBeenCalledTimes(4); }); }); @@ -271,7 +280,7 @@ describe("useChatToolInvalidations", () => { }); await waitFor(() => { - expect(invalidateSpy).toHaveBeenCalledTimes(3); + expect(invalidateSpy).toHaveBeenCalledTimes(4); }); rerender({ @@ -285,7 +294,7 @@ describe("useChatToolInvalidations", () => { }); await waitFor(() => { - expect(invalidateSpy).toHaveBeenCalledTimes(6); + expect(invalidateSpy).toHaveBeenCalledTimes(8); expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: chatEntityKey("chat-2"), exact: true, @@ -311,7 +320,7 @@ describe("useChatToolInvalidations", () => { }); await waitFor(() => { - expect(invalidateSpy).toHaveBeenCalledTimes(3); + expect(invalidateSpy).toHaveBeenCalledTimes(4); }); }); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.ts index d4f2dd0bcca..ddf97e30394 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatToolInvalidations.ts @@ -1,6 +1,9 @@ import { useEffect, useRef } from "react"; import { useQueryClient } from "react-query"; -import { invalidateChatEntity } from "#/api/queries/chats"; +import { + invalidateChatEntity, + invalidateChatsByWorkspace, +} from "#/api/queries/chats"; import { invalidateWorkspaceMutationQueries } from "#/api/queries/workspaces"; import { type ChatStore, useChatSelector } from "./chatStore"; import type { StreamState } from "./types"; @@ -88,6 +91,7 @@ export function useChatToolInvalidations({ if (shouldInvalidateChat) { void invalidateChatEntity(queryClient, chatID); + void invalidateChatsByWorkspace(queryClient); } if (shouldInvalidateWorkspace) {