fix(browser): harden extension relay reconnect race

Co-authored-by: Ho Lim <166576253+HOYALIM@users.noreply.github.com>

src/browser/extension-relay.test.ts

@@ -207,6 +207,51 @@ describe("chrome extension relay server", () => {
     expect(err.message).toContain("401");
   });
 
+  it("rejects a second live extension connection with 409", async () => {
+    const port = await getFreePort();
+    cdpUrl = `http://127.0.0.1:${port}`;
+    await ensureChromeExtensionRelayServer({ cdpUrl });
+
+    const ext1 = new WebSocket(`ws://127.0.0.1:${port}/extension`, {
+      headers: relayAuthHeaders(`ws://127.0.0.1:${port}/extension`),
+    });
+    await waitForOpen(ext1);
+
+    const ext2 = new WebSocket(`ws://127.0.0.1:${port}/extension`, {
+      headers: relayAuthHeaders(`ws://127.0.0.1:${port}/extension`),
+    });
+    const err = await waitForError(ext2);
+    expect(err.message).toContain("409");
+
+    ext1.close();
+  });
+
+  it("allows immediate reconnect when prior extension socket is closing", async () => {
+    const port = await getFreePort();
+    cdpUrl = `http://127.0.0.1:${port}`;
+    await ensureChromeExtensionRelayServer({ cdpUrl });
+
+    const ext1 = new WebSocket(`ws://127.0.0.1:${port}/extension`, {
+      headers: relayAuthHeaders(`ws://127.0.0.1:${port}/extension`),
+    });
+    await waitForOpen(ext1);
+    const ext1Closed = new Promise<void>((resolve) => ext1.once("close", () => resolve()));
+
+    ext1.close();
+    const ext2 = new WebSocket(`ws://127.0.0.1:${port}/extension`, {
+      headers: relayAuthHeaders(`ws://127.0.0.1:${port}/extension`),
+    });
+    await waitForOpen(ext2);
+    await ext1Closed;
+
+    const status = (await fetch(`${cdpUrl}/extension/status`).then((r) => r.json())) as {
+      connected?: boolean;
+    };
+    expect(status.connected).toBe(true);
+
+    ext2.close();
+  });
+
   it("accepts extension websocket access with relay token query param", async () => {
     const port = await getFreePort();
     cdpUrl = `http://127.0.0.1:${port}`;

src/browser/extension-relay.ts

@@ -223,6 +223,7 @@ export async function ensureChromeExtensionRelayServer(opts: {
   let extensionWs: WebSocket | null = null;
   const cdpClients = new Set<WebSocket>();
   const connectedTargets = new Map<string, ConnectedTarget>();
+  const extensionConnected = () => extensionWs?.readyState === WebSocket.OPEN;
 
   const pendingExtension = new Map<
     number,
@@ -386,7 +387,7 @@ export async function ensureChromeExtensionRelayServer(opts: {
 
     if (path === "/extension/status") {
       res.writeHead(200, { "Content-Type": "application/json" });
-      res.end(JSON.stringify({ connected: Boolean(extensionWs) }));
+      res.end(JSON.stringify({ connected: extensionConnected() }));
       return;
     }
 
@@ -403,7 +404,7 @@ export async function ensureChromeExtensionRelayServer(opts: {
         "Protocol-Version": "1.3",
       };
       // Only advertise the WS URL if a real extension is connected.
-      if (extensionWs) {
+      if (extensionConnected()) {
         payload.webSocketDebuggerUrl = cdpWsUrl;
       }
       res.writeHead(200, { "Content-Type": "application/json" });
@@ -504,10 +505,19 @@ export async function ensureChromeExtensionRelayServer(opts: {
         rejectUpgrade(socket, 401, "Unauthorized");
         return;
       }
-      if (extensionWs) {
+      if (extensionConnected()) {
         rejectUpgrade(socket, 409, "Extension already connected");
         return;
       }
+      // MV3 worker reconnect races can leave a stale non-OPEN socket reference.
+      if (extensionWs && extensionWs.readyState !== WebSocket.OPEN) {
+        try {
+          extensionWs.terminate();
+        } catch {
+          // ignore
+        }
+        extensionWs = null;
+      }
       wssExtension.handleUpgrade(req, socket, head, (ws) => {
         wssExtension.emit("connection", ws, req);
       });
@@ -520,7 +530,7 @@ export async function ensureChromeExtensionRelayServer(opts: {
         rejectUpgrade(socket, 401, "Unauthorized");
         return;
       }
-      if (!extensionWs) {
+      if (!extensionConnected()) {
         rejectUpgrade(socket, 503, "Extension not connected");
         return;
       }
@@ -544,6 +554,9 @@ export async function ensureChromeExtensionRelayServer(opts: {
     }, 5000);
 
     ws.on("message", (data) => {
+      if (extensionWs !== ws) {
+        return;
+      }
       let parsed: ExtensionMessage | null = null;
       try {
         parsed = JSON.parse(rawDataToString(data)) as ExtensionMessage;
@@ -645,6 +658,9 @@ export async function ensureChromeExtensionRelayServer(opts: {
 
     ws.on("close", () => {
       clearInterval(ping);
+      if (extensionWs !== ws) {
+        return;
+      }
       extensionWs = null;
       for (const [, pending] of pendingExtension) {
         clearTimeout(pending.timer);
@@ -681,7 +697,7 @@ export async function ensureChromeExtensionRelayServer(opts: {
         return;
       }
 
-      if (!extensionWs) {
+      if (!extensionConnected()) {
         sendResponseToCdp(ws, {
           id: cmd.id,
           sessionId: cmd.sessionId,
@@ -779,7 +795,7 @@ export async function ensureChromeExtensionRelayServer(opts: {
     port,
     baseUrl,
     cdpWsUrl: `ws://${host}:${port}/cdp`,
-    extensionConnected: () => Boolean(extensionWs),
+    extensionConnected,
     stop: async () => {
       relayRuntimeByPort.delete(port);
       try {