fix(tui): narrow streamed-text fallback to tool-block regression shape

src/tui/tui-stream-assembler.test.ts

@@ -98,6 +98,7 @@ describe("TuiStreamAssembler", () => {
         role: "assistant",
         content: [
           { type: "text", text: "Before tool call" },
+          { type: "tool_use", name: "search" },
           { type: "text", text: "After tool call" },
         ],
       },
@@ -108,7 +109,10 @@ describe("TuiStreamAssembler", () => {
       "run-5",
       {
         role: "assistant",
-        content: [{ type: "text", text: "After tool call" }],
+        content: [
+          { type: "tool_use", name: "search" },
+          { type: "text", text: "After tool call" },
+        ],
       },
       false,
     );
@@ -116,6 +120,32 @@ describe("TuiStreamAssembler", () => {
     expect(finalText).toBe("Before tool call\nAfter tool call");
   });
 
+  it("keeps non-empty final text for plain text prefix/suffix updates", () => {
+    const assembler = new TuiStreamAssembler();
+    assembler.ingestDelta(
+      "run-5b",
+      {
+        role: "assistant",
+        content: [
+          { type: "text", text: "Draft line 1" },
+          { type: "text", text: "Draft line 2" },
+        ],
+      },
+      false,
+    );
+
+    const finalText = assembler.finalize(
+      "run-5b",
+      {
+        role: "assistant",
+        content: [{ type: "text", text: "Draft line 1" }],
+      },
+      false,
+    );
+
+    expect(finalText).toBe("Draft line 1");
+  });
+
   it("accepts richer final payload when it extends streamed text", () => {
     const assembler = new TuiStreamAssembler();
     assembler.ingestDelta(

src/tui/tui-stream-assembler.ts

@@ -8,34 +8,69 @@ import {
 type RunStreamState = {
   thinkingText: string;
   contentText: string;
+  contentBlocks: string[];
+  sawNonTextContentBlocks: boolean;
   displayText: string;
 };
 
-function mergeTextPreferRicher(currentText: string, nextText: string): string {
-  const current = currentText.trim();
-  const next = nextText.trim();
-  if (!next) {
-    return current;
+function extractTextBlocksAndSignals(message: unknown): {
+  textBlocks: string[];
+  sawNonTextContentBlocks: boolean;
+} {
+  if (!message || typeof message !== "object") {
+    return { textBlocks: [], sawNonTextContentBlocks: false };
   }
-  if (!current || current === next) {
-    return next;
+  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 (next.includes(current)) {
-    return next;
+  if (!Array.isArray(content)) {
+    return { textBlocks: [], sawNonTextContentBlocks: false };
   }
-  // Keep streamed text only when the newer payload looks like a dropped
-  // leading/trailing block (the known regression shape), not any substring.
-  const currentLines = current.split("\n");
-  const nextLines = next.split("\n");
-  if (currentLines.length > nextLines.length) {
-    const isDroppedLeadingBlock =
-      currentLines.slice(currentLines.length - nextLines.length).join("\n") === next;
-    const isDroppedTrailingBlock = currentLines.slice(0, nextLines.length).join("\n") === next;
-    if (isDroppedLeadingBlock || isDroppedTrailingBlock) {
-      return current;
+
+  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 next;
+  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 {
@@ -47,6 +82,8 @@ export class TuiStreamAssembler {
       state = {
         thinkingText: "",
         contentText: "",
+        contentBlocks: [],
+        sawNonTextContentBlocks: false,
         displayText: "",
       };
       this.runs.set(runId, state);
@@ -57,12 +94,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 = mergeTextPreferRicher(state.thinkingText, thinkingText);
+      state.thinkingText = thinkingText;
     }
     if (contentText) {
-      state.contentText = mergeTextPreferRicher(state.contentText, contentText);
+      state.contentText = contentText;
+      state.contentBlocks = textBlocks.length > 0 ? textBlocks : [contentText];
+    }
+    if (sawNonTextContentBlocks) {
+      state.sawNonTextContentBlocks = true;
     }
 
     const displayText = composeThinkingAndContent({
@@ -89,10 +131,18 @@ 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,
+      finalText: shouldKeepStreamedText ? streamedDisplayText : finalComposed,
       streamedText: streamedDisplayText,
     });