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;