fix(site/src/pages/AgentsPage): clear streamed message when a chat errors (#28277)

site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx

@@ -41,6 +41,7 @@ import { createTestQueryClient } from "#/testHelpers/renderHelpers";
 import type { OneWayMessageEvent } from "#/utils/OneWayWebSocket";
 import {
 	selectChatStatus,
+	selectHasStreamState,
 	selectIsAwaitingFirstStreamChunk,
 	selectMessagesByID,
 	selectOrderedMessageIDs,
@@ -53,6 +54,8 @@ import {
 	useChatSelector,
 	useChatStore,
 } from "./chatStore";
+import { deriveLiveStatus } from "./liveStatusModel";
+import { buildStreamTools } from "./streamState";
 
 vi.mock("#/api/api", () => ({
 	watchChat: vi.fn(),
@@ -2911,6 +2914,474 @@ describe("useChatStore", () => {
 		});
 	});
 
+	// Regression coverage for a chat that stopped with an error but kept the
+	// streamed tool spinning. The error event must clear the partial stream.
+	it("clears the in-flight streamed tool after a terminal error event", async () => {
+		vi.useFakeTimers({ shouldAdvanceTime: true });
+
+		const chatID = "chat-error-clears-running-tool";
+		const existingMessage = buildMessage(
+			chatID,
+			1,
+			"user",
+			"create a workspace",
+		);
+		const mockSocket = createMockSocket();
+		mockWatchChatReturn(mockSocket);
+
+		const queryClient = createTestQueryClient();
+		const wrapper = createWrapper(queryClient);
+		const setChatErrorReason = vi.fn();
+		const clearChatErrorReason = vi.fn();
+
+		const { result } = renderHook(
+			() => {
+				const { store } = useChatStore({
+					chatID,
+					chatMessages: [existingMessage],
+					chatRecord: buildChat(chatID),
+					chatMessagesData: {
+						messages: [existingMessage],
+						queued_messages: [],
+						has_more: false,
+					},
+					chatQueuedMessages: [],
+					setChatErrorReason,
+					clearChatErrorReason,
+				});
+				return {
+					store,
+					hasStreamState: useChatSelector(store, selectHasStreamState),
+					chatStatus: useChatSelector(store, selectChatStatus),
+					streamState: useChatSelector(store, selectStreamState),
+					streamError: useChatSelector(store, selectStreamError),
+				};
+			},
+			{ wrapper },
+		);
+
+		await waitFor(() => {
+			expect(watchChat).toHaveBeenCalledWith(chatID, 1);
+		});
+
+		// A create_workspace tool call begins but never gets a result.
+		act(() => {
+			mockSocket.emitData({
+				type: "message_part",
+				chat_id: chatID,
+				message_part: {
+					part: {
+						type: "tool-call",
+						tool_call_id: "create-workspace-1",
+						tool_name: "create_workspace",
+						args: { name: "dev" },
+					},
+				},
+			});
+		});
+
+		await act(async () => {
+			vi.advanceTimersByTime(1);
+		});
+
+		// A tool call with no result renders as running.
+		await waitFor(() => {
+			expect(
+				result.current.streamState?.toolCalls["create-workspace-1"]?.name,
+			).toBe("create_workspace");
+		});
+		expect(
+			buildStreamTools(
+				result.current.streamState?.toolCalls,
+				result.current.streamState?.toolResults,
+			),
+		).toMatchObject([{ id: "create-workspace-1", status: "running" }]);
+
+		act(() => {
+			mockSocket.emitData({
+				type: "error",
+				chat_id: chatID,
+				error: {
+					message: "The chat session ended unexpectedly.",
+					kind: "generic",
+					retryable: false,
+				},
+			});
+		});
+
+		await waitFor(() => {
+			expect(result.current.chatStatus).toBe("error");
+			expect(result.current.streamError).toMatchObject({
+				kind: "generic",
+				message: "The chat session ended unexpectedly.",
+			});
+		});
+
+		expect(result.current.streamState).toBeNull();
+		expect(result.current.hasStreamState).toBe(false);
+		expect(
+			buildStreamTools(
+				result.current.streamState?.toolCalls,
+				result.current.streamState?.toolResults,
+			),
+		).toEqual([]);
+
+		const liveStatus = deriveLiveStatus({
+			streamState: result.current.streamState,
+			retryState: null,
+			reconnectState: null,
+			streamError: result.current.streamError,
+			persistedError: null,
+			isAwaitingFirstStreamChunk: false,
+			chatStatus: result.current.chatStatus,
+		});
+		expect(liveStatus.phase).toBe("failed");
+		expect(liveStatus.hasAccumulatedOutput).toBe(false);
+	});
+
+	// A part already queued from the errored turn can still arrive after the
+	// error and repopulate the stream. This is accepted because the gate that
+	// would drop it cannot tell it apart from a part belonging to the next
+	// turn. The error status and callout are preserved either way.
+	it("keeps the error status when a late part arrives after a terminal error", async () => {
+		vi.useFakeTimers({ shouldAdvanceTime: true });
+
+		const chatID = "chat-error-late-parts";
+		const existingMessage = buildMessage(
+			chatID,
+			1,
+			"user",
+			"create a workspace",
+		);
+		const mockSocket = createMockSocket();
+		mockWatchChatReturn(mockSocket);
+
+		const queryClient = createTestQueryClient();
+		const wrapper = createWrapper(queryClient);
+		const setChatErrorReason = vi.fn();
+		const clearChatErrorReason = vi.fn();
+
+		const { result } = renderHook(
+			() => {
+				const { store } = useChatStore({
+					chatID,
+					chatMessages: [existingMessage],
+					chatRecord: buildChat(chatID),
+					chatMessagesData: {
+						messages: [existingMessage],
+						queued_messages: [],
+						has_more: false,
+					},
+					chatQueuedMessages: [],
+					setChatErrorReason,
+					clearChatErrorReason,
+				});
+				return {
+					chatStatus: useChatSelector(store, selectChatStatus),
+					streamState: useChatSelector(store, selectStreamState),
+					streamError: useChatSelector(store, selectStreamError),
+				};
+			},
+			{ wrapper },
+		);
+
+		await waitFor(() => {
+			expect(watchChat).toHaveBeenCalledWith(chatID, 1);
+		});
+
+		act(() => {
+			mockSocket.emitData({
+				type: "error",
+				chat_id: chatID,
+				error: {
+					message: "The chat session ended unexpectedly.",
+					kind: "generic",
+					retryable: false,
+				},
+			});
+		});
+
+		await waitFor(() => {
+			expect(result.current.chatStatus).toBe("error");
+			expect(result.current.streamState).toBeNull();
+		});
+
+		// A late part from the errored turn arrives in a later frame.
+		act(() => {
+			mockSocket.emitData({
+				type: "message_part",
+				chat_id: chatID,
+				message_part: {
+					part: { type: "text", text: "late partial output" },
+				},
+			});
+		});
+
+		await act(async () => {
+			vi.advanceTimersByTime(1);
+		});
+
+		// The part is kept and the error status and callout are preserved.
+		await waitFor(() => {
+			expect(result.current.streamState?.blocks).toEqual([
+				{ type: "response", text: "late partial output" },
+			]);
+		});
+		expect(result.current.chatStatus).toBe("error");
+		expect(result.current.streamError).toMatchObject({
+			kind: "generic",
+			message: "The chat session ended unexpectedly.",
+		});
+	});
+
+	// The next turn must stream again after an error. Its first part can
+	// arrive before its running status because the server sends parts and
+	// status on separate paths, so the part must be kept either way.
+	it("re-enables streaming parts on the next turn after a terminal error", async () => {
+		vi.useFakeTimers({ shouldAdvanceTime: true });
+
+		const chatID = "chat-error-next-turn";
+		const existingMessage = buildMessage(
+			chatID,
+			1,
+			"user",
+			"create a workspace",
+		);
+		const mockSocket = createMockSocket();
+		mockWatchChatReturn(mockSocket);
+
+		const queryClient = createTestQueryClient();
+		const wrapper = createWrapper(queryClient);
+		const setChatErrorReason = vi.fn();
+		const clearChatErrorReason = vi.fn();
+
+		const { result } = renderHook(
+			() => {
+				const { store } = useChatStore({
+					chatID,
+					chatMessages: [existingMessage],
+					chatRecord: buildChat(chatID),
+					chatMessagesData: {
+						messages: [existingMessage],
+						queued_messages: [],
+						has_more: false,
+					},
+					chatQueuedMessages: [],
+					setChatErrorReason,
+					clearChatErrorReason,
+				});
+				return {
+					chatStatus: useChatSelector(store, selectChatStatus),
+					streamState: useChatSelector(store, selectStreamState),
+				};
+			},
+			{ wrapper },
+		);
+
+		await waitFor(() => {
+			expect(watchChat).toHaveBeenCalledWith(chatID, 1);
+		});
+
+		act(() => {
+			mockSocket.emitData({
+				type: "error",
+				chat_id: chatID,
+				error: {
+					message: "The chat session ended unexpectedly.",
+					kind: "generic",
+					retryable: false,
+				},
+			});
+		});
+
+		await waitFor(() => {
+			expect(result.current.chatStatus).toBe("error");
+			expect(result.current.streamState).toBeNull();
+		});
+
+		// The next turn's first part arrives before its running status.
+		act(() => {
+			mockSocket.emitData({
+				type: "message_part",
+				chat_id: chatID,
+				message_part: {
+					part: { type: "text", text: "retry output" },
+				},
+			});
+		});
+
+		await act(async () => {
+			vi.advanceTimersByTime(1);
+		});
+
+		act(() => {
+			mockSocket.emitData({
+				type: "status",
+				chat_id: chatID,
+				status: { status: "running" },
+			});
+		});
+
+		await act(async () => {
+			vi.advanceTimersByTime(1);
+		});
+
+		await waitFor(() => {
+			expect(result.current.streamState?.blocks).toEqual([
+				{ type: "response", text: "retry output" },
+			]);
+		});
+	});
+
+	// A non-error status means a new turn began, so the prior error must be
+	// cleared for the new stream to render.
+	it("clears streamError when a non-error status arrives after an error", async () => {
+		vi.useFakeTimers({ shouldAdvanceTime: true });
+
+		const chatID = "chat-error-clears-on-running";
+		const existingMessage = buildMessage(
+			chatID,
+			1,
+			"user",
+			"create a workspace",
+		);
+		const mockSocket = createMockSocket();
+		mockWatchChatReturn(mockSocket);
+
+		const queryClient = createTestQueryClient();
+		const wrapper = createWrapper(queryClient);
+		const setChatErrorReason = vi.fn();
+		const clearChatErrorReason = vi.fn();
+
+		const { result } = renderHook(
+			() => {
+				const { store } = useChatStore({
+					chatID,
+					chatMessages: [existingMessage],
+					chatRecord: buildChat(chatID),
+					chatMessagesData: {
+						messages: [existingMessage],
+						queued_messages: [],
+						has_more: false,
+					},
+					chatQueuedMessages: [],
+					setChatErrorReason,
+					clearChatErrorReason,
+				});
+				return {
+					chatStatus: useChatSelector(store, selectChatStatus),
+					streamError: useChatSelector(store, selectStreamError),
+				};
+			},
+			{ wrapper },
+		);
+
+		await waitFor(() => {
+			expect(watchChat).toHaveBeenCalledWith(chatID, 1);
+		});
+
+		act(() => {
+			mockSocket.emitData({
+				type: "error",
+				chat_id: chatID,
+				error: {
+					message: "The chat session ended unexpectedly.",
+					kind: "generic",
+					retryable: false,
+				},
+			});
+		});
+
+		await waitFor(() => {
+			expect(result.current.streamError).not.toBeNull();
+		});
+
+		act(() => {
+			mockSocket.emitData({
+				type: "status",
+				chat_id: chatID,
+				status: { status: "running" },
+			});
+		});
+
+		await waitFor(() => {
+			expect(result.current.streamError).toBeNull();
+		});
+		expect(result.current.chatStatus).toBe("running");
+	});
+
+	// A status from a turn that was already running must not clear a stored
+	// request failure, because the failed prompt was never retried.
+	it("keeps streamError when a running turn emits a non-error status", async () => {
+		vi.useFakeTimers({ shouldAdvanceTime: true });
+
+		const chatID = "chat-keeps-request-error";
+		const existingMessage = buildMessage(
+			chatID,
+			1,
+			"user",
+			"create a workspace",
+		);
+		const mockSocket = createMockSocket();
+		mockWatchChatReturn(mockSocket);
+
+		const queryClient = createTestQueryClient();
+		const wrapper = createWrapper(queryClient);
+		const setChatErrorReason = vi.fn();
+		const clearChatErrorReason = vi.fn();
+
+		const { result } = renderHook(
+			() => {
+				const { store } = useChatStore({
+					chatID,
+					chatMessages: [existingMessage],
+					chatRecord: { ...buildChat(chatID), status: "running" },
+					chatMessagesData: {
+						messages: [existingMessage],
+						queued_messages: [],
+						has_more: false,
+					},
+					chatQueuedMessages: [],
+					setChatErrorReason,
+					clearChatErrorReason,
+				});
+				return {
+					store,
+					chatStatus: useChatSelector(store, selectChatStatus),
+					streamError: useChatSelector(store, selectStreamError),
+				};
+			},
+			{ wrapper },
+		);
+
+		await waitFor(() => {
+			expect(watchChat).toHaveBeenCalledWith(chatID, 1);
+		});
+
+		// A queued send fails while the prior turn is still running.
+		act(() => {
+			result.current.store.setStreamError({
+				kind: "generic",
+				message: "Prompt could not be sent.",
+			});
+		});
+		expect(result.current.streamError).not.toBeNull();
+
+		// The still-running turn completes and emits a non-error status.
+		act(() => {
+			mockSocket.emitData({
+				type: "status",
+				chat_id: chatID,
+				status: { status: "waiting" },
+			});
+		});
+
+		await waitFor(() => {
+			expect(result.current.chatStatus).toBe("waiting");
+		});
+		expect(result.current.streamError).not.toBeNull();
+	});
+
 	it("uses fallback message when error event has no message", async () => {
 		immediateAnimationFrame();
 

site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.test.ts

@@ -132,13 +132,26 @@ describe("deriveLiveStatus", () => {
 		).toEqual({ phase: "streaming", hasAccumulatedOutput: false });
 	});
 
-	it("tracks accumulated output on failed streams", () => {
+	it("suppresses accumulated output after a terminal stream error", () => {
 		expect(
 			derive({
 				streamState: buildStreamState({
 					blocks: [{ type: "response", text: "Partial response" }],
 				}),
 				streamError: buildStreamError(),
+				chatStatus: "error",
+			}),
+		).toEqual(failedStatus);
+	});
+
+	it("keeps accumulated output when a non-terminal error leaves the stream live", () => {
+		expect(
+			derive({
+				streamState: buildStreamState({
+					blocks: [{ type: "response", text: "Partial response" }],
+				}),
+				streamError: buildStreamError(),
+				chatStatus: "running",
 			}),
 		).toEqual({
 			...failedStatus,
@@ -146,6 +159,17 @@ describe("deriveLiveStatus", () => {
 		});
 	});
 
+	it("does not stream stale output while a retry is in flight after an error", () => {
+		expect(
+			derive({
+				streamState: buildStreamState({
+					blocks: [{ type: "response", text: "Partial response" }],
+				}),
+				chatStatus: "error",
+			}),
+		).toEqual({ phase: "idle", hasAccumulatedOutput: false });
+	});
+
 	it("passes provider detail through failed status", () => {
 		expect(
 			derive({
@@ -173,6 +197,21 @@ describe("deriveLiveStatus", () => {
 		});
 	});
 
+	it("suppresses accumulated output while reconnecting after a terminal error", () => {
+		expect(
+			derive({
+				streamState: buildStreamState({
+					blocks: [{ type: "response", text: "Partial response" }],
+				}),
+				reconnectState: buildReconnectState(),
+				chatStatus: "error",
+			}),
+		).toEqual({
+			...reconnectingStatus,
+			hasAccumulatedOutput: false,
+		});
+	});
+
 	it("prioritizes retrying over failed and reconnecting", () => {
 		expect(
 			derive({

site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.ts

@@ -114,17 +114,25 @@ export const deriveLiveStatus = ({
 	chatStatus,
 }: DeriveLiveStatusParams): LiveStatusModel => {
 	const hasAccumulatedOutput = getHasAccumulatedOutput(streamState);
+	// The stream is cleared on error, so leftover blocks are stale.
+	const showOutput = chatStatus === "error" ? false : hasAccumulatedOutput;
 
 	if (retryState) {
-		return toRetryingLiveStatus(retryState, { hasAccumulatedOutput });
+		return toRetryingLiveStatus(retryState, {
+			hasAccumulatedOutput: showOutput,
+		});
 	}
 
 	if (streamError) {
-		return toFailedLiveStatus(streamError, { hasAccumulatedOutput });
+		return toFailedLiveStatus(streamError, {
+			hasAccumulatedOutput: showOutput,
+		});
 	}
 
 	if (reconnectState) {
-		return toReconnectingLiveStatus(reconnectState, { hasAccumulatedOutput });
+		return toReconnectingLiveStatus(reconnectState, {
+			hasAccumulatedOutput: showOutput,
+		});
 	}
 
 	// The interrupt outranks stream leftovers: while the worker drains and
@@ -138,13 +146,15 @@ export const deriveLiveStatus = ({
 		return { phase: "starting", hasAccumulatedOutput };
 	}
 
-	if (streamState !== null) {
+	if (streamState !== null && chatStatus !== "error") {
 		return { phase: "streaming", hasAccumulatedOutput };
 	}
 
 	if (persistedError) {
-		return toFailedLiveStatus(persistedError, { hasAccumulatedOutput });
+		return toFailedLiveStatus(persistedError, {
+			hasAccumulatedOutput: showOutput,
+		});
 	}
 
-	return { phase: "idle", hasAccumulatedOutput };
+	return { phase: "idle", hasAccumulatedOutput: showOutput };
 };

site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts

@@ -569,17 +569,10 @@ export const useChatStore = (
 						continue;
 					}
 
-					// 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") {
+					// Flush buffered parts before a durable message
+					// commits them. Other events must not flush because
+					// some clear the stream and a flush would restore it.
+					if (streamEvent.type === "message") {
 						flushMessageParts();
 					}
 
@@ -629,12 +622,20 @@ export const useChatStore = (
 							streamReportedWaiting = nextStatus === "waiting";
 							wsStatusReceivedRef.current = true;
 							store.clearRetryState();
+							const prevStatus = store.getSnapshot().chatStatus;
 							store.applyServerChatStatus(nextStatus);
 							if (nextStatus === "waiting") {
 								discardBufferedParts();
 							}
 							if (nextStatus !== "error") {
 								clearChatErrorReasonEvent(chatID);
+								// Only an errored chat starting a new turn
+								// supersedes the stored error. A status from a
+								// turn that was already running must not clear
+								// an unrelated request failure.
+								if (prevStatus === "error") {
+									store.clearStreamError();
+								}
 							}
 							updateSidebarChat((chat) =>
 								chat.status === nextStatus
@@ -648,10 +649,15 @@ export const useChatStore = (
 								kind: "generic",
 								message: "Chat processing failed.",
 							};
-							streamReportedWaiting = false;
+							// An error ends the turn. Clear the partial
+							// stream so no tool keeps spinning. Parts stay
+							// ungated because the next turn can send a part
+							// before its running status.
+							discardBufferedParts();
 							wsStatusReceivedRef.current = true;
 							store.applyServerChatStatus("error");
 							store.setStreamError(reason);
+							store.clearStreamState();
 							store.clearRetryState();
 							setChatErrorReasonEvent(chatID, reason);
 							updateSidebarChat((chat) =>

site/src/pages/AgentsPage/components/ChatPageContent.stories.tsx

@@ -1,7 +1,7 @@
 import { MessageScroller } from "@shadcn/react/message-scroller";
 import type { Meta, StoryObj } from "@storybook/react-vite";
 import type { FC } from "react";
-import { expect, fn, userEvent, within } from "storybook/test";
+import { expect, fn, userEvent, waitFor, within } from "storybook/test";
 import type * as TypesGen from "#/api/typesGenerated";
 import { MockChatQueuedMessage } from "#/testHelpers/chatEntities";
 import { ChatWorkspaceContext } from "../context/ChatWorkspaceContext";
@@ -171,6 +171,51 @@ export const DurableUnresolvedWorkspaceToolRuns: Story = {
 	},
 };
 
+// Matches the fixed terminal error path.
+const errorClearsStreamStore = createChatStore();
+export const ErrorClearsStreamingTool: Story = {
+	render: () => {
+		errorClearsStreamStore.resetTransientState();
+		errorClearsStreamStore.replaceMessages([
+			buildMessage(1, "user", [{ type: "text", text: "Create a workspace" }]),
+		]);
+		errorClearsStreamStore.setChatStatus("running");
+		errorClearsStreamStore.applyMessagePart({
+			type: "tool-call",
+			tool_call_id: "create-workspace-call",
+			tool_name: "create_workspace",
+			args: { name: "dev" },
+		});
+
+		return (
+			<ChatWorkspaceContext value={{ workspaceId: "workspace-1" }}>
+				<StoryChatPageTimeline store={errorClearsStreamStore} />
+			</ChatWorkspaceContext>
+		);
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		expect(canvas.getByText("Creating workspace…")).toBeInTheDocument();
+
+		errorClearsStreamStore.batch(() => {
+			errorClearsStreamStore.applyServerChatStatus("error");
+			errorClearsStreamStore.setStreamError({
+				kind: "generic",
+				message: "The chat session ended unexpectedly.",
+			});
+			errorClearsStreamStore.clearStreamState();
+		});
+
+		await waitFor(() => {
+			expect(canvas.queryByText("Creating workspace…")).toBeNull();
+		});
+		expect(canvas.getByText("Request failed")).toBeInTheDocument();
+		expect(
+			canvas.getByText("The chat session ended unexpectedly."),
+		).toBeInTheDocument();
+	},
+};
+
 export const HiddenAssistantPlaceholderDoesNotRender: Story = {
 	render: () => {
 		const store = createChatStore();