From f31a8277a94cb0c3bc31f725f8bf07dc4441afba Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Thu, 19 Mar 2026 15:27:03 +0200 Subject: [PATCH] fix: show promoted queued message in chat timeline immediately (#23232) Two issues caused the promoted message to never appear: 1. handlePromoteQueuedMessage discarded the ChatMessage returned by the promote API, relying on the WebSocket to deliver it. 2. Even when the WebSocket did deliver it (via upsertDurableMessage), the queue_update event in the same batch called updateChatQueuedMessages, which mutated the React Query cache. This gave chatMessagesList a new reference, triggering the message sync effect. The effect found the promoted message in the store but not in the REST-fetched data, classified it as a stale entry (the path designed for edit truncation), and called replaceMessages, wiping it. Fix (1): capture the ChatMessage from the promote response and upsert it into the store, matching handleSend for non-queued messages. Fix (2): track the fetched message array elements across effect runs using element-level reference comparison. Only run the hasStaleEntries/replaceMessages path when the message objects actually changed (e.g. a refetch producing new objects from the server), not when only an unrelated field like queued_messages caused the query data reference to update. Element references work because useMemo(flatMap) preserves object identity when only non-message fields change in the page data. --- site/src/api/queries/chats.ts | 6 +- site/src/pages/AgentsPage/AgentDetail.tsx | 6 +- .../AgentDetail/ChatContext.test.tsx | 111 ++++++++++++++++++ .../AgentsPage/AgentDetail/ChatContext.ts | 32 ++++- 4 files changed, 147 insertions(+), 8 deletions(-) diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index e61dafb601..9779d64cd6 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -366,9 +366,9 @@ export const promoteChatQueuedMessage = ( ) => ({ mutationFn: (queuedMessageId: number) => API.promoteChatQueuedMessage(chatId, queuedMessageId), - // No onSuccess invalidation needed: the per-chat WebSocket - // delivers the promoted message, queue update, and status - // change in real-time. + // No onSuccess invalidation needed: the caller upserts the + // promoted message from the response, and the per-chat + // WebSocket delivers queue and status updates in real-time. }); export const chatDiffContentsKey = (chatId: string) => diff --git a/site/src/pages/AgentsPage/AgentDetail.tsx b/site/src/pages/AgentsPage/AgentDetail.tsx index 8c0ce3fba4..9914b15404 100644 --- a/site/src/pages/AgentsPage/AgentDetail.tsx +++ b/site/src/pages/AgentsPage/AgentDetail.tsx @@ -709,7 +709,11 @@ const AgentDetail: FC = () => { store.clearStreamError(); store.setChatStatus("pending"); try { - await promoteQueuedMutation.mutateAsync(id); + 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. + store.upsertDurableMessage(promotedMessage); } catch (error) { store.setQueuedMessages(previousQueuedMessages); store.setChatStatus(previousChatStatus); diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx index 793b64fdc2..06c6705521 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx @@ -2338,6 +2338,117 @@ describe("useChatStore", () => { expect(result.current.orderedMessageIDs).toEqual([1]); }); }); + + it("does not wipe WebSocket-delivered message when queue_update triggers cache change", async () => { + immediateAnimationFrame(); + + const chatID = "chat-queue-promote"; + const msg1 = makeMessage(chatID, 1, "user", "hello"); + const msg2 = makeMessage(chatID, 2, "assistant", "hi"); + // The promoted message that will arrive via WebSocket. + const promotedMsg = makeMessage(chatID, 3, "user", "follow-up"); + + const mockSocket = createMockSocket(); + vi.mocked(watchChat).mockReturnValue(mockSocket as never); + + const queryClient = createTestQueryClient(); + const wrapper: FC = ({ children }) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const queuedMsg = makeQueuedMessage(chatID, 10, "follow-up"); + const initialMessages = [msg1, msg2]; + + const initialOptions = { + chatID, + chatMessages: initialMessages, + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: initialMessages, + queued_messages: [queuedMsg], + has_more: false, + }, + chatQueuedMessages: [queuedMsg], + setChatErrorReason, + clearChatErrorReason, + }; + + const { result, rerender } = renderHook( + (options: Parameters[0]) => { + const { store } = useChatStore(options); + return { + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + queuedMessages: useChatSelector(store, selectQueuedMessages), + }; + }, + { initialProps: initialOptions, wrapper }, + ); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2]); + expect(result.current.queuedMessages).toHaveLength(1); + }); + + // Simulate the WebSocket delivering the promoted message + // followed by a queue_update in the same batch (as the server + // does when auto-promoting or when the promote endpoint runs). + act(() => { + mockSocket.emitOpen(); + }); + act(() => { + mockSocket.emitDataBatch([ + { + type: "message", + chat_id: chatID, + message: promotedMsg, + }, + { + type: "queue_update", + chat_id: chatID, + queued_messages: [], + }, + ]); + }); + + // The promoted message should appear in the store and the + // queue should be empty. Before the fix, the queue_update + // caused updateChatQueuedMessages to mutate the React Query + // cache, giving chatMessages a new reference that triggered + // the sync effect. The effect detected the promoted message + // as a "stale entry" (present in store but not in the REST + // data) and called replaceMessages, wiping it. + // + // Now re-render so the updated query cache flows through + // the chatMessages prop (simulating React rerender after + // the query cache mutation). + rerender({ + ...initialOptions, + // chatMessages still comes from REST (not refetched), so + // it only has [msg1, msg2]. The promoted message lives + // only in the store via the WebSocket delivery. + // + // Spread into a new array to simulate what actually + // happens: updateChatQueuedMessages mutates the React + // Query cache (changing queued_messages), which gives + // chatMessagesQuery.data a new reference, causing the + // chatMessagesList useMemo to return a new array with + // the same elements. The new reference triggers the + // sync effect. + chatMessages: [...initialMessages], + chatMessagesData: { + messages: [...initialMessages], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + }); + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + expect(result.current.queuedMessages).toHaveLength(0); + }); + }); }); describe("updateSidebarChat via stream events", () => { diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts index 4be26aabed..6014024a55 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts @@ -448,6 +448,16 @@ export const useChatStore = ( const wsQueueUpdateReceivedRef = useRef(false); const activeChatIDRef = useRef(null); const prevChatIDRef = useRef(chatID); + // Snapshot of the chatMessages elements from the last sync effect + // run. Used to detect whether chatMessages actually changed (e.g. + // after a refetch producing new objects) vs. just getting a new + // array reference because an unrelated field like queued_messages + // was updated in the query cache. Element-level reference + // comparison works because useMemo(flatMap) preserves message + // object references when only non-message fields change in the + // page, while a genuine refetch returns new objects from the + // server. + const lastSyncedMessagesRef = useRef([]); const store = storeRef.current; @@ -544,6 +554,7 @@ export const useChatStore = ( // the new chat's query resolves. if (prevChatIDRef.current !== chatID) { prevChatIDRef.current = chatID; + lastSyncedMessagesRef.current = []; store.replaceMessages([]); } // Merge REST-fetched messages into the store one-by-one instead @@ -554,12 +565,25 @@ export const useChatStore = ( // However, if the fetched set is missing message IDs the store // already has (e.g. after an edit truncation), a full replace // is needed because upsert can only add/update, not remove. + // We must only do this when the fetched messages actually + // changed (new elements from a refetch), not when an + // unrelated field like queued_messages caused the query + // data reference to update. Without this guard, a + // queue_update WebSocket event would trigger + // replaceMessages with the stale REST data, wiping any + // message the WebSocket just delivered. if (chatMessages) { - const fetchedIDs = new Set(chatMessages.map((m) => m.id)); + const prev = lastSyncedMessagesRef.current; + const contentChanged = + chatMessages.length !== prev.length || + chatMessages.some((m, i) => m !== prev[i]); + lastSyncedMessagesRef.current = chatMessages; + const storeSnap = store.getSnapshot(); - const hasStaleEntries = storeSnap.orderedMessageIDs.some( - (id) => !fetchedIDs.has(id), - ); + const fetchedIDs = new Set(chatMessages.map((m) => m.id)); + const hasStaleEntries = + contentChanged && + storeSnap.orderedMessageIDs.some((id) => !fetchedIDs.has(id)); if (hasStaleEntries) { store.replaceMessages(chatMessages); } else {