fix(site): ensure Thinking indicator appears regardless of WebSocket event ordering (#23884)

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).

<details><summary>Implementation plan & decision log</summary>

### 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.

</details>
This commit is contained in:
Kyle Carberry
2026-04-01 07:45:38 -04:00
committed by GitHub
parent f5b98aa12d
commit 4b52656958
2 changed files with 275 additions and 3 deletions
@@ -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) => (
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
);
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) => (
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
);
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) => (
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
);
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();
@@ -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<TypesGen.ServerSentEvent>,
) => {
@@ -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));
}