@aitty/server 0.1.0

package/dist/runtime/websocket-transport.js

@@ -0,0 +1,465 @@
+import { isSessionTokenAuthorized } from "../cli-token.js";
+import { isHostAllowed, isOriginAllowed } from "../network-policy.js";
+import { createStructuredLogger } from "../logging.js";
+import os from "node:os";
+import { createExitControlFrame, createHelloControlFrame, createPingControlFrame, createPongControlFrame, createThemeControlFrame, isAittyResizeControlFrame, parseAittyControlFrame, subscribeThemeSource } from "@aitty/protocol";
+//#region src/runtime/websocket-transport.ts
+const DEFAULT_MAX_BUFFERED_AMOUNT_BYTES = 256 * 1024;
+const DEFAULT_MAX_QUEUED_BYTES = 256 * 1024;
+const DEFAULT_DRAIN_WAIT_MS = 250;
+const DEFAULT_PRIORITY_OUTPUT_RECOVERY_MS = 1e3;
+const DEFAULT_KEEPALIVE_INITIAL_DELAY_MS = 100;
+const DEFAULT_KEEPALIVE_INTERVAL_MS = 3e4;
+const SOCKET_OPEN = 1;
+const SOCKET_CLOSED = 3;
+const NORMAL_CLOSE_CODE = 1e3;
+const SUPERSEDED_CLOSE_CODE = 4001;
+const SUPERSEDED_CLOSE_REASON = "superseded";
+const BACKPRESSURE_CLOSE_CODE = 4002;
+const BACKPRESSURE_CLOSE_REASON = "backpressure";
+function createWebSocketTransport(options) {
+	const server = options.server;
+	const websocketServer = new options.WebSocketServer({ noServer: true });
+	const logger = createStructuredLogger({
+		token: options.token,
+		verbose: options.verbose,
+		writer: options.stderr
+	});
+	const clients = /* @__PURE__ */ new Map();
+	let inputWriteChain = Promise.resolve();
+	const maxBufferedAmountBytes = options.maxBufferedAmountBytes ?? DEFAULT_MAX_BUFFERED_AMOUNT_BYTES;
+	const maxQueuedBytes = options.maxQueuedBytes ?? DEFAULT_MAX_QUEUED_BYTES;
+	const drainWaitMs = options.drainWaitMs ?? DEFAULT_DRAIN_WAIT_MS;
+	const keepaliveInitialDelayMs = Math.max(0, options.keepaliveInitialDelayMs ?? DEFAULT_KEEPALIVE_INITIAL_DELAY_MS);
+	const keepaliveIntervalMs = Math.max(0, options.keepaliveIntervalMs ?? DEFAULT_KEEPALIVE_INTERVAL_MS);
+	let closed = false;
+	let exitNotified = false;
+	let latestThemeFrame = null;
+	const themeSubscription = subscribeThemeSource(options.themeSource, (theme) => {
+		latestThemeFrame = createThemeControlFrame(theme);
+		for (const [websocket, clientState] of clients) {
+			if (!clientState.active || websocket.readyState !== SOCKET_OPEN) continue;
+			enqueueClientFrame(websocket, clientState, latestThemeFrame, false, {
+				drainWaitMs,
+				maxBufferedAmountBytes,
+				maxQueuedBytes,
+				logger
+			});
+		}
+	});
+	const sessionSubscription = options.session.onData((chunk) => {
+		const normalizedChunk = Buffer.from(chunk);
+		const activeClient = getActiveClient(clients);
+		if (!activeClient) return;
+		const [websocket, clientState] = activeClient;
+		if (websocket.readyState !== SOCKET_OPEN) return;
+		enqueueClientFrame(websocket, clientState, normalizedChunk, true, {
+			drainWaitMs,
+			maxBufferedAmountBytes,
+			maxQueuedBytes,
+			logger
+		});
+	});
+	const onUpgrade = (request, socket, head) => {
+		const requestUrl = new URL(request.url ?? "/", "http://127.0.0.1");
+		if (requestUrl.pathname !== "/ws") {
+			rejectUpgrade(socket, 404, "Not Found");
+			return;
+		}
+		if (options.networkPolicy && !isHostAllowed(request.headers.host, options.networkPolicy)) {
+			rejectUpgrade(socket, 403, "Forbidden");
+			return;
+		}
+		if (!isSessionTokenAuthorized(options.token, requestUrl.searchParams.get("t"))) {
+			rejectUpgrade(socket, 403, "Forbidden");
+			return;
+		}
+		if (options.networkPolicy && !isOriginAllowed(request.headers.origin, options.networkPolicy)) {
+			rejectUpgrade(socket, 403, "Forbidden");
+			return;
+		}
+		websocketServer.handleUpgrade(request, socket, head, (websocket) => {
+			websocketServer.emit("connection", websocket, request);
+		});
+	};
+	server.on("upgrade", onUpgrade);
+	websocketServer.on("connection", (websocket) => {
+		if (closed) {
+			websocket.close();
+			return;
+		}
+		logger.info("WebSocket client connected");
+		const clientState = {
+			active: true,
+			keepaliveTimer: null,
+			oversizedAllowanceBytes: 0,
+			pendingBytes: 0,
+			pendingFrames: [],
+			priorityModeUntil: 0,
+			sending: false
+		};
+		const previousClient = getActiveClient(clients);
+		if (previousClient) previousClient[1].active = false;
+		clients.set(websocket, clientState);
+		if (previousClient) supersedeClient(previousClient[0]);
+		const dispose = once(() => {
+			clientState.active = false;
+			clearKeepaliveTimer(clientState);
+			clientState.oversizedAllowanceBytes = 0;
+			clientState.pendingBytes = 0;
+			clientState.pendingFrames.length = 0;
+			clients.delete(websocket);
+		});
+		const replay = options.session.getBufferedOutput();
+		enqueueClientFrame(websocket, clientState, createHelloFrame(options.session, replay), false, {
+			allowOversizedFrame: true,
+			drainWaitMs,
+			maxBufferedAmountBytes,
+			maxQueuedBytes,
+			logger
+		});
+		if (latestThemeFrame) enqueueClientFrame(websocket, clientState, latestThemeFrame, false, {
+			drainWaitMs,
+			maxBufferedAmountBytes,
+			maxQueuedBytes,
+			logger
+		});
+		scheduleKeepaliveFrame(websocket, clientState, {
+			drainWaitMs,
+			keepaliveIntervalMs,
+			maxBufferedAmountBytes,
+			maxQueuedBytes,
+			logger
+		}, keepaliveInitialDelayMs);
+		websocket.on("close", (code, reason) => {
+			const closeReason = Buffer.from(reason).toString("utf8");
+			logger.info(`WebSocket client disconnected code=${code}${closeReason ? ` reason=${closeReason}` : ""}`);
+			dispose();
+		});
+		websocket.on("error", (error) => {
+			const message = error instanceof Error ? error.message : String(error);
+			logger.warn(`WebSocket client error: ${message}`);
+			dispose();
+		});
+		websocket.on("message", (data, isBinary) => {
+			if (!clients.get(websocket)?.active) return;
+			if (isBinary) {
+				const payload = toBuffer(data);
+				scheduleKeepaliveFrame(websocket, clientState, {
+					drainWaitMs,
+					keepaliveIntervalMs,
+					maxBufferedAmountBytes,
+					maxQueuedBytes,
+					logger
+				}, keepaliveIntervalMs);
+				if (containsPriorityControlByte(payload)) prioritizeClientOutput(clientState);
+				inputWriteChain = inputWriteChain.catch(() => void 0).then(() => options.session.write(payload));
+				return;
+			}
+			handleControlFrame(websocket, data, options.session, logger);
+			scheduleKeepaliveFrame(websocket, clientState, {
+				drainWaitMs,
+				keepaliveIntervalMs,
+				maxBufferedAmountBytes,
+				maxQueuedBytes,
+				logger
+			}, keepaliveIntervalMs);
+		});
+	});
+	return {
+		async close() {
+			if (closed) return;
+			closed = true;
+			server.removeListener("upgrade", onUpgrade);
+			sessionSubscription.dispose();
+			themeSubscription.dispose();
+			await Promise.all(Array.from(clients.keys(), (client) => shutdownClient(client)));
+			await closeWebSocketServer(websocketServer);
+		},
+		async notifyExit(event) {
+			if (closed || exitNotified) return;
+			exitNotified = true;
+			const payload = createExitControlFrame({
+				code: event.signal ? null : event.exitCode,
+				signal: resolveSignalName(event.signal)
+			});
+			await Promise.all(Array.from(clients.entries(), ([client, state]) => {
+				if (!state.active) return Promise.resolve();
+				return sendExitAndClose(client, payload);
+			}));
+		}
+	};
+}
+function handleControlFrame(websocket, data, session, logger) {
+	const controlFrame = parseAittyControlFrame(toBuffer(data).toString("utf8"));
+	if (!controlFrame) {
+		logger.warn("Ignoring malformed WebSocket control frame");
+		return;
+	}
+	if (controlFrame.type === "ping") {
+		sendPongFrame(websocket);
+		return;
+	}
+	if (controlFrame.type === "pong") return;
+	if (controlFrame.type === "resize") {
+		handleResizeControlFrame(controlFrame, session, logger);
+		return;
+	}
+	logger.warn(`Ignoring unknown WebSocket control frame type: ${controlFrame.type}`);
+}
+async function sendPongFrame(websocket) {
+	if (websocket.readyState !== SOCKET_OPEN) return;
+	try {
+		await sendFrame(websocket, createPongControlFrame(), false);
+	} catch {}
+}
+function handleResizeControlFrame(frame, session, logger) {
+	if (!isAittyResizeControlFrame(frame)) {
+		logger.warn("Ignoring invalid WebSocket resize frame");
+		return;
+	}
+	if (frame.cols === session.cols && frame.rows === session.rows) return;
+	session.resize(frame.cols, frame.rows);
+}
+function rejectUpgrade(socket, statusCode, statusText) {
+	const body = statusText;
+	socket.write(`HTTP/1.1 ${statusCode} ${statusText}\r\nConnection: close\r
+Content-Type: text/plain; charset=utf-8\r
+Content-Length: ${Buffer.byteLength(body)}\r\n\r
+` + body);
+	socket.destroy();
+}
+function toBuffer(data) {
+	if (typeof data === "string") return Buffer.from(data, "utf8");
+	if (Buffer.isBuffer(data)) return data;
+	if (Array.isArray(data)) return Buffer.concat(data.map((chunk) => toBuffer(chunk)));
+	return Buffer.from(data);
+}
+function createHelloFrame(session, replay) {
+	return createHelloControlFrame({
+		cols: session.cols,
+		replay: replay ? replay.toString("base64") : "",
+		rows: session.rows
+	});
+}
+function enqueueClientFrame(websocket, clientState, payload, binary, options) {
+	if (!clientState.active || websocket.readyState !== SOCKET_OPEN) return;
+	const frame = createOutboundFrame(payload, binary, options.allowOversizedFrame ?? false);
+	if (frame.allowOversized) clientState.oversizedAllowanceBytes = Math.max(clientState.oversizedAllowanceBytes, frame.byteLength);
+	const queuedLimit = Math.max(options.maxQueuedBytes, clientState.oversizedAllowanceBytes);
+	if (clientState.pendingBytes + frame.byteLength > queuedLimit) {
+		if (frame.allowOversized && clientState.pendingBytes === 0 && clientState.pendingFrames.length === 0) {
+			clientState.pendingFrames.push(frame);
+			clientState.pendingBytes += frame.byteLength;
+			pumpClientQueue(websocket, clientState, options.maxBufferedAmountBytes, options.drainWaitMs, options.logger);
+			return;
+		}
+		closeClientForBackpressure(websocket, clientState, options.logger);
+		return;
+	}
+	clientState.pendingFrames.push(frame);
+	clientState.pendingBytes += frame.byteLength;
+	pumpClientQueue(websocket, clientState, options.maxBufferedAmountBytes, options.drainWaitMs, options.logger);
+}
+function createOutboundFrame(payload, binary, allowOversized) {
+	return {
+		allowOversized,
+		binary,
+		byteLength: typeof payload === "string" ? Buffer.byteLength(payload) : payload.length,
+		payload
+	};
+}
+async function pumpClientQueue(websocket, clientState, maxBufferedAmountBytes, drainWaitMs, logger) {
+	if (clientState.sending) return;
+	clientState.sending = true;
+	try {
+		while (clientState.active && websocket.readyState === SOCKET_OPEN) {
+			const frame = clientState.pendingFrames[0];
+			if (!frame) {
+				if (websocket.bufferedAmount <= maxBufferedAmountBytes) clientState.oversizedAllowanceBytes = 0;
+				return;
+			}
+			const bufferedLimit = Math.max(maxBufferedAmountBytes, clientState.oversizedAllowanceBytes);
+			if (websocket.bufferedAmount > bufferedLimit) {
+				await waitForSocketWritable(websocket, clientState.priorityModeUntil > Date.now() ? Math.max(drainWaitMs, DEFAULT_PRIORITY_OUTPUT_RECOVERY_MS) : drainWaitMs);
+				if (!clientState.active || websocket.readyState !== SOCKET_OPEN) return;
+				if (websocket.bufferedAmount > bufferedLimit) {
+					if (clientState.priorityModeUntil > Date.now()) continue;
+					await closeClientForBackpressure(websocket, clientState, logger);
+					return;
+				}
+			}
+			clientState.pendingFrames.shift();
+			clientState.pendingBytes -= frame.byteLength;
+			try {
+				await sendFrame(websocket, frame.payload, frame.binary);
+				if (clientState.pendingFrames.every((pendingFrame) => !pendingFrame.allowOversized) && websocket.bufferedAmount <= maxBufferedAmountBytes) clientState.oversizedAllowanceBytes = 0;
+			} catch {
+				return;
+			}
+		}
+	} finally {
+		clientState.sending = false;
+		if (clientState.active && websocket.readyState === SOCKET_OPEN && clientState.pendingFrames.length > 0) pumpClientQueue(websocket, clientState, maxBufferedAmountBytes, drainWaitMs, logger);
+	}
+}
+function resolveSignalName(signal) {
+	if (!signal) return null;
+	for (const [name, value] of Object.entries(os.constants.signals)) if (value === signal) return name;
+	return `SIG${signal}`;
+}
+async function sendExitAndClose(websocket, payload) {
+	if (websocket.readyState === SOCKET_CLOSED) return;
+	if (websocket.readyState === SOCKET_OPEN) {
+		try {
+			await sendFrame(websocket, payload, false);
+		} catch {}
+		websocket.close(NORMAL_CLOSE_CODE);
+	}
+	if (!await waitForSocketClose(websocket)) {
+		websocket.terminate();
+		await waitForSocketClose(websocket);
+	}
+}
+async function supersedeClient(websocket) {
+	if (websocket.readyState === SOCKET_CLOSED) return;
+	if (websocket.readyState === SOCKET_OPEN) websocket.close(SUPERSEDED_CLOSE_CODE, SUPERSEDED_CLOSE_REASON);
+	if (!await waitForSocketClose(websocket)) {
+		websocket.terminate();
+		await waitForSocketClose(websocket);
+	}
+}
+async function closeClientForBackpressure(websocket, clientState, logger) {
+	if (!clientState.active) return;
+	clientState.active = false;
+	clientState.pendingBytes = 0;
+	clientState.pendingFrames.length = 0;
+	logger.warn("Closing slow WebSocket client after backpressure limit exceeded");
+	if (websocket.readyState === SOCKET_CLOSED) return;
+	if (websocket.readyState === SOCKET_OPEN) websocket.close(BACKPRESSURE_CLOSE_CODE, BACKPRESSURE_CLOSE_REASON);
+	else {
+		websocket.terminate();
+		await waitForSocketClose(websocket);
+		return;
+	}
+	if (!await waitForSocketClose(websocket)) {
+		websocket.terminate();
+		await waitForSocketClose(websocket);
+	}
+}
+async function shutdownClient(websocket) {
+	if (websocket.readyState === SOCKET_CLOSED) return;
+	websocket.terminate();
+	await waitForSocketClose(websocket);
+}
+function closeWebSocketServer(websocketServer) {
+	return new Promise((resolve, reject) => {
+		websocketServer.close((error) => {
+			if (error) {
+				reject(error);
+				return;
+			}
+			resolve();
+		});
+	});
+}
+function sendFrame(websocket, payload, binary = Buffer.isBuffer(payload)) {
+	return new Promise((resolve, reject) => {
+		websocket.send(payload, { binary }, (error) => {
+			if (error) {
+				reject(error);
+				return;
+			}
+			resolve();
+		});
+	});
+}
+function waitForSocketWritable(websocket, timeout) {
+	const socket = getUnderlyingSocket(websocket);
+	if (!socket || websocket.bufferedAmount === 0 && !socket.writableNeedDrain) return Promise.resolve();
+	return new Promise((resolve) => {
+		const finish = once(() => {
+			clearTimeout(timer);
+			socket.removeListener("drain", finish);
+			websocket.removeListener("close", finish);
+			websocket.removeListener("error", finish);
+			resolve();
+		});
+		const timer = setTimeout(() => {
+			finish();
+		}, timeout);
+		socket.once("drain", finish);
+		websocket.once("close", finish);
+		websocket.once("error", finish);
+	});
+}
+function waitForSocketClose(websocket, timeout = 250) {
+	if (websocket.readyState === SOCKET_CLOSED) return Promise.resolve(true);
+	return new Promise((resolve) => {
+		const finish = once((closed) => {
+			clearTimeout(timer);
+			websocket.removeListener("close", onClose);
+			websocket.removeListener("error", onError);
+			resolve(closed);
+		});
+		const timer = setTimeout(() => {
+			finish(websocket.readyState === SOCKET_CLOSED);
+		}, timeout);
+		const onClose = () => {
+			finish(true);
+		};
+		const onError = () => {
+			finish(websocket.readyState === SOCKET_CLOSED);
+		};
+		websocket.once("close", onClose);
+		websocket.once("error", onError);
+	});
+}
+function getUnderlyingSocket(websocket) {
+	return websocket._socket ?? null;
+}
+function once(callback) {
+	let called = false;
+	return (...args) => {
+		if (called) return;
+		called = true;
+		callback(...args);
+	};
+}
+function clearKeepaliveTimer(clientState) {
+	if (clientState.keepaliveTimer === null) return;
+	clearTimeout(clientState.keepaliveTimer);
+	clientState.keepaliveTimer = null;
+}
+function scheduleKeepaliveFrame(websocket, clientState, options, delayMs) {
+	clearKeepaliveTimer(clientState);
+	if (options.keepaliveIntervalMs <= 0) return;
+	clientState.keepaliveTimer = setTimeout(() => {
+		clientState.keepaliveTimer = null;
+		if (!clientState.active || websocket.readyState !== SOCKET_OPEN) return;
+		enqueueClientFrame(websocket, clientState, createPingControlFrame(), false, {
+			drainWaitMs: options.drainWaitMs,
+			maxBufferedAmountBytes: options.maxBufferedAmountBytes,
+			maxQueuedBytes: options.maxQueuedBytes,
+			logger: options.logger
+		});
+		scheduleKeepaliveFrame(websocket, clientState, options, options.keepaliveIntervalMs);
+	}, Math.max(0, delayMs));
+}
+function getActiveClient(clients) {
+	for (const entry of clients.entries()) if (entry[1].active) return entry;
+	return null;
+}
+function containsPriorityControlByte(payload) {
+	return payload.includes(3) || payload.includes(12) || isStandaloneEscapeByte(payload);
+}
+function isStandaloneEscapeByte(payload) {
+	return payload.length === 1 && payload[0] === 27;
+}
+function prioritizeClientOutput(clientState) {
+	clientState.pendingBytes = 0;
+	clientState.pendingFrames.length = 0;
+	clientState.priorityModeUntil = Date.now() + DEFAULT_PRIORITY_OUTPUT_RECOVERY_MS;
+}
+//#endregion
+export { createWebSocketTransport };

package/dist/server.d.ts

@@ -0,0 +1,51 @@
+import { BrowserRuntimeKind, BrowserShellOptions } from "./frontend/browser-shell.js";
+import { PtyModule, PtySession } from "./runtime/pty-session.js";
+import { AittyThemeSource as AittyThemeSource$1 } from "./theme-source.js";
+import { AittyPortlessIdentity, AittyPortlessIdentity as AittyPortlessIdentity$1, AittyPortlessOptions, AittyPortlessOptions as AittyPortlessOptions$1, AittyTheme } from "@aitty/protocol";
+import * as _$ws from "ws";
+
+//#region src/server.d.ts
+interface AittyRuntimeDependencies {
+  WebSocketServer?: typeof _$ws.WebSocketServer;
+  nodePty: PtyModule;
+}
+interface AittyServerOptions {
+  args?: string[];
+  bufferSize?: number;
+  command: string;
+  cwd: string;
+  env?: NodeJS.ProcessEnv;
+  host?: string;
+  port?: number;
+  portless?: AittyPortlessOptions$1;
+  publicHost?: string;
+  publicOrigin?: string;
+  runtimeKind?: BrowserRuntimeKind;
+  shell?: BrowserShellOptions;
+  theme?: AittyTheme;
+  themeSource?: AittyThemeSource$1;
+  verbose?: boolean;
+}
+interface AittyRunningServer {
+  childPid: number;
+  close(signal?: NodeJS.Signals): Promise<number>;
+  closed: Promise<number>;
+  host: string;
+  port: number;
+  portless: AittyPortlessIdentity$1;
+  session: PtySession;
+  token: string;
+  url: string;
+}
+interface Writer {
+  write(chunk: string): boolean;
+}
+interface AittyServerIo {
+  registerSignalHandlers?: boolean;
+  stderr?: Writer;
+}
+declare function createAittyServer(options: AittyServerOptions, dependencies: AittyRuntimeDependencies, io?: AittyServerIo): Promise<AittyRunningServer>;
+declare function buildAittyUrl(origin: string, token: string, pathname: string): string;
+declare function buildAittyOrigin(host: string, port: number): string;
+//#endregion
+export { type AittyPortlessIdentity, type AittyPortlessOptions, AittyRunningServer, AittyRuntimeDependencies, AittyServerIo, AittyServerOptions, Writer, buildAittyOrigin, buildAittyUrl, createAittyServer };