From 4b5265695883cfbe3f4558b0ff0ef345a3791612 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Wed, 1 Apr 2026 07:45:38 -0400 Subject: [PATCH] fix(site): ensure Thinking indicator appears regardless of WebSocket event ordering (#23884) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The "Thinking..." indicator intermittently failed to render after submitting a message. The behavior depended on the order of events within a single WebSocket frame. ## Root Cause `flushMessageParts()` was called before **all** non-`message_part` events in the batch loop. When the server sent `[message_part, status:"running"]` in the same SSE chunk: 1. `message_part` → pushed to `partsBuf` 2. `status:"running"` → `flushMessageParts()` applied parts **first** → `streamState` became non-null → then `chatStatus` set to `"running"` 3. Subscriber saw `streamState != null && chatStatus == "running"` → `selectIsAwaitingFirstStreamChunk` returned `false` → no "Thinking..." When events arrived in the reverse order (`[status:"running", message_part]`), the indicator worked because the status was set before parts were applied. ## Fix - Move `flushMessageParts()` to only fire before `message` and `error` events (which need prior parts visible) - Add `discardBufferedParts()` for events that clear stream state (`status:"pending"/"waiting"`, `retry`) so the deferred `setTimeout(0)` flush doesn't re-populate cleared state - Status changes are now always applied before parts within a batch, and the deferred flush gives React one render cycle to show "Thinking..." | Event | Flush? | Rationale | |---|---|---| | `message` | YES | Durable commit must include all stream parts | | `error` | YES | Partial output should be visible alongside error | | `status` | NO | Status must be set before parts so "starting" phase renders | | `retry` | DISCARD | Retry clears stream state; flushing would re-populate it | | `queue_update` | NO | Doesn't interact with stream state | ## Tests (written first, failing before fix) 1. **"shows starting phase when message_part arrives before status:running in same batch"** — the exact bug scenario 2. **"shows starting phase when status:running arrives before message_part in same batch"** — verifies the "good" order still works 3. **"discards buffered parts when status transitions to pending"** — verifies parts don't leak through pending transitions All tests are deterministic (fake timers, no race conditions).
Implementation plan & decision log ### Why not reorder events within the batch? Reordering would change the semantic ordering of events from the server, which could have subtle side effects. The simpler approach is to be selective about when parts are flushed. ### Why discard (not flush) before pending/waiting/retry? These events clear `streamState`. If parts were flushed before the clear, they'd be visible for one frame then disappear. If the deferred flush ran after the clear, it would re-populate the state. Discarding is the only correct behavior. ### Why keep flush before error? Errors should surface partial output so the user can see what the agent was doing when it failed.
--- .../ChatConversation/chatStore.test.tsx | 244 ++++++++++++++++++ .../ChatConversation/useChatStore.ts | 34 ++- 2 files changed, 275 insertions(+), 3 deletions(-) diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index 692bb3543b..e9c7273b26 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx @@ -37,6 +37,7 @@ import type * as TypesGen from "#/api/typesGenerated"; import type { OneWayMessageEvent } from "#/utils/OneWayWebSocket"; import { selectChatStatus, + selectIsAwaitingFirstStreamChunk, selectOrderedMessageIDs, selectQueuedMessages, selectReconnectState, @@ -2953,6 +2954,249 @@ describe("useChatStore", () => { }); }); +describe("thinking indicator event ordering", () => { + it("shows starting phase when message_part arrives before status:running in same batch", async () => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + immediateAnimationFrame(); + + const chatID = "chat-thinking-parts-before-status"; + const userMsg = makeMessage(chatID, 1, "user", "hello"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [userMsg], + chatRecord: { ...makeChat(chatID), status: "running" }, + chatMessagesData: { + messages: [userMsg], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + chatStatus: useChatSelector(store, selectChatStatus), + isAwaiting: useChatSelector(store, selectIsAwaitingFirstStreamChunk), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + // Server sends message_part BEFORE status:running in the same + // WebSocket frame. This is the event ordering that previously + // caused the "Thinking..." indicator to be skipped. + act(() => { + mockSocket.emitDataBatch([ + { + type: "message_part", + chat_id: chatID, + message_part: { + part: { type: "reasoning", text: "Let me think..." }, + }, + }, + { + type: "status", + chat_id: chatID, + status: { status: "running" }, + }, + ]); + }); + + // After the batch, the status should be "running" but stream + // parts should NOT have been applied yet (deferred to + // setTimeout). This is the window where "Thinking..." shows. + await waitFor(() => { + expect(result.current.chatStatus).toBe("running"); + expect(result.current.streamState).toBeNull(); + expect(result.current.isAwaiting).toBe(true); + }); + + // Let the deferred parts flush fire (setTimeout 0). + await act(async () => { + vi.advanceTimersByTime(1); + }); + + // Now stream state should be populated. + await waitFor(() => { + expect(result.current.streamState).not.toBeNull(); + expect(result.current.isAwaiting).toBe(false); + }); + }); + + it("shows starting phase when status:running arrives before message_part in same batch", async () => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + immediateAnimationFrame(); + + const chatID = "chat-thinking-status-before-parts"; + const userMsg = makeMessage(chatID, 1, "user", "hello"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [userMsg], + chatRecord: { ...makeChat(chatID), status: "running" }, + chatMessagesData: { + messages: [userMsg], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + chatStatus: useChatSelector(store, selectChatStatus), + isAwaiting: useChatSelector(store, selectIsAwaitingFirstStreamChunk), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + // Server sends status:running BEFORE message_part (the "good" order). + act(() => { + mockSocket.emitDataBatch([ + { + type: "status", + chat_id: chatID, + status: { status: "running" }, + }, + { + type: "message_part", + chat_id: chatID, + message_part: { + part: { type: "text", text: "Hello" }, + }, + }, + ]); + }); + + // Same contract: status set, parts deferred. + await waitFor(() => { + expect(result.current.chatStatus).toBe("running"); + expect(result.current.streamState).toBeNull(); + expect(result.current.isAwaiting).toBe(true); + }); + + // Let the deferred parts flush fire. + await act(async () => { + vi.advanceTimersByTime(1); + }); + + await waitFor(() => { + expect(result.current.streamState).not.toBeNull(); + expect(result.current.isAwaiting).toBe(false); + }); + }); + + it("discards buffered parts when status transitions to pending", async () => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + immediateAnimationFrame(); + + const chatID = "chat-thinking-discard-pending"; + const userMsg = makeMessage(chatID, 1, "user", "hello"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [userMsg], + chatRecord: { ...makeChat(chatID), status: "running" }, + chatMessagesData: { + messages: [userMsg], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + chatStatus: useChatSelector(store, selectChatStatus), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + // Server sends message_part then immediately transitions to pending. + // The buffered parts must be discarded (not applied) because + // pending status clears stream state. + act(() => { + mockSocket.emitDataBatch([ + { + type: "message_part", + chat_id: chatID, + message_part: { + part: { type: "text", text: "partial response" }, + }, + }, + { + type: "status", + chat_id: chatID, + status: { status: "pending" }, + }, + ]); + }); + + await waitFor(() => { + expect(result.current.chatStatus).toBe("pending"); + expect(result.current.streamState).toBeNull(); + }); + + // Even after timers fire, parts should not re-appear. + await act(async () => { + vi.advanceTimersByTime(50); + }); + + expect(result.current.streamState).toBeNull(); + }); +}); + describe("updateSidebarChat via stream events", () => { it("updates sidebar chat status on status stream event", async () => { immediateAnimationFrame(); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index cac0a21dab..6721aa67e0 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts @@ -381,8 +381,8 @@ export const useChatStore = ( }; // Immediate flush for non-message_part events that need - // the parts applied before they execute (e.g. a status - // change right after the last part). + // the parts applied before they execute (e.g. a durable + // message commit right after the last part). const flushMessageParts = () => { if (partsBuf.length === 0) { return; @@ -398,6 +398,19 @@ export const useChatStore = ( } store.applyMessageParts(parts); }; + + // Discard buffered parts without applying them. Used when + // stream state is about to be cleared (pending, waiting, + // retry) — flushing would re-populate the state that the + // event is about to clear. + const discardBufferedParts = () => { + partsBuf.length = 0; + if (partsFlushTimer !== null) { + clearTimeout(partsFlushTimer); + partsFlushTimer = null; + } + }; + const handleMessage = ( payload: OneWayMessageEvent, ) => { @@ -444,7 +457,20 @@ export const useChatStore = ( } continue; } - flushMessageParts(); + + // Only flush buffered parts before events that + // need them applied first. `message` events + // commit durable state that must include all + // stream parts. `error` events should surface + // partial output. Other events (status, retry, + // queue_update) must NOT flush — status changes + // need to be visible before parts so the + // "Thinking..." indicator can render, and retry + // clears stream state which a flush would + // re-populate. + if (streamEvent.type === "message" || streamEvent.type === "error") { + flushMessageParts(); + } switch (streamEvent.type) { case "message": { @@ -494,6 +520,7 @@ export const useChatStore = ( store.clearRetryState(); store.setChatStatus(nextStatus); if (nextStatus === "pending" || nextStatus === "waiting") { + discardBufferedParts(); store.clearStreamState(); store.clearRetryState(); } @@ -530,6 +557,7 @@ export const useChatStore = ( } const retry = streamEvent.retry; if (retry) { + discardBufferedParts(); store.clearStreamState(); store.setRetryState(normalizeRetryState(retry)); }