diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index ace1f20e95..3a53f6f5e3 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx @@ -2953,6 +2953,77 @@ describe("useChatStore", () => { expect(textBlock).toBeDefined(); }); }); + + it("does not let a stale REST chatRecord.status override WS-delivered status", async () => { + immediateAnimationFrame(); + + const chatID = "chat-stale-rest-status"; + const userMsg = makeMessage(chatID, 1, "user", "hello"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + + // Start with a "running" chatRecord so the WS opens. + const { result, rerender } = renderHook( + (props: { chatRecord: TypesGen.Chat }) => { + const { store } = useChatStore({ + chatID, + chatMessages: [userMsg], + chatRecord: props.chatRecord, + chatMessagesData: { + messages: [userMsg], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason: vi.fn(), + clearChatErrorReason: vi.fn(), + }); + return { + chatStatus: useChatSelector(store, selectChatStatus), + }; + }, + { + wrapper, + initialProps: { + chatRecord: makeChat(chatID), + }, + }, + ); + + // Wait for WS to connect. + await waitFor(() => { + expect(result.current.chatStatus).toBe("running"); + }); + + // Deliver a status event over WS so wsStatusReceivedRef is set. + act(() => { + mockSocket.emitData({ + type: "status", + chat_id: chatID, + status: { status: "running" }, + }); + }); + + await waitFor(() => { + expect(result.current.chatStatus).toBe("running"); + }); + + // Simulate a stale REST refetch returning "pending". + rerender({ + chatRecord: { ...makeChat(chatID), status: "pending" }, + }); + + // The store must ignore the stale REST value because the + // WS already delivered a status event for this chat. + await waitFor(() => { + expect(result.current.chatStatus).toBe("running"); + }); + }); }); describe("thinking indicator event ordering", () => { diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index 63db88d184..58ec521da2 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts @@ -102,6 +102,14 @@ export const useChatStore = ( // messages are corrected when switching back to a chat whose // queue was drained while the user was away. const wsQueueUpdateReceivedRef = useRef(false); + // Tracks whether the WebSocket has delivered a status event for + // the current chat. Once true, the WS is the authoritative + // source for chatStatus and the REST-fetched chatRecord.status + // must not overwrite it. Without this guard, a React Query + // refetch (e.g. on window focus) can regress chatStatus to a + // stale value like "pending", causing shouldApplyMessagePart() + // to drop all incoming parts. + const wsStatusReceivedRef = useRef(false); const activeChatIDRef = useRef(null); const prevChatIDRef = useRef(chatID); // Snapshot of the chatMessages elements from the last sync effect @@ -185,12 +193,19 @@ export const useChatStore = ( }, [chatID, chatMessages, store]); useEffect(() => { - store.setChatStatus(chatRecord?.status ?? null); + // Only hydrate from REST when the WebSocket hasn't delivered + // a status event yet. Once the WS is the authoritative + // source, a stale REST refetch must not overwrite the + // fresher WS-delivered value. + if (!wsStatusReceivedRef.current) { + store.setChatStatus(chatRecord?.status ?? null); + } }, [chatRecord?.status, store]); useEffect(() => { queuedMessagesHydratedChatIDRef.current = null; wsQueueUpdateReceivedRef.current = false; + wsStatusReceivedRef.current = false; store.setQueuedMessages([]); if (!chatID) { return; @@ -496,6 +511,7 @@ export const useChatStore = ( continue; } + wsStatusReceivedRef.current = true; store.clearRetryState(); store.setChatStatus(nextStatus); if (nextStatus === "pending" || nextStatus === "waiting") {