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