fix(tui): preserve streamed text when final payload regresses (#15452) (#15573)

Merged via /review-pr -> /prepare-pr -> /merge-pr.

Prepared head SHA: e4a5e3c8a6
Co-authored-by: TsekaLuk <79151285+TsekaLuk@users.noreply.github.com>
Co-authored-by: steipete <58493+steipete@users.noreply.github.com>
Reviewed-by: @steipete

src/tui/tui-stream-assembler.ts

@@ -8,9 +8,73 @@ import {
 type RunStreamState = {
   thinkingText: string;
   contentText: string;
+  contentBlocks: string[];
+  sawNonTextContentBlocks: boolean;
   displayText: string;
 };
 
+function extractTextBlocksAndSignals(message: unknown): {
+  textBlocks: string[];
+  sawNonTextContentBlocks: boolean;
+} {
+  if (!message || typeof message !== "object") {
+    return { textBlocks: [], sawNonTextContentBlocks: false };
+  }
+  const record = message as Record<string, unknown>;
+  const content = record.content;
+
+  if (typeof content === "string") {
+    const text = content.trim();
+    return {
+      textBlocks: text ? [text] : [],
+      sawNonTextContentBlocks: false,
+    };
+  }
+  if (!Array.isArray(content)) {
+    return { textBlocks: [], sawNonTextContentBlocks: false };
+  }
+
+  const textBlocks: string[] = [];
+  let sawNonTextContentBlocks = false;
+  for (const block of content) {
+    if (!block || typeof block !== "object") {
+      continue;
+    }
+    const rec = block as Record<string, unknown>;
+    if (rec.type === "text" && typeof rec.text === "string") {
+      const text = rec.text.trim();
+      if (text) {
+        textBlocks.push(text);
+      }
+      continue;
+    }
+    if (typeof rec.type === "string" && rec.type !== "thinking") {
+      sawNonTextContentBlocks = true;
+    }
+  }
+  return { textBlocks, sawNonTextContentBlocks };
+}
+
+function isDroppedBoundaryTextBlockSubset(params: {
+  streamedTextBlocks: string[];
+  finalTextBlocks: string[];
+}): boolean {
+  const { streamedTextBlocks, finalTextBlocks } = params;
+  if (finalTextBlocks.length === 0 || finalTextBlocks.length >= streamedTextBlocks.length) {
+    return false;
+  }
+
+  const prefixMatches = finalTextBlocks.every(
+    (block, index) => streamedTextBlocks[index] === block,
+  );
+  if (prefixMatches) {
+    return true;
+  }
+
+  const suffixStart = streamedTextBlocks.length - finalTextBlocks.length;
+  return finalTextBlocks.every((block, index) => streamedTextBlocks[suffixStart + index] === block);
+}
+
 export class TuiStreamAssembler {
   private runs = new Map<string, RunStreamState>();
 
@@ -20,6 +84,8 @@ export class TuiStreamAssembler {
       state = {
         thinkingText: "",
         contentText: "",
+        contentBlocks: [],
+        sawNonTextContentBlocks: false,
         displayText: "",
       };
       this.runs.set(runId, state);
@@ -30,12 +96,17 @@ export class TuiStreamAssembler {
   private updateRunState(state: RunStreamState, message: unknown, showThinking: boolean) {
     const thinkingText = extractThinkingFromMessage(message);
     const contentText = extractContentFromMessage(message);
+    const { textBlocks, sawNonTextContentBlocks } = extractTextBlocksAndSignals(message);
 
     if (thinkingText) {
       state.thinkingText = thinkingText;
     }
     if (contentText) {
       state.contentText = contentText;
+      state.contentBlocks = textBlocks.length > 0 ? textBlocks : [contentText];
+    }
+    if (sawNonTextContentBlocks) {
+      state.sawNonTextContentBlocks = true;
     }
 
     const displayText = composeThinkingAndContent({
@@ -61,11 +132,20 @@ export class TuiStreamAssembler {
 
   finalize(runId: string, message: unknown, showThinking: boolean): string {
     const state = this.getOrCreateRun(runId);
+    const streamedDisplayText = state.displayText;
+    const streamedTextBlocks = [...state.contentBlocks];
+    const streamedSawNonTextContentBlocks = state.sawNonTextContentBlocks;
     this.updateRunState(state, message, showThinking);
     const finalComposed = state.displayText;
+    const shouldKeepStreamedText =
+      streamedSawNonTextContentBlocks &&
+      isDroppedBoundaryTextBlockSubset({
+        streamedTextBlocks,
+        finalTextBlocks: state.contentBlocks,
+      });
     const finalText = resolveFinalAssistantText({
-      finalText: finalComposed,
-      streamedText: state.displayText,
+      finalText: shouldKeepStreamedText ? streamedDisplayText : finalComposed,
+      streamedText: streamedDisplayText,
     });
 
     this.runs.delete(runId);