Open History to see the last commit
1import {
2 ORCHESTRATION_V2_WS_METHODS,
3 type TerminalAttachStreamEvent,
4 type TerminalEvent,
5 WS_METHODS,
6} from "@t3tools/contracts";
7import type * as Arr from "effect/Array";
8import * as Cause from "effect/Cause";
9import * as Effect from "effect/Effect";
10import * as Latch from "effect/Latch";
11import * as Stream from "effect/Stream";
12import type { RpcServer } from "effect/rpc";
13
14const MAX_PENDING_CHUNKS = 8;
15const MAX_PENDING_BYTES = 64 * 1024;
16// Charge an envelope so tiny events cannot bypass the byte budget.
17const EVENT_OVERHEAD = 128;
18const isFull = (sizes: number[]) =>
19 sizes.length >= MAX_PENDING_CHUNKS ||
20 sizes.reduce((total, size) => total + size, 0) >= MAX_PENDING_BYTES;
21// Terminal output is small chunks, so its window also caps bytes. The other streams
22// window by count only: sizing a multi-MB thread snapshot would serialize it twice.
23const BYTE_WINDOWED_STREAMS: ReadonlySet<string> = new Set([
24 WS_METHODS.terminalAttach,
25 WS_METHODS.terminalObserve,
26]);
27// One ACK round trip per chunk stalls sync on high-latency links.
28const COUNT_WINDOWED_STREAMS: ReadonlySet<string> = new Set([
29 WS_METHODS.subscribeServerConfig,
30 ORCHESTRATION_V2_WS_METHODS.subscribeShell,
31 ORCHESTRATION_V2_WS_METHODS.subscribeThread,
32]);
33
34const byteLengthOf = (value: unknown) => (typeof value === "string" ? Buffer.byteLength(value) : 0);
35
36// Charges every variable-size field a terminal event can carry, so the byte
37// budgets below cannot be bypassed by a large payload in an uncharged field.
38const queuedSize = (event: TerminalEvent | TerminalAttachStreamEvent) =>
39 EVENT_OVERHEAD +
40 ("threadId" in event ? byteLengthOf(event.threadId) + byteLengthOf(event.terminalId) : 0) +
41 (event.type === "output"
42 ? byteLengthOf(event.data)
43 : event.type === "error"
44 ? byteLengthOf(event.message)
45 : event.type === "activity"
46 ? byteLengthOf(event.label)
47 : "snapshot" in event
48 ? byteLengthOf(event.snapshot.threadId) +
49 byteLengthOf(event.snapshot.terminalId) +
50 byteLengthOf(event.snapshot.history) +
51 byteLengthOf(event.snapshot.cwd) +
52 byteLengthOf(event.snapshot.worktreePath) +
53 byteLengthOf(event.snapshot.label)
54 : 0);
55
56// Byte-windowed streams carry TerminalAttachStreamEvent, whose encoded form is the
57// event itself, so charge it like the output queue does instead of serializing twice.
58const chunkSize = (values: ReadonlyArray<unknown>) =>
59 values.reduce<number>(
60 (total, value) => total + queuedSize(value as TerminalAttachStreamEvent),
61 EVENT_OVERHEAD,
62 );
63
64interface Window {
65 readonly bytes: boolean;
66 readonly sizes: number[];
67}
68
69export function withOutputWindow(
70 protocol: RpcServer.Protocol["Service"],
71): RpcServer.Protocol["Service"] {
72 // RpcServer drops a disconnected client before its streams exit, so their Exit never
73 // reaches send; departed clients are pruned when a window opens instead.
74 const clients = new Map<number, Map<string, Window>>();
75 let receive: Parameters<RpcServer.Protocol["Service"]["run"]>[0];
76 const windowFor = (clientId: number, requestId: string) => clients.get(clientId)?.get(requestId);
77 const openWindow = Effect.fnUntraced(function* (
78 clientId: number,
79 requestId: string,
80 bytes: boolean,
81 ) {
82 const connected = yield* protocol.clientIds;
83 for (const id of clients.keys()) if (!connected.has(id)) clients.delete(id);
84 let windows = clients.get(clientId);
85 if (!windows) clients.set(clientId, (windows = new Map()));
86 if (!windows.has(requestId)) windows.set(requestId, { bytes, sizes: [] });
87 });
88 return {
89 ...protocol,
90 run: (write) => {
91 receive = write;
92 return protocol.run((clientId, message) =>
93 Effect.suspend(() => {
94 if (message._tag === "Request") {
95 const bytes = BYTE_WINDOWED_STREAMS.has(message.tag);
96 if (bytes || COUNT_WINDOWED_STREAMS.has(message.tag)) {
97 return openWindow(clientId, String(message.id), bytes).pipe(
98 Effect.andThen(write(clientId, message)),
99 );
100 }
101 } else if (message._tag === "Ack") {
102 const window = windowFor(clientId, String(message.requestId));
103 if (window) {
104 const wasFull = isFull(window.sizes);
105 window.sizes.shift();
106 if (!wasFull || isFull(window.sizes)) return Effect.void;
107 }
108 }
109 return write(clientId, message);
110 }),
111 );
112 },
113 send: (clientId, response, transferables) =>
114 Effect.suspend(() => {
115 const send = protocol.send(clientId, response, transferables);
116 if (response._tag === "Exit") {
117 clients.get(clientId)?.delete(String(response.requestId));
118 } else if (response._tag === "Chunk") {
119 const window = windowFor(clientId, String(response.requestId));
120 if (window) {
121 window.sizes.push(window.bytes ? chunkSize(response.values) : 0);
122 if (!isFull(window.sizes)) {
123 return send.pipe(
124 Effect.andThen(() =>
125 receive(clientId, {
126 _tag: "Ack",
127 requestId: response.requestId,
128 }),
129 ),
130 );
131 }
132 }
133 }
134 return send;
135 }),
136 };
137}
138
139const MAX_QUEUED_OUTPUT = 8 * 1024 * 1024;
140
141type Listener<A> = (event: A) => Effect.Effect<void>;
142// Acquisition must register cleanup synchronously with installing the listener.
143type Subscribe<A, E> = (
144 listener: Listener<A>,
145 registerCleanup: (stop: () => void) => void,
146) => Effect.Effect<() => void, E>;
147
148export function terminalOutputStream<A extends TerminalEvent | TerminalAttachStreamEvent, E>(
149 subscribe: Subscribe<A, E>,
150 resync?: Subscribe<A, E>,
151 maxQueued = MAX_QUEUED_OUTPUT,
152): Stream.Stream<A, E> {
153 return Stream.unwrap(
154 Effect.gen(function* () {
155 const ready = yield* Latch.make();
156 let pending: Array<A | undefined> = [];
157 let sizes: number[] = [];
158 let head = 0;
159 let queued = 0;
160 let overflowed = false;
161 let subscription = 0;
162 let unsubscribe: (() => void) | null = null;
163 const release = () => {
164 subscription += 1;
165 unsubscribe?.();
166 unsubscribe = null;
167 pending = [];
168 sizes = [];
169 head = 0;
170 queued = 0;
171 };
172 yield* Effect.addFinalizer(() => Effect.sync(release));
173
174 const listen = (from: Subscribe<A, E>) =>
175 Effect.suspend(() => {
176 const id = ++subscription;
177 return from(
178 (event) =>
179 Effect.sync(() => {
180 if (id !== subscription) return;
181 const size = queuedSize(event);
182 pending.push(event);
183 sizes.push(size);
184 queued += size;
185 // A single snapshot must fit even when history fills the budget.
186 if (queued - sizes[head]! > maxQueued) {
187 release();
188 overflowed = true;
189 }
190 ready.openUnsafe();
191 }),
192 (stop) => {
193 // Overflow or cancellation may invalidate acquisition before registration.
194 if (id === subscription) unsubscribe = stop;
195 else stop();
196 },
197 ).pipe(
198 Effect.asVoid,
199 Effect.onError(() => Effect.sync(release)),
200 );
201 });
202
203 const take = () => {
204 const event = pending[head]!;
205 pending[head] = undefined;
206 head += 1;
207 return event;
208 };
209 const takeBatch = (): Arr.NonEmptyReadonlyArray<A> => {
210 let bytes = sizes[head]!;
211 const batch: [A, ...A[]] = [take()];
212 while (head < pending.length && bytes + sizes[head]! <= MAX_PENDING_BYTES) {
213 bytes += sizes[head]!;
214 batch.push(take());
215 }
216 queued -= bytes;
217 if (head * 2 >= pending.length) {
218 pending = pending.slice(head);
219 sizes = sizes.slice(head);
220 head = 0;
221 }
222 return batch;
223 };
224
225 const pull: Effect.Effect<Arr.NonEmptyReadonlyArray<A>, E | Cause.Done> = Effect.suspend(
226 () => {
227 if (overflowed) {
228 overflowed = false;
229 return resync ? listen(resync).pipe(Effect.andThen(pull)) : Cause.done();
230 }
231 if (head < pending.length) return Effect.succeed(takeBatch());
232 ready.closeUnsafe();
233 return ready.await.pipe(Effect.andThen(pull));
234 },
235 );
236
237 yield* listen(subscribe);
238 return Stream.fromPull(Effect.succeed(pull));
239 }),
240 );
241}