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

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

@@ -89,4 +89,56 @@ describe("TuiStreamAssembler", () => {
 
     expect(second).toBeNull();
   });
+
+  it("keeps richer streamed text when final payload drops earlier blocks", () => {
+    const assembler = new TuiStreamAssembler();
+    assembler.ingestDelta(
+      "run-5",
+      {
+        role: "assistant",
+        content: [
+          { type: "text", text: "Before tool call" },
+          { type: "text", text: "After tool call" },
+        ],
+      },
+      false,
+    );
+
+    const finalText = assembler.finalize(
+      "run-5",
+      {
+        role: "assistant",
+        content: [{ type: "text", text: "After tool call" }],
+      },
+      false,
+    );
+
+    expect(finalText).toBe("Before tool call\nAfter tool call");
+  });
+
+  it("accepts richer final payload when it extends streamed text", () => {
+    const assembler = new TuiStreamAssembler();
+    assembler.ingestDelta(
+      "run-6",
+      {
+        role: "assistant",
+        content: [{ type: "text", text: "Before tool call" }],
+      },
+      false,
+    );
+
+    const finalText = assembler.finalize(
+      "run-6",
+      {
+        role: "assistant",
+        content: [
+          { type: "text", text: "Before tool call" },
+          { type: "text", text: "After tool call" },
+        ],
+      },
+      false,
+    );
+
+    expect(finalText).toBe("Before tool call\nAfter tool call");
+  });
 });

src/tui/tui-stream-assembler.ts

@@ -11,6 +11,24 @@ type RunStreamState = {
   displayText: string;
 };
 
+function mergeTextPreferRicher(currentText: string, nextText: string): string {
+  const current = currentText.trim();
+  const next = nextText.trim();
+  if (!next) {
+    return current;
+  }
+  if (!current || current === next) {
+    return next;
+  }
+  if (next.includes(current)) {
+    return next;
+  }
+  if (current.includes(next)) {
+    return current;
+  }
+  return next;
+}
+
 export class TuiStreamAssembler {
   private runs = new Map<string, RunStreamState>();
 
@@ -32,10 +50,10 @@ export class TuiStreamAssembler {
     const contentText = extractContentFromMessage(message);
 
     if (thinkingText) {
-      state.thinkingText = thinkingText;
+      state.thinkingText = mergeTextPreferRicher(state.thinkingText, thinkingText);
     }
     if (contentText) {
-      state.contentText = contentText;
+      state.contentText = mergeTextPreferRicher(state.contentText, contentText);
     }
 
     const displayText = composeThinkingAndContent({
@@ -61,11 +79,12 @@ export class TuiStreamAssembler {
 
   finalize(runId: string, message: unknown, showThinking: boolean): string {
     const state = this.getOrCreateRun(runId);
+    const streamedDisplayText = state.displayText;
     this.updateRunState(state, message, showThinking);
     const finalComposed = state.displayText;
     const finalText = resolveFinalAssistantText({
       finalText: finalComposed,
-      streamedText: state.displayText,
+      streamedText: streamedDisplayText,
     });
 
     this.runs.delete(runId);