fix(site): prevent sticky message cycling when submitting edited message (#24292)

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

@@ -823,7 +823,7 @@ describe("mutation invalidation scope", () => {
 		}
 	});
 
-	it("editChatMessage invalidates chat detail, messages, and debug runs", async () => {
+	it("editChatMessage invalidates chat detail and debug runs, not messages", async () => {
 		const queryClient = createTestQueryClient();
 		const chatId = "chat-1";
 		seedAllActiveQueries(queryClient, chatId);
@@ -833,19 +833,22 @@ describe("mutation invalidation scope", () => {
 
 		await new Promise((r) => setTimeout(r, 0));
 
-		// These queries should be invalidated -- editing changes
-		// message content, may update the chat record, and can start
-		// a new debug run.
+		// Chat metadata and debug runs should be invalidated because
+		// editing changes the chat's updated_at and can start a new
+		// debug run.
 		const chatState = queryClient.getQueryState(chatKey(chatId));
 		expect(chatState?.isInvalidated, "chatKey should be invalidated").toBe(
 			true,
 		);
 
+		// Messages are NOT invalidated. The per-chat WebSocket handles
+		// post-edit message delivery, making REST invalidation
+		// unnecessary.
 		const messagesState = queryClient.getQueryState(chatMessagesKey(chatId));
 		expect(
 			messagesState?.isInvalidated,
-			"chatMessagesKey should be invalidated",
-		).toBe(true);
+			"chatMessagesKey should not be invalidated",
+		).not.toBe(true);
 
 		expect(
 			queryClient.getQueryState(chatDebugRunsKey(chatId))?.isInvalidated,
@@ -853,6 +856,37 @@ describe("mutation invalidation scope", () => {
 		).toBe(true);
 	});
 
+	it("editChatMessage onError invalidates messages", async () => {
+		const queryClient = createTestQueryClient();
+		const chatId = "chat-1";
+		const messages = [3, 2, 1].map((id) => makeMsg(chatId, id));
+
+		queryClient.setQueryData<InfMessages>(chatMessagesKey(chatId), {
+			pages: [{ messages, queued_messages: [], has_more: false }],
+			pageParams: [undefined],
+		});
+
+		const mutation = editChatMessage(queryClient, chatId);
+		mutation.onError(
+			new Error("fail"),
+			{ messageId: 2, req: editReq },
+			{
+				previousData: {
+					pages: [{ messages, queued_messages: [], has_more: false }],
+					pageParams: [undefined],
+				},
+			},
+		);
+
+		await new Promise((r) => setTimeout(r, 0));
+
+		const messagesState = queryClient.getQueryState(chatMessagesKey(chatId));
+		expect(
+			messagesState?.isInvalidated,
+			"chatMessagesKey should be invalidated on error",
+		).toBe(true);
+	});
+
 	// Shared type for the infinite messages cache shape used by
 	// editChatMessage tests below.
 	type InfMessages = {
@@ -1083,7 +1117,7 @@ describe("mutation invalidation scope", () => {
 
 		const mutation = editChatMessage(queryClient, chatId);
 
-		// Pass undefined context — simulates onMutate throwing before
+		// Pass undefined context. This simulates onMutate throwing before
 		// it could return a snapshot.
 		mutation.onError(
 			new Error("fail"),
@@ -1091,9 +1125,16 @@ describe("mutation invalidation scope", () => {
 			undefined,
 		);
 
-		// Cache should be untouched — no crash, no corruption.
+		// Cache should be untouched: no crash, no corruption.
 		const data = queryClient.getQueryData<InfMessages>(chatMessagesKey(chatId));
 		expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([3, 2, 1]);
+
+		await new Promise((r) => setTimeout(r, 0));
+		const messagesState = queryClient.getQueryState(chatMessagesKey(chatId));
+		expect(
+			messagesState?.isInvalidated,
+			"chatMessagesKey should be invalidated even without context",
+		).toBe(true);
 	});
 
 	it("editChatMessage onMutate updates the first page and preserves older pages", async () => {

site/src/api/queries/chats.ts

@@ -1010,6 +1010,13 @@ export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({
 		if (context?.previousData) {
 			queryClient.setQueryData(chatMessagesKey(chatId), context.previousData);
 		}
+		// Invalidate messages as a safety net: the restored snapshot
+		// may be missing WebSocket-delivered messages that arrived
+		// during the mutation's flight time.
+		void queryClient.invalidateQueries({
+			queryKey: chatMessagesKey(chatId),
+			exact: true,
+		});
 	},
 	onSuccess: (
 		response: TypesGen.EditChatMessageResponse,
@@ -1026,20 +1033,18 @@ export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({
 		);
 	},
 	onSettled: () => {
-		// Always reconcile with the server regardless of whether
-		// the mutation succeeded or failed. On success this picks
-		// up the replacement message; on failure it confirms the
-		// restore from onError matches the server state. Use exact
-		// matching to avoid cascading to unrelated queries
-		// (diff-status, diff-contents, cost summaries, etc.).
+		// Refresh chat metadata (status, title, etc.). The messages
+		// query is intentionally NOT invalidated here. The per-chat
+		// WebSocket handles post-edit message delivery via
+		// FullRefresh, making REST invalidation unnecessary.
+		// Invalidating chatMessagesKey would trigger a redundant
+		// refetch that causes extra store mutations while the
+		// sticky user message is settling after the optimistic
+		// truncation.
 		void queryClient.invalidateQueries({
 			queryKey: chatKey(chatId),
 			exact: true,
 		});
-		void queryClient.invalidateQueries({
-			queryKey: chatMessagesKey(chatId),
-			exact: true,
-		});
 		void invalidateChatDebugRuns(queryClient, chatId);
 	},
 });

site/src/pages/AgentsPage/AgentChatPage.test.ts

@@ -7,6 +7,7 @@ import {
 	filterWorkspaceOptionsByOrganization,
 	getPersistedDraftInputValue,
 	restoreOptimisticRequestSnapshot,
+	submitEditAndScroll,
 	useConversationEditingState,
 	waitForPendingChatSettingsSyncs,
 } from "./AgentChatPage";
@@ -852,3 +853,62 @@ describe("useConversationEditingState", () => {
 		unmount();
 	});
 });
+
+describe("submitEditAndScroll", () => {
+	const dummyArgs = {
+		messageId: 42,
+		req: { content: [{ type: "text" as const, text: "edited" }] },
+	};
+
+	it("calls scrollToBottom after editMessage resolves", async () => {
+		const callOrder: string[] = [];
+		const editMessage = vi.fn(async () => {
+			callOrder.push("editMessage");
+		});
+		const scrollToBottom = vi.fn(() => {
+			callOrder.push("scrollToBottom");
+		});
+
+		await submitEditAndScroll({
+			editMessage,
+			editArgs: dummyArgs,
+			scrollToBottom,
+			onError: vi.fn(),
+		});
+
+		expect(callOrder).toEqual(["editMessage", "scrollToBottom"]);
+	});
+
+	it("does not call scrollToBottom when editMessage throws", async () => {
+		const scrollToBottom = vi.fn();
+		const onError = vi.fn();
+		const editMessage = vi.fn().mockRejectedValue(new Error("boom"));
+
+		await expect(
+			submitEditAndScroll({
+				editMessage,
+				editArgs: dummyArgs,
+				scrollToBottom,
+				onError,
+			}),
+		).rejects.toThrow("boom");
+
+		expect(scrollToBottom).not.toHaveBeenCalled();
+		expect(onError).toHaveBeenCalledWith(
+			expect.objectContaining({ message: "boom" }),
+		);
+	});
+
+	it("tolerates null scrollToBottom", async () => {
+		const editMessage = vi.fn().mockResolvedValue(undefined);
+
+		await submitEditAndScroll({
+			editMessage,
+			editArgs: dummyArgs,
+			scrollToBottom: null,
+			onError: vi.fn(),
+		});
+
+		expect(editMessage).toHaveBeenCalled();
+	});
+});

site/src/pages/AgentsPage/AgentChatPage.tsx

@@ -140,6 +140,40 @@ export const restoreOptimisticRequestSnapshot = (
 	});
 };
 
+export async function submitEditAndScroll({
+	editMessage,
+	editArgs,
+	scrollToBottom,
+	onError,
+}: {
+	editMessage: (args: {
+		messageId: number;
+		optimisticMessage?: TypesGen.ChatMessage;
+		req: TypesGen.EditChatMessageRequest;
+	}) => Promise<unknown>;
+	editArgs: {
+		messageId: number;
+		optimisticMessage?: TypesGen.ChatMessage;
+		req: TypesGen.EditChatMessageRequest;
+	};
+	scrollToBottom: (() => void) | null | undefined;
+	onError: (error: unknown) => void;
+}): Promise<void> {
+	try {
+		await editMessage(editArgs);
+	} catch (error) {
+		onError(error);
+		throw error;
+	}
+	// Scroll after the mutation resolves so the optimistic
+	// truncation and server reconciliation have already been
+	// applied to the DOM. Scrolling before this point causes
+	// the sticky user message to cycle through prior messages
+	// as the IntersectionObserver reacts to rapid layout
+	// shifts between the old and truncated content.
+	scrollToBottom?.();
+}
+
 /** @internal Exported for testing. */
 export const waitForPendingChatSettingsSyncs = async (
 	pendingSyncs: readonly (Promise<unknown> | null | undefined)[],
@@ -1226,18 +1260,19 @@ const AgentChatPage: FC = () => {
 				store.setChatStatus("running");
 				store.clearStreamState();
 			});
-			scrollToBottomRef.current?.();
-			try {
-				await editMessage({
+			await submitEditAndScroll({
+				editMessage,
+				editArgs: {
 					messageId: editedMessageID,
 					optimisticMessage,
 					req: request,
-				});
-			} catch (error) {
-				restoreOptimisticRequestSnapshot(store, previousSnapshot);
-				handleUsageLimitError(error);
-				throw error;
-			}
+				},
+				scrollToBottom: scrollToBottomRef.current,
+				onError: (error) => {
+					restoreOptimisticRequestSnapshot(store, previousSnapshot);
+					handleUsageLimitError(error);
+				},
+			});
 			return;
 		}
 

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

@@ -1146,3 +1146,63 @@ export const ScrollRepinnedAfterWheelDeferredAppend: Story = {
 		).toBeNull();
 	},
 };
+
+const editSubmitScrollStore = buildStoreWithMessages(buildLongConversation(30));
+
+/**
+ * Verifies that the scroll position settles at the bottom of the
+ * conversation after an optimistic edit truncation removes messages.
+ * The actual scroll-ordering regression (scrollToBottom must fire
+ * after editMessage resolves) is covered by the submitEditAndScroll
+ * unit tests in AgentChatPage.test.ts.
+ */
+export const ScrollStableAfterEditTruncation: Story = {
+	parameters: { chromatic: { disableSnapshot: true } },
+	decorators: scrollStoryDecorators,
+	render: () => <StoryAgentChatPageView store={editSubmitScrollStore} />,
+	play: async ({ canvasElement }) => {
+		// Reset the module-scoped store so interactive re-runs in
+		// Storybook start from the full 30-message conversation.
+		editSubmitScrollStore.replaceMessages(buildLongConversation(30));
+		editSubmitScrollStore.setChatStatus("completed");
+
+		const canvas = within(canvasElement);
+		const scrollContainer = canvas.getByTestId("scroll-container");
+
+		await waitForScrollOverflow(scrollContainer);
+
+		await waitFor(
+			() => {
+				const dist =
+					scrollContainer.scrollHeight -
+					scrollContainer.scrollTop -
+					scrollContainer.clientHeight;
+				expect(dist).toBeLessThan(5);
+			},
+			{ timeout: 2000 },
+		);
+
+		const existing = getStoreMessages(editSubmitScrollStore);
+		const editIndex = 10;
+		const truncated = existing.slice(0, editIndex);
+		truncated.push(
+			buildMessage(existing[editIndex].id, "user", "Edited question"),
+		);
+		editSubmitScrollStore.replaceMessages(truncated);
+
+		await waitFor(
+			() => {
+				const dist =
+					scrollContainer.scrollHeight -
+					scrollContainer.scrollTop -
+					scrollContainer.clientHeight;
+				expect(dist).toBeLessThan(5);
+			},
+			{ timeout: 2000 },
+		);
+
+		expect(
+			canvas.queryByRole("button", { name: "Scroll to bottom" }),
+		).toBeNull();
+	},
+};