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((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((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((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((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((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((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(); } }); });