diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 8fe6e17d09..5774f549ca 100644 --- a/site/src/api/queries/chats.test.ts +++ b/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(chatMessagesKey("chat-1"), data); + const seeded = queryClient.getQueryData( + chatMessagesKey("chat-1"), + ); + if (!seeded) { + throw new Error("failed to seed messages cache"); + } + return seeded; + }; + + const readMessagePages = ( + queryClient: QueryClient, + ): InfMessages | undefined => + queryClient.getQueryData(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(); + }); +}); diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index f81d109c70..83e031392b 100644 --- a/site/src/api/queries/chats.ts +++ b/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 | undefined >(chatMessagesKey(chatId), updater); +const replaceMessagesInPage = ( + page: TypesGen.ChatMessagesResponse, + incomingByID: ReadonlyMap, + foundIDs: Set, +): 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 | undefined, + messages: readonly TypesGen.ChatMessage[], +): InfiniteData | undefined => { + if (!currentData?.pages?.length || messages.length === 0) { + return currentData; + } + + const incomingByID = new Map( + messages.map((message) => [message.id, message]), + ); + const foundIDs = new Set(); + 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 | undefined, + messages: readonly TypesGen.ChatMessage[], +): InfiniteData | 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; diff --git a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx index 4f7ca1c562..581d9b5ba4 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx @@ -313,6 +313,11 @@ const buildQueries = ( ]; }; +const withoutQuery = ( + queries: ReturnType, + 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, - queryKey: readonly unknown[], -) => queries.filter(({ key }) => hashKey(key) !== hashKey(queryKey)); - export const DetailQueryError: Story = { parameters: { queries: withoutQuery( diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 27d2ff021d..f2761de194 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/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()); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index ae9db95cd1..80c4ad6c7e 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/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(); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts index 99d5da05f4..08f4375f0c 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts +++ b/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(); const orderedMessageIDs: number[] = []; @@ -53,29 +54,6 @@ const arraysEqual = (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. diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index 298a8b1d7d..bf77434503 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/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],