Enable BiDi support for connect-existing mode

src/firefox/core.ts

@@ -5,6 +5,7 @@
 import { Builder, Browser } from 'selenium-webdriver';
 import firefox from 'selenium-webdriver/firefox.js';
 import { spawn, type ChildProcess } from 'node:child_process';
+import WebSocket from 'ws';
 import type { FirefoxLaunchOptions } from './types.js';
 import { log, logDebug } from '../utils/logger.js';
 
@@ -107,18 +108,22 @@ class GeckodriverElement implements IElement {
  *
  * This exists because selenium-webdriver's Driver.createSession() tries to
  * auto-upgrade to BiDi WebSocket, which hangs when connecting to an existing
- * Firefox instance. By talking directly to geckodriver's HTTP API we avoid
- * the BiDi issue entirely.
+ * Firefox instance. By talking directly to geckodriver's HTTP API we can
+ * create the session without hanging, then manually open the BiDi WebSocket
+ * using the webSocketUrl returned in the session capabilities.
  */
 class GeckodriverHttpDriver implements IDriver {
   private baseUrl: string;
   private sessionId: string;
   private gdProcess: ChildProcess;
+  private webSocketUrl: string | null;
+  private bidiConnection: { subscribe: (event: string, contexts?: string[]) => Promise<void>; socket: WebSocket } | null = null;
 
-  constructor(baseUrl: string, sessionId: string, gdProcess: ChildProcess) {
+  constructor(baseUrl: string, sessionId: string, gdProcess: ChildProcess, webSocketUrl: string | null) {
     this.baseUrl = baseUrl;
     this.sessionId = sessionId;
     this.gdProcess = gdProcess;
+    this.webSocketUrl = webSocketUrl;
   }
 
   static async connect(marionettePort: number): Promise<GeckodriverHttpDriver> {
@@ -191,11 +196,11 @@ class GeckodriverHttpDriver implements IDriver {
 
     const baseUrl = `http://127.0.0.1:${port}`;
 
-    // Create a WebDriver session
+    // Create a WebDriver session with BiDi opt-in
     const resp = await fetch(`${baseUrl}/session`, {
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },
-      body: JSON.stringify({ capabilities: { alwaysMatch: {} } }),
+      body: JSON.stringify({ capabilities: { alwaysMatch: { webSocketUrl: true } } }),
     });
     const json = (await resp.json()) as {
       value: { sessionId: string; capabilities: Record<string, unknown> };
@@ -204,7 +209,14 @@ class GeckodriverHttpDriver implements IDriver {
       throw new Error(`Failed to create session: ${JSON.stringify(json)}`);
     }
 
-    return new GeckodriverHttpDriver(baseUrl, json.value.sessionId, gd);
+    const wsUrl = json.value.capabilities.webSocketUrl as string | undefined;
+    if (wsUrl) {
+      logDebug(`BiDi WebSocket URL: ${wsUrl}`);
+    } else {
+      logDebug('BiDi WebSocket URL not available (Firefox may not support it or Remote Agent is not running)');
+    }
+
+    return new GeckodriverHttpDriver(baseUrl, json.value.sessionId, gd, wsUrl ?? null);
   }
 
   private async cmd(method: string, path: string, body?: unknown): Promise<unknown> {
@@ -406,7 +418,67 @@ class GeckodriverHttpDriver implements IDriver {
     return builder;
   }
 
+  /**
+   * Return a BiDi handle compatible with the interface used by ConsoleEvents
+   * and NetworkEvents. Opens a WebSocket to Firefox's Remote Agent on first
+   * call, using the webSocketUrl returned in the session capabilities.
+   */
+  async getBidi(): Promise<{ subscribe: (event: string, contexts?: string[]) => Promise<void>; socket: WebSocket }> {
+    if (this.bidiConnection) return this.bidiConnection;
+    if (!this.webSocketUrl) {
+      throw new Error(
+        'BiDi is not available: no webSocketUrl in session capabilities. ' +
+        'Ensure Firefox was started with --remote-debugging-port.'
+      );
+    }
+
+    const ws = new WebSocket(this.webSocketUrl);
+    await new Promise<void>((resolve, reject) => {
+      ws.on('open', resolve);
+      ws.on('error', reject);
+    });
+    logDebug('BiDi WebSocket connected');
+
+    let cmdId = 0;
+    const subscribe = async (event: string, contexts?: string[]): Promise<void> => {
+      const msg: Record<string, unknown> = {
+        id: ++cmdId,
+        method: 'session.subscribe',
+        params: { events: [event] },
+      };
+      if (contexts) msg.params = { events: [event], contexts };
+      ws.send(JSON.stringify(msg));
+      // Wait for the response matching our command id
+      await new Promise<void>((resolve, reject) => {
+        const timeout = setTimeout(() => reject(new Error(`BiDi subscribe timeout for ${event}`)), 5000);
+        const onMsg = (data: WebSocket.Data) => {
+          try {
+            const payload = JSON.parse(data.toString());
+            if (payload.id === cmdId) {
+              clearTimeout(timeout);
+              ws.off('message', onMsg);
+              if (payload.error) {
+                reject(new Error(`BiDi subscribe error: ${payload.error}`));
+              } else {
+                resolve();
+              }
+            }
+          } catch { /* ignore parse errors from event messages */ }
+        };
+        ws.on('message', onMsg);
+      });
+      logDebug(`BiDi subscribed to ${event}`);
+    };
+
+    this.bidiConnection = { subscribe, socket: ws };
+    return this.bidiConnection;
+  }
+
   async quit(): Promise<void> {
+    if (this.bidiConnection) {
+      this.bidiConnection.socket.close();
+      this.bidiConnection = null;
+    }
     try {
       await this.cmd('DELETE', '');
     } catch {
@@ -417,6 +489,10 @@ class GeckodriverHttpDriver implements IDriver {
 
   /** Kill the geckodriver process without closing Firefox */
   kill(): void {
+    if (this.bidiConnection) {
+      this.bidiConnection.socket.close();
+      this.bidiConnection = null;
+    }
     this.gdProcess.kill();
   }
 }