diff --git a/config/assertion-safety-baseline.txt b/config/assertion-safety-baseline.txt index d4f111037920..f762d39d384d 100644 --- a/config/assertion-safety-baseline.txt +++ b/config/assertion-safety-baseline.txt @@ -3653,7 +3653,6 @@ src/tui/tui-session-actions.ts 5 src/tui/tui-session-events.ts 1 src/tui/tui-session-projection.ts 2 src/tui/tui-session-run-coordinator.ts 1 -src/tui/tui-stream-assembler.ts 2 src/tui/tui-task-suggestions.ts 6 src/tui/tui.ts 5 src/utils.ts 1 diff --git a/src/tui/tui-stream-assembler.test.ts b/src/tui/tui-stream-assembler.test.ts index 2573157085e9..330fa4f8a75b 100644 --- a/src/tui/tui-stream-assembler.test.ts +++ b/src/tui/tui-stream-assembler.test.ts @@ -14,8 +14,6 @@ const messageWithContent = (content: readonly Record[]) => content, }) as const; -const TEXT_ONLY_TWO_BLOCKS = messageWithContent([text("Draft line 1"), text("Draft line 2")]); - type FinalizeBoundaryCase = { name: string; streamedContent: readonly Record[]; @@ -24,18 +22,6 @@ type FinalizeBoundaryCase = { }; const FINALIZE_BOUNDARY_CASES: FinalizeBoundaryCase[] = [ - { - name: "preserves streamed text when tool-boundary final payload drops prefix blocks", - streamedContent: [text("Before tool call"), toolUse(), text("After tool call")], - finalContent: [toolUse(), text("After tool call")], - expected: "Before tool call\nAfter tool call", - }, - { - name: "preserves streamed text when streamed run had non-text and final drops suffix blocks", - streamedContent: [text("Before tool call"), toolUse(), text("After tool call")], - finalContent: [text("Before tool call")], - expected: "Before tool call\nAfter tool call", - }, { name: "prefers final text when non-text appears only in final payload", streamedContent: [text("Draft line 1"), text("Draft line 2")], @@ -313,19 +299,6 @@ describe("TuiStreamAssembler", () => { ); }); - it("keeps streamed delta text when incoming tool boundary drops a block", () => { - const assembler = new TuiStreamAssembler(); - const first = assembler.ingestDelta("run-delta-boundary", TEXT_ONLY_TWO_BLOCKS, false); - expect(first).toBe("Draft line 1\nDraft line 2"); - - const second = assembler.ingestDelta( - "run-delta-boundary", - messageWithContent([toolUse(), text("Draft line 2")]), - false, - ); - expect(second).toBeNull(); - }); - for (const testCase of FINALIZE_BOUNDARY_CASES) { it(testCase.name, () => { const assembler = new TuiStreamAssembler(); diff --git a/src/tui/tui-stream-assembler.ts b/src/tui/tui-stream-assembler.ts index 835849b4f10f..e63c72abd3f9 100644 --- a/src/tui/tui-stream-assembler.ts +++ b/src/tui/tui-stream-assembler.ts @@ -12,98 +12,9 @@ const MAX_TRACKED_STREAM_RUNS = 200; type RunStreamState = { thinkingText: string; contentText: string; - contentBlocks: string[]; - sawNonTextContentBlocks: boolean; displayText: string; }; -type BoundaryDropMode = "streamed-only" | "streamed-or-incoming"; - -// Pull text blocks out of provider-style content arrays while remembering non-text blocks. -function extractTextBlocksAndSignals(message: unknown): { - textBlocks: string[]; - sawNonTextContentBlocks: boolean; -} { - if (!message || typeof message !== "object") { - return { textBlocks: [], sawNonTextContentBlocks: false }; - } - const record = message as Record; - 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; - 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 }; -} - -// Detects final messages that dropped streamed boundary text around a non-text block. -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); -} - -// Some providers omit text adjacent to images/files in the final message; preserve streamed text. -function shouldPreserveBoundaryDroppedText(params: { - boundaryDropMode: BoundaryDropMode; - streamedSawNonTextContentBlocks: boolean; - incomingSawNonTextContentBlocks: boolean; - streamedTextBlocks: string[]; - nextContentBlocks: string[]; -}) { - const sawEligibleNonTextContent = - params.boundaryDropMode === "streamed-or-incoming" - ? params.streamedSawNonTextContentBlocks || params.incomingSawNonTextContentBlocks - : params.streamedSawNonTextContentBlocks; - if (!sawEligibleNonTextContent) { - return false; - } - return isDroppedBoundaryTextBlockSubset({ - streamedTextBlocks: params.streamedTextBlocks, - finalTextBlocks: params.nextContentBlocks, - }); -} - /** Assembles assistant stream deltas and final messages into stable TUI display text. */ export class TuiStreamAssembler { private readonly runs = new Map(); @@ -114,8 +25,6 @@ export class TuiStreamAssembler { return { thinkingText: "", contentText: "", - contentBlocks: [], - sawNonTextContentBlocks: false, displayText: "", }; } @@ -146,36 +55,15 @@ export class TuiStreamAssembler { return state; } - private updateRunState( - state: RunStreamState, - message: unknown, - showThinking: boolean, - opts: { boundaryDropMode: BoundaryDropMode }, - ) { + 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) { - const nextContentBlocks = textBlocks.length > 0 ? textBlocks : [contentText]; - const shouldKeepStreamedBoundaryText = shouldPreserveBoundaryDroppedText({ - boundaryDropMode: opts.boundaryDropMode, - streamedSawNonTextContentBlocks: state.sawNonTextContentBlocks, - incomingSawNonTextContentBlocks: sawNonTextContentBlocks, - streamedTextBlocks: state.contentBlocks, - nextContentBlocks, - }); - - if (!shouldKeepStreamedBoundaryText) { - state.contentText = contentText; - state.contentBlocks = nextContentBlocks; - } - } - if (sawNonTextContentBlocks) { - state.sawNonTextContentBlocks = true; + state.contentText = contentText; } const displayText = composeThinkingAndContent({ @@ -191,9 +79,7 @@ export class TuiStreamAssembler { ingestDelta(runId: string, message: unknown, showThinking: boolean): string | null { const state = this.getTrackedRun(runId); const previousDisplayText = state.displayText; - this.updateRunState(state, message, showThinking, { - boundaryDropMode: "streamed-or-incoming", - }); + this.updateRunState(state, message, showThinking); if (!state.displayText || state.displayText === previousDisplayText) { return null; @@ -212,9 +98,7 @@ export class TuiStreamAssembler { // Late finals must not insert an evicted run and displace a live stream. const state = this.runs.get(runId) ?? this.createRunState(); const streamedContentText = state.contentText; - this.updateRunState(state, message, showThinking, { - boundaryDropMode: "streamed-only", - }); + this.updateRunState(state, message, showThinking); const responseText = resolveFinalAssistantText({ finalText: state.contentText, streamedText: streamedContentText,