From beb99c17de2864974620782e1b68c4a6f6fad135 Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Tue, 7 Apr 2026 11:05:40 +0100 Subject: [PATCH] fix(site): prevent chat messages from disappearing and duplicating (#23995) --- site/src/pages/AgentsPage/AgentChatPage.tsx | 22 ++- .../chatStore.createStore.test.ts | 18 ++ .../ChatConversation/chatStore.test.tsx | 175 ++++++++++++++++++ .../components/ChatConversation/chatStore.ts | 14 +- .../ChatConversation/useChatStore.ts | 117 +++++++----- .../AgentsPage/components/ChatPageContent.tsx | 24 ++- 6 files changed, 310 insertions(+), 60 deletions(-) diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index ddf2702c00..1be0dff9d3 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -616,11 +616,17 @@ const AgentChatPage: FC = () => { const chatMessagesList = (() => { const pages = chatMessagesQuery.data?.pages; if (!pages || pages.length === 0) return undefined; - // Collect all messages, then sort chronologically by ID. + // Collect all messages and deduplicate by ID. + // Cross-page duplication can occur when upsertCacheMessages + // writes a message into page 0 while the same ID still + // exists in a later page. Last occurrence wins so the + // most up-to-date content is preserved. const all = pages.flatMap((p) => p.messages); + const byID = new Map(all.map((m) => [m.id, m])); + const deduped = Array.from(byID.values()); // Sort ascending by ID for chronological order. - all.sort((a, b) => a.id - b.id); - return all; + deduped.sort((a, b) => a.id - b.id); + return deduped; })(); // Queued messages are only in the first page (most recent). @@ -654,7 +660,7 @@ const AgentChatPage: FC = () => { promoteChatQueuedMessage(queryClient, agentId ?? ""), ); - const { store, clearStreamError } = useChatStore({ + const { store, clearStreamError, upsertCacheMessages } = useChatStore({ chatID: agentId, chatMessages: chatMessagesList, chatRecord, @@ -891,6 +897,7 @@ const AgentChatPage: FC = () => { store.setChatStatus("running"); if (response.message) { store.upsertDurableMessage(response.message); + upsertCacheMessages([response.message]); } } if (selectedModelConfigID) { @@ -935,10 +942,11 @@ const AgentChatPage: FC = () => { store.setChatStatus("pending"); try { const promotedMessage = await promoteQueuedMutation.mutateAsync(id); - // Insert the promoted message into the store immediately - // so it appears in the timeline without waiting for the - // WebSocket to deliver it. + // Insert the promoted message into the store and cache + // immediately so it appears in the timeline without + // waiting for the WebSocket to deliver it. store.upsertDurableMessage(promotedMessage); + upsertCacheMessages([promotedMessage]); } catch (error) { store.setQueuedMessages(previousQueuedMessages); store.setChatStatus(previousChatStatus); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts index 738eb4f49f..cd4504178e 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts @@ -701,3 +701,21 @@ describe("selectIsAwaitingFirstStreamChunk", () => { expect(selectIsAwaitingFirstStreamChunk(store.getSnapshot())).toBe(true); }); }); + +describe("duplicate message deduplication", () => { + it("replaceMessages deduplicates orderedMessageIDs when input has duplicate IDs", () => { + const store = createChatStore(); + const msg1 = makeMessage(1, "user", "hello"); + const msg2 = makeMessage(2, "assistant", "hi"); + // Simulate cross-page duplication: same ID appears twice. + const msg2Copy = makeMessage(2, "assistant", "hi"); + + store.replaceMessages([msg1, msg2, msg2Copy]); + + const state = store.getSnapshot(); + // Map deduplicates by key — only 2 unique entries. + expect(state.messagesByID.size).toBe(2); + // orderedMessageIDs MUST also have only 2 entries. + expect(state.orderedMessageIDs).toEqual([1, 2]); + }); +}); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index 3a53f6f5e3..c6fd57230b 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx @@ -4033,3 +4033,178 @@ describe("partsBuf cleanup on reconnect (Bug 2)", () => { }); }); }); + +describe("store/cache desync protection", () => { + it("does not wipe a message added via upsertDurableMessage when a genuine refetch follows", async () => { + // RED TEST: Simulates what handleSend does today — it calls + // store.upsertDurableMessage without writing to the React + // Query cache. A subsequent rerender with new message refs + // (genuine refetch) should NOT wipe the store-only message. + immediateAnimationFrame(); + + const chatID = "chat-send-desync"; + const msg1 = makeMessage(chatID, 1, "user", "hello"); + const msg2 = makeMessage(chatID, 2, "assistant", "hi"); + const msg3 = makeMessage(chatID, 3, "user", "follow-up"); + + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatMessagesKey(chatID), { + pages: [ + { + messages: [msg2, msg1], + queued_messages: [], + has_more: false, + }, + ], + pageParams: [undefined], + }); + const wrapper: FC = ({ children }) => ( + {children} + ); + + const initialMessages = [msg1, msg2]; + const initialOptions = { + chatID, + chatMessages: initialMessages, + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: initialMessages, + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [] as TypesGen.ChatQueuedMessage[], + setChatErrorReason: vi.fn(), + clearChatErrorReason: vi.fn(), + }; + + const { result, rerender } = renderHook( + (options: Parameters[0]) => { + const { store } = useChatStore(options); + return { + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + store, + }; + }, + { initialProps: initialOptions, wrapper }, + ); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2]); + }); + + act(() => { + mockSocket.emitOpen(); + }); + + // Simulate handleSend: write to store only, no cache write. + act(() => { + result.current.store.upsertDurableMessage(msg3); + }); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + }); + + // Genuine refetch: new object refs for msg1 and msg2, + // msg3 absent from the fetched set. + const msg1New = makeMessage(chatID, 1, "user", "hello"); + const msg2New = makeMessage(chatID, 2, "assistant", "hi"); + rerender({ + ...initialOptions, + chatMessages: [msg1New, msg2New], + chatMessagesData: { + messages: [msg1New, msg2New], + queued_messages: [], + has_more: false, + }, + }); + + // msg3 was added to the store AFTER the last sync. It + // should NOT be classified as stale — it's new, not + // something the server removed. + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + }); + }); + + it("still removes messages that were in the previous sync but are absent from a refetch (edit truncation)", async () => { + immediateAnimationFrame(); + + const chatID = "chat-edit-truncation"; + const msg1 = makeMessage(chatID, 1, "user", "hello"); + const msg2 = makeMessage(chatID, 2, "assistant", "hi"); + const msg3 = makeMessage(chatID, 3, "user", "more"); + + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + queryClient.setQueryData(chatMessagesKey(chatID), { + pages: [ + { + messages: [msg3, msg2, msg1], + queued_messages: [], + has_more: false, + }, + ], + pageParams: [undefined], + }); + const wrapper: FC = ({ children }) => ( + {children} + ); + + const initialOptions = { + chatID, + chatMessages: [msg1, msg2, msg3], + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: [msg1, msg2, msg3], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [] as TypesGen.ChatQueuedMessage[], + setChatErrorReason: vi.fn(), + clearChatErrorReason: vi.fn(), + }; + + const { result, rerender } = renderHook( + (options: Parameters[0]) => { + const { store } = useChatStore(options); + return { + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { initialProps: initialOptions, wrapper }, + ); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + }); + + act(() => { + mockSocket.emitOpen(); + }); + + // Simulate edit truncation: rerender with only msg1. + const msg1New = makeMessage(chatID, 1, "user", "hello"); + rerender({ + ...initialOptions, + chatMessages: [msg1New], + chatMessagesData: { + messages: [msg1New], + queued_messages: [], + has_more: false, + }, + }); + + // msg2 and msg3 WERE in the previous sync data and are + // now absent — they are genuinely stale (edit truncation) + // and should be removed. + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1]); + }); + }); +}); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts index 6619fc1224..7c71ddf8dd 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts @@ -26,7 +26,19 @@ const buildOrderedMessageIDs = ( ): readonly number[] => { const sorted = [...messages]; sorted.sort(byMessageCreatedAt); - return sorted.map((message) => message.id); + // Deduplicate by ID. The input can contain duplicate IDs when + // cross-page duplication occurs in the React Query cache (e.g. + // upsertCacheMessages writes to page 0 while the same message + // still exists in a later page). The Map-based messagesByID + // already deduplicates, but orderedMessageIDs must match. + const seen = new Set(); + return sorted + .map((message) => message.id) + .filter((id) => { + if (seen.has(id)) return false; + seen.add(id); + return true; + }); }; const mapsEqualByRef = (left: Map, right: Map): boolean => { diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index 58ec521da2..a289fc9650 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts @@ -81,7 +81,11 @@ interface UseChatStoreOptions { export const useChatStore = ( options: UseChatStoreOptions, -): { store: ChatStore; clearStreamError: () => void } => { +): { + store: ChatStore; + clearStreamError: () => void; + upsertCacheMessages: (messages: readonly TypesGen.ChatMessage[]) => void; +} => { const { chatID, chatMessages, @@ -150,6 +154,56 @@ export const useChatStore = ( // its snapshot, defeating pagination. const initialDataLoaded = chatMessages !== undefined; + // 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 = useEffectEvent( + (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), + ], + }; + }); + }, + ); + useEffect(() => { store.batch(() => { // When the active chat changes, clear stale messages @@ -180,9 +234,18 @@ export const useChatStore = ( const storeSnap = store.getSnapshot(); const fetchedIDs = new Set(chatMessages.map((m) => m.id)); + // Only classify a store-held ID as stale if it was + // present in the PREVIOUS sync's fetched data. IDs + // added to the store after the last sync (by the WS + // handler or handleSend) are new, not stale, and + // must not trigger the destructive replaceMessages + // path. + const prevIDs = new Set(prev.map((m) => m.id)); const hasStaleEntries = contentChanged && - storeSnap.orderedMessageIDs.some((id) => !fetchedIDs.has(id)); + storeSnap.orderedMessageIDs.some( + (id) => !fetchedIDs.has(id) && prevIDs.has(id), + ); if (hasStaleEntries) { store.replaceMessages(chatMessages); } else { @@ -286,54 +349,6 @@ 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), - ], - }; - }); - }; - store.resetTransientState(); activeChatIDRef.current = chatID ?? null; @@ -652,11 +667,13 @@ export const useChatStore = ( store, setChatErrorReasonStable, clearChatErrorReasonStable, + upsertCacheMessages, ]); return { store, clearStreamError: () => { store.clearStreamError(); }, + upsertCacheMessages, }; }; diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.tsx index 95058771ed..760e647c5f 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.tsx +++ b/site/src/pages/AgentsPage/components/ChatPageContent.tsx @@ -67,7 +67,17 @@ export const ChatPageTimeline: FC = ({ const orderedMessageIDs = useChatSelector(store, selectOrderedMessageIDs); const messages = orderedMessageIDs - .map((messageID) => messagesByID.get(messageID)) + .map((messageID) => { + const message = messagesByID.get(messageID); + if (!message && process.env.NODE_ENV !== "production") { + console.warn( + `[ChatPageContent] orderedMessageIDs contains ID ${messageID} ` + + "not found in messagesByID. This may indicate a store/cache " + + "desync bug.", + ); + } + return message; + }) .filter(isChatMessage); const parsedMessages = parseMessagesWithMergedTools(messages); const subagentTitles = buildSubagentTitles(parsedMessages); @@ -205,7 +215,17 @@ export const ChatPageInput: FC = ({ const queuedMessages = useChatSelector(store, selectQueuedMessages); const messages = orderedMessageIDs - .map((messageID) => messagesByID.get(messageID)) + .map((messageID) => { + const message = messagesByID.get(messageID); + if (!message && process.env.NODE_ENV !== "production") { + console.warn( + `[ChatPageContent] orderedMessageIDs contains ID ${messageID} ` + + "not found in messagesByID. This may indicate a store/cache " + + "desync bug.", + ); + } + return message; + }) .filter(isChatMessage); let lastEditableUserMessage: TypesGen.ChatMessage | undefined; for (let index = orderedMessageIDs.length - 1; index >= 0; index--) {