fix(agents): write WebSocket cache updates to infinite query key (#22905)

## Problem

Chat sidebar title/status updates from WebSocket events don't take
effect immediately — they only appear after a full server re-fetch.

**Root cause:** All `setQueryData(chatsKey, ...)` calls write to cache
key `["chats"]`, but the rendered chat list reads from
`useInfiniteQuery(infiniteChats())` on key `["chats", undefined]`.
TanStack Query v5 `setQueryData` requires an exact key match, so these
are different cache entries.

WebSocket events (`title_change`, `status_change`, `created`, `deleted`)
and `updateSidebarChat` were all updating a cache entry that nothing
rendered from. The only way changes reached the UI was via
`invalidateQueries` (which prefix-matches), triggering a full server
re-fetch. This caused visible flicker when the re-fetch raced with
subsequent events.

## Fix

Add `updateInfiniteChatsCache()` helper that uses `setQueriesData({
queryKey: chatsKey })` — this **prefix-matches** all infinite query
variants (`["chats", undefined]`, `["chats", { archived: true }]`, etc.)
and correctly updates the `{ pages, pageParams }` structure.

Replace all direct `setQueryData(chatsKey, ...)` calls:
- WebSocket handler in `AgentsPage.tsx` (deleted, created, title_change,
status_change events)
- `updateSidebarChat` in `ChatContext.ts`
- Archive/unarchive optimistic updates in `chats.ts`

Also adds `readInfiniteChatsCache()` helper for reading the flat chat
list from the infinite query (used by the chime status lookup).

## Files changed

| File | Change |
|------|--------|
| `site/src/api/queries/chats.ts` | Added helpers, updated
archive/unarchive mutations |
| `site/src/pages/AgentsPage/AgentsPage.tsx` | WebSocket handler uses
new helpers |
| `site/src/pages/AgentsPage/AgentDetail/ChatContext.ts` |
`updateSidebarChat` uses new helper |
| `site/src/api/queries/chats.test.ts` | Tests seed/read infinite query
format |
| `site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx` | Tests
seed/read infinite query format |

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<InfiniteData>(infiniteChatsTestKey, {
+		pages: [chats],
+		pageParams: [0],
+	});
+};
+
+/** Read chats back from the infinite query cache. */
+const readInfiniteChats = (
+	queryClient: QueryClient,
+): TypesGen.Chat[] | undefined => {
+	const data = queryClient.getQueryData<InfiniteData>(infiniteChatsTestKey);
+	return data?.pages.flat();
+};
+
 const makeChat = (
 	id: string,
 	overrides?: Partial<TypesGen.Chat>,
@@ -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<TypesGen.Chat[]>(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<TypesGen.Chat[]>(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<TypesGen.Chat[]>(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<TypesGen.Chat[]>(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<TypesGen.Chat[]>(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<TypesGen.Chat[]>(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<TypesGen.Chat[]>(chatsKey)?.[0].archived,
-		).toBe(false);
+		expect(readInfiniteChats(queryClient)?.[0].archived).toBe(false);
 		expect(
 			queryClient.getQueryData<TypesGen.ChatWithMessages>(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<TypesGen.Chat[]>(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<TypesGen.ChatWithMessages>(chatKey(chatId))?.chat
 				.archived,

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<TypesGen.Chat[]>(chatsKey);
 		const previousChat = queryClient.getQueryData<TypesGen.ChatWithMessages>(
 			chatKey(chatId),
 		);
-		queryClient.setQueryData<TypesGen.Chat[]>(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<TypesGen.Chat[]>(
-				chatsKey,
-				context.previousChats,
-			);
-		}
+		// Rollback: invalidate to re-fetch the correct state.
+		void queryClient.invalidateQueries({ queryKey: chatsKey });
 		if (context?.previousChat) {
 			queryClient.setQueryData<TypesGen.ChatWithMessages>(
 				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<TypesGen.Chat[]>(chatsKey);
 		const previousChat = queryClient.getQueryData<TypesGen.ChatWithMessages>(
 			chatKey(chatId),
 		);
-		queryClient.setQueryData<TypesGen.Chat[]>(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<TypesGen.Chat[]>(
-				chatsKey,
-				context.previousChats,
-			);
-		}
+		// Rollback: invalidate to re-fetch the correct state.
+		void queryClient.invalidateQueries({ queryKey: chatsKey });
 		if (context?.previousChat) {
 			queryClient.setQueryData<TypesGen.ChatWithMessages>(
 				chatKey(chatId),

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<InfiniteData>(infiniteChatsTestKey, {
+		pages: [chats],
+		pageParams: [0],
+	});
+};
+
+/** Read chats back from the infinite query cache. */
+const readInfiniteChats = (
+	queryClient: QueryClient,
+): TypesGen.Chat[] | undefined => {
+	const data = queryClient.getQueryData<InfiniteData>(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) => (
 			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
@@ -2229,7 +2258,7 @@ describe("updateSidebarChat via stream events", () => {
 		});
 
 		await waitFor(() => {
-			const sidebarChats = queryClient.getQueryData<TypesGen.Chat[]>(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) => (
 			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
@@ -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<TypesGen.Chat[]>(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) => (
 			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
@@ -2366,7 +2395,7 @@ describe("updateSidebarChat via stream events", () => {
 		});
 
 		await waitFor(() => {
-			const sidebarChats = queryClient.getQueryData<TypesGen.Chat[]>(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) => (
 			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
@@ -2432,14 +2461,14 @@ describe("updateSidebarChat via stream events", () => {
 		});
 
 		await waitFor(() => {
-			const sidebarChats = queryClient.getQueryData<TypesGen.Chat[]>(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<TypesGen.Chat[]>(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) => (
 			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
@@ -2512,7 +2541,7 @@ describe("updateSidebarChat via stream events", () => {
 		});
 
 		await waitFor(() => {
-			const sidebarChats = queryClient.getQueryData<TypesGen.Chat[]>(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) => (
 			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
@@ -2576,7 +2605,7 @@ describe("updateSidebarChat via stream events", () => {
 		});
 
 		await waitFor(() => {
-			const sidebarChats = queryClient.getQueryData<TypesGen.Chat[]>(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) => (
 			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
@@ -2642,7 +2671,7 @@ describe("updateSidebarChat via stream events", () => {
 		});
 
 		await waitFor(() => {
-			const sidebarChats = queryClient.getQueryData<TypesGen.Chat[]>(chatsKey);
+			const sidebarChats = readInfiniteChats(queryClient);
 			expect(sidebarChats?.[0].status).toBe("error");
 			expect(sidebarChats?.[0].updated_at).toBe(initialChat.updated_at);
 		});

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<readonly TypesGen.Chat[] | undefined>(
-				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],
 	);

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<TypesGen.Chat[]>(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<TypesGen.ChatWithMessages | undefined>(
 						chatKey(updatedChat.id),
 						(previousChat) => {