fix(site): write WebSocket messages to React Query cache (#23618)

site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx

@@ -963,6 +963,112 @@ describe("useChatStore", () => {
 		expect(cachedData?.pages[0]?.queued_messages).toEqual([]);
 	});
 
+	it("writes WebSocket message events into the chat query cache", async () => {
+		const chatID = "chat-1";
+		const existingMessage = makeMessage(chatID, 1, "user", "hello");
+		const mockSocket = createMockSocket();
+		mockWatchChatReturn(mockSocket);
+
+		const queryClient = new QueryClient({
+			defaultOptions: {
+				queries: {
+					retry: false,
+					gcTime: Number.POSITIVE_INFINITY,
+					refetchOnWindowFocus: false,
+					networkMode: "offlineFirst",
+				},
+			},
+		});
+		const initialChatMessagesData: TypesGen.ChatMessagesResponse = {
+			messages: [existingMessage],
+			queued_messages: [],
+			has_more: false,
+		};
+		queryClient.setQueryData(chatMessagesKey(chatID), {
+			pages: [initialChatMessagesData],
+			pageParams: [undefined],
+		});
+
+		const wrapper = ({ children }: PropsWithChildren) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+		const setChatErrorReason = vi.fn();
+		const clearChatErrorReason = vi.fn();
+
+		const { result } = renderHook(
+			() => {
+				const { store } = useChatStore({
+					chatID,
+					chatMessages: [existingMessage],
+					chatRecord: makeChat(chatID),
+					chatMessagesData: initialChatMessagesData,
+					chatQueuedMessages: [],
+					setChatErrorReason,
+					clearChatErrorReason,
+				});
+				return {
+					orderedIDs: useChatSelector(store, selectOrderedMessageIDs),
+				};
+			},
+			{ wrapper },
+		);
+
+		await waitFor(() => {
+			expect(watchChat).toHaveBeenCalledWith(chatID, 1);
+		});
+
+		const newMessage = makeMessage(chatID, 2, "assistant", "hi there");
+		act(() => {
+			mockSocket.emitData({
+				type: "message",
+				chat_id: chatID,
+				message: newMessage,
+			});
+		});
+
+		await waitFor(() => {
+			expect(result.current.orderedIDs).toContain(2);
+		});
+
+		// The React Query cache should also contain the new message.
+		const cachedData = queryClient.getQueryData<{
+			pages: TypesGen.ChatMessagesResponse[];
+			pageParams: unknown[];
+		}>(chatMessagesKey(chatID));
+		const cachedMessages = cachedData?.pages[0]?.messages ?? [];
+		// Verifies insertion, preservation, and DESC order.
+		expect(cachedMessages.map((m) => m.id)).toEqual([2, 1]);
+		// Emitting the same message again should not change the
+		// cache reference (reference stability).
+		const refBefore = queryClient.getQueryData(chatMessagesKey(chatID));
+		act(() => {
+			mockSocket.emitData({
+				type: "message",
+				chat_id: chatID,
+				message: newMessage,
+			});
+		});
+		const refAfter = queryClient.getQueryData(chatMessagesKey(chatID));
+		expect(refAfter).toBe(refBefore);
+
+		// Emitting the same message ID with different content should
+		// update the cached entry (content-update path).
+		const revised = makeMessage(chatID, 2, "assistant", "revised");
+		act(() => {
+			mockSocket.emitData({
+				type: "message",
+				chat_id: chatID,
+				message: revised,
+			});
+		});
+		const updatedCache = queryClient.getQueryData<{
+			pages: TypesGen.ChatMessagesResponse[];
+			pageParams: unknown[];
+		}>(chatMessagesKey(chatID));
+		const updatedFirst = updatedCache?.pages[0]?.messages[0];
+		expect(updatedFirst?.content).toEqual([{ type: "text", text: "revised" }]);
+	});
+
 	it("closes old WebSocket and resets state when chatID changes", async () => {
 		immediateAnimationFrame();
 

site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts

@@ -64,7 +64,7 @@ const jsonValuesEqual = (left: unknown, right: unknown): boolean => {
 	}
 };
 
-const chatMessagesEqualByValue = (
+export const chatMessagesEqualByValue = (
 	left: TypesGen.ChatMessage,
 	right: TypesGen.ChatMessage,
 ): boolean =>

site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts

@@ -11,6 +11,7 @@ import { asNumber, asString } from "../ChatElements/runtimeTypeUtils";
 import {
 	type ChatStore,
 	type ChatStoreState,
+	chatMessagesEqualByValue,
 	chatQueuedMessagesEqualByID,
 	createChatStore,
 	isActiveChatStatus,
@@ -286,6 +287,55 @@ export const useChatStore = (
 				};
 			});
 		};
+
+		// Write WebSocket-delivered durable messages into the React
+		// Query infinite cache so that navigating away and back
+		// serves up-to-date data instead of the stale REST snapshot.
+		// Without this, the cache only contains messages from the
+		// last REST fetch, and structural sharing can suppress the
+		// refetch-driven store update when no new durable messages
+		// have been committed to the DB yet.
+		const upsertCacheMessages = (messages: readonly TypesGen.ChatMessage[]) => {
+			if (!chatID || messages.length === 0) {
+				return;
+			}
+			queryClient.setQueryData<
+				InfiniteData<TypesGen.ChatMessagesResponse> | undefined
+			>(chatMessagesKey(chatID), (currentData) => {
+				if (!currentData?.pages?.length) {
+					return currentData;
+				}
+				const firstPage = currentData.pages[0];
+				const existingByID = new Map(firstPage.messages.map((m) => [m.id, m]));
+
+				let changed = false;
+				for (const msg of messages) {
+					const existing = existingByID.get(msg.id);
+					if (!existing || !chatMessagesEqualByValue(existing, msg)) {
+						changed = true;
+						existingByID.set(msg.id, msg);
+					}
+				}
+
+				if (!changed) {
+					return currentData;
+				}
+
+				// Sort descending to match the API page order
+				// (newest first).
+				const updatedMessages = Array.from(existingByID.values());
+				updatedMessages.sort((a, b) => b.id - a.id);
+
+				return {
+					...currentData,
+					pages: [
+						{ ...firstPage, messages: updatedMessages },
+						...currentData.pages.slice(1),
+					],
+				};
+			});
+		};
+
 		cancelScheduledStreamReset();
 		store.resetTransientState();
 		activeChatIDRef.current = chatID ?? null;
@@ -499,6 +549,7 @@ export const useChatStore = (
 				// pass: one Map copy + one sort instead of N each.
 				if (pendingMessages.length > 0) {
 					store.upsertDurableMessages(pendingMessages);
+					upsertCacheMessages(pendingMessages);
 				}
 			});
 			if (needsStreamReset) {