From b74d60e88ce3dd5d8304be88a8cd6c82924bfa00 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Tue, 10 Mar 2026 09:11:45 -0700 Subject: [PATCH] fix(site): correct stale queued messages when switching back to a chat (#22911) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem When a user navigates away from a chat and its queued messages are processed server-side, switching back shows stale queued messages until a hard page refresh. The issue is purely frontend state — the backend is correct. ### Root cause Three things conspire to cause the bug: 1. **Stale React Query cache** — the `chatKey(chatId)` cache entry retains the old `queued_messages` from the last fetch. When the user is on a different chat, no refetch or WebSocket updates the cache for the inactive chat. 2. **One-shot hydration guard** — `queuedMessagesHydratedChatIDRef` blocks all REST-sourced re-hydration after the first hydration for a given chat ID. This was designed to prevent a stale REST refetch from overwriting a fresher `queue_update` from the WebSocket, but it also blocks the corrected data that arrives when the query actually refetches from the server. 3. **No unsolicited `queue_update`** — the WebSocket only sends `queue_update` events when the queue changes. If the queue was already drained before the WebSocket connected, no event is ever sent, so the stale data persists. ## Fix Add a `wsQueueUpdateReceivedRef` flag that tracks whether the WebSocket has delivered a `queue_update` for the current chat. The hydration guard now only blocks REST re-hydration **after** a `queue_update` has been received (since the stream is authoritative at that point). Before any `queue_update` arrives, REST refetches are allowed through to correct stale cached data. The flag is reset on chat switch alongside the existing hydration guard reset. ## Changes - **`ChatContext.ts`**: Add `wsQueueUpdateReceivedRef`, update hydration guard condition, set flag on `queue_update` events, reset on chat switch. - **`ChatContext.test.tsx`**: Add test covering the exact scenario — stale cached queued messages are corrected by a REST refetch when no `queue_update` has arrived. --- .../AgentDetail/ChatContext.test.tsx | 74 +++++++++++++++++++ .../AgentsPage/AgentDetail/ChatContext.ts | 19 ++++- 2 files changed, 92 insertions(+), 1 deletion(-) diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx index b4244bb406..7a737e433d 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx @@ -797,6 +797,80 @@ describe("useChatStore", () => { }); }); + it("corrects stale queued messages from cache when switching back to a chat", async () => { + const chatID = "chat-1"; + const existingMessage = makeMessage(chatID, 1, "user", "hello"); + const queuedMessage = makeQueuedMessage(chatID, 10, "queued"); + const mockSocket = createMockSocket(); + vi.mocked(watchChat).mockReturnValue(mockSocket as never); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + // Start with queued messages from a stale React Query cache. + // This simulates coming back to a chat whose queue was drained + // server-side while the user was viewing a different chat. + const staleOptions = { + chatID, + chatMessages: [existingMessage], + chatRecord: makeChat(chatID), + chatData: { + chat: makeChat(chatID), + messages: [existingMessage], + queued_messages: [queuedMessage], + }, + chatQueuedMessages: [queuedMessage], + setChatErrorReason, + clearChatErrorReason, + }; + + const { result, rerender } = renderHook( + (options: Parameters[0]) => { + const { store } = useChatStore(options); + return { + queuedMessages: useChatSelector(store, selectQueuedMessages), + }; + }, + { + initialProps: staleOptions, + wrapper, + }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + // Initially shows the stale queued message from cache. + expect(result.current.queuedMessages.map((m) => m.id)).toEqual([ + queuedMessage.id, + ]); + + // Simulate the REST query refetching and returning fresh + // data with an empty queue (no queue_update from WS yet). + rerender({ + ...staleOptions, + chatData: { + chat: { + ...makeChat(chatID), + updated_at: "2025-01-01T00:00:02.000Z", + }, + messages: [existingMessage], + queued_messages: [], + }, + chatQueuedMessages: [], + }); + + // The store should accept the fresh REST data because the + // WebSocket hasn't sent a queue_update yet. + await waitFor(() => { + expect(result.current.queuedMessages).toEqual([]); + }); + }); + it("writes queue_update snapshots into the chat query cache", async () => { const chatID = "chat-1"; const existingMessage = makeMessage(chatID, 1, "user", "hello"); diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts index 6c571d9b73..6217cd76f4 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts @@ -454,6 +454,13 @@ export const useChatStore = ( const storeRef = useRef(createChatStore()); const streamResetFrameRef = useRef(null); const queuedMessagesHydratedChatIDRef = useRef(null); + // Tracks whether the WebSocket has delivered a queue_update for the + // current chat. When true, the stream is the authoritative source + // and REST re-fetches must not overwrite the store. When false, + // REST data is allowed to re-hydrate so stale cached queued + // messages are corrected when switching back to a chat whose + // queue was drained while the user was away. + const wsQueueUpdateReceivedRef = useRef(false); const activeChatIDRef = useRef(null); const prevChatIDRef = useRef(chatID); @@ -553,6 +560,7 @@ export const useChatStore = ( useEffect(() => { queuedMessagesHydratedChatIDRef.current = null; + wsQueueUpdateReceivedRef.current = false; store.setQueuedMessages([]); if (!chatID) { return; @@ -563,7 +571,15 @@ export const useChatStore = ( if (!chatID || !chatData) { return; } - if (queuedMessagesHydratedChatIDRef.current === chatID) { + // Allow re-hydration from REST as long as the WebSocket hasn't + // delivered a queue_update yet (which would be fresher). This + // ensures that when the user navigates back to a chat whose + // queued messages were drained server-side while they were + // away, the REST refetch corrects the stale cached state. + if ( + queuedMessagesHydratedChatIDRef.current === chatID && + wsQueueUpdateReceivedRef.current + ) { return; } queuedMessagesHydratedChatIDRef.current = chatID; @@ -688,6 +704,7 @@ export const useChatStore = ( continue; } } + wsQueueUpdateReceivedRef.current = true; store.setQueuedMessages(streamEvent.queued_messages); updateChatQueuedMessages(streamEvent.queued_messages); continue;