1import type { ResponseStreamingMode } from "@t3tools/contracts";
2import {
3 advanceScan,
4 EMPTY_SCAN,
5 readyBoundary,
6 type AssistantTextScan,
7} from "@t3tools/shared/assistantStreaming";
8import type * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter";
9
10// Long enough to ride out normal gaps between chunks; short enough that text written before a
11// long tool call or background wait does not look lost.
12const ASSISTANT_IDLE_FLUSH_MS = 800;
13
14type StreamedMessage = {
15 readonly id: string;
16 readonly text: string;
17 readonly streaming: boolean;
18};
19
20interface StreamTrack {
21 readonly text: string;
22 readonly at: number;
23 // Latest snapshot whose text has not all been delivered yet.
24 readonly held: ProviderAdapter.ProviderAdapterV2Event | null;
25 readonly dueAt: number;
26 readonly scan: AssistantTextScan;
27}
28
29function streamedMessage(event: ProviderAdapter.ProviderAdapterV2Event): StreamedMessage | null {
30 if (event.type === "message.updated" && event.message.role === "assistant") return event.message;
31 if (
32 event.type === "turn_item.updated" &&
33 (event.turnItem.type === "assistant_message" || event.turnItem.type === "reasoning")
34 )
35 return event.turnItem;
36 return null;
37}
38
39function withText(
40 event: ProviderAdapter.ProviderAdapterV2Event,
41 text: string,
42): ProviderAdapter.ProviderAdapterV2Event {
43 if (event.type === "message.updated") return { ...event, message: { ...event.message, text } };
44 if (
45 event.type === "turn_item.updated" &&
46 (event.turnItem.type === "assistant_message" || event.turnItem.type === "reasoning")
47 )
48 return { ...event, turnItem: { ...event.turnItem, text } };
49 return event;
50}
51
52/**
53 * Paragraph mode delivers streamed text at Markdown boundaries, at most every 400ms. A snapshot
54 * held back that way is delivered whole once its message has been idle for `idleFlushMs`, so a
55 * provider that pauses mid-turn still shows everything it wrote.
56 */
57export function makeAssistantStreamingFilter(
58 mode: ResponseStreamingMode,
59 idleFlushMs = ASSISTANT_IDLE_FLUSH_MS,
60) {
61 const tracks = new Map<string, StreamTrack>();
62 // Earliest deadline of any held snapshot, or earlier: a newer snapshot only pushes its own
63 // deadline later, so the cache stays O(1) per event and is recomputed when it comes due.
64 let nextDueAt: number | null = null;
65 const track = (key: string, value: StreamTrack) => {
66 tracks.set(key, value);
67 if (value.held !== null && (nextDueAt === null || value.dueAt < nextDueAt))
68 nextDueAt = value.dueAt;
69 };
70 return {
71 filter(
72 event: ProviderAdapter.ProviderAdapterV2Event,
73 now: number,
74 ): ProviderAdapter.ProviderAdapterV2Event | null {
75 if (
76 event.type === "node.updated" &&
77 event.node.kind === "assistant_message" &&
78 event.node.status === "running"
79 )
80 return null;
81 const message = streamedMessage(event);
82 if (!message) return event;
83 const key = `${event.type}:${message.id}`;
84 if (!message.streaming) {
85 tracks.delete(key);
86 return event;
87 }
88 if (mode === "turn") return null;
89 const previous = tracks.get(key);
90 const delivered = previous?.text ?? "";
91 const hold = (scan = previous?.scan ?? EMPTY_SCAN) => {
92 track(key, {
93 text: delivered,
94 at: previous?.at ?? Number.NEGATIVE_INFINITY,
95 held: message.text !== delivered ? event : null,
96 dueAt: now + idleFlushMs,
97 scan,
98 });
99 return null;
100 };
101 if (previous && now - previous.at < 400) return hold();
102 const scan = advanceScan(previous?.scan ?? EMPTY_SCAN, message.text);
103 const boundary = readyBoundary(scan);
104 const ready = boundary === -1 ? "" : message.text.slice(0, boundary);
105 // An idle delivery may already show more than the next boundary; never shrink it. A
106 // provider rewrite or truncation is not a prefix and replaces it.
107 if (delivered.startsWith(ready)) return hold(scan);
108 track(key, {
109 text: ready,
110 at: now,
111 held: ready !== message.text ? event : null,
112 dueAt: now + idleFlushMs,
113 scan,
114 });
115 return withText(event, ready);
116 },
117 nextIdleFlushAt: (): number | null => nextDueAt,
118 /** Returns the held snapshots due by `now` and counts them as delivered. */
119 takeIdleFlushes(now: number): ReadonlyArray<ProviderAdapter.ProviderAdapterV2Event> {
120 const due: Array<ProviderAdapter.ProviderAdapterV2Event> = [];
121 nextDueAt = null;
122 for (const [key, current] of tracks) {
123 const event = current.held;
124 if (event === null) continue;
125 if (current.dueAt > now) {
126 if (nextDueAt === null || current.dueAt < nextDueAt) nextDueAt = current.dueAt;
127 continue;
128 }
129 due.push(event);
130 tracks.set(key, {
131 text: streamedMessage(event)?.text ?? current.text,
132 at: now,
133 held: null,
134 dueAt: now,
135 scan: current.scan,
136 });
137 }
138 return due;
139 },
140 };
141}