authorgravatar for me@t3.ggTheo Browne <me@t3.gg> 2026-09-04 09:47:23-07:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2026-09-04 09:47:23-07:00
log108f295cc3672716c3cb8291ce487846df5ce098
treeea0a6e668c6437bd4b54f482b682f0d14b334d3b
parentd536b0580d044967dc644498d8c5a6e96464f767
signature Signed by PGP key B5690EEEBB952194

fix(server): bound slow-client event buffers (#9715)

6 files changed, 786 insertions(+), 105 deletions(-)

apps/server/src/orchestration/LiveStreamBudget.test.ts created+47
......@@ -0,0 +1,47 @@
1import { it } from "@effect/vitest";
2import * as Deferred from "effect/Deferred";
3import * as Effect from "effect/Effect";
4import * as Queue from "effect/Queue";
5import * as Stream from "effect/Stream";
6import { describe, expect } from "vite-plus/test";
7
8import { makeLiveStreamBudget, type RetainedLiveItem } from "./LiveStreamBudget.ts";
9
10describe("LiveStreamBudget", () => {
11 it.effect("closes the source without releasing a batch still waiting for an ACK", () =>
12 Effect.scoped(
13 Effect.gen(function* () {
14 const budget = yield* makeLiveStreamBudget({ maxItems: 3 });
15 const queue = yield* Queue.unbounded<RetainedLiveItem<{ text: string }>>();
16 const sourceClosed = yield* Deferred.make<void>();
17 const first = yield* budget.retain({ text: "first" });
18 const second = yield* budget.retain({ text: "second" });
19 const third = yield* budget.retain({ text: "third" });
20 yield* Queue.offerAll(queue, [first, second, third]);
21 yield* Effect.scoped(
22 Effect.gen(function* () {
23 const pull = yield* Stream.toPull(
24 budget.deliver(
25 Stream.fromQueue(queue).pipe(
26 Stream.rechunk(1),
27 Stream.ensuring(Deferred.succeed(sourceClosed, undefined)),
28 ),
29 ),
30 );
31 expect(yield* pull).toEqual([{ text: "first" }]);
32 // The other two items are in the source's pull state, not its queue.
33 expect(yield* Queue.size(queue)).toBe(0);
34 expect((yield* budget.usage).retainedItems).toBe(3);
35 const overflow = yield* budget.retain({ text: "fourth" }).pipe(Effect.result);
36 expect(overflow._tag).toBe("Failure");
37 // Do not resume the consumer. Its source scope must close now.
38 yield* Deferred.await(sourceClosed);
39 yield* budget.closed;
40 expect((yield* budget.usage).retainedItems).toBe(1);
41 }),
42 );
43 expect(yield* budget.usage).toEqual({ retainedItems: 0, retainedSerializedBytes: 0 });
44 }),
45 ),
46 );
47});
apps/server/src/orchestration/LiveStreamBudget.ts created+196
......@@ -0,0 +1,196 @@
1import { OrchestrationGetSnapshotError } from "@t3tools/contracts";
2import * as Arr from "effect/Array";
3import * as Deferred from "effect/Deferred";
4import * as Effect from "effect/Effect";
5import * as Exit from "effect/Exit";
6import * as Scope from "effect/Scope";
7import * as Stream from "effect/Stream";
8
9export const LIVE_STREAM_MAX_ITEMS = 1_000;
10export const LIVE_STREAM_MAX_SERIALIZED_BYTES = 8 * 1024 * 1024;
11
12export interface RetainedLiveItem<A> {
13 readonly value: A;
14 readonly serializedBytes: number;
15}
16
17// Published events are immutable and shared across subscriptions. Measure each
18// object once without keeping the event or its serialized copy alive.
19const serializedSizes = new WeakMap<object, number>();
20
21function serializedSize(value: object): number {
22 const cached = serializedSizes.get(value);
23 if (cached !== undefined) {
24 return cached;
25 }
26 const bytes = Buffer.byteLength(JSON.stringify(value));
27 serializedSizes.set(value, bytes);
28 return bytes;
29}
30
31/** One budget covers one subscription and its delivery stream, including the batch waiting for an RPC ACK. */
32export const makeLiveStreamBudget = Effect.fn("makeLiveStreamBudget")(function* (limits?: {
33 readonly maxItems?: number;
34 readonly maxSerializedBytes?: number;
35}) {
36 const maxItems = limits?.maxItems ?? LIVE_STREAM_MAX_ITEMS;
37 const maxSerializedBytes = limits?.maxSerializedBytes ?? LIVE_STREAM_MAX_SERIALIZED_BYTES;
38 const failed = yield* Deferred.make<never, OrchestrationGetSnapshotError>();
39 const cleanupComplete = yield* Deferred.make<void>();
40 let failure: OrchestrationGetSnapshotError | undefined;
41 const retained = new Set<RetainedLiveItem<unknown>>();
42 let retainedSerializedBytes = 0;
43
44 const release = (items: Iterable<RetainedLiveItem<unknown>>) => {
45 for (const item of items) {
46 if (!retained.delete(item)) {
47 continue;
48 }
49 retainedSerializedBytes -= item.serializedBytes;
50 }
51 };
52 yield* Effect.addFinalizer(() =>
53 Effect.sync(() => release(retained)).pipe(
54 Effect.andThen(Deferred.succeed(cleanupComplete, undefined)),
55 ),
56 );
57
58 const check = Effect.suspend(() => (failure ? Effect.fail(failure) : Effect.void));
59
60 const overflow = Effect.fn("LiveStreamBudget.overflow")(function* (
61 nextItems: number,
62 nextSerializedBytes: number,
63 ) {
64 failure ??= new OrchestrationGetSnapshotError({
65 message: "The live event buffer is full. Resume from the last received sequence.",
66 });
67 yield* Deferred.fail(failed, failure);
68 yield* Effect.logWarning("orchestration live event buffer is full", {
69 retainedItems: retained.size,
70 retainedSerializedBytes,
71 nextItems,
72 nextSerializedBytes,
73 maxItems,
74 maxSerializedBytes,
75 });
76 return yield* failure;
77 });
78
79 const retain = <A extends object>(value: A, payload: object = value) =>
80 Effect.suspend(() => {
81 if (failure) {
82 return Effect.fail(failure);
83 }
84 const serializedBytes = serializedSize(payload);
85 const nextItems = retained.size + 1;
86 const nextSerializedBytes = retainedSerializedBytes + serializedBytes;
87 if (nextItems > maxItems || nextSerializedBytes > maxSerializedBytes) {
88 return overflow(nextItems, nextSerializedBytes);
89 }
90 const item = { value, serializedBytes };
91 retained.add(item);
92 retainedSerializedBytes = nextSerializedBytes;
93 return Effect.succeed(item);
94 });
95
96 // Replace one coalescing batch atomically. Both raw and projected payloads
97 // count against the same budget, and discarded updates release their charge.
98 const replace = <A extends object>(
99 previous: ReadonlyArray<RetainedLiveItem<unknown>>,
100 values: ReadonlyArray<A>,
101 payload: (value: A) => object = (value) => value,
102 ) =>
103 Effect.suspend(() => {
104 if (failure) {
105 return Effect.fail(failure);
106 }
107 const next = values.map((value) => ({
108 value,
109 serializedBytes: serializedSize(payload(value)),
110 }));
111 let nextItems = retained.size + next.length;
112 let nextSerializedBytes =
113 retainedSerializedBytes + next.reduce((sum, item) => sum + item.serializedBytes, 0);
114 for (const item of previous) {
115 if (retained.has(item)) {
116 nextItems -= 1;
117 nextSerializedBytes -= item.serializedBytes;
118 }
119 }
120 if (nextItems > maxItems || nextSerializedBytes > maxSerializedBytes) {
121 return overflow(nextItems, nextSerializedBytes);
122 }
123 release(previous);
124 for (const item of next) {
125 retained.add(item);
126 }
127 retainedSerializedBytes = nextSerializedBytes;
128 return Effect.succeed(next);
129 });
130
131 const deliver = <A, E, R>(stream: Stream.Stream<RetainedLiveItem<A>, E, R>) =>
132 Stream.fromPull(
133 Effect.gen(function* () {
134 yield* check;
135 const sourceScope = yield* Scope.fork(yield* Effect.scope);
136 const source = {
137 pull: yield* Stream.toPull(stream).pipe(Scope.provide(sourceScope)),
138 };
139 let inFlight: ReadonlyArray<RetainedLiveItem<A>> = [];
140 yield* Effect.addFinalizer(() =>
141 Effect.sync(() => {
142 release(inFlight);
143 inFlight = [];
144 source.pull = Effect.interrupt;
145 }),
146 );
147 yield* Deferred.await(failed).pipe(
148 Effect.catchTags({
149 OrchestrationGetSnapshotError: (error) =>
150 Effect.sync(() => {
151 // A grouped source can retain its own pending chunk. Close it
152 // and release the pull closure without waiting for an RPC ACK.
153 source.pull = Effect.interrupt;
154 }).pipe(
155 Effect.andThen(Scope.close(sourceScope, Exit.fail(error))),
156 Effect.andThen(
157 Effect.sync(() => {
158 const delivered = new Set<RetainedLiveItem<unknown>>(inFlight);
159 for (const item of retained) {
160 if (!delivered.has(item)) {
161 release([item]);
162 }
163 }
164 }),
165 ),
166 Effect.andThen(Deferred.succeed(cleanupComplete, undefined)),
167 ),
168 }),
169 Effect.forkScoped,
170 );
171 // @effect-diagnostics-next-line returnEffectInGen:off - Stream.fromPull needs the pull effect as its result.
172 return Effect.gen(function* () {
173 // RpcServer requests the next batch only after the client ACKs this
174 // one. Removing items from a queue alone does not mean delivery ended.
175 release(inFlight);
176 inFlight = [];
177 yield* check;
178 const items = yield* Effect.raceFirst(source.pull, Deferred.await(failed));
179 inFlight = items;
180 yield* check;
181 return Arr.map(items, (item) => item.value);
182 });
183 }),
184 );
185
186 return {
187 retain,
188 replace,
189 release,
190 deliver,
191 check,
192 failed: Deferred.await(failed),
193 closed: Deferred.await(cleanupComplete),
194 usage: Effect.sync(() => ({ retainedItems: retained.size, retainedSerializedBytes })),
195 };
196});
apps/server/src/orchestration/ThreadLiveEventCoalescer.test.ts+90-11
......@@ -3,12 +3,14 @@ import {
33 MessageId,
44 ThreadId,
55 TurnId,
6 type OrchestrationEvent,
6 OrchestrationEvent,
77 type OrchestrationThreadActivity,
88} from "@t3tools/contracts";
99import { it } from "@effect/vitest";
1010import * as Clock from "effect/Clock";
1111import * as Effect from "effect/Effect";
12import * as Schema from "effect/Schema";
13import * as Stream from "effect/Stream";
1214import * as TestClock from "effect/testing/TestClock";
1315import { describe, expect } from "vite-plus/test";
1416
......@@ -19,6 +21,7 @@ import {
1921
2022const threadId = ThreadId.make("thread-coalescer-test");
2123const turnId = TurnId.make("turn-coalescer-test");
24const encodeEvent = Schema.encodeSync(Schema.fromJsonString(OrchestrationEvent));
2225
2326function makeToolActivity(
2427 sequence: number,
......@@ -61,7 +64,7 @@ function makeToolActivity(
6164 };
6265}
6366
64function makeMessage(sequence: number): OrchestrationEvent {
67function makeMessage(sequence: number, text = "Still working"): OrchestrationEvent {
6568 return {
6669 sequence,
6770 eventId: EventId.make(`event-${sequence}`),
......@@ -77,7 +80,7 @@ function makeMessage(sequence: number): OrchestrationEvent {
7780 threadId,
7881 messageId: MessageId.make(`message-${sequence}`),
7982 role: "assistant",
80 text: "Still working",
83 text,
8184 turnId,
8285 streaming: false,
8386 createdAt: "2026-01-01T00:00:02.000Z",
......@@ -134,15 +137,14 @@ describe("ThreadLiveEventCoalescer", () => {
134137 const startedAt = yield* Clock.currentTimeMillis;
135138 yield* Effect.forEach(
136139 Array.from({ length: 10 }, (_, index) => index + 2),
137 (sequence) =>
138 coalescer.offerAndWait({ kind: "event", event: makeToolActivity(sequence) }),
140 (sequence) => coalescer.offer({ kind: "event", event: makeToolActivity(sequence) }),
139141 { discard: true },
140142 );
141 yield* coalescer.offerAndWait({ kind: "event", event: makeMessage(12) });
143 yield* coalescer.offer({ kind: "event", event: makeMessage(12) });
142144
143145 expect(yield* Clock.currentTimeMillis).toBe(startedAt);
144146 expect(
145 Array.from(yield* coalescer.takeAll).map((item) =>
147 (yield* coalescer.stream.pipe(Stream.take(2), Stream.runCollect)).map((item) =>
146148 item.kind === "event" ? item.event.sequence : item.kind,
147149 ),
148150 ).toEqual([11, 12]);
......@@ -155,17 +157,94 @@ describe("ThreadLiveEventCoalescer", () => {
155157 Effect.gen(function* () {
156158 const coalescer = yield* makeThreadLiveEventCoalescer({ coalesceWindow: "500 millis" });
157159 const startedAt = yield* Clock.currentTimeMillis;
158 yield* coalescer.offerAndWait({ kind: "event", event: makeToolActivity(2) });
159 yield* coalescer.offerAndWait({ kind: "event", event: makeToolActivity(3) });
160 yield* coalescer.offerAndWait({ kind: "synchronized" });
160 yield* coalescer.offer({ kind: "event", event: makeToolActivity(2) });
161 yield* coalescer.offer({ kind: "event", event: makeToolActivity(3) });
162 yield* coalescer.offer({ kind: "synchronized" });
161163
162164 expect(yield* Clock.currentTimeMillis).toBe(startedAt);
163165 expect(
164 Array.from(yield* coalescer.takeAll).map((item) =>
166 (yield* coalescer.stream.pipe(Stream.take(2), Stream.runCollect)).map((item) =>
165167 item.kind === "event" ? item.event.sequence : item.kind,
166168 ),
167169 ).toEqual([3, "synchronized"]);
168170 }),
169171 ).pipe(Effect.provide(TestClock.layer())),
170172 );
173
174 it.effect("fails and clears pending updates when their serialized payload fills the budget", () =>
175 Effect.scoped(
176 Effect.gen(function* () {
177 const first = makeToolActivity(1);
178 const coalescer = yield* makeThreadLiveEventCoalescer({
179 coalesceWindow: "500 millis",
180 maxSerializedBytes: Buffer.byteLength(encodeEvent(first)),
181 });
182 yield* coalescer.offer({ kind: "event", event: first });
183 expect(yield* coalescer.usage).toEqual({
184 retainedItems: 1,
185 retainedSerializedBytes: Buffer.byteLength(encodeEvent(first)),
186 });
187
188 const overflow = yield* coalescer
189 .offer({ kind: "event", event: makeToolActivity(2) })
190 .pipe(Effect.result);
191 expect(overflow._tag).toBe("Failure");
192 yield* coalescer.closed;
193 expect(yield* coalescer.usage).toEqual({ retainedItems: 0, retainedSerializedBytes: 0 });
194 const marker = yield* coalescer.offer({ kind: "synchronized" }).pipe(Effect.result);
195 expect(marker._tag).toBe("Failure");
196 const delivered = yield* coalescer.stream.pipe(Stream.runCollect, Effect.result);
197 expect(delivered._tag).toBe("Failure");
198 }),
199 ),
200 );
201
202 it.effect("keeps the flush timer alive when an offer's shorter scope closes", () =>
203 Effect.scoped(
204 Effect.gen(function* () {
205 const coalescer = yield* makeThreadLiveEventCoalescer({ coalesceWindow: "50 millis" });
206 yield* Effect.scoped(coalescer.offer({ kind: "event", event: makeToolActivity(1) }));
207 yield* TestClock.adjust("50 millis");
208 const items = yield* coalescer.stream.pipe(Stream.take(1), Stream.runCollect);
209 expect(
210 items.map((item) => (item.kind === "event" ? item.event.sequence : item.kind)),
211 ).toEqual([1]);
212 }),
213 ),
214 );
215
216 it.effect("keeps an unacknowledged batch charged and clears later events on overflow", () =>
217 Effect.scoped(
218 Effect.gen(function* () {
219 const coalescer = yield* makeThreadLiveEventCoalescer({ maxItems: 3 });
220 const first = makeMessage(1, "é".repeat(1_024));
221 yield* coalescer.offer({ kind: "event", event: first });
222
223 yield* Effect.scoped(
224 Effect.gen(function* () {
225 const pull = yield* Stream.toPull(coalescer.stream);
226 const batch = yield* pull;
227 expect(
228 batch.map((item) => (item.kind === "event" ? item.event.sequence : null)),
229 ).toEqual([1]);
230 yield* coalescer.offer({ kind: "event", event: makeMessage(2) });
231 yield* coalescer.offer({ kind: "event", event: makeToolActivity(3) });
232 expect((yield* coalescer.usage).retainedItems).toBe(3);
233
234 const overflow = yield* coalescer
235 .offer({ kind: "event", event: makeToolActivity(4, { kind: "tool.completed" }) })
236 .pipe(Effect.result);
237 expect(overflow._tag).toBe("Failure");
238 // Do not pull or acknowledge the batch. Cleanup must still finish.
239 yield* coalescer.closed;
240 expect(yield* coalescer.usage).toEqual({
241 retainedItems: 1,
242 retainedSerializedBytes: Buffer.byteLength(encodeEvent(first)),
243 });
244 }),
245 );
246 expect(yield* coalescer.usage).toEqual({ retainedItems: 0, retainedSerializedBytes: 0 });
247 }),
248 ),
249 );
171250});
apps/server/src/orchestration/ThreadLiveEventCoalescer.ts+94-61
......@@ -1,4 +1,8 @@
1import type { OrchestrationEvent, OrchestrationThreadStreamItem } from "@t3tools/contracts";
1import type {
2 OrchestrationEvent,
3 OrchestrationGetSnapshotError,
4 OrchestrationThreadStreamItem,
5} from "@t3tools/contracts";
26import * as Deferred from "effect/Deferred";
37import * as Duration from "effect/Duration";
48import * as Effect from "effect/Effect";
......@@ -9,6 +13,7 @@ import * as Semaphore from "effect/Semaphore";
913import * as Stream from "effect/Stream";
1014
1115import { projectActivityEvent } from "./ActivityPayloadProjection.ts";
16import { makeLiveStreamBudget, type RetainedLiveItem } from "./LiveStreamBudget.ts";
1217
1318const COALESCE_WINDOW = Duration.millis(50);
1419const MAX_PENDING_UPDATES = 512;
......@@ -89,17 +94,24 @@ export function coalesceLiveToolUpdatedEvents(
8994}
9095
9196export const makeThreadLiveEventCoalescer = Effect.fn("makeThreadLiveEventCoalescer")(
92 function* (options?: { readonly coalesceWindow?: Duration.Input }) {
93 const output = yield* Queue.unbounded<OrchestrationThreadStreamItem>();
94 const input = yield* Queue.unbounded<{
95 readonly value: ThreadLiveInput;
96 readonly processed?: Deferred.Deferred<void>;
97 }>();
97 function* (options?: {
98 readonly coalesceWindow?: Duration.Input;
99 readonly maxItems?: number;
100 readonly maxSerializedBytes?: number;
101 }) {
102 const coalescerScope = yield* Effect.scope;
103 const budget = yield* makeLiveStreamBudget(options);
104 const cleanupComplete = yield* Deferred.make<void>();
105 const output = yield* Queue.unbounded<
106 RetainedLiveItem<OrchestrationThreadStreamItem>,
107 OrchestrationGetSnapshotError
108 >();
98109 const mutex = yield* Semaphore.make(1);
99110 const coalesceWindow = options?.coalesceWindow ?? COALESCE_WINDOW;
100 let pendingUpdates: Array<OrchestrationEvent> = [];
111 let pendingUpdates: Array<RetainedLiveItem<OrchestrationEvent>> = [];
101112 let windowGeneration = 0;
102113 let windowFiber: Fiber.Fiber<void, never> | null = null;
114 let closed = false;
103115
104116 const cancelWindow = Effect.fn("ThreadLiveEventCoalescer.cancelWindow")(function* () {
105117 const fiber = windowFiber;
......@@ -110,22 +122,21 @@ export const makeThreadLiveEventCoalescer = Effect.fn("makeThreadLiveEventCoales
110122 yield* Fiber.interrupt(fiber);
111123 });
112124
113 const flushPending = Effect.fn("ThreadLiveEventCoalescer.flushPending")(function* (
114 boundary?: OrchestrationEvent,
115 ) {
116 const events = boundary ? [...pendingUpdates, boundary] : pendingUpdates;
117 pendingUpdates = [];
118 if (events.length === 0) {
125 const flushPending = Effect.fn("ThreadLiveEventCoalescer.flushPending")(function* () {
126 if (pendingUpdates.length === 0) {
119127 return;
120128 }
121 yield* Queue.offerAll(
122 output,
123 coalesceLiveToolUpdatedEvents(events).map((event) => ({
129 const items = yield* budget.replace(
130 pendingUpdates,
131 coalesceLiveToolUpdatedEvents(pendingUpdates.map((item) => item.value)).map((event) => ({
124132 kind: "event" as const,
125133 event: projectActivityEvent(event),
126134 })),
135 (item) => item.event,
127136 );
128 });
137 pendingUpdates = [];
138 yield* Queue.offerAll(output, items);
139 }, Effect.uninterruptible);
129140
130141 const flushWindow = (generation: number) =>
131142 Effect.sleep(coalesceWindow).pipe(
......@@ -141,67 +152,89 @@ export const makeThreadLiveEventCoalescer = Effect.fn("makeThreadLiveEventCoales
141152 }
142153 }),
143154 ),
155 Effect.catchTags({ OrchestrationGetSnapshotError: () => Effect.void }),
144156 );
145157
146 const process = Effect.fn("ThreadLiveEventCoalescer.process")(function* (
147 input: ThreadLiveInput,
158 // Keep each source batch together so a synchronization marker cannot pass
159 // events already pulled from PubSub but still being coalesced.
160 const offerAll = Effect.fn("ThreadLiveEventCoalescer.offerAll")(function* (
161 inputs: ReadonlyArray<ThreadLiveInput>,
148162 ) {
149163 yield* mutex.withPermits(1)(
150 Effect.gen(function* () {
151 if (input.kind === "event" && isToolUpdated(input.event)) {
152 pendingUpdates.push(input.event);
153 if (pendingUpdates.length === 1) {
154 const generation = ++windowGeneration;
155 windowFiber = yield* Effect.forkScoped(flushWindow(generation));
156 }
157 if (pendingUpdates.length >= MAX_PENDING_UPDATES) {
164 Effect.forEach(
165 inputs,
166 (input) =>
167 Effect.gen(function* () {
168 yield* budget.check;
169 if (input.kind === "event") {
170 yield* budget.retain(input.event).pipe(
171 Effect.tap((item) => Effect.sync(() => pendingUpdates.push(item))),
172 Effect.uninterruptible,
173 );
174 }
175 if (input.kind === "event" && isToolUpdated(input.event)) {
176 if (pendingUpdates.length === 1) {
177 const generation = ++windowGeneration;
178 windowFiber = yield* Effect.forkIn(flushWindow(generation), coalescerScope);
179 }
180 if (pendingUpdates.length >= MAX_PENDING_UPDATES) {
181 yield* cancelWindow();
182 windowGeneration += 1;
183 yield* flushPending();
184 }
185 return;
186 }
187
158188 yield* cancelWindow();
159189 windowGeneration += 1;
190 // A non-update event closes the run immediately. The coalescer keeps
191 // that boundary after the final update from the run.
160192 yield* flushPending();
161 }
193 if (input.kind === "synchronized") {
194 yield* budget.retain({ kind: "synchronized" as const }).pipe(
195 Effect.flatMap((marker) => Queue.offer(output, marker)),
196 Effect.uninterruptible,
197 );
198 }
199 }),
200 { discard: true },
201 ),
202 );
203 });
204
205 const close = (error?: OrchestrationGetSnapshotError) =>
206 mutex.withPermits(1)(
207 Effect.gen(function* () {
208 if (closed) {
162209 return;
163210 }
164
165 yield* cancelWindow();
211 closed = true;
166212 windowGeneration += 1;
167 // A non-update event closes the run immediately. The coalescer keeps
168 // that boundary after the final update from the run.
169 if (input.kind === "event") {
170 yield* flushPending(input.event);
171 } else {
172 yield* flushPending();
173 yield* Queue.offer(output, { kind: "synchronized" });
213 yield* cancelWindow();
214 budget.release(pendingUpdates);
215 pendingUpdates = [];
216 budget.release(yield* Queue.clear(output).pipe(Effect.orDie));
217 if (error) {
218 yield* Queue.fail(output, error);
174219 }
220 yield* Queue.shutdown(output);
221 yield* Deferred.succeed(cleanupComplete, undefined);
175222 }),
176223 );
177 });
178224
179 yield* Stream.fromQueue(input).pipe(
180 Stream.runForEach(({ value, processed }) =>
181 process(value).pipe(
182 Effect.andThen(processed ? Deferred.succeed(processed, undefined) : Effect.void),
183 ),
184 ),
225 yield* Effect.addFinalizer(() => close());
226 yield* budget.failed.pipe(
227 Effect.catchTags({ OrchestrationGetSnapshotError: close }),
185228 Effect.forkScoped,
186229 );
187230
188 const offer = (value: ThreadLiveInput) => Queue.offer(input, { value }).pipe(Effect.asVoid);
189
190 // Synchronization callers wait for their marker to pass through the same
191 // ordered input queue before draining output produced ahead of it.
192 const offerAndWait = Effect.fn("ThreadLiveEventCoalescer.offerAndWait")(function* (
193 value: ThreadLiveInput,
194 ) {
195 const processed = yield* Deferred.make<void>();
196 yield* Queue.offer(input, { value, processed });
197 yield* Deferred.await(processed);
198 });
199
200231 return {
201 offer,
202 offerAndWait,
203 stream: Stream.fromQueue(output),
204 takeAll: Queue.takeAll(output),
232 offer: (input: ThreadLiveInput) => offerAll([input]),
233 offerAll,
234 stream: budget.deliver(Stream.fromQueue(output)),
235 failed: budget.failed,
236 closed: Deferred.await(cleanupComplete),
237 usage: budget.usage,
205238 } as const;
206239 },
207240);
apps/server/src/server.test.ts+293
......@@ -1194,6 +1194,32 @@ const withWsRpcClient = <A, E, R>(
11941194 onMessage?: (message: string) => void,
11951195) => makeWsRpcClient.pipe(Effect.flatMap(f), Effect.provide(wsRpcProtocolLayer(wsUrl, onMessage)));
11961196
1197const withFirstWsAckHeld = (
1198 wsUrl: string,
1199 held: Deferred.Deferred<void>,
1200 release: Deferred.Deferred<void>,
1201) => {
1202 let holdNextAck = true;
1203 return Layer.effect(RpcClient.Protocol)(
1204 Effect.map(RpcClient.Protocol, (protocol) =>
1205 RpcClient.Protocol.of({
1206 ...protocol,
1207 send: (clientId, request, transferables) => {
1208 const send = protocol.send(clientId, request, transferables);
1209 if (request._tag !== "Ack" || !holdNextAck) {
1210 return send;
1211 }
1212 holdNextAck = false;
1213 return Deferred.succeed(held, undefined).pipe(
1214 Effect.andThen(Deferred.await(release)),
1215 Effect.andThen(send),
1216 );
1217 },
1218 }),
1219 ),
1220 ).pipe(Layer.provide(wsRpcProtocolLayer(wsUrl)));
1221};
1222
11971223const appendSessionCookieToWsUrl = (url: string, sessionCookieHeader: string) => {
11981224 const isAbsoluteUrl = /^[a-zA-Z][a-zA-Z\d+.-]*:/.test(url);
11991225 const next = new URL(url, "http://localhost");
......@@ -8067,6 +8093,273 @@ it.layer(NodeServices.layer)("server router seam", (it) => {
80678093 }).pipe(Effect.provide(NodeHttpServer.layerTest)),
80688094 );
80698095
8096 it.effect(
8097 "stops an overflowing thread producer without an ACK and replays the missing events",
8098 () =>
8099 Effect.gen(function* () {
8100 const liveEvents = yield* PubSub.unbounded<OrchestrationEvent>();
8101 const attached = yield* Deferred.make<void>();
8102 const detached = yield* Deferred.make<void>();
8103 const ackHeld = yield* Deferred.make<void>();
8104 const releaseAck = yield* Deferred.make<void>();
8105 const firstApplied = yield* Deferred.make<void>();
8106 const replayStarted = yield* Deferred.make<void>();
8107 const replayCalls: Array<{ afterSequence: number; limit: number | undefined }> = [];
8108 let headSequence = 0;
8109 let snapshotCalls = 0;
8110 const message = (sequence: number, text: string) =>
8111 ({
8112 sequence,
8113 eventId: EventId.make(`slow-thread-${sequence}`),
8114 aggregateKind: "thread",
8115 aggregateId: defaultThreadId,
8116 occurredAt: "2026-01-01T00:00:01.000Z",
8117 commandId: null,
8118 causationEventId: null,
8119 correlationId: null,
8120 metadata: {},
8121 type: "thread.message-sent",
8122 payload: {
8123 threadId: defaultThreadId,
8124 messageId: MessageId.make(`slow-message-${sequence}`),
8125 role: "assistant",
8126 text,
8127 turnId: TurnId.make("turn-edit"),
8128 streaming: false,
8129 createdAt: "2026-01-01T00:00:01.000Z",
8130 updatedAt: "2026-01-01T00:00:01.000Z",
8131 },
8132 }) satisfies OrchestrationEvent;
8133 const events = [
8134 message(1, "a".repeat(4 * 1024 * 1024)),
8135 message(2, "b".repeat(4 * 1024 * 1024)),
8136 makeLiveToolActivityEvent(3, "tool.completed"),
8137 message(4, "Finished"),
8138 ] as const;
8139
8140 yield* buildAppUnderTest({
8141 layers: {
8142 orchestrationEngine: {
8143 latestSequence: Effect.sync(() => headSequence),
8144 streamDomainEvents: Stream.unwrap(
8145 Effect.gen(function* () {
8146 const subscription = yield* PubSub.subscribe(liveEvents);
8147 yield* Deferred.succeed(attached, undefined);
8148 return Stream.fromSubscription(subscription);
8149 }),
8150 ).pipe(Stream.ensuring(Deferred.succeed(detached, undefined))),
8151 readEvents: (afterSequence, limit) => {
8152 replayCalls.push({ afterSequence, limit });
8153 const range = events.filter(
8154 (event) => event.sequence > afterSequence && event.sequence <= headSequence,
8155 );
8156 return Stream.concat(
8157 Stream.fromEffect(Deferred.succeed(replayStarted, undefined)).pipe(Stream.drain),
8158 Stream.fromIterable(range),
8159 );
8160 },
8161 },
8162 projectionSnapshotQuery: {
8163 getThreadDetailSnapshot: () =>
8164 Effect.sync(() => {
8165 snapshotCalls += 1;
8166 return Option.none();
8167 }),
8168 getEventReplayStats: ({ fromSequenceExclusive, toSequenceInclusive }) => {
8169 const range = events.filter(
8170 (event) =>
8171 event.sequence > fromSequenceExclusive && event.sequence <= toSequenceInclusive,
8172 );
8173 return Effect.succeed({
8174 eventCount: range.length,
8175 payloadBytes: range.reduce(
8176 (bytes, event) => bytes + Buffer.byteLength(jsonRequestBody(event.payload)),
8177 0,
8178 ),
8179 });
8180 },
8181 },
8182 },
8183 });
8184
8185 const wsUrl = yield* getWsServerUrl("/ws");
8186 yield* makeWsRpcClient.pipe(
8187 Effect.flatMap((client) =>
8188 Effect.gen(function* () {
8189 let cursor = 0;
8190 const received: number[] = [];
8191 const attempt = yield* client[ORCHESTRATION_WS_METHODS.subscribeThread]({
8192 threadId: defaultThreadId,
8193 afterSequence: cursor,
8194 }).pipe(
8195 Stream.tap((item) => {
8196 if (item.kind !== "event") return Effect.void;
8197 cursor = item.event.sequence;
8198 received.push(cursor);
8199 return Deferred.succeed(firstApplied, undefined);
8200 }),
8201 Stream.runDrain,
8202 Effect.result,
8203 Effect.forkScoped,
8204 );
8205 yield* Deferred.await(attached);
8206 yield* Deferred.await(replayStarted);
8207 headSequence = 1;
8208 yield* PubSub.publish(liveEvents, events[0]!);
8209 yield* Deferred.await(ackHeld);
8210 yield* Deferred.await(firstApplied);
8211 headSequence = 4;
8212 yield* PubSub.publishAll(liveEvents, events.slice(1));
8213
8214 // This must finish while the server is still waiting for the first
8215 // batch's ACK. A failed output queue alone would leave PubSub live.
8216 yield* Deferred.await(detached);
8217 assert.equal(yield* PubSub.size(liveEvents), 0);
8218 assert.deepEqual(received, [1]);
8219 yield* Deferred.succeed(releaseAck, undefined);
8220 const result = yield* Fiber.join(attempt);
8221 assertTrue(result._tag === "Failure");
8222 assert.equal(result.failure._tag, "OrchestrationGetSnapshotError");
8223
8224 const recovered = yield* client[ORCHESTRATION_WS_METHODS.subscribeThread]({
8225 threadId: defaultThreadId,
8226 afterSequence: cursor,
8227 requestCompletionMarker: true,
8228 }).pipe(
8229 Stream.takeUntil((item) => item.kind === "synchronized"),
8230 Stream.runCollect,
8231 );
8232 assert.deepEqual(
8233 recovered.map((item) => (item.kind === "event" ? item.event.sequence : item.kind)),
8234 [2, 3, 4, "synchronized"],
8235 );
8236 const first = recovered[0];
8237 assertTrue(first?.kind === "event" && first.event.type === "thread.message-sent");
8238 assert.equal(first.event.payload.text, events[1]!.payload.text);
8239 const completed = recovered[1];
8240 assertTrue(
8241 completed?.kind === "event" && completed.event.type === "thread.activity-appended",
8242 );
8243 assert.equal(completed.event.payload.activity.kind, "tool.completed");
8244 assert.deepEqual(replayCalls, [
8245 { afterSequence: 0, limit: 0 },
8246 { afterSequence: 1, limit: 3 },
8247 ]);
8248 assert.equal(snapshotCalls, 0);
8249 }),
8250 ),
8251 Effect.provide(withFirstWsAckHeld(wsUrl, ackHeld, releaseAck)),
8252 );
8253 }).pipe(Effect.scoped, Effect.provide(NodeHttpServer.layerTest)),
8254 );
8255
8256 it.effect("stops an overflowing shell producer without an ACK and recovers deleted entries", () =>
8257 Effect.gen(function* () {
8258 const liveEvents = yield* PubSub.unbounded<OrchestrationEvent>();
8259 const detached = yield* Deferred.make<void>();
8260 const ackHeld = yield* Deferred.make<void>();
8261 const releaseAck = yield* Deferred.make<void>();
8262 let headSequence = 1;
8263 let snapshotCalls = 0;
8264 let replayCalls = 0;
8265 const thread = makeDefaultOrchestrationThreadShell();
8266 const project = makeDefaultOrchestrationReadModel().projects[0]!;
8267 const events: OrchestrationEvent[] = Array.from({ length: 1_001 }, (_, index) =>
8268 makeLiveToolActivityEvent(index + 2, "tool.completed"),
8269 );
8270 events.push(
8271 {
8272 ...events[0]!,
8273 sequence: 1_003,
8274 type: "thread.archived",
8275 payload: {
8276 threadId: defaultThreadId,
8277 archivedAt: "2026-01-01T00:00:02.000Z",
8278 updatedAt: "2026-01-01T00:00:02.000Z",
8279 },
8280 },
8281 {
8282 ...events[0]!,
8283 sequence: 1_004,
8284 aggregateKind: "project",
8285 aggregateId: defaultProjectId,
8286 type: "project.deleted",
8287 payload: { projectId: defaultProjectId, deletedAt: "2026-01-01T00:00:02.000Z" },
8288 },
8289 );
8290 yield* buildAppUnderTest({
8291 layers: {
8292 orchestrationEngine: {
8293 latestSequence: Effect.sync(() => headSequence),
8294 streamDomainEvents: Stream.fromPubSub(liveEvents).pipe(
8295 Stream.ensuring(Deferred.succeed(detached, undefined)),
8296 ),
8297 readEvents: () => {
8298 replayCalls += 1;
8299 return Stream.empty;
8300 },
8301 },
8302 projectionSnapshotQuery: {
8303 getShellSnapshot: () =>
8304 Effect.sync(() => {
8305 snapshotCalls += 1;
8306 return {
8307 snapshotSequence: headSequence,
8308 projects: headSequence === 1 ? [project] : [],
8309 threads: headSequence === 1 ? [thread] : [],
8310 updatedAt: "2026-01-01T00:00:02.000Z",
8311 };
8312 }),
8313 },
8314 },
8315 });
8316 const wsUrl = yield* getWsServerUrl("/ws");
8317 yield* makeWsRpcClient.pipe(
8318 Effect.flatMap((client) =>
8319 Effect.gen(function* () {
8320 const received: OrchestrationShellStreamItem[] = [];
8321 const attempt = yield* client[ORCHESTRATION_WS_METHODS.subscribeShell]({}).pipe(
8322 Stream.tap((item) => Effect.sync(() => received.push(item))),
8323 Stream.runDrain,
8324 Effect.result,
8325 Effect.forkScoped,
8326 );
8327 yield* Deferred.await(ackHeld);
8328 headSequence = 1_004;
8329 yield* PubSub.publishAll(liveEvents, events);
8330 yield* Deferred.await(detached);
8331 assert.equal(yield* PubSub.size(liveEvents), 0);
8332 yield* Deferred.succeed(releaseAck, undefined);
8333 const result = yield* Fiber.join(attempt);
8334 assertTrue(result._tag === "Failure");
8335 assert.equal(result.failure._tag, "OrchestrationGetSnapshotError");
8336 assert.equal(received.length, 1);
8337 const initial = received[0];
8338 assertTrue(initial?.kind === "snapshot");
8339 assert.equal(initial.snapshot.threads.length, 1);
8340
8341 const recovered = yield* client[ORCHESTRATION_WS_METHODS.subscribeShell]({
8342 afterSequence: initial.snapshot.snapshotSequence,
8343 requestCompletionMarker: true,
8344 }).pipe(
8345 Stream.takeUntil((item) => item.kind === "synchronized"),
8346 Stream.runCollect,
8347 );
8348 const snapshot = recovered[0];
8349 assertTrue(snapshot?.kind === "snapshot");
8350 assert.equal(snapshot.snapshot.snapshotSequence, 1_004);
8351 assert.deepEqual(snapshot.snapshot.projects, []);
8352 assert.deepEqual(snapshot.snapshot.threads, []);
8353 assert.deepEqual(recovered[1], { kind: "synchronized" });
8354 assert.equal(snapshotCalls, 2);
8355 assert.equal(replayCalls, 0);
8356 }),
8357 ),
8358 Effect.provide(withFirstWsAckHeld(wsUrl, ackHeld, releaseAck)),
8359 );
8360 }).pipe(Effect.scoped, Effect.provide(NodeHttpServer.layerTest)),
8361 );
8362
80708363 it.effect("subscribeThread replaces a cursor ahead of the authoritative head", () =>
80718364 Effect.gen(function* () {
80728365 let readEventsCalls = 0;
apps/server/src/ws.ts+66-33
......@@ -81,6 +81,7 @@ import {
8181 projectThreadDetailSnapshot,
8282} from "./orchestration/ActivityPayloadProjection.ts";
8383import { makeThreadLiveEventCoalescer } from "./orchestration/ThreadLiveEventCoalescer.ts";
84import { makeLiveStreamBudget, type RetainedLiveItem } from "./orchestration/LiveStreamBudget.ts";
8485import {
8586 cleanupFailedUploadedAttachments,
8687 normalizeDispatchCommand,
......@@ -927,15 +928,6 @@ const makeWsRpcLayer = (
927928 return output;
928929 });
929930
930 const coalesceShellLiveStream = <E, R>(
931 stream: Stream.Stream<ShellLiveInput, E, R>,
932 ): Stream.Stream<OrchestrationShellStreamItem, E, R> =>
933 stream.pipe(
934 Stream.groupedWithin(SHELL_COALESCE_MAX_CHUNK, SHELL_COALESCE_WINDOW),
935 Stream.mapEffect(coalesceShellLiveInputs),
936 Stream.flatMap((items) => Stream.fromIterable(items)),
937 );
938
939931 const dispatchBootstrapTurnStart = (
940932 command: Extract<OrchestrationCommand, { type: "thread.turn.start" }>,
941933 ): Effect.Effect<{ readonly sequence: number }, OrchestrationDispatchCommandError> =>
......@@ -1416,16 +1408,55 @@ const makeWsRpcLayer = (
14161408 // sequence but the live subscription is not attached yet). Every
14171409 // path below emits from this same buffered live tail. Overlapping
14181410 // events are deduped by sequence on the client.
1419 const liveBuffer = yield* Queue.unbounded<ShellLiveInput>();
1411 const liveBudget = yield* makeLiveStreamBudget();
1412 const liveBuffer = yield* Queue.unbounded<
1413 RetainedLiveItem<ShellLiveInput>,
1414 OrchestrationGetSnapshotError
1415 >();
1416 let liveBufferClosed = false;
1417 const closeLiveBuffer = (error?: OrchestrationGetSnapshotError) =>
1418 Effect.gen(function* () {
1419 if (liveBufferClosed) {
1420 return;
1421 }
1422 liveBufferClosed = true;
1423 liveBudget.release(yield* Queue.clear(liveBuffer).pipe(Effect.orDie));
1424 if (error) {
1425 yield* Queue.fail(liveBuffer, error);
1426 }
1427 yield* Queue.shutdown(liveBuffer);
1428 });
1429 yield* Effect.addFinalizer(() => closeLiveBuffer());
1430 yield* liveBudget.failed.pipe(
1431 Effect.catchTags({ OrchestrationGetSnapshotError: closeLiveBuffer }),
1432 Effect.forkScoped,
1433 );
14201434 yield* Effect.forkScoped(
14211435 orchestrationEngine.streamDomainEvents.pipe(
14221436 Stream.runForEach((event) =>
1423 Queue.offer(liveBuffer, { kind: "event" as const, event }),
1437 liveBudget.retain({ kind: "event" as const, event }, event).pipe(
1438 Effect.flatMap((item) => Queue.offer(liveBuffer, item)),
1439 Effect.uninterruptible,
1440 ),
14241441 ),
1442 // Stop the PubSub consumer even if RPC delivery is waiting
1443 // for an ACK and never pulls the failed buffer again.
1444 Effect.raceFirst(liveBudget.failed),
1445 Effect.catchTags({ OrchestrationGetSnapshotError: () => Effect.void }),
14251446 ),
14261447 { startImmediately: true },
14271448 );
1428 const bufferedLiveStream = coalesceShellLiveStream(Stream.fromQueue(liveBuffer));
1449 const coalesceRetainedInputs = (
1450 items: ReadonlyArray<RetainedLiveItem<ShellLiveInput>>,
1451 ) =>
1452 coalesceShellLiveInputs(items.map((item) => item.value)).pipe(
1453 Effect.flatMap((output) => liveBudget.replace(items, output)),
1454 );
1455 const bufferedLiveStream = Stream.fromQueue(liveBuffer).pipe(
1456 Stream.groupedWithin(SHELL_COALESCE_MAX_CHUNK, SHELL_COALESCE_WINDOW),
1457 Stream.mapEffect(coalesceRetainedInputs),
1458 Stream.flatMap((items) => Stream.fromIterable(items)),
1459 );
14291460
14301461 const loadSnapshot = projectionSnapshotQuery.getShellSnapshot().pipe(
14311462 Effect.tapError((cause) =>
......@@ -1443,18 +1474,21 @@ const makeWsRpcLayer = (
14431474 // Offer the completion marker into the same queue as live events.
14441475 // Anything buffered while snapshot/replay work was in flight is
14451476 // therefore delivered before the client is told it is synchronized.
1446 const synchronizedThenLive =
1477 const synchronizedThenLive = liveBudget.deliver(
14471478 input.requestCompletionMarker === true
14481479 ? Stream.concat(
14491480 Stream.fromEffect(
1450 Queue.offer(liveBuffer, { kind: "synchronized" as const }).pipe(
1481 liveBudget.retain({ kind: "synchronized" as const }).pipe(
1482 Effect.flatMap((item) => Queue.offer(liveBuffer, item)),
1483 Effect.uninterruptible,
14511484 Effect.andThen(Queue.takeAll(liveBuffer)),
1452 Effect.flatMap(coalesceShellLiveInputs),
1485 Effect.flatMap(coalesceRetainedInputs),
14531486 ),
14541487 ).pipe(Stream.flatMap((items) => Stream.fromIterable(items))),
14551488 bufferedLiveStream,
14561489 )
1457 : bufferedLiveStream;
1490 : bufferedLiveStream,
1491 );
14581492
14591493 // When the client already holds a shell snapshot (cached, or loaded
14601494 // over HTTP) it passes that snapshot's sequence, and we resume by
......@@ -1551,9 +1585,14 @@ const makeWsRpcLayer = (
15511585 // Attach live delivery before reading either replay or snapshot state.
15521586 // Otherwise an event published while the snapshot is loading is lost.
15531587 const liveBuffer = yield* makeThreadLiveEventCoalescer();
1554 yield* Effect.forkScoped(liveStream.pipe(Stream.runForEach(liveBuffer.offer)), {
1555 startImmediately: true,
1556 });
1588 yield* Effect.forkScoped(
1589 liveStream.pipe(
1590 Stream.runForEachArray(liveBuffer.offerAll),
1591 Effect.raceFirst(liveBuffer.failed),
1592 Effect.catchTags({ OrchestrationGetSnapshotError: () => Effect.void }),
1593 ),
1594 { startImmediately: true },
1595 );
15571596 const bufferedLiveStream = liveBuffer.stream;
15581597
15591598 // When the client already loaded the snapshot over HTTP it passes
......@@ -1603,13 +1642,10 @@ const makeWsRpcLayer = (
16031642 );
16041643 const afterCatchUp =
16051644 input.requestCompletionMarker === true
1606 ? Stream.concat(
1607 Stream.fromEffect(
1608 liveBuffer
1609 .offerAndWait({ kind: "synchronized" as const })
1610 .pipe(Effect.andThen(liveBuffer.takeAll)),
1611 ).pipe(Stream.flatMap((items) => Stream.fromIterable(items))),
1612 bufferedLiveStream,
1645 ? Stream.unwrap(
1646 liveBuffer
1647 .offer({ kind: "synchronized" as const })
1648 .pipe(Effect.as(bufferedLiveStream)),
16131649 )
16141650 : bufferedLiveStream;
16151651 return Stream.concat(catchUpStream, afterCatchUp);
......@@ -1647,13 +1683,10 @@ const makeWsRpcLayer = (
16471683
16481684 const afterSnapshot =
16491685 input.requestCompletionMarker === true
1650 ? Stream.concat(
1651 Stream.fromEffect(
1652 liveBuffer
1653 .offerAndWait({ kind: "synchronized" as const })
1654 .pipe(Effect.andThen(liveBuffer.takeAll)),
1655 ).pipe(Stream.flatMap((items) => Stream.fromIterable(items))),
1656 bufferedLiveStream,
1686 ? Stream.unwrap(
1687 liveBuffer
1688 .offer({ kind: "synchronized" as const })
1689 .pipe(Effect.as(bufferedLiveStream)),
16571690 )
16581691 : bufferedLiveStream;
16591692 return Stream.concat(