diff --git a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx index d124f03114..41168f472b 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx @@ -2402,6 +2402,65 @@ export const DurableUpdateFansOutToOlderPage: Story = { }, }; +/** + * The streamed thinking block is asserted via its disclosure header + * because smoothed body text never reveals in the test iframe + * (requestAnimationFrame is suspended). + */ +export const StreamedPartWhileStatusWaiting: Story = { + parameters: { + queries: buildQueries( + { + id: CHAT_ID, + ...baseChatFields, + title: "Stale status stream", + status: "waiting", + }, + { + messages: [ + { + id: 1, + chat_id: CHAT_ID, + created_at: "2026-02-18T00:05:00.000Z", + role: "user", + content: [{ type: "text", text: "Start the next turn" }], + }, + ], + queued_messages: [], + has_more: false, + }, + { diffUrl: undefined }, + ), + webSocket: { + "/chats/": [ + { + event: "message", + data: JSON.stringify([ + { + type: "message_part", + chat_id: CHAT_ID, + message_part: { + part: { + type: "reasoning", + text: "Streaming while the chat still reads waiting", + }, + }, + }, + ] satisfies TypesGen.ChatStreamEvent[]), + }, + ], + }, + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await canvas.findByText("Start the next turn"); + // The disclosure header is the only "Thinking" text in the DOM + // at this point: the generic indicator is suppressed at status + // "waiting". + expect(await canvas.findByText("Thinking")).toBeVisible(); + }, +}; + /** * Live agent turn with streaming reasoning and a back-to-back flurry of * in-progress file tool calls. The persisted history establishes context diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index 80c4ad6c7e..4de73c5731 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx @@ -1409,7 +1409,7 @@ describe("useChatStore", () => { }); }); - it("ignores message_part updates while chat is waiting", async () => { + it("applies message_part updates while a REST-hydrated chat status reads waiting", async () => { immediateAnimationFrame(); const chatID = "chat-1"; @@ -1427,7 +1427,10 @@ describe("useChatStore", () => { const { store } = useChatStore({ chatID, chatMessages: [existingMessage], - chatRecord: buildChat(chatID), + // REST hydrates the store with a waiting status, but the + // stream has not delivered any status event, so the + // server's view may already be ahead. + chatRecord: { ...buildChat(chatID), status: "waiting" }, chatMessagesData: { messages: [existingMessage], queued_messages: [], @@ -1456,18 +1459,89 @@ describe("useChatStore", () => { role: "assistant", part: { type: "text", - text: "first", + text: "live output", }, }, }); }); + await waitFor(() => { + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "live output" }, + ]); + }); + }); + + it("drops message_part updates after the stream reports waiting", async () => { + immediateAnimationFrame(); + + const chatID = "chat-1"; + const existingMessage = buildMessage(chatID, 1, "user", "hello"); + 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: "waiting" }, + chatMessagesData: { + messages: [existingMessage], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + chatStatus: useChatSelector(store, selectChatStatus), + store, + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + act(() => { + mockSocket.emitData({ + type: "status", + chat_id: chatID, + status: { status: "running" }, + }); + }); + + act(() => { + mockSocket.emitData({ + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "first" }, + }, + }); + }); + await waitFor(() => { expect(result.current.streamState?.blocks).toEqual([ { type: "response", text: "first" }, ]); }); + // The stream reports waiting: the turn is over server-side. + // A part arriving now comes from the closed episode draining + // (the server keeps closed episodes subscribed for replay for + // up to 15s) and must not repopulate the stream. act(() => { mockSocket.emitData({ type: "status", @@ -1477,14 +1551,7 @@ describe("useChatStore", () => { }); await waitFor(() => { - // Stream state is preserved after status=waiting (the - // durable message event handles cleanup via - // needsStreamReset). Only new message_parts should be - // blocked by the shouldApplyMessagePart gate. - expect(result.current.streamState).not.toBeNull(); - expect(result.current.streamState?.blocks).toEqual([ - { type: "response", text: "first" }, - ]); + expect(result.current.chatStatus).toBe("waiting"); }); act(() => { @@ -1493,20 +1560,70 @@ describe("useChatStore", () => { chat_id: chatID, message_part: { role: "assistant", - part: { - type: "text", - text: "late", - }, + part: { type: "text", text: "late drain" }, + }, + }); + }); + + // Wait past the coalesced flush window so the drop is + // observable rather than "not yet flushed". + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 10)); + }); + + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "first" }, + ]); + + // An optimistic send status must not reopen the window; drain + // parts from the closed episode are still dropped. + act(() => { + result.current.store.setChatStatus("running"); + }); + + act(() => { + mockSocket.emitData({ + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "drain after send" }, + }, + }); + }); + + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 10)); + }); + + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "first" }, + ]); + + // The stream reporting running reopens the window: the next + // turn's parts must flow again. + act(() => { + mockSocket.emitData({ + type: "status", + chat_id: chatID, + status: { status: "running" }, + }); + }); + + act(() => { + mockSocket.emitData({ + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "next turn" }, }, }); }); await waitFor(() => { - // The late message_part should not be applied because - // shouldApplyMessagePart gates on waiting. - // Stream state still shows the original "first". expect(result.current.streamState?.blocks).toEqual([ - { type: "response", text: "first" }, + { type: "response", text: "firstnext turn" }, ]); }); }); @@ -4384,7 +4501,7 @@ describe("thinking indicator event ordering", () => { }); }); - it("discards buffered parts when status transitions to pending", async () => { + it("discards buffered parts when status transitions to waiting", async () => { vi.useFakeTimers({ shouldAdvanceTime: true }); immediateAnimationFrame(); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index b60efc0dec..06da6a8bf8 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts @@ -143,8 +143,8 @@ export const useChatStore = ( // source for chatStatus and the REST-fetched chatRecord.status // must not overwrite it. Without this guard, a React Query // refetch (e.g. on window focus) can regress chatStatus to a - // stale value like "waiting", causing shouldApplyMessagePart() - // to drop all incoming parts. + // stale value like "waiting", hiding live status from the user + // until the next server event arrives. const wsStatusReceivedRef = useRef(false); const [pendingStatusResync, setPendingStatusResync] = useState(false); const pendingStatusResyncUpdatedAtRef = useRef(null); @@ -423,9 +423,13 @@ export const useChatStore = ( let historyResetPending = false; const historyReplacementBuf: TypesGen.ChatMessage[] = []; - const shouldApplyMessagePart = (): boolean => { - return store.getSnapshot().chatStatus !== "waiting"; - }; + // Set when the stream reports "waiting", cleared by any other + // stream status. While set, parts are dropped: they are late + // leftovers from the finished turn. REST and optimistic + // statuses never set it, because they can lag a live turn. + let streamReportedWaiting = false; + + const shouldKeepMessagePart = (): boolean => !streamReportedWaiting; const schedulePartsFlush = () => { if (partsFlushTimer !== null || partsBuf.length === 0) { @@ -436,11 +440,11 @@ export const useChatStore = ( if (disposed || activeChatIDRef.current !== chatID) { return; } - const parts = partsBuf.splice(0); - if (parts.length === 0 || !shouldApplyMessagePart()) { + if (!shouldKeepMessagePart()) { + partsBuf.length = 0; return; } - store.applyMessageParts(parts); + store.applyMessageParts(partsBuf.splice(0)); }, 0); }; @@ -455,11 +459,11 @@ export const useChatStore = ( clearTimeout(partsFlushTimer); partsFlushTimer = null; } - const parts = partsBuf.splice(0); - if (activeChatIDRef.current !== chatID || !shouldApplyMessagePart()) { + if (activeChatIDRef.current !== chatID || !shouldKeepMessagePart()) { + partsBuf.length = 0; return; } - store.applyMessageParts(parts); + store.applyMessageParts(partsBuf.splice(0)); }; // Discard buffered parts without applying them. Used when @@ -521,7 +525,7 @@ export const useChatStore = ( continue; } commitHistoryReplacement(); - if (!shouldApplyMessagePart()) { + if (!shouldKeepMessagePart()) { continue; } const part = streamEvent.message_part?.part; @@ -622,6 +626,7 @@ export const useChatStore = ( continue; } + streamReportedWaiting = nextStatus === "waiting"; wsStatusReceivedRef.current = true; store.clearRetryState(); store.applyServerChatStatus(nextStatus); @@ -643,6 +648,7 @@ export const useChatStore = ( kind: "generic", message: "Chat processing failed.", }; + streamReportedWaiting = false; wsStatusReceivedRef.current = true; store.applyServerChatStatus("error"); store.setStreamError(reason); @@ -698,9 +704,8 @@ export const useChatStore = ( clearTimeout(partsFlushTimer); partsFlushTimer = null; } - const nextParts = partsBuf.splice(0); - if (shouldApplyMessagePart()) { - store.applyMessageParts(nextParts); + if (shouldKeepMessagePart()) { + store.applyMessageParts(partsBuf.splice(0)); } } }