From a16755dd66aa5204a02dd133e9652cb9c498a281 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Sun, 5 Apr 2026 14:10:26 -0400 Subject: [PATCH] fix: prevent stale REST status from dropping streamed parts (#24040) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `useEffect` that syncs `chatRecord.status` from React Query unconditionally overwrites the store's `chatStatus`. The `chat(chatId)` query has no `staleTime` (defaults to 0), so it refetches on window focus, remount, etc. If the REST response catches a transient `"pending"` status (e.g. between multi-step tool-call cycles), it regresses `chatStatus` from `"running"` to `"pending"`. Since `shouldApplyMessagePart()` drops ALL parts when status is `"pending"` or `"waiting"`, every incoming `message_part` event is silently discarded — not even buffered. Parts are visible on the WebSocket but nothing renders, and the UI shows "Response is taking longer than expected". A page reload fixes it because a fresh REST fetch returns the current status. **Fix:** Add `wsStatusReceivedRef` — once the WebSocket delivers a status event, it becomes the authoritative source and REST refetches can no longer overwrite it. This mirrors the existing `wsQueueUpdateReceivedRef` pattern already used for queued messages. The ref resets on chat change. > Generated with [Coder Agents](https://coder.com/agents) --- .../ChatConversation/chatStore.test.tsx | 71 +++++++++++++++++++ .../ChatConversation/useChatStore.ts | 18 ++++- 2 files changed, 88 insertions(+), 1 deletion(-) 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") {