fix(site): prevent chat messages from disappearing and duplicating (#23995)

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);

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]);
+	});
+});

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<PropsWithChildren> = ({ children }) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+
+		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<typeof useChatStore>[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<PropsWithChildren> = ({ children }) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+
+		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<typeof useChatStore>[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]);
+		});
+	});
+});

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<number>();
+	return sorted
+		.map((message) => message.id)
+		.filter((id) => {
+			if (seen.has(id)) return false;
+			seen.add(id);
+			return true;
+		});
 };
 
 const mapsEqualByRef = <K, V>(left: Map<K, V>, right: Map<K, V>): boolean => {

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<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),
+					],
+				};
+			});
+		},
+	);
+
 	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<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),
-					],
-				};
-			});
-		};
-
 		store.resetTransientState();
 		activeChatIDRef.current = chatID ?? null;
 
@@ -652,11 +667,13 @@ export const useChatStore = (
 		store,
 		setChatErrorReasonStable,
 		clearChatErrorReasonStable,
+		upsertCacheMessages,
 	]);
 	return {
 		store,
 		clearStreamError: () => {
 			store.clearStreamError();
 		},
+		upsertCacheMessages,
 	};
 };

site/src/pages/AgentsPage/components/ChatPageContent.tsx

@@ -67,7 +67,17 @@ export const ChatPageTimeline: FC<ChatPageTimelineProps> = ({
 	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<ChatPageInputProps> = ({
 	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--) {