pi-gui/src/protocol/index.js

import path from "node:path";

export const PROTOCOL_VERSION = "v1";
export const MAX_FRAME_BYTES = 64 * 1024;

const requestOperations = new Map([
	["list_agents", { agent: false, payload: "none" }],
	["list_directories", { agent: false, payload: "none" }],
	["select_agent", { agent: false, payload: "worktree" }],
	["forget_directory", { agent: false, payload: "worktree" }],
	["get_state", { agent: true, payload: "none" }],
	["get_session_stats", { agent: true, payload: "none" }],
	["list_sessions", { agent: true, payload: "none" }],
	["new_session", { agent: true, payload: "none" }],
	["switch_session", { agent: true, payload: "session" }],
	["get_transcript", { agent: true, payload: "none" }],
	["subscribe", { agent: true, payload: "cursor" }],
	["prompt", { agent: true, payload: "message" }],
	["submit_prompt", { agent: true, payload: "message" }],
	["steer", { agent: true, payload: "message" }],
	["follow_up", { agent: true, payload: "message" }],
	["abort", { agent: true, payload: "none" }],
	["extension_response", { agent: true, payload: "extensionResponse" }],
	["retry", { agent: true, payload: "none" }],
	["restart", { agent: true, payload: "none" }],
	["get_available_models", { agent: true, payload: "none" }],
	["get_commands", { agent: true, payload: "none" }],
	["set_model", { agent: true, payload: "model" }],
	["set_thinking_level", { agent: true, payload: "thinking" }],
]);

const eventTypes = new Set([
	"agent_state",
	"transcript",
	"stream",
	"tool",
	"queue",
	"recovery",
	"extension_ui_request",
]);

const thinkingLevels = new Set([
	"off",
	"minimal",
	"low",
	"medium",
	"high",
	"xhigh",
	"max",
]);
const idPattern = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/;

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

function isRecord(value) {
	return value !== null && typeof value === "object" && !Array.isArray(value);
}

function assertRecord(value, field) {
	if (!isRecord(value))
		throw new ProtocolError("invalid_message", `${field} must be an object`);
	return value;
}

function assertAllowedKeys(value, allowed, field) {
	for (const key of Object.keys(value)) {
		if (!allowed.has(key))
			throw new ProtocolError(
				"invalid_message",
				`${field}.${key} is not supported`,
			);
	}
}

function assertString(value, field, { maxLength = 4096, pattern } = {}) {
	if (
		typeof value !== "string" ||
		value.length === 0 ||
		value.length > maxLength
	) {
		throw new ProtocolError(
			"invalid_message",
			`${field} must be a non-empty string up to ${maxLength} characters`,
		);
	}
	if (pattern && !pattern.test(value))
		throw new ProtocolError(
			"invalid_message",
			`${field} has an invalid format`,
		);
	return value;
}

function assertOptionalCursor(value) {
	if (value === undefined) return undefined;
	if (Number.isSafeInteger(value) && value >= 0) return value;
	return assertString(value, "payload.cursor", {
		maxLength: 256,
		pattern: idPattern,
	});
}

function validatePayload(kind, value) {
	if (kind === "none") {
		if (value !== undefined)
			throw new ProtocolError(
				"invalid_message",
				"payload is not allowed for this operation",
			);
		return undefined;
	}

	const payload = assertRecord(value, "payload");
	switch (kind) {
		case "worktree": {
			assertAllowedKeys(payload, new Set(["worktreePath"]), "payload");
			const worktreePath = assertString(
				payload.worktreePath,
				"payload.worktreePath",
				{ maxLength: 4096 },
			);
			if (!path.isAbsolute(worktreePath))
				throw new ProtocolError(
					"invalid_message",
					"payload.worktreePath must be absolute",
				);
			return { worktreePath };
		}
		case "session": {
			assertAllowedKeys(payload, new Set(["sessionPath"]), "payload");
			const sessionPath = assertString(
				payload.sessionPath,
				"payload.sessionPath",
				{
					maxLength: 4096,
				},
			);
			if (!path.isAbsolute(sessionPath))
				throw new ProtocolError(
					"invalid_message",
					"payload.sessionPath must be absolute",
				);
			return { sessionPath };
		}
		case "cursor": {
			assertAllowedKeys(payload, new Set(["cursor"]), "payload");
			return {
				...(payload.cursor === undefined
					? {}
					: { cursor: assertOptionalCursor(payload.cursor) }),
			};
		}
		case "message": {
			assertAllowedKeys(payload, new Set(["message"]), "payload");
			return {
				message: assertString(payload.message, "payload.message", {
					maxLength: 32 * 1024,
				}),
			};
		}
		case "extensionResponse": {
			assertAllowedKeys(payload, new Set(["requestId", "response"]), "payload");
			const response = assertRecord(payload.response, "payload.response");
			if (
				!Object.hasOwn(response, "value") &&
				!Object.hasOwn(response, "confirmed") &&
				!Object.hasOwn(response, "cancelled")
			) {
				throw new ProtocolError(
					"invalid_message",
					"payload.response must contain value, confirmed, or cancelled",
				);
			}
			return {
				requestId: assertString(payload.requestId, "payload.requestId", {
					maxLength: 128,
					pattern: idPattern,
				}),
				response,
			};
		}
		case "model": {
			assertAllowedKeys(payload, new Set(["provider", "modelId"]), "payload");
			return {
				provider: assertString(payload.provider, "payload.provider", {
					maxLength: 128,
				}),
				modelId: assertString(payload.modelId, "payload.modelId", {
					maxLength: 512,
				}),
			};
		}
		case "thinking": {
			assertAllowedKeys(payload, new Set(["level"]), "payload");
			const level = assertString(payload.level, "payload.level", {
				maxLength: 16,
			});
			if (!thinkingLevels.has(level))
				throw new ProtocolError(
					"invalid_message",
					"payload.level is not supported",
				);
			return { level };
		}
		default:
			throw new ProtocolError("invalid_message", "unsupported payload shape");
	}
}

export function validateRequest(value) {
	const request = assertRecord(value, "request");
	assertAllowedKeys(
		request,
		new Set(["version", "id", "op", "agentId", "payload"]),
		"request",
	);
	if (request.version !== PROTOCOL_VERSION) {
		throw new ProtocolError(
			"unsupported_version",
			`version must be ${PROTOCOL_VERSION}`,
		);
	}

	const id = assertString(request.id, "request.id", {
		maxLength: 128,
		pattern: idPattern,
	});
	const op = assertString(request.op, "request.op", { maxLength: 64 });
	const operation = requestOperations.get(op);
	if (!operation)
		throw new ProtocolError(
			"unsupported_operation",
			`operation ${op} is not supported`,
		);

	let agentId;
	if (operation.agent) {
		agentId = assertString(request.agentId, "request.agentId", {
			maxLength: 128,
			pattern: idPattern,
		});
	} else if (request.agentId !== undefined) {
		throw new ProtocolError(
			"invalid_message",
			"request.agentId is not allowed for this operation",
		);
	}

	const payload = validatePayload(operation.payload, request.payload);
	return {
		version: PROTOCOL_VERSION,
		id,
		op,
		...(agentId ? { agentId } : {}),
		...(payload === undefined ? {} : { payload }),
	};
}

export function parseRequestFrame(frame, maxFrameBytes = MAX_FRAME_BYTES) {
	const bytes = Buffer.isBuffer(frame)
		? frame.length
		: Buffer.byteLength(frame, "utf8");
	if (bytes > maxFrameBytes)
		throw new ProtocolError(
			"frame_too_large",
			`frame exceeds ${maxFrameBytes} bytes`,
		);
	const line = String(frame).endsWith("\r")
		? String(frame).slice(0, -1)
		: String(frame);
	if (line.length === 0)
		throw new ProtocolError("invalid_message", "frame cannot be empty");

	let parsed;
	try {
		parsed = JSON.parse(line);
	} catch {
		throw new ProtocolError("invalid_json", "frame must contain valid JSON");
	}
	return validateRequest(parsed);
}

export function validateEvent(value) {
	const event = assertRecord(value, "event");
	assertAllowedKeys(
		event,
		new Set(["version", "seq", "type", "agentId", "data"]),
		"event",
	);
	if (event.version !== PROTOCOL_VERSION)
		throw new ProtocolError(
			"unsupported_version",
			`version must be ${PROTOCOL_VERSION}`,
		);
	if (!Number.isSafeInteger(event.seq) || event.seq < 1)
		throw new ProtocolError(
			"invalid_message",
			"event.seq must be a positive integer",
		);
	const type = assertString(event.type, "event.type", { maxLength: 64 });
	if (!eventTypes.has(type))
		throw new ProtocolError(
			"unsupported_event",
			`event ${type} is not supported`,
		);
	return {
		version: PROTOCOL_VERSION,
		seq: event.seq,
		type,
		agentId: assertString(event.agentId, "event.agentId", {
			maxLength: 128,
			pattern: idPattern,
		}),
		data: assertRecord(event.data, "event.data"),
	};
}

export function successResponse(id, result = {}) {
	return { version: PROTOCOL_VERSION, id, ok: true, result };
}

export function errorResponse(id, error) {
	const code = error instanceof ProtocolError ? error.code : "internal_error";
	const message =
		error instanceof Error ? error.message : "internal bridge error";
	return {
		version: PROTOCOL_VERSION,
		...(id ? { id } : {}),
		ok: false,
		error: { code, message },
	};
}

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