From bb7c5f93f31c1f56289df48788316b4c86ad2948 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Thu, 19 Mar 2026 18:59:28 +0200
Subject: [PATCH] fix(site): don't wipe stream state when non-assistant message
 arrives (#23295)

scheduleStreamReset() fired for every durable message event with
changed=true, including user messages (e.g. promoted queued messages).
When a batch contained trailing message_parts followed by a user
message event, the batch loop flushed the parts (building stream
state), then scheduleStreamReset cleared it immediately.

Restrict the reset to assistant messages, which are the only role
that ends a streaming turn.
---
 .../AgentDetail/ChatContext.test.tsx          | 105 ++++++++++++++++++
 .../AgentsPage/AgentDetail/ChatContext.ts     |   2 +-
 2 files changed, 106 insertions(+), 1 deletion(-)

diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
index c429ff5e8e..8730423572 100644
--- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
+++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
@@ -2450,6 +2450,111 @@ describe("useChatStore", () => {
 			expect(result.current.queuedMessages).toHaveLength(0);
 		});
 	});
+
+	it("does not wipe in-progress stream state when user message arrives in batch", async () => {
+		immediateAnimationFrame();
+
+		const chatID = "chat-promote-stream";
+		const msg1 = makeMessage(chatID, 1, "user", "hello");
+		const msg2 = makeMessage(chatID, 2, "assistant", "hi");
+
+		const mockSocket = createMockSocket();
+		vi.mocked(watchChat).mockReturnValue(mockSocket as never);
+
+		const queryClient = createTestQueryClient();
+		const wrapper: FC<PropsWithChildren> = ({ children }) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+		const setChatErrorReason = vi.fn();
+		const clearChatErrorReason = vi.fn();
+
+		const queuedMsg = makeQueuedMessage(chatID, 10, "follow-up");
+		const initialMessages = [msg1, msg2];
+
+		const initialOptions = {
+			chatID,
+			chatMessages: initialMessages,
+			chatRecord: makeChat(chatID),
+			chatMessagesData: {
+				messages: initialMessages,
+				queued_messages: [queuedMsg],
+				has_more: false,
+			},
+			chatQueuedMessages: [queuedMsg],
+			setChatErrorReason,
+			clearChatErrorReason,
+		};
+
+		const { result } = renderHook(
+			(options: Parameters<typeof useChatStore>[0]) => {
+				const { store } = useChatStore(options);
+				return {
+					store,
+					streamState: useChatSelector(store, selectStreamState),
+					chatStatus: useChatSelector(store, selectChatStatus),
+					orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs),
+				};
+			},
+			{ initialProps: initialOptions, wrapper },
+		);
+
+		await waitFor(() => {
+			expect(result.current.orderedMessageIDs).toEqual([1, 2]);
+		});
+
+		// Open the WebSocket and set the chat to running.
+		act(() => {
+			mockSocket.emitOpen();
+		});
+		act(() => {
+			mockSocket.emitData({
+				type: "status",
+				chat_id: chatID,
+				status: { status: "running" },
+			});
+		});
+
+		await waitFor(() => {
+			expect(result.current.chatStatus).toBe("running");
+		});
+
+		// Deliver a batch containing trailing message_parts for
+		// the current response followed by the promoted user
+		// message. The batch loop flushes pending parts when it
+		// hits the message event (building stream state). Before
+		// the fix, scheduleStreamReset would fire for the user
+		// message because it only checked `changed`, and with
+		// immediateAnimationFrame the RAF fires synchronously,
+		// wiping the stream state that was just built.
+		const promotedUser = makeMessage(chatID, 3, "user", "follow-up");
+
+		act(() => {
+			mockSocket.emitDataBatch([
+				{
+					type: "message_part",
+					chat_id: chatID,
+					message_part: {
+						part: { type: "text", text: "I am helping you" },
+					},
+				},
+				{
+					type: "message",
+					chat_id: chatID,
+					message: promotedUser,
+				},
+			]);
+		});
+
+		// The stream state must survive: the promoted user message
+		// should not wipe the in-progress assistant stream.
+		await waitFor(() => {
+			expect(result.current.orderedMessageIDs).toContain(3);
+			expect(result.current.streamState).not.toBeNull();
+			const blocks = result.current.streamState?.blocks ?? [];
+			const textBlock = blocks.find((b) => b.type === "response");
+			expect(textBlock).toBeDefined();
+		});
+	});
 });
 
 describe("updateSidebarChat via stream events", () => {
diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
index 6014024a55..84d446107b 100644
--- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
+++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
@@ -725,7 +725,7 @@ export const useChatStore = (
 						) {
 							lastMessageIdRef.current = message.id;
 						}
-						if (changed) {
+						if (changed && message.role === "assistant") {
 							scheduleStreamReset();
 						}
 						// Do not update updated_at here. The global