1import { MessageId, ProviderDriverKind, ThreadId } from "@t3tools/contracts";
2import type * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter";
3import * as DateTime from "effect/DateTime";
4import { assert, describe, expect, it, vi } from "vite-plus/test";
5import { makeAssistantStreamingFilter, splitBufferedAssistantText } from "./assistantStreaming.ts";
6
7describe("V2 assistant streaming", () => {
8 it("keeps code fences intact", () => {
9 expect(splitBufferedAssistantText("Intro\n\n```ts\nx()\n\n")).toEqual({
10 ready: "Intro\n\n",
11 rest: "```ts\nx()\n\n",
12 });
13 expect(splitBufferedAssistantText("```ts\nx()\n```\nrest")).toEqual({
14 ready: "```ts\nx()\n```\n",
15 rest: "rest",
16 });
17 });
18});
19
20it("holds a streamed section heading until its content has a boundary", () => {
21 expect(splitBufferedAssistantText("Intro\n\n## Results\n\n")).toEqual({
22 ready: "Intro\n\n",
23 rest: "## Results\n\n",
24 });
25 expect(splitBufferedAssistantText("**Results**\n\nBody\n\nNext")).toEqual({
26 ready: "**Results**\n\nBody\n\n",
27 rest: "Next",
28 });
29});
30
31type Step =
32 // A snapshot at `at` ms, and the text it delivers at once (null when held back).
33 | { readonly at: number; readonly text: string; readonly shows: string | null }
34 | { readonly at: number; readonly final: string }
35 | { readonly id: string; readonly at: number; readonly text: string; readonly shows: string }
36 // Held snapshots due by `idle` ms, delivered whole.
37 | { readonly idle: number; readonly shows: ReadonlyArray<string> }
38 | { readonly dueAt: number | null };
39
40const createdAt = DateTime.makeUnsafe(0);
41const snapshot = (id: string, text: string, streaming: boolean) =>
42 ({
43 type: "message.updated",
44 driver: ProviderDriverKind.make("codex"),
45 message: {
46 createdBy: "agent",
47 creationSource: "provider",
48 id: MessageId.make(id),
49 threadId: ThreadId.make("thread"),
50 runId: null,
51 nodeId: null,
52 role: "assistant",
53 text,
54 attachments: [],
55 streaming,
56 createdAt,
57 updatedAt: createdAt,
58 },
59 }) satisfies ProviderAdapter.ProviderAdapterV2Event;
60const shown = (event: ProviderAdapter.ProviderAdapterV2Event | null) =>
61 event?.type === "message.updated" ? event.message.text : null;
62
63it.each<{ readonly name: string; readonly mode?: "turn"; readonly steps: ReadonlyArray<Step> }>([
64 {
65 name: "an unfinished code fence",
66 steps: [
67 { at: 0, text: "A\n\n```ts\nx()", shows: "A\n\n" },
68 { dueAt: 800 },
69 { idle: 799, shows: [] },
70 { idle: 800, shows: ["A\n\n```ts\nx()"] },
71 { dueAt: null },
72 { at: 1_300, text: "A\n\n```ts\nx()\n```\n\nB", shows: "A\n\n```ts\nx()\n```\n\n" },
73 ],
74 },
75 {
76 name: "a list's partial last item",
77 steps: [
78 { at: 0, text: "Steps:\n\n- one\n- tw", shows: "Steps:\n\n- one\n" },
79 { idle: 800, shows: ["Steps:\n\n- one\n- tw"] },
80 { at: 1_300, text: "Steps:\n\n- one\n- two\n- th", shows: "Steps:\n\n- one\n- two\n" },
81 ],
82 },
83 {
84 name: "an incomplete table row",
85 steps: [
86 { at: 0, text: "Intro.\n\n| a | b |\n|---|---|\n| 1 |", shows: "Intro.\n\n" },
87 { idle: 800, shows: ["Intro.\n\n| a | b |\n|---|---|\n| 1 |"] },
88 {
89 at: 1_300,
90 text: "Intro.\n\n| a | b |\n|---|---|\n| 1 | 2 |\n\nAfter",
91 shows: "Intro.\n\n| a | b |\n|---|---|\n| 1 | 2 |\n\n",
92 },
93 ],
94 },
95 {
96 name: "more of the same paragraph after an idle delivery",
97 steps: [
98 { at: 0, text: "A\n\nB", shows: "A\n\n" },
99 { idle: 800, shows: ["A\n\nB"] },
100 { at: 900, text: "A\n\nB more", shows: null },
101 { idle: 1_699, shows: [] },
102 { idle: 1_700, shows: ["A\n\nB more"] },
103 { idle: 5_000, shows: [] },
104 // Its only boundary is behind what the idle delivery showed.
105 { at: 5_100, text: "A\n\nB more and", shows: null },
106 { at: 5_200, text: "A\n\nB more and.\n\nC", shows: "A\n\nB more and.\n\n" },
107 { at: 5_300, final: "A\n\nB more and.\n\nC." },
108 { idle: 10_000, shows: [] },
109 ],
110 },
111 {
112 name: "completion at the idle deadline",
113 steps: [
114 { at: 0, text: "A\n\nB", shows: "A\n\n" },
115 { at: 800, final: "A\n\nB." },
116 { idle: 800, shows: [] },
117 ],
118 },
119 {
120 name: "completion right after an idle delivery at the same instant",
121 steps: [
122 { at: 0, text: "A\n\nB", shows: "A\n\n" },
123 { idle: 800, shows: ["A\n\nB"] },
124 { at: 800, final: "A\n\nB." },
125 { idle: 5_000, shows: [] },
126 ],
127 },
128 {
129 name: "several fast paragraphs inside the throttle",
130 steps: [
131 { at: 0, text: "A\n\nB", shows: "A\n\n" },
132 { at: 100, text: "A\n\nB\n\nC", shows: null },
133 { at: 200, text: "A\n\nB\n\nC\n\nD", shows: null },
134 { at: 300, text: "A\n\nB\n\nC\n\nD\n\nE", shows: null },
135 { idle: 1_099, shows: [] },
136 { idle: 1_100, shows: ["A\n\nB\n\nC\n\nD\n\nE"] },
137 ],
138 },
139 {
140 name: "two messages due at different times",
141 steps: [
142 { id: "a", at: 0, text: "A\n\nx", shows: "A\n\n" },
143 { id: "b", at: 300, text: "B\n\ny", shows: "B\n\n" },
144 { dueAt: 800 },
145 { idle: 800, shows: ["A\n\nx"] },
146 { dueAt: 1_100 },
147 { idle: 1_100, shows: ["B\n\ny"] },
148 ],
149 },
150 {
151 name: "empty text",
152 steps: [{ at: 0, text: "", shows: null }, { dueAt: null }],
153 },
154 {
155 name: "whitespace-only text",
156 steps: [
157 { at: 0, text: " \n", shows: null },
158 { idle: 800, shows: [" \n"] },
159 { idle: 5_000, shows: [] },
160 ],
161 },
162 {
163 // Rewrites replace delivered text: shorter at once, same length after the throttle.
164 name: "provider rewrites and truncation",
165 steps: [
166 { at: 0, text: "Original long paragraph.\n\n", shows: "Original long paragraph.\n\n" },
167 { at: 1_000, text: "New.\n\n", shows: "New.\n\n" },
168 { at: 1_100, text: "Old.\n\n", shows: null },
169 { idle: 1_900, shows: ["Old.\n\n"] },
170 { at: 2_500, text: "Old", shows: null },
171 { idle: 3_300, shows: ["Old"] },
172 { at: 4_000, text: "", shows: null },
173 { idle: 4_800, shows: [""] },
174 ],
175 },
176 {
177 name: "a rewrite that drops an open code fence",
178 steps: [
179 { at: 0, text: "Intro\n\n```ts\nx()\n", shows: "Intro\n\n" },
180 { at: 400, text: "Retry\n\n- a\n- b", shows: "Retry\n\n- a\n" },
181 ],
182 },
183 {
184 name: "turn mode",
185 mode: "turn",
186 steps: [{ at: 0, text: "A\n\nB", shows: null }, { dueAt: null }],
187 },
188])("paragraph streaming holds and idle-delivers $name", ({ mode, steps }) => {
189 const filter = makeAssistantStreamingFilter(mode ?? "paragraph", 800);
190 for (const step of steps) {
191 if ("dueAt" in step) assert.equal(filter.nextIdleFlushAt(), step.dueAt);
192 else if ("idle" in step)
193 assert.deepEqual(filter.takeIdleFlushes(step.idle).map(shown), [...step.shows]);
194 else if ("final" in step)
195 assert.equal(shown(filter.filter(snapshot("m", step.final, false), step.at)), step.final);
196 else
197 assert.equal(
198 shown(filter.filter(snapshot("id" in step ? step.id : "m", step.text, true), step.at)),
199 step.shows,
200 );
201 }
202});
203
204it("scans each streamed line once while a long code fence is still open", () => {
205 const filter = makeAssistantStreamingFilter("paragraph", 800);
206 const lineCount = 400;
207 let text = "Intro\n\n```ts\n";
208 const indexOf = vi.spyOn(String.prototype, "indexOf");
209 try {
210 for (let line = 0; line < lineCount; line++) {
211 text += `const value${line} = ${line};\n`;
212 filter.filter(snapshot("m", text, true), line * 400);
213 }
214 const newlineSearches = indexOf.mock.calls.filter(([search]) => search === "\n").length;
215 expect(newlineSearches).toBeLessThan(lineCount * 4);
216 } finally {
217 indexOf.mockRestore();
218 }
219 text += "```\n";
220 assert.equal(shown(filter.filter(snapshot("m", text, true), lineCount * 400)), text);
221});