diff --git a/src/server/sessiond/sessionProxyRoutes.test.ts b/src/server/sessiond/sessionProxyRoutes.test.ts index 7911741..32e0974 100644 --- a/src/server/sessiond/sessionProxyRoutes.test.ts +++ b/src/server/sessiond/sessionProxyRoutes.test.ts @@ -1,5 +1,6 @@ import Fastify, { type FastifyInstance } from "fastify"; import fastifyWebsocket from "@fastify/websocket"; +import { WebSocket, WebSocketServer } from "ws"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { registerSessionProxyRoutes } from "./sessionProxyRoutes"; @@ -9,12 +10,13 @@ let daemon: FakeSessionDaemon; beforeEach(async () => { app = Fastify({ logger: false }); await app.register(fastifyWebsocket); - daemon = new FakeSessionDaemon(); + daemon = await FakeSessionDaemon.create(); registerSessionProxyRoutes(app, daemon, "/api/machines/local"); }); afterEach(async () => { await app.close(); + await daemon.close(); }); describe("machine-scoped session proxy routes", () => { @@ -33,17 +35,86 @@ describe("machine-scoped session proxy routes", () => { expect(response.json()).toEqual({ ok: true }); expect(daemon.requests).toEqual([{ method: "POST", path: "/auth/api-key", body: { providerId: "p", key: "k" } }]); }); + + it("preserves cwd query context when forwarding session event websockets", async () => { + await app.listen({ host: "127.0.0.1", port: 0 }); + const socket = new WebSocket(`${serverUrl(app)}/api/machines/local/sessions/session-1/events?cwd=${encodeURIComponent("/repo")}`); + + try { + await waitForOpen(socket); + expect(daemon.websocketPaths).toEqual(["/sessions/session-1/events?cwd=%2Frepo"]); + } finally { + socket.close(); + } + }); }); class FakeSessionDaemon { readonly requests: { method: string; path: string; body: unknown }[] = []; + readonly websocketPaths: string[] = []; + private readonly sockets = new Set(); + + private constructor(private readonly upstream: WebSocketServer) { + this.upstream.on("connection", (socket) => { + this.sockets.add(socket); + socket.on("close", () => { this.sockets.delete(socket); }); + }); + } + + static async create(): Promise { + const upstream = new WebSocketServer({ host: "127.0.0.1", port: 0 }); + await waitForListening(upstream); + return new FakeSessionDaemon(upstream); + } request(method: string, path: string, body?: unknown): Promise<{ statusCode: number; headers: Record; body: string }> { this.requests.push({ method, path, body }); return Promise.resolve({ statusCode: 200, headers: { "content-type": "application/json" }, body: JSON.stringify({ ok: true }) }); } - connectWebSocket(): never { - throw new Error("not implemented"); + connectWebSocket(path: string): WebSocket { + this.websocketPaths.push(path); + return new WebSocket(`${webSocketServerUrl(this.upstream)}${path}`); + } + + async close(): Promise { + for (const socket of this.sockets) socket.terminate(); + await closeWebSocketServer(this.upstream); } } + +function serverUrl(instance: FastifyInstance): string { + const address = instance.server.address(); + if (address === null || typeof address === "string") throw new Error("Expected TCP server address"); + return `ws://127.0.0.1:${String(address.port)}`; +} + +function webSocketServerUrl(server: WebSocketServer): string { + const address = server.address(); + if (address === null || typeof address === "string") throw new Error("Expected TCP server address"); + return `ws://127.0.0.1:${String(address.port)}`; +} + +function waitForListening(server: WebSocketServer): Promise { + return new Promise((resolve) => { + server.once("listening", () => { resolve(); }); + }); +} + +function closeWebSocketServer(server: WebSocketServer): Promise { + return new Promise((resolve, reject) => { + server.close((error) => { + if (error !== undefined) reject(error); + else resolve(); + }); + }); +} + +function waitForOpen(socket: WebSocket): Promise { + if (socket.readyState === WebSocket.OPEN) return Promise.resolve(); + return new Promise((resolve, reject) => { + socket.once("open", () => { resolve(); }); + socket.once("error", reject); + socket.once("close", () => { reject(new Error("WebSocket closed before opening")); }); + }); +}