1import * as Duration from "effect/Duration";
2import * as Effect from "effect/Effect";
3import * as Ref from "effect/Ref";
4import * as Scope from "effect/Scope";
5import * as Semaphore from "effect/Semaphore";
6
7export interface ProviderTextDeltaUpdate {
8 readonly turnId: string;
9 readonly itemId: string;
10 readonly text: string;
11 readonly completed: boolean;
12}
13
14export interface ProviderTextDeltaCoalescer {
15 readonly append: (input: {
16 readonly turnId: string;
17 readonly itemId: string;
18 readonly delta: string;
19 }) => Effect.Effect<void>;
20 readonly complete: (input: {
21 readonly turnId: string;
22 readonly itemId: string;
23 readonly finalText?: string;
24 readonly emitEmpty?: boolean;
25 }) => Effect.Effect<string>;
26 readonly flushTurn: (turnId: string) => Effect.Effect<void>;
27}
28
29interface BufferedProviderText {
30 readonly turnId: string;
31 readonly itemId: string;
32 readonly text: string;
33 readonly dirty: boolean;
34}
35
36function providerTextBufferKey(turnId: string, itemId: string): string {
37 return `${turnId}\u0000${itemId}`;
38}
39
40export const makeProviderTextDeltaCoalescer = Effect.fn("makeProviderTextDeltaCoalescer")(
41 function* (input: {
42 readonly flushIntervalMs: number;
43 readonly emit: (update: ProviderTextDeltaUpdate) => Effect.Effect<void>;
44 }): Effect.fn.Return<ProviderTextDeltaCoalescer, never, Scope.Scope> {
45 const buffered = yield* Ref.make(new Map<string, BufferedProviderText>());
46 const flushScheduled = yield* Ref.make(false);
47 const flushLock = yield* Semaphore.make(1);
48 const coalescerScope = yield* Effect.scope;
49
50 const drain = (options: {
51 readonly predicate: (message: BufferedProviderText) => boolean;
52 readonly completed: boolean;
53 readonly onlyDirty: boolean;
54 readonly releaseSchedule?: boolean;
55 }) =>
56 flushLock.withPermit(
57 Effect.gen(function* () {
58 const current = yield* Ref.get(buffered);
59 const updates: Array<ProviderTextDeltaUpdate> = [];
60 for (const message of current.values()) {
61 if (!options.predicate(message) || (options.onlyDirty && !message.dirty)) {
62 continue;
63 }
64 updates.push({
65 turnId: message.turnId,
66 itemId: message.itemId,
67 text: message.text,
68 completed: options.completed,
69 });
70 }
71 const emitUpdates = Effect.forEach(updates, input.emit, { discard: true });
72 yield* options.releaseSchedule === true
73 ? emitUpdates.pipe(Effect.ensuring(Ref.set(flushScheduled, false)))
74 : emitUpdates;
75 yield* Ref.update(buffered, (current) => {
76 const next = new Map(current);
77 for (const [key, message] of current) {
78 if (!options.predicate(message) || (options.onlyDirty && !message.dirty)) {
79 continue;
80 }
81 if (options.completed) {
82 next.delete(key);
83 } else {
84 next.set(key, { ...message, dirty: false });
85 }
86 }
87 return next;
88 });
89 }),
90 );
91
92 const flushDirty = drain({
93 predicate: () => true,
94 completed: false,
95 onlyDirty: true,
96 releaseSchedule: true,
97 });
98
99 return {
100 append: ({ turnId, itemId, delta }) =>
101 delta.length === 0
102 ? Effect.void
103 : Effect.uninterruptible(
104 Effect.gen(function* () {
105 const shouldSchedule = yield* flushLock.withPermit(
106 Effect.gen(function* () {
107 yield* Ref.update(buffered, (current) => {
108 const key = providerTextBufferKey(turnId, itemId);
109 const existing = current.get(key);
110 const next = new Map(current);
111 next.set(key, {
112 turnId,
113 itemId,
114 text: `${existing?.text ?? ""}${delta}`,
115 dirty: true,
116 });
117 return next;
118 });
119 return yield* Ref.modify(flushScheduled, (scheduled) => [!scheduled, true]);
120 }),
121 );
122 if (shouldSchedule) {
123 yield* Effect.sleep(Duration.millis(Math.max(1, input.flushIntervalMs))).pipe(
124 Effect.andThen(flushDirty),
125 Effect.interruptible,
126 Effect.forkIn(coalescerScope),
127 );
128 }
129 }),
130 ),
131 complete: ({ turnId, itemId, finalText, emitEmpty = true }) =>
132 flushLock.withPermit(
133 Effect.gen(function* () {
134 const key = providerTextBufferKey(turnId, itemId);
135 const existing = (yield* Ref.get(buffered)).get(key);
136 const text = finalText !== undefined ? finalText : (existing?.text ?? "");
137 if (emitEmpty || text.length > 0) {
138 yield* input.emit({ turnId, itemId, text, completed: true });
139 }
140 yield* Ref.update(buffered, (current) => {
141 const next = new Map(current);
142 next.delete(key);
143 return next;
144 });
145 return text;
146 }),
147 ),
148 flushTurn: (turnId) =>
149 drain({
150 predicate: (message) => message.turnId === turnId,
151 completed: true,
152 onlyDirty: false,
153 }),
154 };
155 },
156);