From 3f8e3007d8098768154251493d5f4ad045e46bd9 Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Mon, 30 Mar 2026 15:56:08 +0100 Subject: [PATCH] fix(site): write WebSocket messages to React Query cache (#23618) --- .../ChatConversation/chatStore.test.tsx | 106 ++++++++++++++++++ .../components/ChatConversation/chatStore.ts | 2 +- .../ChatConversation/useChatStore.ts | 51 +++++++++ 3 files changed, 158 insertions(+), 1 deletion(-) diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index 59a46c953c..692bb3543b 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/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) => ( + {children} + ); + 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(); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts index 978ffd87b6..0558d68245 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts +++ b/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 => diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index 93117ccc59..cac0a21dab 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/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 | 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) {