fix(chat): fix stream state discrepancies between frontend and backend (#22437)

## Summary

Fixes four frontend↔backend discrepancies in chat stream state
management that could cause duplicate content, UI flicker, and stale
stream state.

### Backend fixes (`coderd/chatd/chatd.go`)

**1. No-pubsub path double-replayed message_part events**

`Subscribe()` built an `initialSnapshot` containing `message_part`
events from `localSnapshot`, then the no-pubsub goroutine replayed the
same `localSnapshot` into the `mergedEvents` channel. Since `streamChat`
sends the snapshot first then reads the channel, the frontend received
every `message_part` twice. `applyMessagePartToStreamState` doesn't
deduplicate — text gets concatenated, so content appeared doubled.

Fix: Only forward live `localParts` in the no-pubsub goroutine; the
snapshot already contains the historical events.

**2. Snapshot missing status event**

The initial snapshot never included a `status` event. The frontend's
`shouldApplyMessagePart()` gates on status (`pending`/`waiting`), but
the initial status came from a separate REST query via `useEffect`.
During the race window between snapshot arrival and REST resolution,
`message_part` events could be incorrectly accepted or rejected.

Fix: Prepend a `status` event to the snapshot after loading the chat
from DB, so the frontend has the authoritative status from the very
first batch.

### Frontend fixes (`ChatContext.ts`)

**3. Scheduled stream reset not canceled by subsequent message_parts**

When a `message` event arrived, `scheduleStreamReset()` queued
`clearStreamState` via `requestAnimationFrame`. If new `message_part`
events arrived in the next WebSocket frame before the rAF fired, they
were pushed to `pendingMessageParts` without canceling the scheduled
reset. The rAF would fire between frames, clearing stream state, then
the next flush would re-populate it — causing a visible flash.

Fix: Call `cancelScheduledStreamReset()` when accumulating
`message_part` events.

**4. startTransition race with synchronous clearStreamState**

`flushMessageParts` wrapped `applyMessageParts` in `startTransition`,
which React can defer. If a `status: "waiting"` event arrived in the
same batch after `message_part` events, the status handler cleared
stream state synchronously, but the deferred `applyMessageParts`
callback could fire afterward and re-populate it.

Fix: Re-check `shouldApplyMessagePart()` inside the `startTransition`
callback at execution time.

### Tests added

- **Go**: `TestSubscribeSnapshotIncludesStatusEvent` — asserts the first
snapshot event is a status event
- **Go**: `TestSubscribeNoPubsubNoDuplicateMessageParts` — asserts the
events channel doesn't replay snapshot events
- **TS**: `cancels scheduled stream reset when message_part arrives
after message` — verifies stream state survives a [message,
message_part] batch
- **TS**: `does not apply message parts after status changes to waiting`
— verifies deferred applyMessageParts respects status transitions

coderd/chatd/chatd.go

@@ -1086,6 +1086,23 @@ func (p *Server) Subscribe(
 		}
 	}
 
+	// Include the current chat status in the snapshot so the
+	// frontend can gate message_part processing correctly from
+	// the very first batch, without waiting for a separate REST
+	// query.
+	if err == nil {
+		statusEvent := codersdk.ChatStreamEvent{
+			Type:   codersdk.ChatStreamEventTypeStatus,
+			ChatID: chatID,
+			Status: &codersdk.ChatStreamStatus{
+				Status: codersdk.ChatStatus(chat.Status),
+			},
+		}
+		// Prepend so the frontend sees the status before any
+		// message_part events.
+		initialSnapshot = append([]codersdk.ChatStreamEvent{statusEvent}, initialSnapshot...)
+	}
+
 	// Track the last message ID we've seen for DB queries
 	var lastMessageID int64
 	if len(messages) > 0 {
@@ -1307,16 +1324,11 @@ func (p *Server) Subscribe(
 			}
 		}()
 	} else {
-		// No pubsub, just merge local parts
+		// No pubsub, just merge local parts.
+		// localSnapshot was already included in initialSnapshot,
+		// so only forward new events here.
 		go func() {
 			defer close(mergedEvents)
-			for _, event := range localSnapshot {
-				select {
-				case <-mergedCtx.Done():
-					return
-				case mergedEvents <- event:
-				}
-			}
 			for event := range localParts {
 				select {
 				case <-mergedCtx.Done():

coderd/chatd/chatd_test.go

@@ -632,6 +632,74 @@ func TestUpdateChatStatusPersistsLastError(t *testing.T) {
 	require.False(t, fromDB.LastError.Valid)
 }
 
+func TestSubscribeSnapshotIncludesStatusEvent(t *testing.T) {
+	t.Parallel()
+
+	db, ps := dbtestutil.NewDB(t)
+	replica := newTestServer(t, db, ps, uuid.New())
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	user, model := seedChatDependencies(ctx, t, db)
+
+	chat, err := replica.CreateChat(ctx, chatd.CreateOptions{
+		OwnerID:            user.ID,
+		Title:              "status-snapshot",
+		ModelConfigID:      model.ID,
+		InitialUserContent: []fantasy.Content{fantasy.TextContent{Text: "hello"}},
+	})
+	require.NoError(t, err)
+
+	snapshot, _, cancel, ok := replica.Subscribe(ctx, chat.ID, nil)
+	require.True(t, ok)
+	t.Cleanup(cancel)
+
+	// The first event in the snapshot must be a status event.
+	require.NotEmpty(t, snapshot)
+	require.Equal(t, codersdk.ChatStreamEventTypeStatus, snapshot[0].Type)
+	require.NotNil(t, snapshot[0].Status)
+	require.Equal(t, codersdk.ChatStatusPending, snapshot[0].Status.Status)
+}
+
+func TestSubscribeNoPubsubNoDuplicateMessageParts(t *testing.T) {
+	t.Parallel()
+
+	// Use nil pubsub to force the no-pubsub path.
+	db, _ := dbtestutil.NewDB(t)
+	replica := newTestServer(t, db, nil, uuid.New())
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	user, model := seedChatDependencies(ctx, t, db)
+
+	chat, err := replica.CreateChat(ctx, chatd.CreateOptions{
+		OwnerID:            user.ID,
+		Title:              "no-dup-parts",
+		ModelConfigID:      model.ID,
+		InitialUserContent: []fantasy.Content{fantasy.TextContent{Text: "hello"}},
+	})
+	require.NoError(t, err)
+
+	snapshot, events, cancel, ok := replica.Subscribe(ctx, chat.ID, nil)
+	require.True(t, ok)
+	t.Cleanup(cancel)
+
+	// Snapshot should have events (at minimum: status + message).
+	require.NotEmpty(t, snapshot)
+
+	// The events channel should NOT immediately produce any
+	// events — the snapshot already contained everything. Before
+	// the fix, localSnapshot was replayed into the channel,
+	// causing duplicates.
+	select {
+	case event, ok := <-events:
+		if ok {
+			t.Fatalf("unexpected event from channel (would be a duplicate): type=%s", event.Type)
+		}
+		// Channel closed without events is fine.
+	case <-time.After(200 * time.Millisecond):
+		// No events — correct behavior.
+	}
+}
+
 func newTestServer(
 	t *testing.T,
 	db database.Store,

site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx

@@ -1044,6 +1044,95 @@ describe("useChatStore", () => {
 		});
 	});
 
+	it("cancels scheduled stream reset when message_part arrives after message", async () => {
+		immediateAnimationFrame();
+
+		const chatID = "chat-raf";
+		const existingMessage = makeMessage(chatID, 1, "user", "hello");
+		const mockSocket = createMockSocket();
+		vi.mocked(watchChat).mockReturnValue(mockSocket as never);
+
+		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: [existingMessage],
+					chatRecord: makeChat(chatID),
+					chatData: {
+						chat: makeChat(chatID),
+						messages: [existingMessage],
+						queued_messages: [],
+					},
+					chatQueuedMessages: [],
+					setChatErrorReason,
+					clearChatErrorReason,
+				});
+				return {
+					streamState: useChatSelector(store, selectStreamState),
+				};
+			},
+			{ wrapper },
+		);
+
+		await waitFor(() => {
+			expect(watchChat).toHaveBeenCalledWith(chatID);
+		});
+
+		// Build up stream state first.
+		act(() => {
+			mockSocket.emitData({
+				type: "message_part",
+				chat_id: chatID,
+				message_part: {
+					role: "assistant",
+					part: { type: "text", text: "working" },
+				},
+			});
+		});
+
+		await waitFor(() => {
+			expect(result.current.streamState?.blocks).toEqual([
+				{ type: "response", text: "working" },
+			]);
+		});
+
+		// Emit a durable message followed by a message_part in the
+		// same batch. The message handler calls scheduleStreamReset
+		// (via rAF), and the subsequent message_part handler calls
+		// cancelScheduledStreamReset to prevent a flash. The final
+		// flushMessageParts re-populates stream state.
+		act(() => {
+			mockSocket.emitDataBatch([
+				{
+					type: "message",
+					chat_id: chatID,
+					message: makeMessage(chatID, 2, "assistant", "done"),
+				},
+				{
+					type: "message_part",
+					chat_id: chatID,
+					message_part: {
+						role: "assistant",
+						part: { type: "text", text: " more" },
+					},
+				},
+			]);
+		});
+
+		// Stream state should be non-null because the message_part
+		// after the message kept it populated.
+		await waitFor(() => {
+			expect(result.current.streamState).not.toBeNull();
+		});
+	});
+
 	it("startTransition deferred parts are discarded after chat switch", async () => {
 		immediateAnimationFrame();
 
@@ -1214,4 +1303,74 @@ describe("useChatStore", () => {
 		});
 		expect(result.current.queuedMessages).toEqual([]);
 	});
+
+	it("does not apply message parts after status changes to waiting", async () => {
+		immediateAnimationFrame();
+
+		const chatID = "chat-status-guard";
+		const mockSocket = createMockSocket();
+		vi.mocked(watchChat).mockReturnValue(mockSocket as never);
+
+		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: [],
+					chatRecord: makeChat(chatID),
+					chatData: {
+						chat: makeChat(chatID),
+						messages: [],
+						queued_messages: [],
+					},
+					chatQueuedMessages: [],
+					setChatErrorReason,
+					clearChatErrorReason,
+				});
+				return {
+					streamState: useChatSelector(store, selectStreamState),
+				};
+			},
+			{ wrapper },
+		);
+
+		await waitFor(() => {
+			expect(watchChat).toHaveBeenCalledWith(chatID);
+		});
+
+		// Emit a batch with message_parts followed by a status change
+		// to "waiting". The status handler clears stream state
+		// synchronously, and the startTransition guard should prevent
+		// the deferred applyMessageParts from re-populating it.
+		act(() => {
+			mockSocket.emitDataBatch([
+				{
+					type: "message_part",
+					chat_id: chatID,
+					message_part: {
+						role: "assistant",
+						part: { type: "text", text: "should be discarded" },
+					},
+				},
+				{
+					type: "status",
+					chat_id: chatID,
+					status: { status: "waiting" },
+				},
+			]);
+		});
+
+		// Stream state should be null — the status change cleared it,
+		// and the deferred applyMessageParts should not have
+		// re-populated it.
+		await waitFor(() => {
+			expect(result.current.streamState).toBeNull();
+		});
+	});
 });

site/src/pages/AgentsPage/AgentDetail/ChatContext.ts

@@ -584,6 +584,12 @@ export const useChatStore = (
 					if (activeChatIDRef.current !== currentChatID) {
 						return;
 					}
+					// Re-check status at execution time. A status
+					// event processed between scheduling and running
+					// this callback may have cleared stream state.
+					if (!shouldApplyMessagePart()) {
+						return;
+					}
 					store.applyMessageParts(parts);
 				});
 			};
@@ -599,6 +605,7 @@ export const useChatStore = (
 					}
 					const part = asRecord(streamEvent.message_part?.part);
 					if (part) {
+						cancelScheduledStreamReset();
 						pendingMessageParts.push(part);
 					}
 					continue;