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 isToolUpdated(input: ThreadLiveEvent): boolean {
25 if (input.event.type !== "turn-item.updated" || input.event.payload.status !== "running")
26 return false;
27 switch (input.event.payload.type) {
28 case "command_execution":
29 case "file_change":
30 case "file_search":
31 case "web_search":
32 case "dynamic_tool":
33 return true;
34 default:
35 return false;
36 }
37}
38
39function stableToolCallIdentity(input: ThreadLiveEvent): string | null {
40 return input.event.type === "turn-item.updated" ? input.event.payload.id : null;
41}
42
43/**
44 * Retain only the latest in-flight update for each stable tool-call id in a
45 * live run. Item IDs distinguish parallel calls with the same label.
46 * Survivors remain in sequence order; terminal updates are never discarded.
47 */
48export function coalesceLiveToolUpdatedEvents(
49 events: ReadonlyArray<ThreadLiveEvent>,
50): ReadonlyArray<ThreadLiveEvent> {
51 const survivors: Array<ThreadLiveEvent> = [];
52 let pendingUpdates: Array<ThreadLiveEvent> = [];
53
54 const flushUpdates = () => {
55 const seen = new Set<string>();
56 const latestUpdates: Array<ThreadLiveEvent> = [];
57 for (let index = pendingUpdates.length - 1; index >= 0; index -= 1) {
58 const event = pendingUpdates[index]!;
59 const identity = stableToolCallIdentity(event);
60 const key = identity
61 ? `${event.event.threadId}\u0000${event.event.runId ?? ""}\u0000${identity}`
62 : null;
63 if (key && seen.has(key)) {
64 continue;
65 }
66 if (key) {
67 seen.add(key);
68 }
69 latestUpdates.push(event);
70 }
71 latestUpdates.reverse();
72 survivors.push(...latestUpdates);
73 pendingUpdates = [];
74 };
75
76 for (const event of events) {
77 if (isToolUpdated(event)) {
78 pendingUpdates.push(event);
79 continue;
80 }
81 flushUpdates();
82 survivors.push(event);
83 }
84 flushUpdates();
85 return survivors;
86}
87
88export const makeThreadLiveEventCoalescer = <E = never>(options?: {
89 readonly coalesceWindow?: Duration.Input;
90 readonly maxItems?: number;
91 readonly maxSerializedBytes?: number;
92}) =>
93 Effect.gen(function* () {
94 const coalescerScope = yield* Effect.scope;
95 const budget = yield* makeLiveStreamBudget(options);
96 const cleanupComplete = yield* Deferred.make<void>();
97 const output = yield* Queue.unbounded<
98 RetainedLiveItem<ThreadLiveInput>,
99 E | LiveStreamBufferError | Cause.Done
100 >();
101 const mutex = yield* Semaphore.make(1);
102 const coalesceWindow = options?.coalesceWindow ?? COALESCE_WINDOW;
103 let pendingUpdates: Array<RetainedLiveItem<ThreadLiveEvent>> = [];
104 let windowGeneration = 0;
105 let windowFiber: Fiber.Fiber<void, never> | null = null;
106 let closed = false;
107
108 const cancelWindow = Effect.fn("ThreadLiveEventCoalescer.cancelWindow")(function* () {
109 const fiber = windowFiber;
110 if (!fiber) {
111 return;
112 }
113 windowFiber = null;
114 yield* Fiber.interrupt(fiber);
115 });
116
117 const flushPending = Effect.fn("ThreadLiveEventCoalescer.flushPending")(function* () {
118 if (pendingUpdates.length === 0) {
119 return;
120 }
121 const items = yield* budget.replace(
122 pendingUpdates,
123 coalesceLiveToolUpdatedEvents(pendingUpdates.map((item) => item.value)),
124 );
125 pendingUpdates = [];
126 yield* Queue.offerAll(output, items);
127 }, Effect.uninterruptible);
128
129 const flushWindow = (generation: number) =>
130 Effect.sleep(coalesceWindow).pipe(
131 Effect.andThen(
132 mutex.withPermits(1)(
133 Effect.suspend(() => (generation === windowGeneration ? flushPending() : Effect.void)),
134 ),
135 ),
136 Effect.ensuring(
137 Effect.sync(() => {
138 if (generation === windowGeneration) {
139 windowFiber = null;
140 }
141 }),
142 ),
143 Effect.catchTags({ LiveStreamBufferError: () => Effect.void }),
144 );
145
146 // Keep each source batch together so a synchronization marker cannot pass
147 // events already pulled from PubSub but still being coalesced.
148 const offerAll = Effect.fn("ThreadLiveEventCoalescer.offerAll")(function* (
149 inputs: ReadonlyArray<ThreadLiveInput>,
150 ) {
151 yield* mutex.withPermits(1)(
152 Effect.forEach(
153 inputs,
154 (input) =>
155 Effect.gen(function* () {
156 yield* budget.check;
157 if (input.kind === "event") {
158 yield* budget.retain(input).pipe(
159 Effect.tap((item) => Effect.sync(() => pendingUpdates.push(item))),
160 Effect.uninterruptible,
161 );
162 }
163 if (input.kind === "event" && isToolUpdated(input)) {
164 if (pendingUpdates.length === 1) {
165 const generation = ++windowGeneration;
166 windowFiber = yield* Effect.forkIn(flushWindow(generation), coalescerScope);
167 }
168 if (pendingUpdates.length >= MAX_PENDING_UPDATES) {
169 yield* cancelWindow();
170 windowGeneration += 1;
171 yield* flushPending();
172 }
173 return;
174 }
175
176 yield* cancelWindow();
177 windowGeneration += 1;
178 // A non-update event closes the run immediately. The coalescer keeps
179 // that boundary after the final update from the run.
180 yield* flushPending();
181 if (input.kind === "synchronized") {
182 yield* budget.retain({ kind: "synchronized" as const }).pipe(
183 Effect.flatMap((marker) => Queue.offer(output, marker)),
184 Effect.uninterruptible,
185 );
186 }
187 }),
188 { discard: true },
189 ),
190 );
191 });
192
193 const close = (cause?: Cause.Cause<E | LiveStreamBufferError>) =>
194 mutex.withPermits(1)(
195 Effect.gen(function* () {
196 if (closed) {
197 return;
198 }
199 closed = true;
200 windowGeneration += 1;
201 yield* cancelWindow();
202 budget.release(pendingUpdates);
203 pendingUpdates = [];
204 budget.release(yield* Queue.clear(output).pipe(Effect.orElseSucceed(() => [])));
205 if (cause) {
206 yield* Queue.failCause(output, cause);
207 }
208 yield* Queue.shutdown(output);
209 yield* Deferred.succeed(cleanupComplete, undefined);
210 }),
211 );
212
213 yield* Effect.addFinalizer(() => close());
214 yield* budget.failed.pipe(
215 Effect.catchTags({ LiveStreamBufferError: (error) => close(Cause.fail(error)) }),
216 Effect.forkScoped,
217 );
218
219 return {
220 offer: (input: ThreadLiveInput) => offerAll([input]),
221 close,
222 end: mutex.withPermits(1)(
223 Effect.gen(function* () {
224 yield* cancelWindow();
225 windowGeneration += 1;
226 yield* flushPending();
227 yield* Queue.end(output);
228 }),
229 ),
230 offerAll,
231 stream: budget.deliver(Stream.fromQueue(output)),
232 failed: budget.failed,
233 closed: Deferred.await(cleanupComplete),
234 usage: budget.usage,
235 } as const;
236 });
237
238/** Keep tool updates within a bounded live window; lifecycle events flush it immediately. */
239export const coalesceThreadLiveStream = <E, R>(source: Stream.Stream<ThreadLiveEvent, E, R>) =>
240 Stream.unwrap(
241 Effect.gen(function* () {
242 const coalescer = yield* makeThreadLiveEventCoalescer<E>();
243 yield* source.pipe(
244 Stream.runForEachArray(coalescer.offerAll),
245 Effect.andThen(coalescer.end),
246 Effect.raceFirst(coalescer.failed),
247 Effect.exit,
248 Effect.flatMap((exit) =>
249 Exit.isFailure(exit) ? coalescer.close(exit.cause) : Effect.void,
250 ),
251 Effect.forkScoped({ startImmediately: true }),
252 );
253 return coalescer.stream;
254 }),
255 );