Skip to content
Source
Sign inSign up

assistantStreaming.ts

242 lines8.9 KB

1import type { ResponseStreamingMode } from "@t3tools/contracts";
2import type * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter";
3​
4// List fences can open at any indentation; closers may be at most three spaces deeper.
5const MARKDOWN_FENCE_PATTERN = /^( *)(`{3,}|~{3,})/;
6// CommonMark blank lines contain only spaces and tabs.
7const BLANK_LINE_PATTERN = /^[ \t]*$/;
8// Require trailing space so partial list markers cannot match during streaming.
9const LIST_ITEM_START_PATTERN = /^[ \t]*(?:[-*+]|\d{1,9}[.)])[ \t]/;
10const SECTION_TITLE_PATTERN = /^ {0,3}(?:#{1,6}(?:[ \t]|$)|\*\*(?:[^*]|\*(?!\*))+\*\*:?$)/;
11// Unindented ATX headings end preceding paragraphs and lists without blank lines; bold lines do not.
12const TOP_LEVEL_HEADING_PATTERN = /^#{1,6}(?:[ \t]|$)/;
13​
14// Matches LIST_ITEM_START_PATTERN on the trimmed partial line without copying it.
15const PARTIAL_LIST_ITEM_START_PATTERN = /[ \t]*(?:[-*+]|\d{1,9}[.)])[ \t][ \t\r]*[^ \t\r]/y;
16​
17interface AssistantTextScan {
18 readonly text: string;
19 // Start of the partial line; every line before it has been scanned.
20 readonly lineStart: number;
21 readonly openFence: { readonly marker: string; readonly indent: number } | null;
22 readonly boundary: number;
23 readonly titleAwaitingContent: boolean;
25​
26const EMPTY_SCAN: AssistantTextScan = {
27 text: "",
28 lineStart: 0,
29 openFence: null,
30 boundary: -1,
31 titleAwaitingContent: false,
32};
33​
34function trimLineEnd(text: string, start: number, end: number): string {
35 while (end > start) {
36 const char = text.charCodeAt(end - 1);
37 if (char !== 32 && char !== 9 && char !== 13) break;
38 end--;
39 }
40 return text.slice(start, end);
42​
43/** Scans the terminated lines of `text`, resuming after `previous` when `text` extends it. */
44function advanceScan(previous: AssistantTextScan, text: string): AssistantTextScan {
45 const resumed = text.startsWith(previous.text);
46 let { openFence, boundary, lineStart, titleAwaitingContent } = resumed ? previous : EMPTY_SCAN;
47 let searchFrom = resumed ? previous.text.length : 0;
48 for (;;) {
49 const newline = text.indexOf("\n", searchFrom);
50 if (newline === -1) break;
51 const line = trimLineEnd(text, lineStart, newline);
52 if (
53 openFence === null &&
54 lineStart > 0 &&
55 !titleAwaitingContent &&
56 LIST_ITEM_START_PATTERN.test(line)
57 ) {
58 boundary = lineStart;
59 }
60 const fenceMatch = MARKDOWN_FENCE_PATTERN.exec(line);
61 if (fenceMatch) {
62 const indent = fenceMatch[1]!.length;
63 const marker = fenceMatch[2]!;
64 if (openFence === null) {
65 openFence = { marker, indent };
66 titleAwaitingContent = false;
67 } else if (
68 marker[0] === openFence.marker[0] &&
69 marker.length >= openFence.marker.length &&
70 indent <= openFence.indent + 3 &&
71 line.length === indent + marker.length
72 ) {
73 // CommonMark: a closing fence carries no info string.
74 openFence = null;
75 boundary = newline + 1;
76 }
77 } else if (openFence === null && BLANK_LINE_PATTERN.test(line) && lineStart > 0) {
78 if (!titleAwaitingContent) {
79 boundary = newline + 1;
80 }
81 } else if (openFence === null) {
82 if (lineStart > 0 && !titleAwaitingContent && TOP_LEVEL_HEADING_PATTERN.test(line)) {
83 boundary = lineStart;
84 }
85 titleAwaitingContent = SECTION_TITLE_PATTERN.test(line);
86 }
87 lineStart = newline + 1;
88 searchFrom = lineStart;
89 }
90 return { text, lineStart, openFence, boundary, titleAwaitingContent };
92​
93// The partial line may still change, so its list item start is checked here, not kept in the scan.
94function readyBoundary(scan: AssistantTextScan): number {
95 if (scan.openFence === null && scan.lineStart > 0 && !scan.titleAwaitingContent) {
96 PARTIAL_LIST_ITEM_START_PATTERN.lastIndex = scan.lineStart;
97 if (PARTIAL_LIST_ITEM_START_PATTERN.test(scan.text)) return scan.lineStart;
98 }
99 return scan.boundary;
101​
102/** Emit terminated lines only, except partial list-start lookahead for tight lists. Hold section titles until content follows. */
103export function splitBufferedAssistantText(text: string): { ready: string; rest: string } {
104 const boundary = readyBoundary(advanceScan(EMPTY_SCAN, text));
105 if (boundary === -1) {
106 return { ready: "", rest: text };
107 }
108 return { ready: text.slice(0, boundary), rest: text.slice(boundary) };
110​
111// Long enough to ride out normal gaps between chunks; short enough that text written before a
112// long tool call or background wait does not look lost.
113const ASSISTANT_IDLE_FLUSH_MS = 800;
114​
115type StreamedMessage = {
116 readonly id: string;
117 readonly text: string;
118 readonly streaming: boolean;
119};
120​
121interface StreamTrack {
122 readonly text: string;
123 readonly at: number;
124 // Latest snapshot whose text has not all been delivered yet.
125 readonly held: ProviderAdapter.ProviderAdapterV2Event | null;
126 readonly dueAt: number;
127 readonly scan: AssistantTextScan;
129​
130function streamedMessage(event: ProviderAdapter.ProviderAdapterV2Event): StreamedMessage | null {
131 if (event.type === "message.updated" && event.message.role === "assistant") return event.message;
132 if (
133 event.type === "turn_item.updated" &&
134 (event.turnItem.type === "assistant_message" || event.turnItem.type === "reasoning")
135 )
136 return event.turnItem;
137 return null;
139​
140function withText(
141 event: ProviderAdapter.ProviderAdapterV2Event,
142 text: string,
143): ProviderAdapter.ProviderAdapterV2Event {
144 if (event.type === "message.updated") return { ...event, message: { ...event.message, text } };
145 if (
146 event.type === "turn_item.updated" &&
147 (event.turnItem.type === "assistant_message" || event.turnItem.type === "reasoning")
148 )
149 return { ...event, turnItem: { ...event.turnItem, text } };
150 return event;
152​
153/**
154 * Paragraph mode delivers streamed text at Markdown boundaries, at most every 400ms. A snapshot
155 * held back that way is delivered whole once its message has been idle for `idleFlushMs`, so a
156 * provider that pauses mid-turn still shows everything it wrote.
157 */
158export function makeAssistantStreamingFilter(
159 mode: ResponseStreamingMode,
160 idleFlushMs = ASSISTANT_IDLE_FLUSH_MS,
161) {
162 const tracks = new Map<string, StreamTrack>();
163 // Earliest deadline of any held snapshot, or earlier: a newer snapshot only pushes its own
164 // deadline later, so the cache stays O(1) per event and is recomputed when it comes due.
165 let nextDueAt: number | null = null;
166 const track = (key: string, value: StreamTrack) => {
167 tracks.set(key, value);
168 if (value.held !== null && (nextDueAt === null || value.dueAt < nextDueAt))
169 nextDueAt = value.dueAt;
170 };
171 return {
172 filter(
173 event: ProviderAdapter.ProviderAdapterV2Event,
174 now: number,
175 ): ProviderAdapter.ProviderAdapterV2Event | null {
176 if (
177 event.type === "node.updated" &&
178 event.node.kind === "assistant_message" &&
179 event.node.status === "running"
180 )
181 return null;
182 const message = streamedMessage(event);
183 if (!message) return event;
184 const key = `${event.type}:${message.id}`;
185 if (!message.streaming) {
186 tracks.delete(key);
187 return event;
188 }
189 if (mode === "turn") return null;
190 const previous = tracks.get(key);
191 const delivered = previous?.text ?? "";
192 const hold = (scan = previous?.scan ?? EMPTY_SCAN) => {
193 track(key, {
194 text: delivered,
195 at: previous?.at ?? Number.NEGATIVE_INFINITY,
196 held: message.text !== delivered ? event : null,
197 dueAt: now + idleFlushMs,
198 scan,
199 });
200 return null;
201 };
202 if (previous && now - previous.at < 400) return hold();
203 const scan = advanceScan(previous?.scan ?? EMPTY_SCAN, message.text);
204 const boundary = readyBoundary(scan);
205 const ready = boundary === -1 ? "" : message.text.slice(0, boundary);
206 // An idle delivery may already show more than the next boundary; never shrink it. A
207 // provider rewrite or truncation is not a prefix and replaces it.
208 if (delivered.startsWith(ready)) return hold(scan);
209 track(key, {
210 text: ready,
211 at: now,
212 held: ready !== message.text ? event : null,
213 dueAt: now + idleFlushMs,
214 scan,
215 });
216 return withText(event, ready);
217 },
218 nextIdleFlushAt: (): number | null => nextDueAt,
219 /** Returns the held snapshots due by `now` and counts them as delivered. */
220 takeIdleFlushes(now: number): ReadonlyArray<ProviderAdapter.ProviderAdapterV2Event> {
221 const due: Array<ProviderAdapter.ProviderAdapterV2Event> = [];
222 nextDueAt = null;
223 for (const [key, current] of tracks) {
224 const event = current.held;
225 if (event === null) continue;
226 if (current.dueAt > now) {
227 if (nextDueAt === null || current.dueAt < nextDueAt) nextDueAt = current.dueAt;
228 continue;
229 }
230 due.push(event);
231 tracks.set(key, {
232 text: streamedMessage(event)?.text ?? current.text,
233 at: now,
234 held: null,
235 dueAt: now,
236 scan: current.scan,
237 });
238 }
239 return due;
240 },
241 };