test: cover session websocket cwd proxying

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<WebSocket>();
+
+  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<FakeSessionDaemon> {
+    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<string, string>; 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<void> {
+    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<void> {
+  return new Promise((resolve) => {
+    server.once("listening", () => { resolve(); });
+  });
+}
+
+function closeWebSocketServer(server: WebSocketServer): Promise<void> {
+  return new Promise((resolve, reject) => {
+    server.close((error) => {
+      if (error !== undefined) reject(error);
+      else resolve();
+    });
+  });
+}
+
+function waitForOpen(socket: WebSocket): Promise<void> {
+  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")); });
+  });
+}