1import { WS_METHODS, type TerminalAttachStreamEvent, type TerminalEvent } from "@t3tools/contracts";
2import * as Effect from "effect/Effect";
3import * as Stream from "effect/Stream";
4import type { RpcServer } from "effect/rpc";
5
6const MAX_PENDING_CHUNKS = 8;
7const MAX_PENDING_BYTES = 64 * 1024;
8const isFull = (sizes: number[]) =>
9 sizes.length >= MAX_PENDING_CHUNKS ||
10 sizes.reduce((total, size) => total + size, 0) >= MAX_PENDING_BYTES;
11
12/** Keeps sequenced events intact, including snapshots larger than the batch budget. */
13export function batchTerminalOutput<A extends TerminalAttachStreamEvent | TerminalEvent, E, R>(
14 stream: Stream.Stream<A, E, R>,
15): Stream.Stream<A, E, R> {
16 return Stream.transformPull(stream, (pull) =>
17 Effect.sync(() => {
18 let pending: ReadonlyArray<A> = [];
19 let index = 0;
20
21 return Effect.gen(function* () {
22 if (index === pending.length) {
23 pending = yield* pull;
24 index = 0;
25 }
26
27 const first = pending[index++]!;
28 const batch: [A, ...A[]] = [first];
29 // @effect-diagnostics-next-line preferSchemaOverJson:off -- Match the RPC JSON wire encoding.
30 let bytes = Buffer.byteLength(JSON.stringify(first)) + 2;
31
32 // Reserve 1 KiB for RPC framing; oversized events remain indivisible.
33 while (index < pending.length && bytes < MAX_PENDING_BYTES - 1024) {
34 const next = pending[index]!;
35 // @effect-diagnostics-next-line preferSchemaOverJson:off -- Match the RPC JSON wire encoding.
36 const size = Buffer.byteLength(JSON.stringify(next)) + 1;
37 if (bytes + size > MAX_PENDING_BYTES - 1024) break;
38
39 batch.push(next);
40 bytes += size;
41 index += 1;
42 }
43
44 if (index === pending.length) {
45 pending = [];
46 index = 0;
47 }
48 return batch;
49 });
50 }),
51 );
52}
53
54export function withTerminalOutputWindow(
55 protocol: RpcServer.Protocol["Service"],
56): RpcServer.Protocol["Service"] {
57 const windows = new Map<string, number[]>();
58 let receive: Parameters<RpcServer.Protocol["Service"]["run"]>[0];
59 return {
60 ...protocol,
61 run: (write) => {
62 receive = write;
63 return protocol.run((clientId, message) =>
64 Effect.suspend(() => {
65 if (
66 message._tag === "Request" &&
67 (message.tag === WS_METHODS.terminalAttach ||
68 message.tag === WS_METHODS.terminalObserve ||
69 message.tag === WS_METHODS.subscribeTerminalEvents)
70 ) {
71 const key = `${clientId}:${message.id}`;
72 if (!windows.has(key)) windows.set(key, []);
73 } else if (message._tag === "Ack") {
74 const window = windows.get(`${clientId}:${message.requestId}`);
75 if (window) {
76 const wasFull = isFull(window);
77 window.shift();
78 if (!wasFull || isFull(window)) return Effect.void;
79 }
80 }
81 return write(clientId, message);
82 }),
83 );
84 },
85 send: (clientId, response, transferables) =>
86 Effect.suspend(() => {
87 const send = protocol.send(clientId, response, transferables);
88 if (response._tag === "Exit") {
89 windows.delete(`${clientId}:${response.requestId}`);
90 } else if (response._tag === "Chunk") {
91 const window = windows.get(`${clientId}:${response.requestId}`);
92 if (window) {
93 // @effect-diagnostics-next-line preferSchemaOverJson:off
94 const size = Buffer.byteLength(JSON.stringify(response));
95 window.push(size);
96 if (!isFull(window)) {
97 return send.pipe(
98 Effect.andThen(() =>
99 receive(clientId, {
100 _tag: "Ack",
101 requestId: response.requestId,
102 }),
103 ),
104 );
105 }
106 }
107 }
108 return send;
109 }),
110 };
111}