Repository Public
1/**2 * Bidirectional WebSocket pipe browser <-> runtime with backpressure in both directions.3 * Frame types (binary/text) are preserved. Frame contents are never logged or inspected,4 * except client TEXT frames which must validate as `TermClientMsg` (invalid ones are dropped).5 */6import type { Socket } from 'node:net';7import type WebSocket from 'ws';8import { TermClientMsg, WS_CLOSE } from '@cr/shared';910const rawSocket = (ws: WebSocket): Socket | null => (ws as unknown as { _socket?: Socket | null })._socket ?? null;11const SATURATED_MS = 30_000;12const PRE_OPEN_QUEUE_BYTES = 2 * 1024 * 1024;1314export interface PipeOptions {15 /** High-water mark of either outbound buffer (CR_TERM_BUFFER_BYTES). */16 bufferBytes: number;17 /** false ⇒ client frames are ignored (read-only streams such as task logs). */18 clientToUpstream: boolean;19 onClose?: () => void;20 /**21 * Called for every BINARY client frame relayed upstream (terminal input). The frame is not22 * passed: its content is never inspected. Used by agent profiles without a prompt-submit hook.23 */24 onClientInput?: () => void;25}2627/**28 * Returns a sender source→sink that pauses reading `source` (its underlying stream) while29 * `sink` has more than `limit` bytes queued, and resumes once it drained below half of it.30 * A sink that stays saturated for 30 s triggers `onStuck`.31 */32function makeForwarder(source: WebSocket, sink: WebSocket, limit: number, onStuck: () => void): (data: WebSocket.RawData, isBinary: boolean) => void {33 let paused = false;34 let stuck: NodeJS.Timeout | null = null;35 const resume = (): void => {36 if (!paused) return;37 paused = false;38 if (stuck) clearTimeout(stuck);39 stuck = null;40 rawSocket(source)?.resume();41 };42 sink.once('close', () => stuck && clearTimeout(stuck));43 return (data, isBinary) => {44 if (sink.readyState !== sink.OPEN) return;45 sink.send(data, { binary: isBinary }, () => {46 if (paused && sink.bufferedAmount <= limit / 2) resume();47 });48 if (!paused && sink.bufferedAmount > limit) {49 paused = true;50 rawSocket(source)?.pause();51 stuck = setTimeout(onStuck, SATURATED_MS);52 stuck.unref();53 }54 };55}5657export function pipeWebSockets(client: WebSocket, upstream: WebSocket, opts: PipeOptions): void {58 let closed = false;59 let queue: { data: WebSocket.RawData; isBinary: boolean }[] | null = [];60 let queued = 0;61 let lastPong = Date.now();6263 const finish = (clientCode: number, reason: string): void => {64 if (closed) return;65 closed = true;66 clearInterval(heartbeat);67 // A socket that is already CLOSING is finishing its own close handshake (e.g. `ws` refusing68 // an oversized frame with 1009): leave it alone instead of cutting it with terminate().69 if (client.readyState === client.OPEN) client.close(clientCode, reason);70 else if (client.readyState === client.CONNECTING) client.terminate();71 if (upstream.readyState === upstream.OPEN) upstream.close(1000);72 else if (upstream.readyState === upstream.CONNECTING) upstream.terminate();73 opts.onClose?.();74 };75 const saturated = (): void => finish(WS_CLOSE.BACKPRESSURE, 'backpressure');7677 const heartbeat = setInterval(() => {78 if (Date.now() - lastPong > 70_000) return finish(1001, 'heartbeat timeout');79 if (client.readyState === client.OPEN) client.ping();80 }, 25_000);81 heartbeat.unref();82 client.on('pong', () => (lastPong = Date.now()));8384 const up = makeForwarder(client, upstream, opts.bufferBytes, saturated);85 const down = makeForwarder(upstream, client, opts.bufferBytes, saturated);86 const toUpstream = (data: WebSocket.RawData, isBinary: boolean): void => {87 if (!isBinary) {88 // Control frames: only well-formed resize/ping messages are relayed.89 let ok = false;90 try {91 ok = TermClientMsg.safeParse(JSON.parse(data.toString())).success;92 } catch {93 ok = false;94 }95 if (!ok) return;96 } else {97 opts.onClientInput?.();98 }99 up(data, isBinary);100 };101102 client.on('message', (data, isBinary) => {103 if (closed || !opts.clientToUpstream) return;104 if (queue) {105 queued += Array.isArray(data) ? data.reduce((n, b) => n + b.length, 0) : (data as Buffer).byteLength;106 if (queued > PRE_OPEN_QUEUE_BYTES) return finish(1009, 'too much data before ready');107 queue.push({ data, isBinary });108 return;109 }110 toUpstream(data, isBinary);111 });112 upstream.on('open', () => {113 const pending = queue ?? [];114 queue = null;115 for (const m of pending) toUpstream(m.data, m.isBinary);116 });117 upstream.on('message', (data, isBinary) => {118 if (closed) return;119 down(data, isBinary);120 });121122 client.on('close', () => finish(1000, 'client closed'));123 client.on('error', () => finish(1011, 'client error'));124 upstream.on('close', (code) => {125 // Runtime close codes in the application range are meaningful to the browser (4410, 4429…).126 finish(code >= 4000 && code < 5000 ? code : code === 1000 ? 1000 : WS_CLOSE.RUNTIME_UNAVAILABLE, 'runtime closed');127 });128 upstream.on('error', () => finish(WS_CLOSE.RUNTIME_UNAVAILABLE, 'runtime unavailable'));129 upstream.on('unexpected-response', (_req, res) => {130 res.resume();131 finish(res.statusCode === 404 ? WS_CLOSE.SESSION_ENDED : WS_CLOSE.RUNTIME_UNAVAILABLE, 'runtime refused');132 });133}134