fix(site/src): restart first chat list fetches after a move and keep the drop story reachable

site/src/api/queries/chats.test.ts

@@ -106,6 +106,7 @@ import {
 	pinChat,
 	planModeFieldsForCreateMessage,
 	prependToInfiniteChatsCache,
+	projectChats,
 	projectChatsKey,
 	promoteChatQueuedMessage,
 	proposeChatTitle,
@@ -115,6 +116,7 @@ import {
 	reorderPinnedChat,
 	replaceChatMessagesHistory,
 	resetUnloadedChatEntity,
+	resetUnloadedChatLists,
 	setChatGroupRole,
 	setChatUserRole,
 	shouldInvalidateChatSearches,
@@ -268,6 +270,34 @@ const observeChatWithDeferredFirstFetch = (
 	};
 };
 
+/** resetQueries refetches only observed queries. */
+const observeChatLists = (
+	queryClient: QueryClient,
+	options: { staleTime?: number } = {},
+) => {
+	const observers = [
+		new InfiniteQueryObserver(queryClient, { ...infiniteChats(), ...options }),
+		new InfiniteQueryObserver(queryClient, {
+			...projectChats("project-1"),
+			...options,
+		}),
+	];
+	const unsubscribes = observers.map((observer) =>
+		observer.subscribe(() => {}),
+	);
+	return {
+		chats: () =>
+			observers.map((observer) =>
+				observer.getCurrentResult().data?.pages.flat(),
+			),
+		unsubscribe: () => {
+			for (const unsubscribe of unsubscribes) {
+				unsubscribe();
+			}
+		},
+	};
+};
+
 describe("chat model query factories", () => {
 	const organizationId = "organization-1";
 	const otherOrganizationId = "organization-2";
@@ -4895,6 +4925,62 @@ describe("semantic cache operations: cancellation", () => {
 		expect(resetSpy).not.toHaveBeenCalled();
 	});
 
+	it("resetUnloadedChatLists restarts first fetches of sidebar and project lists", async () => {
+		const queryClient = createTestQueryClient();
+		const staleChat = makeChat("chat-1");
+		const movedChat = makeChat("chat-1", { project_id: "project-1" });
+		const firstFetches = [
+			createDeferred<TypesGen.Chat[]>(),
+			createDeferred<TypesGen.Chat[]>(),
+		];
+		const getChats = vi
+			.mocked(API.experimental.getChats)
+			.mockReturnValueOnce(firstFetches[0].promise)
+			.mockReturnValueOnce(firstFetches[1].promise)
+			.mockResolvedValueOnce([movedChat])
+			.mockResolvedValueOnce([movedChat]);
+		const lists = observeChatLists(queryClient);
+		try {
+			expect(getChats).toHaveBeenCalledTimes(2);
+
+			const reset = resetUnloadedChatLists(queryClient);
+			// Both first responses read the lists before the move.
+			for (const firstFetch of firstFetches) {
+				firstFetch.resolve([staleChat]);
+			}
+			await reset;
+
+			expect(getChats).toHaveBeenCalledTimes(4);
+			expect(lists.chats()).toEqual([[movedChat], [movedChat]]);
+		} finally {
+			lists.unsubscribe();
+		}
+	});
+
+	it("resetUnloadedChatLists keeps loaded lists and sends no request", async () => {
+		const queryClient = createTestQueryClient();
+		// Answers any request, so a refetch fails below instead of timing out.
+		const getChats = vi.mocked(API.experimental.getChats).mockResolvedValue([]);
+		const loadedChats = [makeChat("chat-1")];
+		seedInfiniteChats(queryClient, loadedChats);
+		queryClient.setQueryData<InfiniteData>(projectChatsKey("project-1"), {
+			pages: [loadedChats],
+			pageParams: [0],
+		});
+		// Fresh data keeps the observers from fetching on mount.
+		const lists = observeChatLists(queryClient, {
+			staleTime: Number.POSITIVE_INFINITY,
+		});
+		try {
+			await resetUnloadedChatLists(queryClient);
+
+			expect(getChats).not.toHaveBeenCalled();
+			expect(lists.chats()).toEqual([loadedChats, loadedChats]);
+		} finally {
+			lists.unsubscribe();
+		}
+	});
+
 	it("cancelChatMessages cancels the exact messages entry", async () => {
 		const queryClient = createTestQueryClient();
 		const cancelSpy = vi.spyOn(queryClient, "cancelQueries");

site/src/api/queries/chats.ts

@@ -1136,6 +1136,17 @@ export const resetUnloadedChatEntity = (
 	});
 };
 
+/** As resetUnloadedChatEntity, for every chat list that has not loaded. */
+export const resetUnloadedChatLists = (queryClient: QueryClient) =>
+	Promise.all(
+		[chatListFamilyKey, chatProjectListFamilyKey].map((queryKey) =>
+			queryClient.resetQueries({
+				queryKey,
+				predicate: (query) => query.state.data === undefined,
+			}),
+		),
+	);
+
 export const cancelChatMessages = (queryClient: QueryClient, chatId: string) =>
 	queryClient.cancelQueries({
 		queryKey: chatMessagesKey(chatId),

site/src/pages/AgentsPage/AgentsPageLayout.test.tsx

@@ -943,4 +943,31 @@ describe("AgentsPageLayout project moves", () => {
 
 		await expectOpenChatProject(queryClient, "project-b");
 	});
+
+	it("ends on the moved project when a move overtakes the first chat list fetch", async () => {
+		const firstFetch = createDeferred<Chat[]>();
+		vi.mocked(API.experimental.getChats).mockReturnValueOnce(
+			firstFetch.promise,
+		);
+		const { queryClient } = renderLayout();
+		await waitFor(() =>
+			expect(API.experimental.getChats).toHaveBeenCalledTimes(1),
+		);
+
+		commitMove(mockChatInProjectB);
+		publishChatWatchEvent({ kind: "project_change", chat: mockChatInProjectB });
+		// The first fetch read the list before the move committed.
+		await act(async () => firstFetch.resolve([mockChatInProjectA]));
+
+		await waitFor(() =>
+			expect(
+				queryClient
+					.getQueryData<InfiniteData<Chat[]>>(
+						infiniteChats({ sources: DEFAULT_AGENT_SIDEBAR_FILTERS.sources })
+							.queryKey,
+					)
+					?.pages.flat(),
+			).toEqual([mockChatInProjectB]),
+		);
+	});
 });

site/src/pages/AgentsPage/AgentsPageLayout.tsx

@@ -35,6 +35,7 @@ import {
 	prependToInfiniteChatsCache,
 	proposeChatTitle,
 	resetUnloadedChatEntity,
+	resetUnloadedChatLists,
 	shouldInvalidateChatSearches,
 	shouldInvalidateChatsByWorkspace,
 	updateChatTitle,
@@ -439,6 +440,11 @@ const AgentsPageLayout: React.FC = () => {
 							// and restart a first fetch that may predate the move.
 							void resetUnloadedChatEntity(queryClient, updatedChat.id);
 							void invalidateChatEntity(queryClient, updatedChat.id);
+							// First list fetches may predate the move too. Every move
+							// sends its root's event, so one reset per move is enough.
+							if (!updatedChat.parent_chat_id) {
+								void resetUnloadedChatLists(queryClient);
+							}
 						}
 					}
 				});

site/src/pages/AgentsPage/components/ChatsSidebar/ChatsSidebar.stories.tsx

@@ -1,5 +1,6 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
 import { useEffect, useState } from "react";
+import { useInfiniteQuery } from "react-query";
 import { useLocation } from "react-router";
 import {
 	expect,
@@ -14,7 +15,12 @@ import {
 import { reactRouterParameters } from "storybook-addon-remix-react-router";
 import { API } from "#/api/api";
 import { chatProjectsKey } from "#/api/queries/chatProjects";
-import { userChatProviderConfigsKey } from "#/api/queries/chats";
+import {
+	chatListKey,
+	infiniteChats,
+	toChatListParams,
+	userChatProviderConfigsKey,
+} from "#/api/queries/chats";
 import type * as TypesGen from "#/api/typesGenerated";
 import type { Chat } from "#/api/typesGenerated";
 import { MockChat, MockChatDiffStatus } from "#/testHelpers/chatEntities";
@@ -2941,16 +2947,33 @@ export const ChatDraggedOverProjectFolder: Story = {
 	},
 };
 
-/** The args are static, so the dropped chat stays under Chats. */
+/** Reads the cached chat list so the drop moves the chat. */
 export const ChatDroppedOnProjectFolder: Story = {
-	args: { chats: mockProjectChats },
 	parameters: {
 		experiments: ["chat-projects"],
-		queries: [{ key: chatProjectsKey, data: [MockChatProject] }],
+		queries: [
+			{ key: chatProjectsKey, data: [MockChatProject] },
+			{
+				key: chatListKey(
+					toChatListParams({ sources: defaultSidebarFilters.sources }),
+				),
+				data: { pages: [mockProjectChats], pageParams: [0] },
+			},
+		],
 	},
 	beforeEach: () => {
 		spyOn(API.experimental, "updateChat").mockResolvedValue(undefined);
 	},
+	render: function ChatDroppedOnProjectFolderRender(args) {
+		// Stories have no API server, so the list must not fetch.
+		const chatsQuery = useInfiniteQuery({
+			...infiniteChats({ sources: defaultSidebarFilters.sources }),
+			enabled: false,
+		});
+		return (
+			<ChatsSidebar {...args} chats={chatsQuery.data?.pages.flat() ?? []} />
+		);
+	},
 	play: async ({ canvasElement }) => {
 		await startKeyboardDrag(canvasElement, /Loose chat/);
 		await userEvent.keyboard("[ArrowUp]");