fix: narrow finalize boundary-drop guard (#27711) (thanks @scz2011)

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

@@ -152,6 +152,35 @@ describe("TuiStreamAssembler", () => {
     expect(finalText).toBe("Draft line 1");
   });
 
+  it("prefers final text when non-text blocks appear only in final payload", () => {
+    const assembler = new TuiStreamAssembler();
+    assembler.ingestDelta(
+      "run-5c",
+      {
+        role: "assistant",
+        content: [
+          { type: "text", text: "Draft line 1" },
+          { type: "text", text: "Draft line 2" },
+        ],
+      },
+      false,
+    );
+
+    const finalText = assembler.finalize(
+      "run-5c",
+      {
+        role: "assistant",
+        content: [
+          { type: "tool_use", name: "search" },
+          { type: "text", text: "Draft line 2" },
+        ],
+      },
+      false,
+    );
+
+    expect(finalText).toBe("Draft line 2");
+  });
+
   it("accepts richer final payload when it extends streamed text", () => {
     const assembler = new TuiStreamAssembler();
     assembler.ingestDelta(

src/tui/tui-stream-assembler.ts

@@ -97,7 +97,10 @@ export class TuiStreamAssembler {
     state: RunStreamState,
     message: unknown,
     showThinking: boolean,
-    opts?: { protectBoundaryDrops?: boolean },
+    opts?: {
+      protectBoundaryDrops?: boolean;
+      useIncomingNonTextForBoundaryDrops?: boolean;
+    },
   ) {
     const thinkingText = extractThinkingFromMessage(message);
     const contentText = extractContentFromMessage(message);
@@ -108,9 +111,11 @@ export class TuiStreamAssembler {
     }
     if (contentText) {
       const nextContentBlocks = textBlocks.length > 0 ? textBlocks : [contentText];
+      const useIncomingNonTextForBoundaryDrops = opts?.useIncomingNonTextForBoundaryDrops !== false;
       const shouldPreserveBoundaryDroppedText =
         opts?.protectBoundaryDrops === true &&
-        (state.sawNonTextContentBlocks || sawNonTextContentBlocks) &&
+        (state.sawNonTextContentBlocks ||
+          (useIncomingNonTextForBoundaryDrops && sawNonTextContentBlocks)) &&
         isDroppedBoundaryTextBlockSubset({
           streamedTextBlocks: state.contentBlocks,
           finalTextBlocks: nextContentBlocks,
@@ -151,7 +156,10 @@ export class TuiStreamAssembler {
     const streamedDisplayText = state.displayText;
     const streamedTextBlocks = [...state.contentBlocks];
     const streamedSawNonTextContentBlocks = state.sawNonTextContentBlocks;
-    this.updateRunState(state, message, showThinking, { protectBoundaryDrops: true });
+    this.updateRunState(state, message, showThinking, {
+      protectBoundaryDrops: true,
+      useIncomingNonTextForBoundaryDrops: false,
+    });
     const finalComposed = state.displayText;
     const shouldKeepStreamedText =
       streamedSawNonTextContentBlocks &&