pi-gui/src/client/local-client.js

import { randomUUID } from "node:crypto";
import { createConnection } from "node:net";
import path from "node:path";
import { StringDecoder } from "node:string_decoder";
import {
	MAX_FRAME_BYTES,
	ProtocolError,
	validateRequest,
} from "../protocol/index.js";

export class LocalClientError extends Error {
	constructor(code, message) {
		super(message);
		this.name = "LocalClientError";
		this.code = code;
	}
}

function frame(value) {
	return `${JSON.stringify(value)}\n`;
}

export async function connectLocalClient({
	socketPath,
	maxFrameBytes = MAX_FRAME_BYTES,
}) {
	if (typeof socketPath !== "string" || !path.isAbsolute(socketPath)) {
		throw new TypeError("socketPath must be an absolute Unix-socket path");
	}

	const socket = await new Promise((resolve, reject) => {
		const connection = createConnection(socketPath);
		connection.once("connect", () => resolve(connection));
		connection.once("error", reject);
	});
	const decoder = new StringDecoder("utf8");
	const pending = new Map();
	const subscriptions = new Map();
	let buffer = "";
	let closed = false;

	const rejectPending = (error) => {
		for (const { reject } of pending.values()) reject(error);
		pending.clear();
	};
	const dispatchEvent = (event) => {
		const subscription = subscriptions.get(event.agentId);
		if (!subscription) return;
		if (subscription.ready) subscription.listener(event);
		else subscription.queued.push(event);
	};
	const dispatchLine = (line) => {
		let message;
		try {
			message = JSON.parse(line.endsWith("\r") ? line.slice(0, -1) : line);
		} catch {
			rejectPending(
				new LocalClientError("invalid_json", "bridge emitted invalid JSON"),
			);
			return;
		}
		if (message?.type === "event" && message.event) {
			dispatchEvent(message.event);
			return;
		}
		if (typeof message?.id !== "string") return;
		const request = pending.get(message.id);
		if (!request) return;
		pending.delete(message.id);
		if (message.ok) request.resolve(message.result);
		else
			request.reject(
				new LocalClientError(
					message.error?.code ?? "bridge_error",
					message.error?.message ?? "bridge request failed",
				),
			);
	};

	socket.on("data", (chunk) => {
		if (closed) return;
		buffer += decoder.write(chunk);
		let index;
		while ((index = buffer.indexOf("\n")) !== -1) {
			const line = buffer.slice(0, index);
			buffer = buffer.slice(index + 1);
			dispatchLine(line);
		}
		if (Buffer.byteLength(buffer, "utf8") > maxFrameBytes) {
			rejectPending(
				new LocalClientError(
					"frame_too_large",
					`bridge frame exceeds ${maxFrameBytes} bytes`,
				),
			);
			socket.destroy();
		}
	});
	socket.on("error", (error) =>
		rejectPending(new LocalClientError("socket_error", error.message)),
	);
	socket.on("close", () => {
		closed = true;
		rejectPending(
			new LocalClientError("socket_closed", "bridge socket closed"),
		);
	});

	const request = (op, { agentId, payload } = {}) => {
		if (closed || socket.destroyed)
			return Promise.reject(
				new LocalClientError("socket_closed", "bridge socket is closed"),
			);
		const id = `client-${randomUUID()}`;
		let requestFrame;
		try {
			requestFrame = validateRequest({
				version: "v1",
				id,
				op,
				...(agentId ? { agentId } : {}),
				...(payload === undefined ? {} : { payload }),
			});
		} catch (error) {
			return Promise.reject(
				error instanceof ProtocolError
					? error
					: new LocalClientError("invalid_request", "invalid bridge request"),
			);
		}
		return new Promise((resolve, reject) => {
			pending.set(id, { resolve, reject });
			socket.write(frame(requestFrame));
		});
	};

	return {
		request,
		async subscribe(agentId, cursor, listener) {
			if (typeof listener !== "function")
				throw new TypeError("listener must be a function");
			const subscription = { listener, queued: [], ready: false };
			subscriptions.set(agentId, subscription);
			try {
				const result = await request("subscribe", {
					agentId,
					payload: cursor === undefined ? {} : { cursor },
				});
				for (const event of result.events) listener(event);
				subscription.ready = true;
				for (const event of subscription.queued) listener(event);
				subscription.queued = [];
			} catch (error) {
				subscriptions.delete(agentId);
				throw error;
			}
			return () => subscriptions.delete(agentId);
		},
		async close() {
			if (closed) return;
			closed = true;
			socket.end();
			socket.destroy();
		},
	};
}