bridge.mjs

import http from "node:http";
import { spawn } from "node:child_process";
import process from "node:process";
import { WebSocketServer, WebSocket } from "ws";
import { buildPiArgs } from "./pi-args.mjs";
import { BridgeTurnFence } from "./bridge-turn-fence.mjs";

const PORT = parsePositiveInt(process.env.BRIDGE_PORT, 4777);
const HOST = process.env.BRIDGE_HOST || "127.0.0.1";
const WORKSPACE = process.env.WORKSPACE_DIR || "/workspace";
const PI_BIN = process.env.PI_BIN || "pi";
const MAX_EVENTS = parsePositiveInt(process.env.BRIDGE_MAX_EVENTS, 5_000);
const MAX_BUFFER_BYTES = parsePositiveInt(process.env.BRIDGE_MAX_BUFFER_BYTES, 8 * 1024 * 1024);
const MAX_LINE_BYTES = parsePositiveInt(process.env.BRIDGE_MAX_LINE_BYTES, 1024 * 1024);
const MAX_COMMANDS = parsePositiveInt(process.env.BRIDGE_MAX_COMMANDS, 20_000);
const WS_HIGH_WATER_BYTES = parsePositiveInt(
	process.env.BRIDGE_WS_HIGH_WATER_BYTES,
	4 * 1024 * 1024
);

const piPrefixArgs = parseStringArray(process.env.PI_BIN_PREFIX_ARGS_JSON);
const piArgs = buildPiArgs(piPrefixArgs, process.env);

let nextSequence = 1;
let bufferedBytes = 0;
let stdoutBuffer = "";
const replayBuffer = [];
const sockets = new Set();
// The bridge outlives chat-ui pods. Stable command ids turn a server retry into
// an idempotent acknowledgement instead of a second pi prompt.
const commands = new Map();
const turnFence = new BridgeTurnFence();

const pi = spawn(PI_BIN, piArgs, {
	cwd: WORKSPACE,
	env: {
		...process.env,
		PI_TELEMETRY: process.env.PI_TELEMETRY ?? "0",
	},
	stdio: ["pipe", "pipe", "pipe"],
});

const server = http.createServer((request, response) => {
	if (request.url === "/health") {
		response.writeHead(pi.exitCode === null ? 200 : 503, { "Content-Type": "application/json" });
		response.end(
			JSON.stringify({
				status: pi.exitCode === null ? "ready" : "stopped",
				pid: pi.pid,
				sequence: nextSequence - 1,
			})
		);
		return;
	}
	response.writeHead(404).end();
});

const webSockets = new WebSocketServer({ server, path: "/ws", maxPayload: MAX_LINE_BYTES });

webSockets.on("connection", (socket, request) => {
	sockets.add(socket);
	const requestUrl = new URL(request.url ?? "/ws", "http://bridge.local");
	const fromSequence = parseNonNegativeInt(requestUrl.searchParams.get("fromSeq"), 0);

	send(socket, {
		type: "bridge_ready",
		pid: pi.pid,
		fromSeq: replayBuffer.at(0)?.seq ?? nextSequence,
		toSeq: nextSequence - 1,
		piExited: pi.exitCode !== null,
	});
	for (const record of replayBuffer) {
		if (record.seq > fromSequence) sendRecord(socket, record);
	}

	socket.on("message", (raw, isBinary) => {
		if (isBinary) {
			send(socket, { type: "bridge_error", message: "Binary RPC commands are not supported" });
			return;
		}
		const text = raw.toString("utf8");
		if (Buffer.byteLength(text) > MAX_LINE_BYTES) {
			send(socket, { type: "bridge_error", message: "RPC command exceeds bridge limit" });
			return;
		}
		let command;
		try {
			command = JSON.parse(text);
		} catch {
			send(socket, { type: "bridge_error", message: "RPC command must be valid JSON" });
			return;
		}
		if (!command || typeof command !== "object" || Array.isArray(command)) {
			send(socket, { type: "bridge_error", message: "RPC command must be an object" });
			return;
		}
		if (command.type === "bridge_ping") {
			send(socket, { type: "bridge_pong", sequence: nextSequence - 1 });
			return;
		}
		if (command.type === "bridge_state") {
			send(socket, {
				type: "bridge_state",
				pid: pi.pid,
				fromSeq: replayBuffer.at(0)?.seq ?? nextSequence,
				toSeq: nextSequence - 1,
				piExited: pi.exitCode !== null,
			});
			return;
		}
		if (typeof command.id === "string" && command.id.length > 0) {
			const fingerprint = JSON.stringify(command);
			const previous = commands.get(command.id);
			if (previous) {
				if (previous.fingerprint !== fingerprint) {
					send(socket, {
						type: "bridge_error",
						message: `RPC command id '${command.id}' was reused with a different payload`,
					});
					return;
				}
				// An in-flight retry simply waits for the original broadcast response.
				// A completed retry gets a fresh replay sequence so a client whose
				// cursor already covered the original response can correlate it.
				if (previous.response) record(previous.response);
				return;
			}
		}
		const prepared = turnFence.prepare(command);
		if (prepared.error) {
			rejectCommand(socket, command, prepared.error);
			return;
		}
		if (!pi.stdin.writable || pi.exitCode !== null) {
			send(socket, { type: "bridge_error", message: "pi is not running" });
			return;
		}
		const commandEntry =
			typeof command.id === "string" && command.id.length > 0
				? { fingerprint: JSON.stringify(command), type: command.type }
				: undefined;
		if (commandEntry) {
			commands.set(command.id, commandEntry);
			pruneCommands();
		}
		if (prepared.promptId) turnFence.activate(prepared.promptId);
		try {
			pi.stdin.write(`${JSON.stringify(prepared.command)}\n`, (error) => {
				if (!error && command.type === "bridge_control" && commandEntry) {
					if (!commandEntry.response) {
						const response = {
							id: command.id,
							type: "response",
							command: command.type,
							success: true,
						};
						commandEntry.response = response;
						record(response);
					}
					return;
				}
				if (!error || !commandEntry || commandEntry.response) return;
				if (commands.get(command.id) === commandEntry) commands.delete(command.id);
				if (prepared.promptId) turnFence.failPrompt(prepared.promptId);
				send(socket, {
					type: "bridge_error",
					message: `Unable to send RPC command: ${error.message}`,
				});
			});
		} catch (error) {
			if (commandEntry && commands.get(command.id) === commandEntry) commands.delete(command.id);
			if (prepared.promptId) turnFence.failPrompt(prepared.promptId);
			send(socket, {
				type: "bridge_error",
				message: `Unable to send RPC command: ${error instanceof Error ? error.message : String(error)}`,
			});
		}
	});

	socket.on("close", () => {
		sockets.delete(socket);
	});
});

pi.stdout.on("data", (chunk) => {
	stdoutBuffer += chunk.toString("utf8");
	let newline = stdoutBuffer.indexOf("\n");
	while (newline !== -1) {
		const line = stdoutBuffer.slice(0, newline).replace(/\r$/, "");
		stdoutBuffer = stdoutBuffer.slice(newline + 1);
		handlePiLine(line);
		newline = stdoutBuffer.indexOf("\n");
	}
	if (Buffer.byteLength(stdoutBuffer) > MAX_LINE_BYTES) {
		record({ type: "bridge_error", message: "pi emitted an oversized JSONL record" });
		stdoutBuffer = "";
	}
});

pi.stderr.on("data", (chunk) => {
	const data = chunk.toString("utf8");
	record({
		type: "bridge_stderr",
		data: truncateUtf8(data, Math.min(MAX_LINE_BYTES, 64 * 1024)),
	});
});

pi.on("error", (error) => {
	turnFence.settle();
	record({ type: "bridge_error", message: `Unable to start pi: ${error.message}` });
});

pi.on("exit", (code, signal) => {
	turnFence.settle();
	record({ type: "bridge_exit", exitCode: code, signal });
});

server.listen(PORT, HOST);

function handlePiLine(line) {
	if (!line) return;
	if (Buffer.byteLength(line) > MAX_LINE_BYTES) {
		record({ type: "bridge_error", message: "pi emitted an oversized JSONL record" });
		return;
	}
	try {
		const payload = JSON.parse(line);
		if (payload?.type === "agent_settled" || payload?.type === "bridge_exit") {
			turnFence.settle();
		}
		if (payload?.type === "response" && typeof payload.id === "string") {
			const command = commands.get(payload.id);
			if (command) command.response = payload;
			if (command?.type === "prompt" && payload.success === false) {
				turnFence.failPrompt(payload.id);
			}
		}
		record(payload);
	} catch {
		record({
			type: "bridge_error",
			message: "pi emitted invalid JSONL",
			preview: truncateUtf8(line, 2_000),
		});
	}
}

function rejectCommand(socket, command, message) {
	if (typeof command.id === "string" && command.id) {
		record({
			id: command.id,
			type: "response",
			command: command.type,
			success: false,
			error: message,
		});
		return;
	}
	send(socket, { type: "bridge_error", message });
}

function pruneCommands() {
	if (commands.size <= MAX_COMMANDS) return;
	for (const [id, command] of commands) {
		// Prompt ids are the cross-server exactly-once boundary. Keep them for
		// the lifetime of this pi process even if a very chatty extension causes
		// ordinary diagnostic/control command entries to roll over.
		if (!command.response || command.type === "prompt") continue;
		commands.delete(id);
		if (commands.size <= MAX_COMMANDS) return;
	}
}

function record(payload) {
	const record = { seq: nextSequence++, payload };
	const bytes = Buffer.byteLength(JSON.stringify(record));
	replayBuffer.push({ ...record, bytes });
	bufferedBytes += bytes;
	while (replayBuffer.length > MAX_EVENTS || bufferedBytes > MAX_BUFFER_BYTES) {
		const removed = replayBuffer.shift();
		bufferedBytes -= removed?.bytes ?? 0;
	}
	for (const socket of sockets) {
		if (socket.readyState === WebSocket.OPEN) sendRecord(socket, record);
	}
}

function sendRecord(socket, record) {
	if (socket.bufferedAmount > WS_HIGH_WATER_BYTES) {
		// The replay buffer remains authoritative. Closing forces a reconnect with
		// fromSeq instead of letting unbounded WebSocket backpressure consume RAM.
		socket.close(1013, "Controller is too slow; reconnect to replay");
		return;
	}
	send(socket, { type: "bridge_event", seq: record.seq, payload: record.payload });
}

function send(socket, value) {
	if (socket.readyState === WebSocket.OPEN) socket.send(JSON.stringify(value));
}

function truncateUtf8(value, maxBytes) {
	const bytes = Buffer.from(value);
	if (bytes.length <= maxBytes) return value;
	return `${bytes.subarray(0, maxBytes).toString("utf8")}\n… output truncated by bridge`;
}

function parsePositiveInt(value, fallback) {
	const parsed = Number.parseInt(value ?? "", 10);
	return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback;
}

function parseNonNegativeInt(value, fallback) {
	const parsed = Number.parseInt(value ?? "", 10);
	return Number.isFinite(parsed) && parsed >= 0 ? parsed : fallback;
}

function parseStringArray(value) {
	if (!value) return [];
	try {
		const parsed = JSON.parse(value);
		return Array.isArray(parsed) && parsed.every((entry) => typeof entry === "string")
			? parsed
			: [];
	} catch {
		return [];
	}
}

function shutdown() {
	webSockets.close();
	server.close();
	if (pi.exitCode === null) pi.kill("SIGTERM");
	setTimeout(() => {
		if (pi.exitCode === null) pi.kill("SIGKILL");
	}, 5_000).unref();
}

process.once("SIGTERM", shutdown);
process.once("SIGINT", shutdown);