src/lib/server/agent/PiRpcClient.spec.ts

import { createServer, type Server } from "node:http";
import { afterEach, describe, expect, it } from "vitest";
import { WebSocketServer } from "ws";
import { PI_RPC_REPLAY_GAP_MESSAGE, PiRpcClient, PiRpcReplayGapError } from "./PiRpcClient";

const servers: Server[] = [];

afterEach(async () => {
	await Promise.all(
		servers.splice(0).map(
			(server) =>
				new Promise<void>((resolve) => {
					server.close(() => resolve());
				})
		)
	);
});

describe("PiRpcClient", () => {
	it("retries a transient HTTP 400 upgrade within one connect deadline", async () => {
		const server = createServer();
		servers.push(server);
		const webSockets = new WebSocketServer({ noServer: true });
		let upgrades = 0;
		server.on("upgrade", (request, socket, head) => {
			upgrades += 1;
			if (upgrades === 1) {
				socket.end(
					"HTTP/1.1 400 Bad Request\r\nConnection: close\r\nContent-Length: 11\r\n\r\nBad Request"
				);
				return;
			}
			webSockets.handleUpgrade(request, socket, head, (webSocket) => {
				webSockets.emit("connection", webSocket, request);
			});
		});
		webSockets.on("connection", (socket) => {
			socket.send(JSON.stringify({ type: "bridge_ready", fromSeq: 1, toSeq: 0 }));
		});
		await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
		const address = server.address();
		if (!address || typeof address === "string") throw new Error("Test server has no port");

		const statuses: string[] = [];
		const client = new PiRpcClient(`ws://127.0.0.1:${address.port}/ws`, {
			headers: {},
			connectTimeoutMs: 2_000,
			minReconnectDelayMs: 10,
			maxReconnectDelayMs: 10,
			onStatus: (status) => statuses.push(status),
		});
		try {
			await Promise.all([client.connect(), client.connect()]);
			expect(client.connected).toBe(true);
			expect(upgrades).toBe(2);
			expect(statuses.slice(0, 3)).toEqual(["connecting", "reconnecting", "connected"]);
		} finally {
			client.close();
			webSockets.close();
		}
	});

	it("keeps reconnecting when an upgraded socket never becomes bridge-ready", async () => {
		const server = createServer();
		servers.push(server);
		const webSockets = new WebSocketServer({ noServer: true });
		let upgrades = 0;
		let firstSocket: import("ws").WebSocket | undefined;
		server.on("upgrade", (request, socket, head) => {
			upgrades += 1;
			webSockets.handleUpgrade(request, socket, head, (webSocket) => {
				webSockets.emit("connection", webSocket, request);
			});
		});
		webSockets.on("connection", (socket) => {
			if (upgrades === 1) firstSocket = socket;
			if (upgrades !== 2) {
				socket.send(JSON.stringify({ type: "bridge_ready", fromSeq: 1, toSeq: 0 }));
			}
		});
		await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
		const address = server.address();
		if (!address || typeof address === "string") throw new Error("Test server has no port");

		let connectedCount = 0;
		let resolveReconnected: (() => void) | undefined;
		const reconnected = new Promise<void>((resolve) => {
			resolveReconnected = resolve;
		});
		const client = new PiRpcClient(`ws://127.0.0.1:${address.port}/ws`, {
			headers: {},
			connectTimeoutMs: 100,
			minReconnectDelayMs: 5,
			maxReconnectDelayMs: 5,
			onStatus: (status) => {
				if (status === "connected" && ++connectedCount === 2) resolveReconnected?.();
			},
		});
		try {
			await client.connect();
			firstSocket?.terminate();
			await reconnected;
			expect(client.connected).toBe(true);
			expect(upgrades).toBe(3);
		} finally {
			client.close();
			webSockets.close();
		}
	});

	it("fails an attach gap before dispatching a retry prompt", async () => {
		const server = createServer();
		servers.push(server);
		const webSockets = new WebSocketServer({ server });
		const receivedTypes: string[] = [];
		webSockets.on("connection", (socket) => {
			socket.on("message", (raw) => {
				const message = JSON.parse(raw.toString()) as { type?: string };
				if (message.type) receivedTypes.push(message.type);
			});
			socket.send(JSON.stringify({ type: "bridge_ready", fromSeq: 5, toSeq: 7 }));
		});
		await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
		const address = server.address();
		if (!address || typeof address === "string") throw new Error("Test server has no port");

		const client = new PiRpcClient(`ws://127.0.0.1:${address.port}/ws`, {
			headers: {},
			lastSequence: 1,
			connectTimeoutMs: 1_000,
		});
		try {
			const failure = await client
				.command({ type: "prompt", message: "do not dispatch twice" })
				.catch((error: unknown) => error);
			expect(failure).toBeInstanceOf(PiRpcReplayGapError);
			expect(failure).toMatchObject({
				message: PI_RPC_REPLAY_GAP_MESSAGE,
				gap: { expected: 2, received: 5, earliestAvailable: 5, latestAvailable: 7 },
			});
			expect(client.lastSequence).toBe(1);
			await new Promise((resolve) => setImmediate(resolve));
			expect(receivedTypes).not.toContain("prompt");
		} finally {
			client.close();
			webSockets.close();
		}
	});

	it("does not advance its cursor across a hole inside the initial replay", async () => {
		const server = createServer();
		servers.push(server);
		const webSockets = new WebSocketServer({ server });
		const receivedTypes: string[] = [];
		webSockets.on("connection", (socket) => {
			socket.on("message", (raw) => {
				const message = JSON.parse(raw.toString()) as { type?: string };
				if (message.type) receivedTypes.push(message.type);
			});
			socket.send(JSON.stringify({ type: "bridge_ready", fromSeq: 1, toSeq: 3 }));
			socket.send(
				JSON.stringify({ type: "bridge_event", seq: 1, payload: { type: "agent_start" } })
			);
			socket.send(
				JSON.stringify({ type: "bridge_event", seq: 3, payload: { type: "agent_settled" } })
			);
		});
		await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
		const address = server.address();
		if (!address || typeof address === "string") throw new Error("Test server has no port");

		const client = new PiRpcClient(`ws://127.0.0.1:${address.port}/ws`, {
			headers: {},
			connectTimeoutMs: 1_000,
		});
		const sequences: number[] = [];
		client.subscribe((_event, sequence) => sequences.push(sequence));
		try {
			await expect(
				client.command({ type: "prompt", message: "wait for replay first" })
			).rejects.toMatchObject({
				name: "PiRpcReplayGapError",
				gap: { expected: 2, received: 3, latestAvailable: 3 },
			});
			expect(sequences).toEqual([1]);
			expect(client.lastSequence).toBe(1);
			await new Promise((resolve) => setImmediate(resolve));
			expect(receivedTypes).not.toContain("prompt");
		} finally {
			client.close();
			webSockets.close();
		}
	});
});