diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index ba7022714b..2501b600f1 100644 --- a/site/src/api/queries/chats.test.ts +++ b/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(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(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 () => { diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 0f206f14a7..fdf33f66c6 100644 --- a/site/src/api/queries/chats.ts +++ b/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); }, }); diff --git a/site/src/pages/AgentsPage/AgentChatPage.test.ts b/site/src/pages/AgentsPage/AgentChatPage.test.ts index 99055dec52..a7f0be7658 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.test.ts +++ b/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(); + }); +}); diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 5bd26cc0ec..98bbb6cdda 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/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; + editArgs: { + messageId: number; + optimisticMessage?: TypesGen.ChatMessage; + req: TypesGen.EditChatMessageRequest; + }; + scrollToBottom: (() => void) | null | undefined; + onError: (error: unknown) => void; +}): Promise { + 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 | 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; } diff --git a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx index b290e7eb7c..81ebbab4ab 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx +++ b/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: () => , + 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(); + }, +};