pi-mono/packages/protocol/test/protocol.test.ts

import { describe, expect, test } from "vitest";
import {
	type ClientHello,
	type ClientMessage,
	ClientMessageDecoder,
	decodeCbor,
	encodeCbor,
	encodeClientMessage,
	encodeFrame,
	encodeServerMessage,
	FrameDecoder,
	isSupportedProtocolVersion,
	PROTOCOL_VERSION,
	ProtocolValidationError,
	parseClientMessage,
	parseServerMessage,
	type ServerHello,
	type ServerMessage,
	ServerMessageDecoder,
	type ServerSnapshot,
} from "../src/index.ts";

const emptyServerSnapshot: ServerSnapshot = {
	serverId: "server-1",
	protocolVersion: PROTOCOL_VERSION,
	revision: 0,
	sessions: [],
	models: [],
};

const clientHello: ClientHello = {
	type: "hello",
	version: PROTOCOL_VERSION,
	token: "secret",
};

const serverHello: ServerHello = {
	type: "hello",
	version: PROTOCOL_VERSION,
	connectionId: "connection-1",
	snapshot: emptyServerSnapshot,
};

describe("protocol validation", () => {
	test("uses protocol version 2", () => {
		expect(PROTOCOL_VERSION).toBe(2);
		expect(isSupportedProtocolVersion(2)).toBe(true);
		expect(isSupportedProtocolVersion(1)).toBe(false);
		expect(isSupportedProtocolVersion(2.5)).toBe(false);
	});

	test("accepts an integer version and bearer token in the client hello", () => {
		expect(parseClientMessage(clientHello)).toEqual(clientHello);
	});

	test.each([
		{ type: "hello", version: "2", token: "secret" },
		{ type: "hello", version: 2.5, token: "secret" },
		{ type: "hello", version: 2, token: "" },
		{ type: "hello", version: 2, token: "secret", extra: true },
	])("rejects an invalid strict handshake: $version", (message) => {
		expect(() => parseClientMessage(message)).toThrow(ProtocolValidationError);
	});

	test("does not parse JSON strings as wire messages", () => {
		expect(() => parseClientMessage(JSON.stringify(clientHello))).toThrow(ProtocolValidationError);
		expect(() => parseServerMessage(JSON.stringify(serverHello))).toThrow(ProtocolValidationError);
	});

	test("rejects image input while the MVP remains text-only", () => {
		expect(() =>
			parseClientMessage({
				type: "request",
				id: "request-1",
				request: {
					command: "prompt",
					sessionId: "session-1",
					text: "inspect",
					images: [{ type: "image", data: "abc", mimeType: "image/png" }],
				},
			}),
		).toThrow(ProtocolValidationError);
	});

	test("parses a server handshake snapshot", () => {
		expect(parseServerMessage(serverHello)).toEqual(serverHello);
	});

	test.each([
		{ type: "hello", version: 1, connectionId: "connection-1", snapshot: emptyServerSnapshot },
		{ type: "response", id: "request-1", ok: true, result: { command: "unknown" } },
		{ type: "event", event: { type: "session_removed", sessionId: 42 } },
	])("rejects invalid server messages", (wire) => {
		expect(() => parseServerMessage(wire)).toThrow(ProtocolValidationError);
	});

	test("validates nested JSON tool details", () => {
		const message = {
			type: "event",
			event: {
				type: "session_progress",
				sessionId: "session-1",
				progress: {
					type: "item_finished",
					item: {
						id: "tool-1",
						role: "tool",
						toolCallId: "call-1",
						toolName: "read",
						input: { path: "/tmp/file" },
						content: [{ type: "text", text: "done" }],
						details: { lines: [1, 2, 3], cached: false },
						status: "complete",
						isError: false,
						timestamp: 1,
					},
				},
			},
		};
		expect(parseServerMessage(message)).toEqual(message);
	});

	test("validation errors do not retain rejected payloads", () => {
		let thrown: unknown;
		try {
			parseClientMessage({ type: "hello", version: "2", token: "x".repeat(2_000_000) });
		} catch (error) {
			thrown = error;
		}
		expect(thrown).toBeInstanceOf(ProtocolValidationError);
		expect(Object.hasOwn(thrown as object, "value")).toBe(false);
		expect((thrown as Error).message.length).toBeLessThan(1_000);
	});
});

describe("validated framed protocol APIs", () => {
	test("encodes complete client and server frames", () => {
		const clientFrames = new FrameDecoder().push(encodeClientMessage(clientHello));
		expect(clientFrames).toHaveLength(1);
		expect(parseClientMessage(decodeCbor(clientFrames[0]!))).toEqual(clientHello);

		const serverFrames = new FrameDecoder().push(encodeServerMessage(serverHello));
		expect(serverFrames).toHaveLength(1);
		expect(parseServerMessage(decodeCbor(serverFrames[0]!))).toEqual(serverHello);
	});

	test("enforces an outbound frame limit before returning encoded bytes", () => {
		expect(() => encodeClientMessage(clientHello, { maxFrameLength: 8 })).toThrow(ProtocolValidationError);
		expect(() => encodeServerMessage(serverHello, { maxFrameLength: 8 })).toThrow(ProtocolValidationError);
	});

	test("validates messages before encoding", () => {
		expect(() => encodeClientMessage({ type: "hello", version: 2, token: "" })).toThrow(ProtocolValidationError);
	});

	test("omits explicit undefined optional properties on the wire", () => {
		const message: ClientMessage = {
			type: "request",
			id: "request-1",
			request: { command: "create", cwd: undefined, name: undefined },
		};
		const [payload] = new FrameDecoder().push(encodeClientMessage(message));
		expect(decodeCbor(payload!)).toEqual({
			type: "request",
			id: "request-1",
			request: { command: "create" },
		});
	});

	test("incrementally decodes fragmented and coalesced client messages", () => {
		const request: ClientMessage = {
			type: "request",
			id: "request-1",
			request: { command: "list" },
		};
		const first = encodeClientMessage(clientHello);
		const second = encodeClientMessage(request);
		const wire = new Uint8Array(first.byteLength + second.byteLength);
		wire.set(first);
		wire.set(second, first.byteLength);

		for (let split = 0; split <= wire.byteLength; split++) {
			const decoder = new ClientMessageDecoder();
			const messages = [...decoder.push(wire.subarray(0, split)), ...decoder.push(wire.subarray(split))];
			decoder.end();
			expect(messages).toEqual([clientHello, request]);
		}
	});

	test("incrementally decodes server messages", () => {
		const errorMessage: ServerMessage = {
			type: "hello_error",
			error: { code: "auth", message: "Invalid token" },
		};
		const decoder = new ServerMessageDecoder();
		expect(decoder.push(encodeServerMessage(errorMessage))).toEqual([errorMessage]);
		decoder.end();
	});

	test.each([
		["empty CBOR payload", encodeFrame(new Uint8Array())],
		["malformed CBOR", encodeFrame(new Uint8Array([0xff]))],
		["schema-invalid CBOR", encodeFrame(encodeCbor({ type: "hello", version: 2, token: "", extra: true }))],
	] as const)("rejects invalid framed client input: %s", (_label, wire) => {
		const decoder = new ClientMessageDecoder();
		expect(() => decoder.push(wire)).toThrow(ProtocolValidationError);
		expect(() => decoder.push(encodeClientMessage(clientHello))).toThrow(/failed/i);
	});

	test("rejects CBOR byte strings nested in JSON-valued fields", () => {
		const wire = encodeFrame(
			encodeCbor({
				type: "response",
				id: "request-1",
				ok: false,
				error: {
					code: "invalid_request",
					message: "invalid",
					details: { nested: new Uint8Array([1, 2, 3]) },
				},
			}),
		);
		expect(() => new ServerMessageDecoder().push(wire)).toThrow(ProtocolValidationError);
	});

	test("rejects truncated and oversized framing through the validated decoder", () => {
		const truncated = new ServerMessageDecoder();
		expect(truncated.push(new Uint8Array([0, 0, 0, 2, 1]))).toEqual([]);
		expect(() => truncated.end()).toThrow(ProtocolValidationError);

		const oversized = new ClientMessageDecoder({ maxFrameLength: 3 });
		expect(() => oversized.push(new Uint8Array([0, 0, 0, 4]))).toThrow(ProtocolValidationError);
	});
});