diff --git a/src/tui/tui-stream-assembler.test.ts b/src/tui/tui-stream-assembler.test.ts index 34476446c0e..ed3aab34788 100644 --- a/src/tui/tui-stream-assembler.test.ts +++ b/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( diff --git a/src/tui/tui-stream-assembler.ts b/src/tui/tui-stream-assembler.ts index dccc4e4c37d..0475e47b974 100644 --- a/src/tui/tui-stream-assembler.ts +++ b/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; + 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; + 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, });