fix(site/src): fan chat message upserts out to every containing page (#27912)

site/src/api/queries/chats.test.ts

@@ -7,6 +7,7 @@ import {
 	ERROR_STATUSES,
 	SUCCESS_STATUSES,
 } from "#/pages/AgentsPage/components/RightPanel/DebugPanel/debugPanelUtils";
+import { MockChatMessage } from "#/testHelpers/chatEntities";
 import { buildOptimisticEditedMessage } from "./chatMessageEdits";
 import {
 	addChildToParentInCache,
@@ -61,6 +62,7 @@ import {
 	removeChatFromChatsByWorkspace,
 	removeChildFromParentInCache,
 	reorderPinnedChat,
+	replaceChatMessagesHistory,
 	setChatGroupRole,
 	setChatUserRole,
 	shouldInvalidateChatSearches,
@@ -75,6 +77,7 @@ import {
 	updateChatWorkspace,
 	updateChildInParentCache,
 	updateInfiniteChatsCache,
+	upsertChatMessages,
 } from "./chats";
 
 vi.mock("#/api/api", () => ({
@@ -3459,3 +3462,280 @@ describe("semantic cache operations: removal and patching", () => {
 		);
 	});
 });
+
+describe("message upsert fan-out and history replacement", () => {
+	type InfMessages = {
+		pages: TypesGen.ChatMessagesResponse[];
+		pageParams: (number | undefined)[];
+	};
+
+	const mockChatMessage = (
+		id: number,
+		text = `msg ${id}`,
+	): TypesGen.ChatMessage => ({
+		...MockChatMessage,
+		id,
+		content: [{ type: "text", text }],
+	});
+
+	/** Seed and read back the canonical stored object so reference
+	 *  assertions compare against what the cache actually holds. */
+	const seedMessagePages = (
+		queryClient: QueryClient,
+		data: InfMessages,
+	): InfMessages => {
+		queryClient.setQueryData<InfMessages>(chatMessagesKey("chat-1"), data);
+		const seeded = queryClient.getQueryData<InfMessages>(
+			chatMessagesKey("chat-1"),
+		);
+		if (!seeded) {
+			throw new Error("failed to seed messages cache");
+		}
+		return seeded;
+	};
+
+	const readMessagePages = (
+		queryClient: QueryClient,
+	): InfMessages | undefined =>
+		queryClient.getQueryData<InfMessages>(chatMessagesKey("chat-1"));
+
+	// pages[0] is the newest page and every page is DESC by ID,
+	// matching chatMessagesForInfiniteScroll.
+	const twoPageFixture = (): InfMessages => ({
+		pages: [
+			{
+				messages: [mockChatMessage(60), mockChatMessage(55)],
+				queued_messages: [],
+				has_more: true,
+			},
+			{
+				messages: [mockChatMessage(50), mockChatMessage(45)],
+				queued_messages: [],
+				has_more: false,
+			},
+		],
+		pageParams: [undefined, 55],
+	});
+
+	it("upsertChatMessages replaces a message found only in an older page in place", () => {
+		const queryClient = createTestQueryClient();
+		const before = seedMessagePages(queryClient, twoPageFixture());
+
+		upsertChatMessages(queryClient, "chat-1", [mockChatMessage(45, "updated")]);
+
+		const after = readMessagePages(queryClient);
+		expect(after?.pages[0]).toBe(before.pages[0]);
+		expect(after?.pages[0]?.messages.map((m) => m.id)).toEqual([60, 55]);
+		expect(after?.pages[1]?.messages.map((m) => m.id)).toEqual([50, 45]);
+		expect(after?.pages[1]?.messages[1]?.content).toEqual([
+			{ type: "text", text: "updated" },
+		]);
+	});
+
+	it("upsertChatMessages gives every containing page the same fresh value for a duplicated ID", () => {
+		const queryClient = createTestQueryClient();
+		seedMessagePages(queryClient, {
+			pages: [
+				{
+					messages: [mockChatMessage(60), mockChatMessage(45)],
+					queued_messages: [],
+					has_more: true,
+				},
+				{
+					messages: [mockChatMessage(50), mockChatMessage(45)],
+					queued_messages: [],
+					has_more: false,
+				},
+			],
+			pageParams: [undefined, 45],
+		});
+
+		upsertChatMessages(queryClient, "chat-1", [mockChatMessage(45, "fresh")]);
+
+		const after = readMessagePages(queryClient);
+		expect(after?.pages[0]?.messages.map((m) => m.id)).toEqual([60, 45]);
+		expect(after?.pages[1]?.messages.map((m) => m.id)).toEqual([50, 45]);
+		expect(after?.pages[0]?.messages[1]?.content).toEqual([
+			{ type: "text", text: "fresh" },
+		]);
+		expect(after?.pages[1]?.messages[1]?.content).toEqual([
+			{ type: "text", text: "fresh" },
+		]);
+	});
+
+	it("upsertChatMessages preserves the previous reference for a found-but-equal ID in an older page", () => {
+		const queryClient = createTestQueryClient();
+		const before = seedMessagePages(queryClient, twoPageFixture());
+
+		upsertChatMessages(queryClient, "chat-1", [mockChatMessage(45)]);
+
+		const after = readMessagePages(queryClient);
+		expect(after).toBe(before);
+		expect(after?.pages[0]?.messages.map((m) => m.id)).toEqual([60, 55]);
+	});
+
+	it("upsertChatMessages prepends an unknown newest message to page 0 in descending order", () => {
+		const queryClient = createTestQueryClient();
+		const before = seedMessagePages(queryClient, twoPageFixture());
+
+		upsertChatMessages(queryClient, "chat-1", [mockChatMessage(70)]);
+
+		const after = readMessagePages(queryClient);
+		expect(after?.pages[0]?.messages.map((m) => m.id)).toEqual([70, 60, 55]);
+		expect(after?.pages[1]).toBe(before.pages[1]);
+	});
+
+	it("upsertChatMessages inserts an unknown interleaving message mid-page-0 in descending order", () => {
+		const queryClient = createTestQueryClient();
+		const before = seedMessagePages(queryClient, twoPageFixture());
+
+		upsertChatMessages(queryClient, "chat-1", [mockChatMessage(57)]);
+
+		const after = readMessagePages(queryClient);
+		expect(after?.pages[0]?.messages.map((m) => m.id)).toEqual([60, 57, 55]);
+		expect(after?.pages[1]).toBe(before.pages[1]);
+	});
+
+	it("upsertChatMessages inserts each unseen ID once when the batch carries two revisions of it", () => {
+		const queryClient = createTestQueryClient();
+		seedMessagePages(queryClient, twoPageFixture());
+
+		upsertChatMessages(queryClient, "chat-1", [
+			mockChatMessage(70, "revision 1"),
+			mockChatMessage(70, "revision 2"),
+		]);
+
+		const after = readMessagePages(queryClient);
+		expect(after?.pages[0]?.messages.map((m) => m.id)).toEqual([70, 60, 55]);
+		expect(after?.pages[0]?.messages[0]?.content).toEqual([
+			{ type: "text", text: "revision 2" },
+		]);
+	});
+
+	it("upsertChatMessages applies a mixed replace-and-insert batch in a single call", () => {
+		const queryClient = createTestQueryClient();
+		seedMessagePages(queryClient, twoPageFixture());
+
+		upsertChatMessages(queryClient, "chat-1", [
+			mockChatMessage(45, "updated"),
+			mockChatMessage(70),
+		]);
+
+		const after = readMessagePages(queryClient);
+		expect(after?.pages[0]?.messages.map((m) => m.id)).toEqual([70, 60, 55]);
+		expect(after?.pages[1]?.messages.map((m) => m.id)).toEqual([50, 45]);
+		expect(after?.pages[1]?.messages[1]?.content).toEqual([
+			{ type: "text", text: "updated" },
+		]);
+	});
+
+	it("upsertChatMessages never changes the pageParams reference", () => {
+		const queryClient = createTestQueryClient();
+		const before = seedMessagePages(queryClient, twoPageFixture());
+
+		upsertChatMessages(queryClient, "chat-1", [
+			mockChatMessage(45, "updated"),
+			mockChatMessage(70),
+		]);
+
+		const after = readMessagePages(queryClient);
+		expect(after?.pageParams).toBe(before.pageParams);
+		expect(after?.pageParams).toEqual([undefined, 55]);
+	});
+
+	it("upsertChatMessages returns the previous reference for a same-value batch", () => {
+		const queryClient = createTestQueryClient();
+		const before = seedMessagePages(queryClient, twoPageFixture());
+
+		upsertChatMessages(queryClient, "chat-1", [
+			mockChatMessage(60),
+			mockChatMessage(45),
+		]);
+
+		expect(readMessagePages(queryClient)).toBe(before);
+	});
+
+	it("upsertChatMessages is a no-op on an absent cache and creates no entry", () => {
+		const queryClient = createTestQueryClient();
+
+		upsertChatMessages(queryClient, "chat-1", [mockChatMessage(70)]);
+
+		expect(
+			queryClient.getQueryCache().find({ queryKey: chatMessagesKey("chat-1") }),
+		).toBeUndefined();
+	});
+
+	it("replaceChatMessagesHistory collapses to one page and one pageParam, preserving queued messages", () => {
+		const queryClient = createTestQueryClient();
+		const queuedMessage: TypesGen.ChatQueuedMessage = {
+			id: 1,
+			chat_id: "chat-1",
+			created_at: "2025-01-01T00:10:00.000Z",
+			content: [{ type: "text", text: "queued" }],
+		};
+		seedMessagePages(queryClient, {
+			pages: [
+				{
+					messages: [mockChatMessage(60), mockChatMessage(55)],
+					queued_messages: [queuedMessage],
+					has_more: true,
+				},
+				{
+					messages: [mockChatMessage(50), mockChatMessage(45)],
+					queued_messages: [],
+					has_more: true,
+				},
+				{
+					messages: [mockChatMessage(40)],
+					queued_messages: [],
+					has_more: false,
+				},
+			],
+			pageParams: [undefined, 55, 45],
+		});
+
+		replaceChatMessagesHistory(queryClient, "chat-1", [
+			mockChatMessage(55),
+			mockChatMessage(60, "rewritten"),
+		]);
+
+		const after = readMessagePages(queryClient);
+		expect(after?.pages).toHaveLength(1);
+		expect(after?.pageParams).toHaveLength(1);
+		expect(after?.pageParams).toEqual([undefined]);
+		expect(after?.pages[0]?.messages.map((m) => m.id)).toEqual([60, 55]);
+		expect(after?.pages[0]?.has_more).toBe(false);
+		expect(after?.pages[0]?.queued_messages).toEqual([queuedMessage]);
+	});
+
+	it("replaceChatMessagesHistory preserves the previous reference when the replacement equals the current single page", () => {
+		const queryClient = createTestQueryClient();
+		const before = seedMessagePages(queryClient, {
+			pages: [
+				{
+					messages: [mockChatMessage(60), mockChatMessage(55)],
+					queued_messages: [],
+					has_more: false,
+				},
+			],
+			pageParams: [undefined],
+		});
+
+		replaceChatMessagesHistory(queryClient, "chat-1", [
+			mockChatMessage(55),
+			mockChatMessage(60),
+		]);
+
+		expect(readMessagePages(queryClient)).toBe(before);
+	});
+
+	it("replaceChatMessagesHistory is a no-op on an absent cache and creates no entry", () => {
+		const queryClient = createTestQueryClient();
+
+		replaceChatMessagesHistory(queryClient, "chat-1", [mockChatMessage(60)]);
+
+		expect(
+			queryClient.getQueryCache().find({ queryKey: chatMessagesKey("chat-1") }),
+		).toBeUndefined();
+	});
+});

site/src/api/queries/chats.ts

@@ -1,3 +1,4 @@
+import isEqual from "lodash/isEqual";
 import {
 	type InfiniteData,
 	type QueryClient,
@@ -816,6 +817,120 @@ export const patchChatMessages = (
 		InfiniteData<TypesGen.ChatMessagesResponse> | undefined
 	>(chatMessagesKey(chatId), updater);
 
+const replaceMessagesInPage = (
+	page: TypesGen.ChatMessagesResponse,
+	incomingByID: ReadonlyMap<number, TypesGen.ChatMessage>,
+	foundIDs: Set<number>,
+): TypesGen.ChatMessagesResponse => {
+	let pageChanged = false;
+
+	const nextMessages = page.messages.map((existing) => {
+		const incoming = incomingByID.get(existing.id);
+		if (!incoming) {
+			return existing;
+		}
+
+		foundIDs.add(existing.id);
+		if (isEqual(existing, incoming)) {
+			return existing;
+		}
+
+		pageChanged = true;
+		return incoming;
+	});
+
+	return pageChanged ? { ...page, messages: nextMessages } : page;
+};
+
+const upsertMessagesAcrossPages = (
+	currentData: InfiniteData<TypesGen.ChatMessagesResponse> | undefined,
+	messages: readonly TypesGen.ChatMessage[],
+): InfiniteData<TypesGen.ChatMessagesResponse> | undefined => {
+	if (!currentData?.pages?.length || messages.length === 0) {
+		return currentData;
+	}
+
+	const incomingByID = new Map(
+		messages.map((message) => [message.id, message]),
+	);
+	const foundIDs = new Set<number>();
+	const nextPages = currentData.pages.map((page) =>
+		replaceMessagesInPage(page, incomingByID, foundIDs),
+	);
+	const pagesChanged = nextPages.some(
+		(page, index) => page !== currentData.pages[index],
+	);
+
+	const messagesToInsert = [...incomingByID.values()].filter(
+		(message) => !foundIDs.has(message.id),
+	);
+	if (messagesToInsert.length === 0) {
+		return pagesChanged ? { ...currentData, pages: nextPages } : currentData;
+	}
+
+	const firstPage = nextPages[0];
+	const firstPageMessages = [...firstPage.messages, ...messagesToInsert].sort(
+		(a, b) => b.id - a.id,
+	);
+
+	return {
+		...currentData,
+		pages: [
+			{ ...firstPage, messages: firstPageMessages },
+			...nextPages.slice(1),
+		],
+	};
+};
+
+const replaceMessagesHistory = (
+	currentData: InfiniteData<TypesGen.ChatMessagesResponse> | undefined,
+	messages: readonly TypesGen.ChatMessage[],
+): InfiniteData<TypesGen.ChatMessagesResponse> | undefined => {
+	if (!currentData?.pages?.length) {
+		return currentData;
+	}
+
+	const firstPage = currentData.pages[0];
+	const nextMessages = [...messages].sort((a, b) => b.id - a.id);
+	const alreadyReplaced =
+		currentData.pages.length === 1 &&
+		!firstPage.has_more &&
+		firstPage.messages.length === nextMessages.length &&
+		firstPage.messages.every((existing, index) =>
+			isEqual(existing, nextMessages[index]),
+		);
+
+	if (alreadyReplaced) {
+		return currentData;
+	}
+
+	return {
+		...currentData,
+		pages: [{ ...firstPage, messages: nextMessages, has_more: false }],
+		pageParams: currentData.pageParams.slice(0, 1),
+	};
+};
+
+export const upsertChatMessages = (
+	queryClient: QueryClient,
+	chatId: string,
+	messages: readonly TypesGen.ChatMessage[],
+) => {
+	return patchChatMessages(queryClient, chatId, (currentData) =>
+		upsertMessagesAcrossPages(currentData, messages),
+	);
+};
+
+export const replaceChatMessagesHistory = (
+	queryClient: QueryClient,
+	chatId: string,
+	messages: readonly TypesGen.ChatMessage[],
+) => {
+	return patchChatMessages(queryClient, chatId, (currentData) =>
+		replaceMessagesHistory(currentData, messages),
+	);
+};
+
 const DEFAULT_CHAT_PAGE_LIMIT = 50;
 export const CHAT_SEARCH_LIMIT = 50;
 

site/src/pages/AgentsPage/AgentChatPage.stories.tsx

@@ -313,6 +313,11 @@ const buildQueries = (
 	];
 };
 
+const withoutQuery = (
+	queries: ReturnType<typeof buildQueries>,
+	queryKey: readonly unknown[],
+) => queries.filter(({ key }) => hashKey(key) !== hashKey(queryKey));
+
 // ---------------------------------------------------------------------------
 // Every-tool showcase: a single completed assistant turn that exercises
 // every tool renderer registered in Tool.tsx, plus the SubagentRenderer
@@ -2247,6 +2252,84 @@ export const StreamedReasoning: Story = {
 // This made the stories render empty chats and fail interaction
 // tests in both local and CI environments.
 
+const mockNewestMessage: TypesGen.ChatMessage = {
+	...MockChatMessage,
+	id: 30,
+	role: "assistant",
+	content: [{ type: "text", text: "Newest message" }],
+};
+
+const mockOlderRevision: TypesGen.ChatMessage = {
+	...MockChatMessage,
+	id: 20,
+	role: "assistant",
+	content: [{ type: "text", text: "Old revision" }],
+};
+
+const mockFreshRevision: TypesGen.ChatMessage = {
+	...mockOlderRevision,
+	content: [{ type: "text", text: "Fresh revision" }],
+};
+
+export const DurableUpdateFansOutToOlderPage: Story = {
+	parameters: {
+		queries: [
+			...withoutQuery(
+				buildQueries(
+					{
+						id: CHAT_ID,
+						...baseChatFields,
+						title: "Fan-out chat",
+						status: "waiting",
+					},
+					{ messages: [], queued_messages: [], has_more: false },
+				),
+				chatMessagesKey(CHAT_ID),
+			),
+			{
+				key: chatMessagesKey(CHAT_ID),
+				data: {
+					pages: [
+						{
+							messages: [mockNewestMessage],
+							queued_messages: [],
+							has_more: true,
+						},
+						{
+							messages: [mockOlderRevision],
+							queued_messages: [],
+							has_more: false,
+						},
+					],
+					pageParams: [undefined, 30],
+				},
+			},
+		],
+		webSocket: {
+			"/chats/": [
+				{
+					event: "message",
+					data: JSON.stringify([
+						{
+							type: "message",
+							chat_id: CHAT_ID,
+							message: mockFreshRevision,
+						},
+					] satisfies TypesGen.ChatStreamEvent[]),
+				},
+			],
+		},
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		expect(await canvas.findByText("Fresh revision")).toBeVisible();
+		await waitFor(() => {
+			expect(canvas.getAllByText("Fresh revision")).toHaveLength(1);
+		});
+		expect(canvas.queryByText("Old revision")).not.toBeInTheDocument();
+	},
+};
+
 /**
  * Live agent turn with streaming reasoning and a back-to-back flurry of
  * in-progress file tool calls. The persisted history establishes context
@@ -3078,11 +3161,6 @@ const mockServerError = {
 	status: 500,
 };
 
-const withoutQuery = (
-	queries: ReturnType<typeof buildQueries>,
-	queryKey: readonly unknown[],
-) => queries.filter(({ key }) => hashKey(key) !== hashKey(queryKey));
-
 export const DetailQueryError: Story = {
 	parameters: {
 		queries: withoutQuery(

site/src/pages/AgentsPage/AgentChatPage.tsx

@@ -1088,11 +1088,11 @@ const AgentChatPage: FC = () => {
 	const chatMessagesList = (() => {
 		const pages = chatMessagesQuery.data?.pages;
 		if (!pages || pages.length === 0) return undefined;
-		// 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.
+		// Collect all messages and deduplicate by ID as a defense
+		// against cross-page duplicates. Cache upserts fan the same
+		// fresh value out to every page containing an ID, so any
+		// surviving duplicates are value-identical and either
+		// occurrence is safe to render.
 		const all = pages.flatMap((p) => p.messages);
 		const byID = new Map(all.map((m) => [m.id, m]));
 		const deduped = Array.from(byID.values());

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

@@ -910,6 +910,192 @@ describe("useChatStore", () => {
 		]);
 	});
 
+	it("collapses a multi-page cache to one page after history_reset", async () => {
+		const chatID = "chat-history-reset-multipage";
+		const initialMessages = [
+			buildMessage(chatID, 1, "user", "old prompt"),
+			buildMessage(chatID, 2, "assistant", "old answer"),
+			buildMessage(chatID, 3, "user", "stale prompt"),
+		];
+		const replacementMessage = buildMessage(chatID, 1, "user", "new prompt");
+		const mockSocket = createMockSocket();
+		mockWatchChatReturn(mockSocket);
+
+		const queryClient = new QueryClient({
+			defaultOptions: {
+				queries: {
+					retry: false,
+					gcTime: Number.POSITIVE_INFINITY,
+					refetchOnWindowFocus: false,
+					networkMode: "offlineFirst",
+				},
+			},
+		});
+		queryClient.setQueryData(chatMessagesKey(chatID), {
+			pages: [
+				{
+					messages: [initialMessages[2], initialMessages[1]],
+					queued_messages: [],
+					has_more: true,
+				},
+				{
+					messages: [initialMessages[0]],
+					queued_messages: [],
+					has_more: false,
+				},
+			],
+			pageParams: [undefined, 2],
+		});
+		const wrapper = createWrapper(queryClient);
+		const setChatErrorReason = vi.fn();
+		const clearChatErrorReason = vi.fn();
+
+		const { result } = renderHook(
+			() => {
+				const { store } = useChatStore({
+					chatID,
+					chatMessages: initialMessages,
+					chatRecord: buildChat(chatID),
+					chatMessagesData: {
+						messages: initialMessages,
+						queued_messages: [],
+						has_more: false,
+					},
+					chatQueuedMessages: [],
+					setChatErrorReason,
+					clearChatErrorReason,
+				});
+				return {
+					orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs),
+				};
+			},
+			{ wrapper },
+		);
+
+		await waitFor(() => {
+			expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]);
+		});
+
+		act(() => {
+			mockSocket.emitDataBatch([
+				{ type: "history_reset", chat_id: chatID },
+				{ type: "message", chat_id: chatID, message: replacementMessage },
+				{ type: "preview_reset", chat_id: chatID },
+			]);
+		});
+
+		await waitFor(() => {
+			expect(result.current.orderedMessageIDs).toEqual([1]);
+		});
+
+		const cached = queryClient.getQueryData<{
+			pages: TypesGen.ChatMessagesResponse[];
+			pageParams: unknown[];
+		}>(chatMessagesKey(chatID));
+		expect(cached?.pages).toHaveLength(1);
+		expect(cached?.pageParams).toHaveLength(1);
+		expect(cached?.pages[0]?.messages.map((message) => message.id)).toEqual([
+			1,
+		]);
+		expect(cached?.pages[0]?.has_more).toBe(false);
+	});
+
+	it("updates an older page in place for a WS durable message without duplicating it into page 0", async () => {
+		const chatID = "chat-fanout-older-page";
+		const initialMessages = [
+			buildMessage(chatID, 1, "user", "old prompt"),
+			buildMessage(chatID, 2, "assistant", "old answer"),
+			buildMessage(chatID, 3, "user", "newest prompt"),
+		];
+		const updatedMessage = buildMessage(chatID, 1, "user", "edited prompt");
+		const mockSocket = createMockSocket();
+		mockWatchChatReturn(mockSocket);
+
+		const queryClient = new QueryClient({
+			defaultOptions: {
+				queries: {
+					retry: false,
+					gcTime: Number.POSITIVE_INFINITY,
+					refetchOnWindowFocus: false,
+					networkMode: "offlineFirst",
+				},
+			},
+		});
+		queryClient.setQueryData(chatMessagesKey(chatID), {
+			pages: [
+				{
+					messages: [initialMessages[2], initialMessages[1]],
+					queued_messages: [],
+					has_more: true,
+				},
+				{
+					messages: [initialMessages[0]],
+					queued_messages: [],
+					has_more: false,
+				},
+			],
+			pageParams: [undefined, 2],
+		});
+		const wrapper = createWrapper(queryClient);
+		const setChatErrorReason = vi.fn();
+		const clearChatErrorReason = vi.fn();
+
+		const { result } = renderHook(
+			() => {
+				const { store } = useChatStore({
+					chatID,
+					chatMessages: initialMessages,
+					chatRecord: buildChat(chatID),
+					chatMessagesData: {
+						messages: initialMessages,
+						queued_messages: [],
+						has_more: false,
+					},
+					chatQueuedMessages: [],
+					setChatErrorReason,
+					clearChatErrorReason,
+				});
+				return {
+					messagesByID: useChatSelector(store, selectMessagesByID),
+					orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs),
+				};
+			},
+			{ wrapper },
+		);
+
+		await waitFor(() => {
+			expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]);
+		});
+
+		act(() => {
+			mockSocket.emitData({
+				type: "message",
+				chat_id: chatID,
+				message: updatedMessage,
+			});
+		});
+
+		await waitFor(() => {
+			expect(result.current.messagesByID.get(1)?.content).toEqual(
+				updatedMessage.content,
+			);
+		});
+
+		const cached = queryClient.getQueryData<{
+			pages: TypesGen.ChatMessagesResponse[];
+			pageParams: unknown[];
+		}>(chatMessagesKey(chatID));
+		expect(cached?.pages[0]?.messages.map((message) => message.id)).toEqual([
+			3, 2,
+		]);
+		expect(cached?.pages[1]?.messages.map((message) => message.id)).toEqual([
+			1,
+		]);
+		expect(cached?.pages[1]?.messages[0]?.content).toEqual(
+			updatedMessage.content,
+		);
+	});
+
 	it("clears stream state when a new durable message arrives", async () => {
 		immediateAnimationFrame();
 

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

@@ -1,3 +1,4 @@
+import isEqual from "lodash/isEqual";
 import { useSyncExternalStore } from "react";
 import type * as TypesGen from "#/api/typesGenerated";
 import { type ChatDetailError, chatDetailErrorsEqual } from "./chatError";
@@ -14,10 +15,10 @@ const buildOrderedMessageIDs = (
 ): readonly number[] => {
 	// created_at is shared across an insert batch, so only id tracks append order.
 	const sorted = messages.toSorted((left, right) => left.id - right.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
+	// Deduplicate by ID as a defense against duplicate IDs in the
+	// input. Cache upserts fan the same fresh value out to every
+	// page containing an ID, so cross-page duplicates in the React
+	// Query cache are value-identical. The Map-based messagesByID
 	// already deduplicates, but orderedMessageIDs must match.
 	const seen = new Set<number>();
 	const orderedMessageIDs: number[] = [];
@@ -53,29 +54,6 @@ const arraysEqual = <T>(left: readonly T[], right: readonly T[]): boolean => {
 	return true;
 };
 
-const jsonValuesEqual = (left: unknown, right: unknown): boolean => {
-	if (left === right) {
-		return true;
-	}
-	try {
-		return JSON.stringify(left) === JSON.stringify(right);
-	} catch {
-		return false;
-	}
-};
-
-export const chatMessagesEqualByValue = (
-	left: TypesGen.ChatMessage,
-	right: TypesGen.ChatMessage,
-): boolean =>
-	left.id === right.id &&
-	left.chat_id === right.chat_id &&
-	left.model_config_id === right.model_config_id &&
-	left.created_at === right.created_at &&
-	left.role === right.role &&
-	jsonValuesEqual(left.content, right.content) &&
-	jsonValuesEqual(left.usage, right.usage);
-
 export const chatQueuedMessagesEqualByID = (
 	left: readonly TypesGen.ChatQueuedMessage[],
 	right: readonly TypesGen.ChatQueuedMessage[],
@@ -311,7 +289,7 @@ export const createChatStore = (): ChatStore => {
 		// concurrent state change (TOCTOU).
 		const existing = state.messagesByID.get(message.id);
 		const isDuplicate = state.messagesByID.has(message.id);
-		if (existing && chatMessagesEqualByValue(existing, message)) {
+		if (existing && isEqual(existing, message)) {
 			return { isDuplicate, changed: false };
 		}
 
@@ -320,7 +298,7 @@ export const createChatStore = (): ChatStore => {
 			// Re-check inside the updater: another call may have
 			// already applied this exact message.
 			const curExisting = current.messagesByID.get(message.id);
-			if (curExisting && chatMessagesEqualByValue(curExisting, message)) {
+			if (curExisting && isEqual(curExisting, message)) {
 				return current;
 			}
 
@@ -358,7 +336,7 @@ export const createChatStore = (): ChatStore => {
 			for (const message of messages) {
 				const map = nextMessagesByID ?? current.messagesByID;
 				const existing = map.get(message.id);
-				if (existing && chatMessagesEqualByValue(existing, message)) {
+				if (existing && isEqual(existing, message)) {
 					continue;
 				}
 				// Lazily copy the map on first actual change.

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

@@ -16,7 +16,9 @@ import {
 	invalidateChatPrompts,
 	invalidateChatSearches,
 	patchChatMessages,
+	replaceChatMessagesHistory,
 	updateInfiniteChatsCache,
+	upsertChatMessages,
 } from "#/api/queries/chats";
 import type * as TypesGen from "#/api/typesGenerated";
 import type { OneWayMessageEvent } from "#/utils/OneWayWebSocket";
@@ -25,7 +27,6 @@ import { type ChatDetailError, normalizeChatErrorPayload } from "./chatError";
 import {
 	type ChatStore,
 	type ChatStoreState,
-	chatMessagesEqualByValue,
 	chatQueuedMessagesEqualByID,
 	createChatStore,
 	isActiveChatStatus,
@@ -197,39 +198,7 @@ export const useChatStore = (
 			if (!chatID || messages.length === 0) {
 				return;
 			}
-			patchChatMessages(queryClient, 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),
-					],
-				};
-			});
+			upsertChatMessages(queryClient, chatID, messages);
 			// Refresh the dedicated prompt-history cache when a user message arrives.
 			const hasNewUserPrompt = messages.some((msg) => msg.role === "user");
 			if (hasNewUserPrompt) {
@@ -245,18 +214,7 @@ export const useChatStore = (
 			if (!chatID) {
 				return;
 			}
-			patchChatMessages(queryClient, chatID, (currentData) => {
-				if (!currentData?.pages?.length) {
-					return currentData;
-				}
-				const firstPage = currentData.pages[0];
-				const updatedMessages = [...messages].sort((a, b) => b.id - a.id);
-				return {
-					...currentData,
-					pages: [{ ...firstPage, messages: updatedMessages, has_more: false }],
-					pageParams: currentData.pageParams.slice(0, 1),
-				};
-			});
+			replaceChatMessagesHistory(queryClient, chatID, messages);
 			void invalidateChatSearches(queryClient);
 		},
 		[chatID, queryClient],