diff --git a/ui/src/pages/chat/chat-gateway.test.ts b/ui/src/pages/chat/chat-gateway.test.ts index 58819234f17..ab04be0d298 100644 --- a/ui/src/pages/chat/chat-gateway.test.ts +++ b/ui/src/pages/chat/chat-gateway.test.ts @@ -1821,6 +1821,37 @@ describe("handleChatGatewayEvent", () => { expect(state.chatRunError).toEqual({ summary: "Error: gateway disconnected" }); }); + it("preserves streamed text before a differing terminal assistant message", () => { + const terminalMessage = { + role: "assistant", + content: [{ type: "text", text: "Configure provider auth, then try again." }], + timestamp: 101, + }; + const state = createState({ + sessionKey: "main", + chatRunId: "run-1", + chatStream: "Partial answer before gateway error.", + chatStreamStartedAt: 100, + }); + + expect( + handleChatGatewayEvent(state, { + runId: "run-1", + sessionKey: "main", + state: "error", + errorMessage: "gateway disconnected", + message: terminalMessage, + }), + ).toBe("error"); + expect(state.chatMessages).toHaveLength(2); + expectTextChatMessage( + state.chatMessages[0], + "assistant", + "Partial answer before gateway error.", + ); + expect(state.chatMessages[1]).toEqual(terminalMessage); + }); + it("preserves terminal extensions after a tool splits the stream", () => { const message = { role: "assistant", @@ -1855,6 +1886,37 @@ describe("handleChatGatewayEvent", () => { expect(state.chatRunError).toEqual({ summary: "Error: gateway disconnected" }); }); + it("preserves a split stream when the terminal message only overlaps its prefix", () => { + const terminalMessage = { + role: "assistant", + content: [{ type: "text", text: "First thought. Configure provider auth." }], + timestamp: 101, + }; + const state = createState({ + sessionKey: "main", + chatRunId: "run-1", + chatStream: "After tool.", + chatStreamStartedAt: 100, + }) as ChatState & { + chatStreamSegments: Array<{ text: string; ts: number; toolCallId: string }>; + }; + state.chatStreamSegments = [{ text: "First thought.", ts: 90, toolCallId: "call-1" }]; + + expect( + handleChatGatewayEvent(state, { + runId: "run-1", + sessionKey: "main", + state: "error", + errorMessage: "gateway disconnected", + message: terminalMessage, + }), + ).toBe("error"); + expect(state.chatMessages).toHaveLength(3); + expectTextChatMessage(state.chatMessages[0], "assistant", "First thought."); + expectTextChatMessage(state.chatMessages[1], "assistant", "After tool."); + expect(state.chatMessages[2]).toEqual(terminalMessage); + }); + it("preserves terminal extensions when split stream punctuation is adjacent", () => { const state = createState({ sessionKey: "main", diff --git a/ui/src/pages/chat/chat-gateway.ts b/ui/src/pages/chat/chat-gateway.ts index 815b6faed05..d2b40f428c9 100644 --- a/ui/src/pages/chat/chat-gateway.ts +++ b/ui/src/pages/chat/chat-gateway.ts @@ -20,6 +20,7 @@ import { appendTerminalAssistantMessage, clearToolStreamSegments, hasVisibleStreamParts, + terminalMessageReplacesVisibleStream, } from "./stream-reconciliation.ts"; import { authoritativeHistoryAppliedForRun, @@ -310,23 +311,41 @@ function handleChatEvent(state: ChatState, payload?: ChatEventPayload) { ); if (hadActiveRunBeforeEvent) { if (visiblePayloadMessage && !projectedErrorMessage) { - if ( - hasVisibleStreamParts(state, { - includeCurrent: false, + const replacesVisibleStream = terminalMessageReplacesVisibleStream( + visiblePayloadMessage, + state, + { isHiddenStreamText: isHiddenAssistantStreamText, - }) - ) { + persistCommentary: state.settings?.chatPersistCommentary === true, + }, + ); + if (replacesVisibleStream) { + if ( + hasVisibleStreamParts(state, { + includeCurrent: false, + isHiddenStreamText: isHiddenAssistantStreamText, + }) + ) { + state.chatMessages = materializeVisibleAssistantStreamMessages( + state.chatMessages, + state, + { includeCurrent: false }, + ); + clearToolStreamSegments(state); + } + state.chatMessages = appendTerminalAssistantMessage( + state.chatMessages, + visiblePayloadMessage, + ); + } else { state.chatMessages = materializeVisibleAssistantStreamMessages( state.chatMessages, state, - { includeCurrent: false }, + { includeCurrent: true }, ); clearToolStreamSegments(state); + state.chatMessages = [...state.chatMessages, visiblePayloadMessage]; } - state.chatMessages = appendTerminalAssistantMessage( - state.chatMessages, - visiblePayloadMessage, - ); } else { state.chatMessages = materializeVisibleAssistantStreamMessages(state.chatMessages, state, { includeCurrent: true, diff --git a/ui/src/pages/chat/stream-reconciliation.ts b/ui/src/pages/chat/stream-reconciliation.ts index a0bdb1a9728..bcb16a2bf61 100644 --- a/ui/src/pages/chat/stream-reconciliation.ts +++ b/ui/src/pages/chat/stream-reconciliation.ts @@ -462,6 +462,37 @@ export function hasVisibleStreamParts( return visibleAssistantStreamParts(state, opts).length > 0; } +export function terminalMessageReplacesVisibleStream( + message: unknown, + state: StreamReconciliationState, + opts: Pick, +): boolean { + const terminalText = extractText(message)?.trim(); + if (!terminalText) { + return false; + } + const parts = visibleAssistantStreamParts(state, { + includeCurrent: true, + isHiddenStreamText: opts.isHiddenStreamText, + }).filter((part) => opts.persistCommentary === true || !part.itemId); + if (parts.length === 0) { + return false; + } + + let searchStart = 0; + for (const [index, part] of parts.entries()) { + const text = part.text.trim(); + const matchIndex = terminalText.indexOf(text, searchStart); + // A terminal replacement must preserve the complete visible stream in order, + // beginning with its first part. Otherwise both outputs are user-visible. + if (matchIndex < 0 || (index === 0 && matchIndex !== 0)) { + return false; + } + searchStart = matchIndex + text.length; + } + return true; +} + function currentToolStreamMessageIndex( messages: unknown[], state: StreamReconciliationState,