refactor(tui): simplify stream boundary-drop modes

This commit is contained in:
Peter Steinberger
2026-02-26 20:54:29 +01:00
parent b01273cfc6
commit 675764e866
2 changed files with 125 additions and 213 deletions

View File

@@ -1,232 +1,122 @@
import { describe, expect, it } from "vitest"; import { describe, expect, it } from "vitest";
import { TuiStreamAssembler } from "./tui-stream-assembler.js"; import { TuiStreamAssembler } from "./tui-stream-assembler.js";
const STREAM_WITH_TOOL_BLOCKS = { const text = (value: string) => ({ type: "text", text: value }) as const;
role: "assistant", const thinking = (value: string) => ({ type: "thinking", thinking: value }) as const;
content: [ const toolUse = () => ({ type: "tool_use", name: "search" }) as const;
{ type: "text", text: "Before tool call" },
{ type: "tool_use", name: "search" },
{ type: "text", text: "After tool call" },
],
} as const;
const STREAM_AFTER_TOOL_BLOCKS = { const messageWithContent = (content: readonly Record<string, unknown>[]) =>
role: "assistant", ({
content: [ role: "assistant",
{ type: "tool_use", name: "search" }, content,
{ type: "text", text: "After tool call" }, }) as const;
],
} as const; const TEXT_ONLY_TWO_BLOCKS = messageWithContent([text("Draft line 1"), text("Draft line 2")]);
type FinalizeBoundaryCase = {
name: string;
streamedContent: readonly Record<string, unknown>[];
finalContent: readonly Record<string, unknown>[];
expected: string;
};
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")],
finalContent: [toolUse(), text("Draft line 2")],
expected: "Draft line 2",
},
{
name: "keeps non-empty final text for plain text boundary drops",
streamedContent: [text("Draft line 1"), text("Draft line 2")],
finalContent: [text("Draft line 1")],
expected: "Draft line 1",
},
{
name: "prefers final replacement text when payload is not a boundary subset",
streamedContent: [text("Before tool call"), toolUse(), text("After tool call")],
finalContent: [toolUse(), text("Replacement")],
expected: "Replacement",
},
{
name: "accepts richer final payload when it extends streamed text",
streamedContent: [text("Before tool call")],
finalContent: [text("Before tool call"), text("After tool call")],
expected: "Before tool call\nAfter tool call",
},
];
describe("TuiStreamAssembler", () => { describe("TuiStreamAssembler", () => {
it("keeps thinking before content even when thinking arrives later", () => { it("keeps thinking before content even when thinking arrives later", () => {
const assembler = new TuiStreamAssembler(); const assembler = new TuiStreamAssembler();
const first = assembler.ingestDelta( const first = assembler.ingestDelta("run-1", messageWithContent([text("Hello")]), true);
"run-1",
{
role: "assistant",
content: [{ type: "text", text: "Hello" }],
},
true,
);
expect(first).toBe("Hello"); expect(first).toBe("Hello");
const second = assembler.ingestDelta( const second = assembler.ingestDelta("run-1", messageWithContent([thinking("Brain")]), true);
"run-1",
{
role: "assistant",
content: [{ type: "thinking", thinking: "Brain" }],
},
true,
);
expect(second).toBe("[thinking]\nBrain\n\nHello"); expect(second).toBe("[thinking]\nBrain\n\nHello");
}); });
it("omits thinking when showThinking is false", () => { it("omits thinking when showThinking is false", () => {
const assembler = new TuiStreamAssembler(); const assembler = new TuiStreamAssembler();
const text = assembler.ingestDelta( const output = assembler.ingestDelta(
"run-2", "run-2",
{ messageWithContent([thinking("Hidden"), text("Visible")]),
role: "assistant",
content: [
{ type: "thinking", thinking: "Hidden" },
{ type: "text", text: "Visible" },
],
},
false, false,
); );
expect(output).toBe("Visible");
expect(text).toBe("Visible");
}); });
it("falls back to streamed text on empty final payload", () => { it("falls back to streamed text on empty final payload", () => {
const assembler = new TuiStreamAssembler(); const assembler = new TuiStreamAssembler();
assembler.ingestDelta( assembler.ingestDelta("run-3", messageWithContent([text("Streamed")]), false);
"run-3", const finalText = assembler.finalize("run-3", { role: "assistant", content: [] }, false);
{
role: "assistant",
content: [{ type: "text", text: "Streamed" }],
},
false,
);
const finalText = assembler.finalize(
"run-3",
{
role: "assistant",
content: [],
},
false,
);
expect(finalText).toBe("Streamed"); expect(finalText).toBe("Streamed");
}); });
it("returns null when delta text is unchanged", () => { it("returns null when delta text is unchanged", () => {
const assembler = new TuiStreamAssembler(); const assembler = new TuiStreamAssembler();
const first = assembler.ingestDelta( const first = assembler.ingestDelta("run-4", messageWithContent([text("Repeat")]), false);
"run-4",
{
role: "assistant",
content: [{ type: "text", text: "Repeat" }],
},
false,
);
expect(first).toBe("Repeat"); expect(first).toBe("Repeat");
const second = assembler.ingestDelta("run-4", messageWithContent([text("Repeat")]), false);
expect(second).toBeNull();
});
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( const second = assembler.ingestDelta(
"run-4", "run-delta-boundary",
{ messageWithContent([toolUse(), text("Draft line 2")]),
role: "assistant",
content: [{ type: "text", text: "Repeat" }],
},
false, false,
); );
expect(second).toBeNull(); expect(second).toBeNull();
}); });
it("keeps richer streamed text when final payload drops earlier blocks", () => { for (const testCase of FINALIZE_BOUNDARY_CASES) {
const assembler = new TuiStreamAssembler(); it(testCase.name, () => {
assembler.ingestDelta("run-5", STREAM_WITH_TOOL_BLOCKS, false); const assembler = new TuiStreamAssembler();
assembler.ingestDelta("run-boundary", messageWithContent(testCase.streamedContent), false);
const finalText = assembler.finalize("run-5", STREAM_AFTER_TOOL_BLOCKS, false); const finalText = assembler.finalize(
"run-boundary",
expect(finalText).toBe("Before tool call\nAfter tool call"); messageWithContent(testCase.finalContent),
}); false,
);
it("does not regress streamed text when a delta drops boundary blocks after tool calls", () => { expect(finalText).toBe(testCase.expected);
const assembler = new TuiStreamAssembler(); });
const first = assembler.ingestDelta("run-5-stream", STREAM_WITH_TOOL_BLOCKS, false); }
expect(first).toBe("Before tool call\nAfter tool call");
const second = assembler.ingestDelta("run-5-stream", STREAM_AFTER_TOOL_BLOCKS, false);
expect(second).toBeNull();
});
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("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(
"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");
});
it("prefers non-empty final payload when it is not a dropped block regression", () => {
const assembler = new TuiStreamAssembler();
assembler.ingestDelta(
"run-7",
{
role: "assistant",
content: [{ type: "text", text: "NOT OK" }],
},
false,
);
const finalText = assembler.finalize(
"run-7",
{
role: "assistant",
content: [{ type: "text", text: "OK" }],
},
false,
);
expect(finalText).toBe("OK");
});
}); });

View File

@@ -13,6 +13,8 @@ type RunStreamState = {
displayText: string; displayText: string;
}; };
type BoundaryDropMode = "off" | "streamed-only" | "streamed-or-incoming";
function extractTextBlocksAndSignals(message: unknown): { function extractTextBlocksAndSignals(message: unknown): {
textBlocks: string[]; textBlocks: string[];
sawNonTextContentBlocks: boolean; sawNonTextContentBlocks: boolean;
@@ -75,6 +77,29 @@ function isDroppedBoundaryTextBlockSubset(params: {
return finalTextBlocks.every((block, index) => streamedTextBlocks[suffixStart + index] === block); return finalTextBlocks.every((block, index) => streamedTextBlocks[suffixStart + index] === block);
} }
function shouldPreserveBoundaryDroppedText(params: {
boundaryDropMode: BoundaryDropMode;
streamedSawNonTextContentBlocks: boolean;
incomingSawNonTextContentBlocks: boolean;
streamedTextBlocks: string[];
nextContentBlocks: string[];
}) {
if (params.boundaryDropMode === "off") {
return false;
}
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,
});
}
export class TuiStreamAssembler { export class TuiStreamAssembler {
private runs = new Map<string, RunStreamState>(); private runs = new Map<string, RunStreamState>();
@@ -97,10 +122,7 @@ export class TuiStreamAssembler {
state: RunStreamState, state: RunStreamState,
message: unknown, message: unknown,
showThinking: boolean, showThinking: boolean,
opts?: { opts?: { boundaryDropMode?: BoundaryDropMode },
protectBoundaryDrops?: boolean;
useIncomingNonTextForBoundaryDrops?: boolean;
},
) { ) {
const thinkingText = extractThinkingFromMessage(message); const thinkingText = extractThinkingFromMessage(message);
const contentText = extractContentFromMessage(message); const contentText = extractContentFromMessage(message);
@@ -111,17 +133,16 @@ export class TuiStreamAssembler {
} }
if (contentText) { if (contentText) {
const nextContentBlocks = textBlocks.length > 0 ? textBlocks : [contentText]; const nextContentBlocks = textBlocks.length > 0 ? textBlocks : [contentText];
const useIncomingNonTextForBoundaryDrops = opts?.useIncomingNonTextForBoundaryDrops !== false; const boundaryDropMode = opts?.boundaryDropMode ?? "off";
const shouldPreserveBoundaryDroppedText = const shouldKeepStreamedBoundaryText = shouldPreserveBoundaryDroppedText({
opts?.protectBoundaryDrops === true && boundaryDropMode,
(state.sawNonTextContentBlocks || streamedSawNonTextContentBlocks: state.sawNonTextContentBlocks,
(useIncomingNonTextForBoundaryDrops && sawNonTextContentBlocks)) && incomingSawNonTextContentBlocks: sawNonTextContentBlocks,
isDroppedBoundaryTextBlockSubset({ streamedTextBlocks: state.contentBlocks,
streamedTextBlocks: state.contentBlocks, nextContentBlocks,
finalTextBlocks: nextContentBlocks, });
});
if (!shouldPreserveBoundaryDroppedText) { if (!shouldKeepStreamedBoundaryText) {
state.contentText = contentText; state.contentText = contentText;
state.contentBlocks = nextContentBlocks; state.contentBlocks = nextContentBlocks;
} }
@@ -142,7 +163,9 @@ export class TuiStreamAssembler {
ingestDelta(runId: string, message: unknown, showThinking: boolean): string | null { ingestDelta(runId: string, message: unknown, showThinking: boolean): string | null {
const state = this.getOrCreateRun(runId); const state = this.getOrCreateRun(runId);
const previousDisplayText = state.displayText; const previousDisplayText = state.displayText;
this.updateRunState(state, message, showThinking, { protectBoundaryDrops: true }); this.updateRunState(state, message, showThinking, {
boundaryDropMode: "streamed-or-incoming",
});
if (!state.displayText || state.displayText === previousDisplayText) { if (!state.displayText || state.displayText === previousDisplayText) {
return null; return null;
@@ -157,8 +180,7 @@ export class TuiStreamAssembler {
const streamedTextBlocks = [...state.contentBlocks]; const streamedTextBlocks = [...state.contentBlocks];
const streamedSawNonTextContentBlocks = state.sawNonTextContentBlocks; const streamedSawNonTextContentBlocks = state.sawNonTextContentBlocks;
this.updateRunState(state, message, showThinking, { this.updateRunState(state, message, showThinking, {
protectBoundaryDrops: true, boundaryDropMode: "streamed-only",
useIncomingNonTextForBoundaryDrops: false,
}); });
const finalComposed = state.displayText; const finalComposed = state.displayText;
const shouldKeepStreamedText = const shouldKeepStreamedText =