diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index 4de73c5731f..a2c0fffb3d5 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/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(); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.test.ts b/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.test.ts index 535bb900844..be23e3508e6 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.test.ts +++ b/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({ diff --git a/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.ts b/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.ts index d9acdd1af82..6fb1005036e 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/liveStatusModel.ts +++ b/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 }; }; diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index 06da6a8bf87..60ee1b12e1c 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/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) => diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.stories.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.stories.tsx index 8030066aa45..cda4b231765 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.stories.tsx +++ b/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 ( + + + + ); + }, + 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();