test-proxy-recorder 0.3.4 → 0.3.6

package/dist/index.cjs

@@ -1,20 +1,20 @@
 'use strict';
 
 var fs = require('fs/promises');
-var http = require('http');
-var https = require('https');
+var http2 = require('http');
 var httpProxy = require('http-proxy');
-var ws = require('ws');
+var https = require('https');
 var crypto = require('crypto');
 var path2 = require('path');
 var filenamify2 = require('filenamify');
+var ws = require('ws');
 
 function _interopDefault (e) { return e && e.__esModule ? e : { default: e }; }
 
 var fs__default = /*#__PURE__*/_interopDefault(fs);
-var http__default = /*#__PURE__*/_interopDefault(http);
-var https__default = /*#__PURE__*/_interopDefault(https);
+var http2__default = /*#__PURE__*/_interopDefault(http2);
 var httpProxy__default = /*#__PURE__*/_interopDefault(httpProxy);
+var https__default = /*#__PURE__*/_interopDefault(https);
 var crypto__default = /*#__PURE__*/_interopDefault(crypto);
 var path2__default = /*#__PURE__*/_interopDefault(path2);
 var filenamify2__default = /*#__PURE__*/_interopDefault(filenamify2);
@@ -28,6 +28,244 @@ var HTTP_STATUS_NOT_FOUND = 404;
 var CONTROL_ENDPOINT = "/__control";
 var RECORDING_ID_HEADER = "x-test-rcrd-id";
 
+// src/utils/cors.ts
+function getCorsHeaders(req) {
+  const origin = req.headers.origin;
+  return {
+    "access-control-allow-origin": origin || "*",
+    "access-control-allow-credentials": "true",
+    "access-control-allow-headers": req.headers["access-control-request-headers"] || `Origin, X-Requested-With, Content-Type, Accept, Authorization, ${RECORDING_ID_HEADER}`,
+    "access-control-allow-methods": "GET, POST, PUT, DELETE, PATCH, OPTIONS",
+    "access-control-expose-headers": "*"
+  };
+}
+function addCorsHeaders(proxyRes, req) {
+  Object.assign(proxyRes.headers, getCorsHeaders(req));
+}
+
+// src/httpRecorder.ts
+async function bufferRequestBody(req) {
+  const chunks = [];
+  req.on("data", (chunk) => {
+    chunks.push(chunk);
+  });
+  try {
+    await new Promise((resolveBuffer, rejectBuffer) => {
+      req.on("end", () => resolveBuffer());
+      req.on("error", (err) => rejectBuffer(err));
+      setTimeout(
+        () => rejectBuffer(new Error("Request buffering timeout")),
+        3e4
+      );
+    });
+  } catch (error) {
+    console.error("Error buffering request:", error);
+  }
+  return chunks;
+}
+function handleProxyResponse(proxyRes, context) {
+  const { options, requestBody, resolve } = context;
+  const { req, res, key, recordingId, sequence, onProxyError } = options;
+  addCorsHeaders(proxyRes, req);
+  const responseChunks = [];
+  proxyRes.on("data", (chunk) => {
+    responseChunks.push(chunk);
+  });
+  proxyRes.on("end", () => {
+    try {
+      const responseBody = Buffer.concat(responseChunks);
+      const responseBodyStr = responseBody.toString("utf8");
+      const recording = {
+        request: {
+          method: req.method,
+          url: req.url,
+          headers: req.headers,
+          body: requestBody || null
+        },
+        response: {
+          statusCode: proxyRes.statusCode,
+          headers: proxyRes.headers,
+          body: responseBodyStr || null
+        },
+        timestamp: (/* @__PURE__ */ new Date()).toISOString(),
+        key,
+        recordingId,
+        sequence
+      };
+      const responseHeaders = {
+        ...proxyRes.headers,
+        ...getCorsHeaders(req)
+      };
+      res.writeHead(proxyRes.statusCode || 200, responseHeaders);
+      res.end(responseBody);
+      console.log(
+        `Recorded: ${req.method} ${req.url} (recordingId: ${recordingId}, sequence: ${sequence})`
+      );
+      resolve(recording);
+    } catch (error) {
+      console.error("Error completing recording:", error);
+      resolve(null);
+    }
+  });
+  proxyRes.on("error", (err) => {
+    console.error("Proxy response error:", err);
+    if (!res.headersSent) {
+      onProxyError(err, req, res);
+    }
+    resolve(null);
+  });
+}
+function proxyWithBufferedBody(options, chunks) {
+  const { req, res, target, onProxyError } = options;
+  const requestBody = Buffer.concat(chunks).toString("utf8");
+  const targetUrl = new URL(target);
+  const isHttps = targetUrl.protocol === "https:";
+  const requestModule = isHttps ? https__default.default : http2__default.default;
+  const defaultPort = isHttps ? 443 : 80;
+  return new Promise((resolve) => {
+    const proxyReq = requestModule.request(
+      {
+        hostname: targetUrl.hostname,
+        port: targetUrl.port || defaultPort,
+        path: req.url,
+        method: req.method,
+        headers: req.headers
+      },
+      (proxyRes) => {
+        handleProxyResponse(proxyRes, { options, requestBody, resolve });
+      }
+    );
+    proxyReq.on("error", (err) => {
+      onProxyError(err, req, res);
+      resolve(null);
+    });
+    if (chunks.length > 0) {
+      proxyReq.write(Buffer.concat(chunks));
+    }
+    proxyReq.end();
+  });
+}
+async function recordAndProxyRequest(options) {
+  const { req, res, onProxyError } = options;
+  try {
+    const chunks = await bufferRequestBody(req);
+    return await proxyWithBufferedBody(options, chunks);
+  } catch (error) {
+    console.error("Error in recordAndProxyRequest:", error);
+    try {
+      onProxyError(error, req, res);
+    } catch (error_) {
+      console.error("Failed to handle proxy error:", error_);
+    }
+    return null;
+  }
+}
+
+// src/replaySessions.ts
+var ReplaySessionManager = class {
+  sessions = /* @__PURE__ */ new Map();
+  evictionTimer = null;
+  timeoutMs;
+  constructor(timeoutMs) {
+    this.timeoutMs = timeoutMs;
+  }
+  get size() {
+    return this.sessions.size;
+  }
+  keys() {
+    return this.sessions.keys();
+  }
+  /**
+   * Get or create a replay session state for a given recording ID
+   */
+  getOrCreate(recordingId) {
+    let session = this.sessions.get(recordingId);
+    if (session) {
+      session.lastAccessTime = Date.now();
+    } else {
+      session = {
+        recordingId,
+        servedRecordingIdsByKey: /* @__PURE__ */ new Map(),
+        loadedSession: null,
+        lastAccessTime: Date.now(),
+        sortedRecordingsByKey: /* @__PURE__ */ new Map()
+      };
+      this.sessions.set(recordingId, session);
+      this.startEvictionTimer();
+      console.log(
+        `[CONCURRENT REPLAY] Created new session for recording: ${recordingId}`
+      );
+    }
+    return session;
+  }
+  delete(sessionId) {
+    this.sessions.delete(sessionId);
+    if (this.sessions.size === 0) {
+      this.stopEvictionTimer();
+    }
+  }
+  startEvictionTimer() {
+    if (this.evictionTimer) {
+      return;
+    }
+    const CHECK_INTERVAL_MS = 3e4;
+    this.evictionTimer = setInterval(() => {
+      const now = Date.now();
+      for (const [id, session] of this.sessions) {
+        if (now - session.lastAccessTime >= this.timeoutMs) {
+          console.log(
+            `[EVICTION] Evicting idle replay session: ${id} (idle for ${Math.round((now - session.lastAccessTime) / 1e3)}s)`
+          );
+          this.sessions.delete(id);
+        }
+      }
+      if (this.sessions.size === 0) {
+        this.stopEvictionTimer();
+      }
+    }, CHECK_INTERVAL_MS);
+    this.evictionTimer.unref();
+  }
+  stopEvictionTimer() {
+    if (this.evictionTimer) {
+      clearInterval(this.evictionTimer);
+      this.evictionTimer = null;
+    }
+  }
+};
+function getServedTracker(sessionState, key) {
+  if (!sessionState.servedRecordingIdsByKey.has(key)) {
+    sessionState.servedRecordingIdsByKey.set(key, /* @__PURE__ */ new Set());
+  }
+  return sessionState.servedRecordingIdsByKey.get(key);
+}
+function getSortedRecordings(sessionState, key) {
+  if (sessionState.sortedRecordingsByKey.has(key)) {
+    return sessionState.sortedRecordingsByKey.get(key);
+  }
+  const session = sessionState.loadedSession;
+  const sortedRecords = session.recordings.filter((r) => r.key === key && r.response).toSorted((a, b) => {
+    const aSeq = a.sequence !== void 0 ? a.sequence : a.recordingId;
+    const bSeq = b.sequence !== void 0 ? b.sequence : b.recordingId;
+    return aSeq - bSeq;
+  });
+  sessionState.sortedRecordingsByKey.set(key, sortedRecords);
+  return sortedRecords;
+}
+function selectReplayRecord(recordsWithKey, servedForThisKey, key, recordingId) {
+  for (const rec of recordsWithKey) {
+    if (!servedForThisKey.has(rec.recordingId)) {
+      return rec;
+    }
+  }
+  if (recordsWithKey.length > 0) {
+    console.log(
+      `[REPLAY WARNING] All ${recordsWithKey.length} recordings already served for ${key} (session: ${recordingId}), reusing last one`
+    );
+    return recordsWithKey[recordsWithKey.length - 1];
+  }
+  return null;
+}
+
 // src/types.ts
 var Modes = {
   transparent: "transparent",
@@ -71,9 +309,9 @@ function processRecordings(recordings) {
   const processedRecordings = [];
   for (const [_key, keyRecordings] of recordingsByKey) {
     keyRecordings.sort((a, b) => a.recordingId - b.recordingId);
-    keyRecordings.forEach((recording, index) => {
+    for (const [index, recording] of keyRecordings.entries()) {
       processedRecordings.push({ ...recording, sequence: index });
-    });
+    }
   }
   processedRecordings.sort((a, b) => a.recordingId - b.recordingId);
   return processedRecordings;
@@ -129,6 +367,180 @@ function sendJsonResponse(res, statusCode, data) {
   res.end(JSON.stringify(data));
 }
 
+// src/utils/recordingId.ts
+function getRecordingIdFromHeader(req) {
+  const headerValue = req.headers[RECORDING_ID_HEADER];
+  if (!headerValue) {
+    return null;
+  }
+  return Array.isArray(headerValue) ? headerValue[0] : headerValue;
+}
+function getRecordingIdFromCookie(req) {
+  const cookies = req.headers.cookie;
+  if (!cookies) {
+    return null;
+  }
+  const match = cookies.match(/proxy-recording-id=([^;]+)/);
+  return match ? decodeURIComponent(match[1]) : null;
+}
+function getRecordingIdFromRequest(req) {
+  const fromHeader = getRecordingIdFromHeader(req);
+  const fromCookie = getRecordingIdFromCookie(req);
+  return fromHeader ?? fromCookie ?? null;
+}
+function getWsRecordingKey(url) {
+  return `WS_${url.replaceAll("/", "_")}`;
+}
+var WS_INTERNAL_HEADERS = /* @__PURE__ */ new Set([
+  "host",
+  "connection",
+  "upgrade",
+  "sec-websocket-key",
+  "sec-websocket-version",
+  "sec-websocket-extensions"
+]);
+function getRecordableWsHeaders(req) {
+  const headers = {};
+  for (const [name, value] of Object.entries(req.headers)) {
+    if (!WS_INTERNAL_HEADERS.has(name) && value !== void 0) {
+      headers[name] = value;
+    }
+  }
+  return headers;
+}
+function getForwardableWsHeaders(req) {
+  const headers = {};
+  for (const [name, value] of Object.entries(getRecordableWsHeaders(req))) {
+    if (name !== "sec-websocket-protocol" && value !== void 0) {
+      headers[name] = Array.isArray(value) ? value.join(", ") : value;
+    }
+  }
+  return headers;
+}
+function getClientSubprotocols(req) {
+  const header = req.headers["sec-websocket-protocol"];
+  if (!header) {
+    return [];
+  }
+  const raw = Array.isArray(header) ? header.join(",") : header;
+  return raw.split(",").map((p) => p.trim()).filter(Boolean);
+}
+function recordWebSocket(req, clientSocket, head, target, session) {
+  const url = req.url || "/";
+  const wsRecording = {
+    url,
+    messages: [],
+    timestamp: (/* @__PURE__ */ new Date()).toISOString(),
+    key: getWsRecordingKey(url),
+    headers: getRecordableWsHeaders(req)
+  };
+  if (session) {
+    session.websocketRecordings.push(wsRecording);
+  }
+  const backendWsUrl = `${target.replace("http", "ws")}${url}`;
+  const backendWs = new ws.WebSocket(backendWsUrl, getClientSubprotocols(req), {
+    headers: getForwardableWsHeaders(req)
+  });
+  const wss = new ws.WebSocketServer({
+    noServer: true,
+    handleProtocols: (protocols) => backendWs.protocol && protocols.has(backendWs.protocol) ? backendWs.protocol : protocols.values().next().value ?? false
+  });
+  backendWs.on("open", () => {
+    console.log(`WebSocket recording: connected to backend ${backendWsUrl}`);
+    if (backendWs.protocol) {
+      wsRecording.protocol = backendWs.protocol;
+    }
+    wss.handleUpgrade(req, clientSocket, head, (clientWs) => {
+      clientWs.on("message", (data) => {
+        const message = data.toString();
+        wsRecording.messages.push({
+          direction: "client-to-server",
+          data: message,
+          timestamp: (/* @__PURE__ */ new Date()).toISOString()
+        });
+        if (backendWs.readyState === ws.WebSocket.OPEN) {
+          backendWs.send(message);
+        }
+      });
+      backendWs.on("message", (data) => {
+        const message = data.toString();
+        wsRecording.messages.push({
+          direction: "server-to-client",
+          data: message,
+          timestamp: (/* @__PURE__ */ new Date()).toISOString()
+        });
+        if (clientWs.readyState === ws.WebSocket.OPEN) {
+          clientWs.send(message);
+        }
+      });
+      clientWs.on("error", (err) => {
+        console.error("Client WebSocket error:", err);
+      });
+      backendWs.on("error", (err) => {
+        console.error("Backend WebSocket error:", err);
+      });
+      clientWs.on("close", () => {
+        backendWs.close();
+        console.log("Client WebSocket closed");
+      });
+      backendWs.on("close", () => {
+        clientWs.close();
+        console.log("Backend WebSocket closed");
+      });
+    });
+  });
+  backendWs.on("error", (err) => {
+    console.error("Backend WebSocket connection error:", err);
+    clientSocket.write("HTTP/1.1 502 Bad Gateway\r\n\r\n");
+    clientSocket.destroy();
+  });
+  wss.on("error", (err) => {
+    console.error("WebSocket server error:", err);
+  });
+}
+function replayWebSocket(req, socket, wsRecording, recordingId) {
+  const url = req.url || "/";
+  const wss = new ws.WebSocketServer({
+    noServer: true,
+    handleProtocols: (protocols) => wsRecording.protocol && protocols.has(wsRecording.protocol) ? wsRecording.protocol : protocols.values().next().value ?? false
+  });
+  const fakeReq = Object.assign(req, {
+    headers: {
+      ...req.headers,
+      "sec-websocket-key": req.headers["sec-websocket-key"] || "replay-key",
+      "sec-websocket-version": "13"
+    }
+  });
+  wss.handleUpgrade(fakeReq, socket, Buffer.alloc(0), (ws$1) => {
+    console.log(`Replaying WebSocket: ${url} (session: ${recordingId})`);
+    const messages = wsRecording.messages;
+    let cursor = 0;
+    const flushServerMessages = () => {
+      while (cursor < messages.length && messages[cursor].direction === "server-to-client") {
+        const msg = messages[cursor];
+        cursor++;
+        if (ws$1.readyState === ws.WebSocket.OPEN) {
+          ws$1.send(msg.data);
+        }
+      }
+    };
+    ws$1.on("message", (data) => {
+      console.log(`Replay: Client sent: ${data.toString()}`);
+      if (cursor < messages.length && messages[cursor].direction === "client-to-server") {
+        cursor++;
+      }
+      flushServerMessages();
+    });
+    ws$1.on("error", (err) => {
+      console.error("Replay WebSocket error:", err);
+    });
+    ws$1.on("close", () => {
+      console.log("Replay WebSocket closed");
+    });
+    flushServerMessages();
+  });
+}
+
 // src/ProxyServer.ts
 var ProxyServer = class {
   target;
@@ -139,6 +551,7 @@ var ProxyServer = class {
   proxy;
   currentSession;
   recordingsDir;
+  timeoutMs;
   recordingIdCounter;
   // Unique ID for each recording entry
   sequenceCounterByKey;
@@ -149,8 +562,9 @@ var ProxyServer = class {
   // Stack of promises that resolve to completed recordings
   flushPromise;
   // Promise for in-progress flush operation
-  constructor(target, recordingsDir) {
+  constructor(target, recordingsDir, timeoutMs) {
     this.target = target;
+    this.timeoutMs = timeoutMs ?? DEFAULT_TIMEOUT_MS;
     this.mode = Modes.transparent;
     this.recordingId = null;
     this.recordingIdCounter = 0;
@@ -159,7 +573,7 @@ var ProxyServer = class {
     this.modeTimeout = null;
     this.currentSession = null;
     this.recordingsDir = recordingsDir;
-    this.replaySessions = /* @__PURE__ */ new Map();
+    this.replaySessions = new ReplaySessionManager(this.timeoutMs);
     this.recordingPromises = [];
     this.flushPromise = null;
     this.proxy = httpProxy__default.default.createProxyServer({
@@ -173,7 +587,7 @@ var ProxyServer = class {
     await fs__default.default.mkdir(this.recordingsDir, { recursive: true });
   }
   listen(port) {
-    const server = http__default.default.createServer((req, res) => {
+    const server = http2__default.default.createServer((req, res) => {
       this.handleRequest(req, res);
     });
     server.on("upgrade", (req, socket, head) => {
@@ -187,15 +601,15 @@ var ProxyServer = class {
   }
   setupProxyEventHandlers() {
     this.proxy.on("error", this.handleProxyError.bind(this));
-    this.proxy.on("proxyRes", this.addCorsHeaders.bind(this));
+    this.proxy.on("proxyRes", addCorsHeaders);
   }
   handleProxyError(err, req, res) {
     console.error("Proxy error:", err);
-    if (!(res instanceof http__default.default.ServerResponse)) {
+    if (!(res instanceof http2__default.default.ServerResponse)) {
       return;
     }
     if (!res.headersSent) {
-      const corsHeaders = this.getCorsHeaders(req);
+      const corsHeaders = getCorsHeaders(req);
       res.writeHead(HTTP_STATUS_BAD_GATEWAY, {
         "Content-Type": "application/json",
         ...corsHeaders
@@ -203,103 +617,12 @@ var ProxyServer = class {
     }
     res.end(JSON.stringify({ error: "Proxy error", message: err.message }));
   }
-  /**
-   * Get CORS headers for a given request
-   * @param req The incoming HTTP request
-   * @returns An object containing CORS headers
-   */
-  getCorsHeaders(req) {
-    const origin = req.headers.origin;
-    return {
-      "access-control-allow-origin": origin || "*",
-      "access-control-allow-credentials": "true",
-      "access-control-allow-headers": req.headers["access-control-request-headers"] || `Origin, X-Requested-With, Content-Type, Accept, Authorization, ${RECORDING_ID_HEADER}`,
-      "access-control-allow-methods": "GET, POST, PUT, DELETE, PATCH, OPTIONS",
-      "access-control-expose-headers": "*"
-    };
-  }
-  addCorsHeaders(proxyRes, req) {
-    const corsHeaders = this.getCorsHeaders(req);
-    Object.assign(proxyRes.headers, corsHeaders);
-  }
-  getTarget() {
-    return this.target;
-  }
-  /**
-   * Extract recording ID from custom HTTP header
-   * Used for concurrent replay session routing, especially with Next.js
-   * @param req The incoming HTTP request
-   * @returns The recording ID from header, or null if not found
-   */
-  getRecordingIdFromHeader(req) {
-    const headerValue = req.headers[RECORDING_ID_HEADER];
-    if (!headerValue) {
-      return null;
-    }
-    return Array.isArray(headerValue) ? headerValue[0] : headerValue;
-  }
-  /**
-   * Extract recording ID from request cookie
-   * Used for concurrent replay session routing (fallback method)
-   * @param req The incoming HTTP request
-   * @returns The recording ID from cookie, or null if not found
-   */
-  getRecordingIdFromCookie(req) {
-    const cookies = req.headers.cookie;
-    if (!cookies) {
-      return null;
-    }
-    const match = cookies.match(/proxy-recording-id=([^;]+)/);
-    return match ? decodeURIComponent(match[1]) : null;
-  }
-  /**
-   * Extract recording ID from request using custom header (preferred) or cookie (fallback)
-   * @param req The incoming HTTP request
-   * @returns The recording ID, or null if not found
-   */
-  getRecordingIdFromRequest(req) {
-    const fromHeader = this.getRecordingIdFromHeader(req);
-    const fromCookie = this.getRecordingIdFromCookie(req);
-    if (fromHeader) {
-      return fromHeader;
-    }
-    if (fromCookie) {
-      return fromCookie;
-    }
-    return null;
-  }
-  /**
-   * Get or create a replay session state for a given recording ID
-   * @param recordingId The recording ID to get/create session for
-   * @returns The replay session state
-   */
-  getOrCreateReplaySession(recordingId) {
-    let session = this.replaySessions.get(recordingId);
-    if (session) {
-      session.lastAccessTime = Date.now();
-    } else {
-      session = {
-        recordingId,
-        servedRecordingIdsByKey: /* @__PURE__ */ new Map(),
-        loadedSession: null,
-        lastAccessTime: Date.now(),
-        sortedRecordingsByKey: /* @__PURE__ */ new Map()
-      };
-      this.replaySessions.set(recordingId, session);
-      console.log(
-        `[CONCURRENT REPLAY] Created new session for recording: ${recordingId}`
-      );
-    }
-    return session;
-  }
   /**
    * Clean up a session - removes it from memory and resets counters
    * @param sessionId The session ID to clean up
    */
   async cleanupSession(sessionId) {
-    if (this.replaySessions.has(sessionId)) {
-      this.replaySessions.delete(sessionId);
-    }
+    this.replaySessions.delete(sessionId);
     if (this.recordingId === sessionId) {
       await this.saveCurrentSession();
       this.currentSession = null;
@@ -310,29 +633,17 @@ var ProxyServer = class {
     }
     console.log(`[CLEANUP] Session ${sessionId} cleaned up successfully`);
   }
-  parseGetParams(req) {
-    const url = new URL(req.url || "", `http://${req.headers.host}`);
-    const mode = url.searchParams.get("mode");
-    const id = url.searchParams.get("id") || void 0;
-    const timeoutParam = url.searchParams.get("timeout");
-    const timeout = timeoutParam ? Number.parseInt(timeoutParam, 10) : void 0;
-    if (!mode) {
-      throw new Error("Mode parameter is required");
-    }
-    return { mode, id, timeout };
-  }
-  async parseControlRequest(req) {
-    if (req.method === "GET") {
-      return this.parseGetParams(req);
-    }
-    if (req.method === "POST") {
-      const body = await readRequestBody(req);
-      console.log(`MODE CHANGE (${req.method})`, body);
-      return JSON.parse(body);
-    }
-    throw new Error("Unsupported control method");
+  async parseControlBody(req) {
+    const body = await readRequestBody(req);
+    console.log(`MODE CHANGE (${req.method})`, body);
+    return JSON.parse(body);
   }
   async handleControlRequest(req, res) {
+    if (req.method === "HEAD") {
+      res.writeHead(HTTP_STATUS_OK);
+      res.end();
+      return;
+    }
     if (req.method === "GET") {
       sendJsonResponse(res, HTTP_STATUS_OK, {
         recordingsDir: this.recordingsDir,
@@ -341,8 +652,11 @@ var ProxyServer = class {
       });
       return;
     }
+    await this.handleControlPost(req, res);
+  }
+  async handleControlPost(req, res) {
     try {
-      const data = await this.parseControlRequest(req);
+      const data = await this.parseControlBody(req);
       const { mode, id, timeout: requestTimeout, cleanup } = data;
       if (cleanup && id) {
         await this.cleanupSession(id);
@@ -353,29 +667,7 @@ var ProxyServer = class {
         });
         return;
       }
-      if (!mode) {
-        throw new Error(
-          "Mode parameter is required when cleanup is not specified"
-        );
-      }
-      const timeout = requestTimeout ?? DEFAULT_TIMEOUT_MS;
-      this.clearModeTimeout();
-      await this.switchMode(mode, id);
-      this.setupModeTimeout(timeout);
-      if (mode === Modes.replay && id) {
-        res.setHeader(
-          "Set-Cookie",
-          `proxy-recording-id=${encodeURIComponent(id)}; HttpOnly; Path=/; SameSite=Lax`
-        );
-        console.log(`[CONCURRENT REPLAY] Set cookie for recording: ${id}`);
-      }
-      sendJsonResponse(res, HTTP_STATUS_OK, {
-        success: true,
-        mode: this.mode,
-        id: this.recordingId || this.replayId,
-        timeout,
-        recordingsDir: this.recordingsDir
-      });
+      await this.applyModeChange(res, mode, id, requestTimeout);
     } catch (error) {
       console.error("Control request error:", error);
       sendJsonResponse(res, HTTP_STATUS_BAD_REQUEST, {
@@ -383,6 +675,31 @@ var ProxyServer = class {
       });
     }
   }
+  async applyModeChange(res, mode, id, requestTimeout) {
+    if (!mode) {
+      throw new Error(
+        "Mode parameter is required when cleanup is not specified"
+      );
+    }
+    const timeout = requestTimeout ?? this.timeoutMs;
+    this.clearModeTimeout();
+    await this.switchMode(mode, id);
+    this.setupModeTimeout(timeout);
+    if (mode === Modes.replay && id) {
+      res.setHeader(
+        "Set-Cookie",
+        `proxy-recording-id=${encodeURIComponent(id)}; HttpOnly; Path=/; SameSite=Lax`
+      );
+      console.log(`[CONCURRENT REPLAY] Set cookie for recording: ${id}`);
+    }
+    sendJsonResponse(res, HTTP_STATUS_OK, {
+      success: true,
+      mode: this.mode,
+      id: this.recordingId || this.replayId,
+      timeout,
+      recordingsDir: this.recordingsDir
+    });
+  }
   clearModeTimeout() {
     clearTimeout(this.modeTimeout || 0);
     this.modeTimeout = null;
@@ -422,7 +739,7 @@ var ProxyServer = class {
     this.recordingId = null;
     this.replayId = null;
     this.currentSession = null;
-    clearTimeout(this.modeTimeout || 0);
+    this.clearModeTimeout();
     console.log("Switched to transparent mode");
   }
   switchToRecordMode(id) {
@@ -439,7 +756,7 @@ var ProxyServer = class {
     this.replayId = id;
     this.recordingId = null;
     this.currentSession = null;
-    const sessionState = this.getOrCreateReplaySession(id);
+    const sessionState = this.replaySessions.getOrCreate(id);
     sessionState.servedRecordingIdsByKey.clear();
     sessionState.sortedRecordingsByKey.clear();
     const filePath = getRecordingPath(this.recordingsDir, id);
@@ -453,7 +770,7 @@ var ProxyServer = class {
     console.log(`Switched to replay mode with ID: ${id}`);
   }
   setupModeTimeout(timeout) {
-    clearTimeout(this.modeTimeout || 0);
+    this.clearModeTimeout();
     this.modeTimeout = setTimeout(async () => {
       console.log("Timeout reached, switching back to transparent mode");
       await this.saveCurrentSession();
@@ -500,7 +817,7 @@ var ProxyServer = class {
     await saveRecordingSession(this.recordingsDir, this.currentSession);
   }
   getRecordingIdOrError(req, res) {
-    const recordingIdFromRequest = this.getRecordingIdFromRequest(req);
+    const recordingIdFromRequest = getRecordingIdFromRequest(req);
     if (recordingIdFromRequest) {
       return recordingIdFromRequest;
     }
@@ -508,7 +825,7 @@ var ProxyServer = class {
       console.warn(
         `[CONCURRENT REPLAY WARNING] Request to ${req.method} ${req.url} is missing ${RECORDING_ID_HEADER} header/cookie. Active sessions: ${[...this.replaySessions.keys()].join(", ")}. this.replayId fallback would be: ${this.replayId} (NOT USING - could be wrong session)`
       );
-      const corsHeaders = this.getCorsHeaders(req);
+      const corsHeaders = getCorsHeaders(req);
       res.writeHead(HTTP_STATUS_BAD_REQUEST, {
         "Content-Type": "application/json",
         ...corsHeaders
@@ -524,7 +841,7 @@ var ProxyServer = class {
     }
     const recordingId = this.replayId;
     if (!recordingId) {
-      const corsHeaders = this.getCorsHeaders(req);
+      const corsHeaders = getCorsHeaders(req);
       res.writeHead(HTTP_STATUS_BAD_REQUEST, {
         "Content-Type": "application/json",
         ...corsHeaders
@@ -537,56 +854,22 @@ var ProxyServer = class {
     );
     return recordingId;
   }
-  getServedTracker(sessionState, key) {
-    if (!sessionState.servedRecordingIdsByKey.has(key)) {
-      sessionState.servedRecordingIdsByKey.set(key, /* @__PURE__ */ new Set());
-    }
-    return sessionState.servedRecordingIdsByKey.get(key);
-  }
-  getSortedRecordings(sessionState, key) {
-    if (sessionState.sortedRecordingsByKey.has(key)) {
-      return sessionState.sortedRecordingsByKey.get(key);
-    }
-    const session = sessionState.loadedSession;
-    const sortedRecords = session.recordings.filter((r) => r.key === key && r.response).toSorted((a, b) => {
-      const aSeq = a.sequence !== void 0 ? a.sequence : a.recordingId;
-      const bSeq = b.sequence !== void 0 ? b.sequence : b.recordingId;
-      return aSeq - bSeq;
-    });
-    sessionState.sortedRecordingsByKey.set(key, sortedRecords);
-    return sortedRecords;
-  }
-  selectReplayRecord(recordsWithKey, servedForThisKey, key, recordingId) {
-    for (const rec of recordsWithKey) {
-      if (!servedForThisKey.has(rec.recordingId)) {
-        return rec;
-      }
-    }
-    if (recordsWithKey.length > 0) {
-      console.log(
-        `[REPLAY WARNING] All ${recordsWithKey.length} recordings already served for ${key} (session: ${recordingId}), reusing last one`
-      );
-      return recordsWithKey[recordsWithKey.length - 1];
-    }
-    return null;
-  }
   async handleReplayRequest(req, res) {
     const recordingId = this.getRecordingIdOrError(req, res);
     if (!recordingId) return;
     const key = getReqID(req);
     const filePath = getRecordingPath(this.recordingsDir, recordingId);
     try {
-      const sessionState = this.getOrCreateReplaySession(recordingId);
+      const sessionState = this.replaySessions.getOrCreate(recordingId);
       if (!sessionState.loadedSession) {
-        const error = new Error(
-          `Recording session file not found: ${filePath}`
+        throw Object.assign(
+          new Error(`Recording session file not found: ${filePath}`),
+          { code: "ENOENT" }
         );
-        error.code = "ENOENT";
-        throw error;
       }
-      const servedForThisKey = this.getServedTracker(sessionState, key);
+      const servedForThisKey = getServedTracker(sessionState, key);
       const host = req.headers.host || "unknown";
-      const recordsWithKey = this.getSortedRecordings(sessionState, key);
+      const recordsWithKey = getSortedRecordings(sessionState, key);
       if (recordsWithKey.length === 0) {
         const errorMsg = `No recording found for ${key} at ${req.method} ${host}${req.url}`;
         console.error(`[REPLAY ERROR] ${errorMsg} (session: ${recordingId})`);
@@ -599,7 +882,7 @@ var ProxyServer = class {
           key,
           sessionId: recordingId
         };
-        const corsHeaders = this.getCorsHeaders(req);
+        const corsHeaders = getCorsHeaders(req);
         res.writeHead(HTTP_STATUS_NOT_FOUND, {
           "Content-Type": "application/json",
           ...corsHeaders
@@ -611,7 +894,7 @@ var ProxyServer = class {
       console.log(
         `[replay request #${requestCount}] ${req.method} ${req.url} (key: ${key}, session: ${recordingId}, total: ${recordsWithKey.length}, served: ${servedForThisKey.size})`
       );
-      const record = this.selectReplayRecord(
+      const record = selectReplayRecord(
         recordsWithKey,
         servedForThisKey,
         key,
@@ -629,7 +912,7 @@ var ProxyServer = class {
       const { statusCode, headers, body } = record.response;
       const responseHeaders = {
         ...headers,
-        ...this.getCorsHeaders(req)
+        ...getCorsHeaders(req)
       };
       res.writeHead(statusCode, responseHeaders);
       res.end(body);
@@ -640,7 +923,7 @@ var ProxyServer = class {
   handleReplayError(req, res, err, key, filePath) {
     const isFileNotFound = err instanceof Error && "code" in err && err.code === "ENOENT";
     console.error("Replay error:", err);
-    const corsHeaders = this.getCorsHeaders(req);
+    const corsHeaders = getCorsHeaders(req);
     res.writeHead(HTTP_STATUS_NOT_FOUND, {
       "Content-Type": "application/json",
       ...corsHeaders
@@ -668,7 +951,7 @@ var ProxyServer = class {
     await this.handleProxyRequest(req, res);
   }
   handleCorsPreflightRequest(req, res) {
-    const corsHeaders = this.getCorsHeaders(req);
+    const corsHeaders = getCorsHeaders(req);
     res.writeHead(HTTP_STATUS_OK, {
       ...corsHeaders,
       "Access-Control-Max-Age": "86400"
@@ -677,16 +960,15 @@ var ProxyServer = class {
     res.end();
   }
   async handleProxyRequest(req, res) {
-    const target = this.getTarget();
+    const target = this.target;
     console.log(`[${this.mode}] ${req.method} ${req.url} -> ${target}`);
     if (this.mode === Modes.record) {
-      await this.recordAndProxyRequest(req, res, target);
+      this.recordAndProxy(req, res, target);
     } else {
       this.proxy.web(req, res, { target });
     }
   }
-  // Note: streaming requests are buffered before proxying; streaming passthrough is not yet implemented
-  async recordAndProxyRequest(req, res, target) {
+  recordAndProxy(req, res, target) {
     if (!this.currentSession) {
       return;
     }
@@ -694,194 +976,66 @@ var ProxyServer = class {
     const recordingId = this.recordingIdCounter++;
     const sequence = this.sequenceCounterByKey.get(key) || 0;
     this.sequenceCounterByKey.set(key, sequence + 1);
-    const recordingPromise = new Promise((resolve) => {
-      (async () => {
-        try {
-          const chunks = [];
-          req.on("data", (chunk) => {
-            chunks.push(chunk);
-          });
-          try {
-            await new Promise((resolveBuffer, rejectBuffer) => {
-              req.on("end", () => resolveBuffer());
-              req.on("error", (err) => rejectBuffer(err));
-              setTimeout(
-                () => rejectBuffer(new Error("Request buffering timeout")),
-                3e4
-              );
-            });
-          } catch (error) {
-            console.error("Error buffering request:", error);
-          }
-          const requestBody = Buffer.concat(chunks).toString("utf8");
-          const targetUrl = new URL(target);
-          const isHttps = targetUrl.protocol === "https:";
-          const requestModule = isHttps ? https__default.default : http__default.default;
-          const defaultPort = isHttps ? 443 : 80;
-          const proxyReq = requestModule.request(
-            {
-              hostname: targetUrl.hostname,
-              port: targetUrl.port || defaultPort,
-              path: req.url,
-              method: req.method,
-              headers: req.headers
-            },
-            (proxyRes) => {
-              this.addCorsHeaders(proxyRes, req);
-              const responseChunks = [];
-              proxyRes.on("data", (chunk) => {
-                responseChunks.push(chunk);
-              });
-              proxyRes.on("end", async () => {
-                try {
-                  const responseBody = Buffer.concat(responseChunks);
-                  const responseBodyStr = responseBody.toString("utf8");
-                  const recording = {
-                    request: {
-                      method: req.method,
-                      url: req.url,
-                      headers: req.headers,
-                      body: requestBody || null
-                    },
-                    response: {
-                      statusCode: proxyRes.statusCode,
-                      headers: proxyRes.headers,
-                      body: responseBodyStr || null
-                    },
-                    timestamp: (/* @__PURE__ */ new Date()).toISOString(),
-                    key,
-                    recordingId,
-                    sequence
-                  };
-                  const responseHeaders = {
-                    ...proxyRes.headers,
-                    ...this.getCorsHeaders(req)
-                  };
-                  res.writeHead(proxyRes.statusCode || 200, responseHeaders);
-                  res.end(responseBody);
-                  console.log(
-                    `Recorded: ${req.method} ${req.url} (recordingId: ${recordingId}, sequence: ${sequence})`
-                  );
-                  resolve(recording);
-                } catch (error) {
-                  console.error("Error completing recording:", error);
-                  resolve(null);
-                }
-              });
-              proxyRes.on("error", (err) => {
-                console.error("Proxy response error:", err);
-                if (!res.headersSent) {
-                  this.handleProxyError(err, req, res);
-                }
-                resolve(null);
-              });
-            }
-          );
-          proxyReq.on("error", (err) => {
-            this.handleProxyError(err, req, res);
-            resolve(null);
-          });
-          if (chunks.length > 0) {
-            proxyReq.write(Buffer.concat(chunks));
-          }
-          proxyReq.end();
-        } catch (error) {
-          console.error("Error in recordAndProxyRequest:", error);
-          try {
-            this.handleProxyError(error, req, res);
-          } catch (error_) {
-            console.error("Failed to handle proxy error:", error_);
-          }
-          resolve(null);
-        }
-      })();
-    });
-    this.recordingPromises.push(recordingPromise);
+    this.recordingPromises.push(
+      recordAndProxyRequest({
+        req,
+        res,
+        target,
+        key,
+        recordingId,
+        sequence,
+        onProxyError: this.handleProxyError.bind(this)
+      })
+    );
   }
   handleUpgrade(req, socket, head) {
     if (this.mode === Modes.replay) {
       this.handleReplayWebSocket(req, socket);
       return;
     }
-    const target = this.getTarget();
+    const target = this.target;
     console.log(`[${this.mode}] WebSocket upgrade ${req.url} -> ${target}`);
     if (this.mode === Modes.record) {
-      this.handleRecordWebSocket(req, socket, head, target);
+      recordWebSocket(req, socket, head, target, this.currentSession);
     } else {
       this.proxy.ws(req, socket, head, { target });
     }
   }
-  handleRecordWebSocket(req, clientSocket, head, target) {
-    const url = req.url || "/";
-    const key = `WS_${url.replaceAll("/", "_")}`;
-    const wsRecording = {
-      url,
-      messages: [],
-      timestamp: (/* @__PURE__ */ new Date()).toISOString(),
-      key
-    };
-    if (this.currentSession) {
-      this.currentSession.websocketRecordings.push(wsRecording);
+  /**
+   * Resolve the recording ID for a WebSocket upgrade request.
+   * Mirrors getRecordingIdOrError(): prefer the header/cookie from the request,
+   * fall back to this.replayId only when there is at most one active session.
+   * Browsers cannot set custom headers on WebSocket handshakes from JS, but
+   * Playwright's setExtraHTTPHeaders / cookies still reach the upgrade request.
+   */
+  getWsRecordingId(req) {
+    const fromRequest = getRecordingIdFromRequest(req);
+    if (fromRequest) {
+      return fromRequest;
     }
-    const backendWsUrl = `${target.replace("http", "ws")}${url}`;
-    const backendWs = new ws.WebSocket(backendWsUrl);
-    const wss = new ws.WebSocketServer({ noServer: true });
-    backendWs.on("open", () => {
-      console.log(`WebSocket recording: connected to backend ${backendWsUrl}`);
-      wss.handleUpgrade(req, clientSocket, head, (clientWs) => {
-        clientWs.on("message", (data) => {
-          const message = data.toString();
-          wsRecording.messages.push({
-            direction: "client-to-server",
-            data: message,
-            timestamp: (/* @__PURE__ */ new Date()).toISOString()
-          });
-          if (backendWs.readyState === ws.WebSocket.OPEN) {
-            backendWs.send(message);
-          }
-        });
-        backendWs.on("message", (data) => {
-          const message = data.toString();
-          wsRecording.messages.push({
-            direction: "server-to-client",
-            data: message,
-            timestamp: (/* @__PURE__ */ new Date()).toISOString()
-          });
-          if (clientWs.readyState === ws.WebSocket.OPEN) {
-            clientWs.send(message);
-          }
-        });
-        clientWs.on("error", (err) => {
-          console.error("Client WebSocket error:", err);
-        });
-        backendWs.on("error", (err) => {
-          console.error("Backend WebSocket error:", err);
-        });
-        clientWs.on("close", () => {
-          backendWs.close();
-          console.log("Client WebSocket closed");
-        });
-        backendWs.on("close", () => {
-          clientWs.close();
-          console.log("Backend WebSocket closed");
-        });
-      });
-    });
-    backendWs.on("error", (err) => {
-      console.error("Backend WebSocket connection error:", err);
-      clientSocket.write("HTTP/1.1 502 Bad Gateway\r\n\r\n");
-      clientSocket.destroy();
-    });
-    wss.on("error", (err) => {
-      console.error("WebSocket server error:", err);
-    });
+    if (this.replaySessions.size > 1) {
+      console.warn(
+        `[CONCURRENT REPLAY WARNING] WebSocket upgrade ${req.url} is missing ${RECORDING_ID_HEADER} header/cookie. Active sessions: ${[...this.replaySessions.keys()].join(", ")}. this.replayId fallback would be: ${this.replayId} (NOT USING - could be wrong session)`
+      );
+      return null;
+    }
+    return this.replayId;
   }
-  handleReplayWebSocket(req, socket) {
-    const url = req.url || "/";
-    const key = `WS_${url.replaceAll("/", "_")}`;
-    const filePath = getRecordingPath(this.recordingsDir, this.replayId);
-    loadRecordingSession(filePath).then((session) => {
-      const wsRecording = session.websocketRecordings.find(
+  async handleReplayWebSocket(req, socket) {
+    const key = getWsRecordingKey(req.url || "/");
+    const recordingId = this.getWsRecordingId(req);
+    if (!recordingId) {
+      socket.write("HTTP/1.1 400 Bad Request\r\n\r\n");
+      socket.destroy();
+      return;
+    }
+    try {
+      const sessionState = this.replaySessions.getOrCreate(recordingId);
+      if (!sessionState.loadedSession) {
+        const filePath = getRecordingPath(this.recordingsDir, recordingId);
+        sessionState.loadedSession = await loadRecordingSession(filePath);
+      }
+      const wsRecording = sessionState.loadedSession.websocketRecordings.find(
         (r) => r.key === key
       );
       if (!wsRecording) {
@@ -890,67 +1044,12 @@ var ProxyServer = class {
         console.log(`No WebSocket recording found for ${key}`);
         return;
       }
-      const wss = new ws.WebSocketServer({ noServer: true });
-      const fakeReq = Object.assign(req, {
-        headers: {
-          ...req.headers,
-          "sec-websocket-key": req.headers["sec-websocket-key"] || "replay-key",
-          "sec-websocket-version": "13"
-        }
-      });
-      wss.handleUpgrade(fakeReq, socket, Buffer.alloc(0), (ws$1) => {
-        console.log(`Replaying WebSocket: ${url}`);
-        const serverMessages = wsRecording.messages.filter(
-          (m) => m.direction === "server-to-client"
-        );
-        let messageIndex = 0;
-        ws$1.on("message", (data) => {
-          const clientMessage = data.toString();
-          console.log(`Replay: Client sent: ${clientMessage}`);
-          if (messageIndex < serverMessages.length) {
-            setTimeout(() => {
-              if (ws$1.readyState === ws.WebSocket.OPEN) {
-                ws$1.send(serverMessages[messageIndex].data);
-                console.log(`Replay: Sent server message ${messageIndex}`);
-                messageIndex++;
-              }
-            }, 10);
-          }
-        });
-        let initialMessagesSent = 0;
-        for (let i = 0; i < wsRecording.messages.length; i++) {
-          const msg = wsRecording.messages[i];
-          if (msg.direction === "client-to-server") {
-            break;
-          }
-          if (msg.direction === "server-to-client") {
-            setTimeout(
-              () => {
-                if (ws$1.readyState === ws.WebSocket.OPEN) {
-                  ws$1.send(msg.data);
-                  console.log(
-                    `Replay: Sent initial server message: ${msg.data}`
-                  );
-                  messageIndex++;
-                  initialMessagesSent++;
-                }
-              },
-              10 * (initialMessagesSent + 1)
-            );
-          }
-        }
-        ws$1.on("error", (err) => {
-          console.error("Replay WebSocket error:", err);
-        });
-        ws$1.on("close", () => {
-          console.log("Replay WebSocket closed");
-        });
-      });
-    }).catch((error) => {
+      replayWebSocket(req, socket, wsRecording, recordingId);
+    } catch (error) {
       console.error("Replay error:", error);
       socket.write("HTTP/1.1 404 Not Found\r\n\r\n");
       socket.destroy();
-    });
+    }
   }
   logServerStartup(port) {
     console.log(`Proxy server running on http://localhost:${port}`);
@@ -961,6 +1060,7 @@ var ProxyServer = class {
     );
   }
 };
+var registeredContexts = /* @__PURE__ */ new WeakSet();
 function getProxyPort() {
   const envPort = process.env.TEST_PROXY_RECORDER_PORT;
   if (envPort) {
@@ -979,7 +1079,7 @@ async function setProxyMode(mode, sessionId, timeout) {
       id: sessionId,
       ...timeout && { timeout }
     };
-    const response = await fetch(`http://127.0.0.1:${proxyPort}/__control`, {
+    const response = await fetch(`http://localhost:${proxyPort}/__control`, {
       method: "POST",
       headers: { "Content-Type": "application/json" },
       body: JSON.stringify(body)
@@ -1002,7 +1102,7 @@ async function cleanupSession(sessionId) {
       cleanup: true,
       id: sessionId
     };
-    const response = await fetch(`http://127.0.0.1:${proxyPort}/__control`, {
+    const response = await fetch(`http://localhost:${proxyPort}/__control`, {
       method: "POST",
       headers: { "Content-Type": "application/json" },
       body: JSON.stringify(body)
@@ -1066,7 +1166,7 @@ async function getRecordingsDir() {
   }
   const proxyPort = getProxyPort();
   try {
-    const response = await fetch(`http://127.0.0.1:${proxyPort}/__control`);
+    const response = await fetch(`http://localhost:${proxyPort}/__control`);
     if (response.ok) {
       const data = await response.json();
       if (data.recordingsDir) {
@@ -1119,6 +1219,14 @@ var playwrightProxy = {
     await page.setExtraHTTPHeaders({
       [RECORDING_ID_HEADER]: sessionId
     });
+    const fallbackProxyPort = getProxyPort();
+    await page.context().addCookies([
+      {
+        name: "proxy-recording-id",
+        value: encodeURIComponent(sessionId),
+        url: `http://localhost:${fallbackProxyPort}`
+      }
+    ]);
     await setProxyMode(mode, sessionId, timeout);
     if (clientSideOptions?.url) {
       await setupClientSideRecording(
@@ -1153,10 +1261,8 @@ var playwrightProxy = {
       // Ensure the handler applies to all matching requests
     );
     const context = page.context();
-    const contextId = context._guid || "default";
-    const handlerKey = `cleanup_${contextId}`;
-    if (!globalThis[handlerKey]) {
-      globalThis[handlerKey] = true;
+    if (!registeredContexts.has(context)) {
+      registeredContexts.add(context);
       context.on("close", async () => {
         try {
           await cleanupSession(sessionId);
@@ -1166,7 +1272,7 @@ var playwrightProxy = {
             error
           );
         } finally {
-          delete globalThis[handlerKey];
+          registeredContexts.delete(context);
         }
       });
     }