Open History to see the last commit
1import type { OrchestrationV2ThreadStreamItem } from "@t3tools/contracts";
2import * as Cause from "effect/Cause";
3import * as Deferred from "effect/Deferred";
4import * as Duration from "effect/Duration";
5import * as Effect from "effect/Effect";
6import * as Exit from "effect/Exit";
7import * as Fiber from "effect/Fiber";
8import * as Queue from "effect/Queue";
9import * as Semaphore from "effect/Semaphore";
10import * as Stream from "effect/Stream";
11
12import {
13 LiveStreamBufferError,
14 makeLiveStreamBudget,
15 type RetainedLiveItem,
16} from "./LiveStreamBudget.ts";
17
18const COALESCE_WINDOW = Duration.millis(50);
19const MAX_PENDING_UPDATES = 512;
20
21type ThreadLiveEvent = Extract<OrchestrationV2ThreadStreamItem, { readonly kind: "event" }>;
22export type ThreadLiveInput = ThreadLiveEvent | { readonly kind: "synchronized" };
23
24function coalescibleUpdateKey(input: ThreadLiveEvent): string | null {
25 const { event } = input;
26 const prefix = `${event.threadId}\u0000${event.runId ?? ""}\u0000${event.type}`;
27 if (event.type === "message.updated") {
28 return event.payload.role === "assistant" && event.payload.streaming
29 ? `${prefix}\u0000${event.payload.id}`
30 : null;
31 }
32 if (event.type !== "turn-item.updated" || event.payload.status !== "running") return null;
33 switch (event.payload.type) {
34 case "assistant_message":
35 case "reasoning":
36 if (!event.payload.streaming) return null;
37 break;
38 case "command_execution":
39 case "file_change":
40 case "file_search":
41 case "web_search":
42 case "dynamic_tool":
43 break;
44 default:
45 return null;
46 }
47 return `${prefix}\u0000${event.payload.type}\u0000${event.payload.id}`;
48}
49
50// These events replace whole entities. Keep each survivor at its original
51// position/sequence so a resume cursor cannot skip an unrelated earlier update.
52export function coalesceLiveUpdatedEvents(
53 events: ReadonlyArray<ThreadLiveEvent>,
54): ReadonlyArray<ThreadLiveEvent> {
55 const latestUpdates = new Map<string, number>();
56 const superseded = new Set<number>();
57 const pendingKeys = new Map<ThreadLiveEvent["event"]["type"], string>();
58
59 for (const [index, event] of events.entries()) {
60 const key = coalescibleUpdateKey(event);
61 if (key === null) {
62 latestUpdates.clear();
63 pendingKeys.clear();
64 continue;
65 }
66 // Both entity families use insertion order; partial timelines also reject
67 // unseen items below their watermark. Keep identity switches within each
68 // family, while allowing paired message/item updates to coalesce.
69 const pendingKey = pendingKeys.get(event.event.type);
70 if (pendingKey !== undefined && pendingKey !== key) latestUpdates.clear();
71 pendingKeys.set(event.event.type, key);
72 const previous = latestUpdates.get(key);
73 if (previous !== undefined) superseded.add(previous);
74 latestUpdates.set(key, index);
75 }
76 return events.filter((_, index) => !superseded.has(index));
77}
78
79export const makeThreadLiveEventCoalescer = Effect.fn(
80 "ThreadLiveEventCoalescer.makeThreadLiveEventCoalescer",
81)(function* <E = never>(options?: {
82 readonly maxItems?: number;
83 readonly maxSerializedBytes?: number;
84}) {
85 const coalescerScope = yield* Effect.scope;
86 const budget = yield* makeLiveStreamBudget(options);
87 const cleanupComplete = yield* Deferred.make<void>();
88 const output = yield* Queue.unbounded<
89 RetainedLiveItem<ThreadLiveInput>,
90 E | LiveStreamBufferError | Cause.Done
91 >();
92 const mutex = yield* Semaphore.make(1);
93 const pendingUpdates = new Set<RetainedLiveItem<ThreadLiveEvent>>();
94 const latestUpdates = new Map<string, RetainedLiveItem<ThreadLiveEvent>>();
95 const pendingKeys = new Map<ThreadLiveEvent["event"]["type"], string>();
96 let windowGeneration = 0;
97 let windowFiber: Fiber.Fiber<void, never> | null = null;
98 let closed = false;
99
100 // cancelWindow/flushPending/offerAll run per live event or batch; sub-ms spans here would flood the rotating trace.
101 const cancelWindow = Effect.fnUntraced(function* () {
102 const fiber = windowFiber;
103 if (!fiber) {
104 return;
105 }
106 windowFiber = null;
107 yield* Fiber.interrupt(fiber);
108 });
109
110 const flushPending = Effect.fnUntraced(function* () {
111 if (pendingUpdates.size === 0) {
112 return;
113 }
114 const items = Array.from(pendingUpdates);
115 pendingUpdates.clear();
116 latestUpdates.clear();
117 pendingKeys.clear();
118 yield* Queue.offerAll(output, items);
119 }, Effect.uninterruptible);
120
121 const flushWindow = (generation: number) =>
122 Effect.sleep(COALESCE_WINDOW).pipe(
123 Effect.andThen(
124 mutex.withPermits(1)(
125 Effect.suspend(() => (generation === windowGeneration ? flushPending() : Effect.void)),
126 ),
127 ),
128 Effect.ensuring(
129 Effect.sync(() => {
130 if (generation === windowGeneration) {
131 windowFiber = null;
132 }
133 }),
134 ),
135 );
136
137 // Keep each source batch together so a synchronization marker cannot pass
138 // events already pulled from PubSub but still being coalesced.
139 const offerAll = Effect.fnUntraced(function* (inputs: ReadonlyArray<ThreadLiveInput>) {
140 yield* mutex.withPermits(1)(
141 Effect.forEach(
142 inputs,
143 (input) =>
144 Effect.gen(function* () {
145 yield* budget.check;
146 const startsWindow = pendingUpdates.size === 0;
147 const key = input.kind === "event" ? coalescibleUpdateKey(input) : null;
148 if (input.kind === "event" && key !== null) {
149 const pendingKey = pendingKeys.get(input.event.type);
150 if (pendingKey !== undefined && pendingKey !== key) latestUpdates.clear();
151 pendingKeys.set(input.event.type, key);
152 const previous = latestUpdates.get(key);
153 yield* budget.replace(previous ? [previous] : [], [input]).pipe(
154 Effect.tap(([item]) =>
155 Effect.sync(() => {
156 if (previous) pendingUpdates.delete(previous);
157 pendingUpdates.add(item!);
158 latestUpdates.set(key, item!);
159 }),
160 ),
161 Effect.uninterruptible,
162 );
163 if (startsWindow) {
164 const generation = ++windowGeneration;
165 windowFiber = yield* Effect.forkIn(flushWindow(generation), coalescerScope);
166 }
167 if (pendingUpdates.size >= MAX_PENDING_UPDATES) {
168 yield* cancelWindow();
169 windowGeneration += 1;
170 yield* flushPending();
171 }
172 return;
173 }
174
175 if (input.kind === "event") {
176 yield* budget.retain(input).pipe(
177 Effect.tap((item) => Effect.sync(() => pendingUpdates.add(item))),
178 Effect.uninterruptible,
179 );
180 }
181 yield* cancelWindow();
182 windowGeneration += 1;
183 yield* flushPending();
184 if (input.kind === "synchronized") {
185 yield* budget.retain({ kind: "synchronized" as const }).pipe(
186 Effect.flatMap((marker) => Queue.offer(output, marker)),
187 Effect.uninterruptible,
188 );
189 }
190 }),
191 { discard: true },
192 ),
193 );
194 });
195
196 const close = (cause?: Cause.Cause<E | LiveStreamBufferError>) =>
197 mutex.withPermits(1)(
198 Effect.gen(function* () {
199 if (closed) {
200 return;
201 }
202 closed = true;
203 windowGeneration += 1;
204 yield* cancelWindow();
205 budget.release(pendingUpdates);
206 pendingUpdates.clear();
207 latestUpdates.clear();
208 pendingKeys.clear();
209 budget.release(yield* Queue.clear(output).pipe(Effect.orElseSucceed(() => [])));
210 if (cause) {
211 yield* Queue.failCause(output, cause);
212 }
213 yield* Queue.shutdown(output);
214 yield* Deferred.succeed(cleanupComplete, undefined);
215 }),
216 );
217
218 yield* Effect.addFinalizer(() => close());
219 yield* budget.failed.pipe(
220 Effect.catchTags({ LiveStreamBufferError: (error) => close(Cause.fail(error)) }),
221 Effect.forkScoped,
222 );
223
224 return {
225 offer: (input: ThreadLiveInput) => offerAll([input]),
226 close,
227 end: mutex.withPermits(1)(
228 Effect.gen(function* () {
229 yield* cancelWindow();
230 windowGeneration += 1;
231 yield* flushPending();
232 yield* Queue.end(output);
233 }),
234 ),
235 offerAll,
236 stream: budget.deliver(Stream.fromQueue(output)),
237 failed: budget.failed,
238 closed: Deferred.await(cleanupComplete),
239 usage: budget.usage,
240 } as const;
241});
242
243export const coalesceThreadLiveStream = <E, R>(source: Stream.Stream<ThreadLiveEvent, E, R>) =>
244 Stream.unwrap(
245 Effect.gen(function* () {
246 const coalescer = yield* makeThreadLiveEventCoalescer<E>();
247 yield* source.pipe(
248 Stream.runForEachArray(coalescer.offerAll),
249 Effect.andThen(coalescer.end),
250 Effect.raceFirst(coalescer.failed),
251 Effect.exit,
252 Effect.flatMap((exit) =>
253 Exit.isFailure(exit) ? coalescer.close(exit.cause) : Effect.void,
254 ),
255 Effect.forkScoped({ startImmediately: true }),
256 );
257 return coalescer.stream;
258 }),
259 );