fix: preserve streamed text before terminal errors (#108509)

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",

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,

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<MaterializeVisibleStreamOptions, "isHiddenStreamText" | "persistCommentary">,
+): 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,