File size: 5,162 Bytes
fc93158 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | import { describe, expect, it } from "vitest";
import { TuiStreamAssembler } from "./tui-stream-assembler.js";
const text = (value: string) => ({ type: "text", text: value }) as const;
const thinking = (value: string) => ({ type: "thinking", thinking: value }) as const;
const toolUse = () => ({ type: "tool_use", name: "search" }) as const;
const messageWithContent = (content: readonly Record<string, unknown>[]) =>
({
role: "assistant",
content,
}) 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", () => {
it("keeps thinking before content even when thinking arrives later", () => {
const assembler = new TuiStreamAssembler();
const first = assembler.ingestDelta("run-1", messageWithContent([text("Hello")]), true);
expect(first).toBe("Hello");
const second = assembler.ingestDelta("run-1", messageWithContent([thinking("Brain")]), true);
expect(second).toBe("[thinking]\nBrain\n\nHello");
});
it("omits thinking when showThinking is false", () => {
const assembler = new TuiStreamAssembler();
const output = assembler.ingestDelta(
"run-2",
messageWithContent([thinking("Hidden"), text("Visible")]),
false,
);
expect(output).toBe("Visible");
});
it("falls back to streamed text on empty final payload", () => {
const assembler = new TuiStreamAssembler();
assembler.ingestDelta("run-3", messageWithContent([text("Streamed")]), false);
const finalText = assembler.finalize("run-3", { role: "assistant", content: [] }, false);
expect(finalText).toBe("Streamed");
});
it("falls back to event error message when final payload has no renderable text", () => {
const assembler = new TuiStreamAssembler();
const finalText = assembler.finalize(
"run-3-error",
{ role: "assistant", content: [] },
false,
'401 {"error":{"message":"Missing scopes: model.request"}}',
);
expect(finalText).toContain("HTTP 401");
expect(finalText).toContain("Missing scopes: model.request");
});
it("returns null when delta text is unchanged", () => {
const assembler = new TuiStreamAssembler();
const first = assembler.ingestDelta("run-4", messageWithContent([text("Repeat")]), false);
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(
"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();
assembler.ingestDelta("run-boundary", messageWithContent(testCase.streamedContent), false);
const finalText = assembler.finalize(
"run-boundary",
messageWithContent(testCase.finalContent),
false,
);
expect(finalText).toBe(testCase.expected);
});
}
});
|