chore(extension): support reconnect, implement relay-extension protocol (#602)

extension/background.js

@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+import { Connection } from './connection.js';
+
 /**
  * Simple Chrome Extension that pumps CDP messages between chrome.debugger and WebSocket
  */
@@ -21,7 +23,7 @@
 // @ts-check
 
 function debugLog(...args) {
-  const enabled = false;
+  const enabled = true;
   if (enabled) {
     console.log('[Extension]', ...args);
   }
@@ -116,173 +118,49 @@ class TabShareExtension {
   async connectTab(tabId, bridgeUrl) {
     try {
       debugLog(`Connecting tab ${tabId} to bridge at ${bridgeUrl}`);
-
-      // Attach chrome debugger
-      const debuggee = { tabId };
-      await chrome.debugger.attach(debuggee, '1.3');
-
-      if (chrome.runtime.lastError)
-        throw new Error(chrome.runtime.lastError.message);
-      const targetInfo = /** @type {any} */ (await chrome.debugger.sendCommand(debuggee, 'Target.getTargetInfo'));
-      debugLog('Target info:', targetInfo);
-
       // Connect to bridge server
       const socket = new WebSocket(bridgeUrl);
-
-      const connection = {
-        debuggee,
-        socket,
-        tabId,
-        sessionId: `pw-tab-${tabId}`
-      };
-
       await new Promise((resolve, reject) => {
-        socket.onopen = () => {
+        socket.onopen = () => resolve(undefined);
-          debugLog(`WebSocket connected for tab ${tabId}`);
-          // Send initial connection info to bridge
-          socket.send(JSON.stringify({
-            type: 'connection_info',
-            sessionId: connection.sessionId,
-            targetInfo: targetInfo?.targetInfo
-          }));
-          resolve(undefined);
-        };
         socket.onerror = reject;
         setTimeout(() => reject(new Error('Connection timeout')), 5000);
       });
 
-      // Set up message handling
+      const info = this._createConnection(tabId, socket);
-      this.setupMessageHandling(connection);
-
       // Store connection
-      this.activeConnections.set(tabId, connection);
+      this.activeConnections.set(tabId, info);
-
-      // Update UI
-      chrome.action.setBadgeText({ tabId, text: '●' });
-      chrome.action.setBadgeBackgroundColor({ tabId, color: '#4CAF50' });
-      chrome.action.setTitle({ tabId, title: 'Disconnect from Playwright MCP' });
 
+      this._updateUI(tabId, { text: '●', color: '#4CAF50', title: 'Disconnect from Playwright MCP' });
       debugLog(`Tab ${tabId} connected successfully`);
-
     } catch (error) {
       debugLog(`Failed to connect tab ${tabId}:`, error.message);
-      await this.cleanupConnection(tabId);
+      await this._cleanupConnection(tabId);
 
       // Show error to user
-      chrome.action.setBadgeText({ tabId, text: '!' });
+      this._updateUI(tabId, { text: '!', color: '#F44336', title: `Connection failed: ${error.message}` });
-      chrome.action.setBadgeBackgroundColor({ tabId, color: '#F44336' });
-      chrome.action.setTitle({ tabId, title: `Connection failed: ${error.message}` });
 
       throw error; // Re-throw for popup to handle
     }
   }
 
-  /**
+  _updateUI(tabId, { text, color, title }) {
-   * Set up bidirectional message handling between debugger and WebSocket
+    chrome.action.setBadgeText({ tabId, text });
-   * @param {Object} connection
+    if (color)
-   */
+      chrome.action.setBadgeBackgroundColor({ tabId, color });
-  setupMessageHandling(connection) {
+    chrome.action.setTitle({ tabId, title });
-    const { debuggee, socket, tabId, sessionId: rootSessionId } = connection;
+  }
 
-    // WebSocket -> chrome.debugger
+  _createConnection(tabId, socket) {
-    socket.onmessage = async (event) => {
+    const connection = new Connection(tabId, socket);
-      let message;
-      try {
-        message = JSON.parse(event.data);
-      } catch (error) {
-        debugLog('Error parsing message:', error);
-        socket.send(JSON.stringify({
-          error: {
-            code: -32700,
-            message: `Error parsing message: ${error.message}`
-          }
-        }));
-        return;
-      }
-
-      try {
-        debugLog('Received from bridge:', message);
-
-        const debuggerSession = { ...debuggee };
-        const sessionId = message.sessionId;
-        // Pass session id, unless it's the root session.
-        if (sessionId && sessionId !== rootSessionId)
-          debuggerSession.sessionId = sessionId;
-
-        // Forward CDP command to chrome.debugger
-        const result = await chrome.debugger.sendCommand(
-          debuggerSession,
-          message.method,
-          message.params || {}
-        );
-
-        // Send response back to bridge
-        const response = {
-          id: message.id,
-          sessionId,
-          result
-        };
-
-        if (chrome.runtime.lastError) {
-          response.error = {
-            code: -32000,
-            message: chrome.runtime.lastError.message,
-          };
-        }
-
-        socket.send(JSON.stringify(response));
-      } catch (error) {
-        debugLog('Error processing WebSocket message:', error);
-        const response = {
-          id: message.id,
-          sessionId: message.sessionId,
-          error: {
-            code: -32000,
-            message: error.message,
-          },
-        };
-        socket.send(JSON.stringify(response));
-      }
-    };
-
-    // chrome.debugger events -> WebSocket
-    const eventListener = (source, method, params) => {
-      if (source.tabId === tabId && socket.readyState === WebSocket.OPEN) {
-        // If the sessionId is not provided, use the root sessionId.
-        const event = {
-          sessionId: source.sessionId || rootSessionId,
-          method,
-          params,
-        };
-        debugLog('Forwarding CDP event:', event);
-        socket.send(JSON.stringify(event));
-      }
-    };
-
-    const detachListener = (source, reason) => {
-      if (source.tabId === tabId) {
-        debugLog(`Debugger detached from tab ${tabId}, reason: ${reason}`);
-        this.disconnectTab(tabId);
-      }
-    };
-
-    // Store listeners for cleanup
-    connection.eventListener = eventListener;
-    connection.detachListener = detachListener;
-
-    chrome.debugger.onEvent.addListener(eventListener);
-    chrome.debugger.onDetach.addListener(detachListener);
-
-    // Handle WebSocket close
     socket.onclose = () => {
       debugLog(`WebSocket closed for tab ${tabId}`);
       this.disconnectTab(tabId);
     };
-
     socket.onerror = (error) => {
       debugLog(`WebSocket error for tab ${tabId}:`, error);
       this.disconnectTab(tabId);
     };
+    return { connection };
   }
 
   /**
@@ -290,12 +168,8 @@ class TabShareExtension {
    * @param {number} tabId
    */
   async disconnectTab(tabId) {
-    await this.cleanupConnection(tabId);
+    await this._cleanupConnection(tabId);
-
+    this._updateUI(tabId, { text: '', color: null, title: 'Share tab with Playwright MCP' });
-    // Update UI
-    chrome.action.setBadgeText({ tabId, text: '' });
-    chrome.action.setTitle({ tabId, title: 'Share tab with Playwright MCP' });
-
     debugLog(`Tab ${tabId} disconnected`);
   }
 
@@ -303,31 +177,21 @@ class TabShareExtension {
    * Clean up connection resources
    * @param {number} tabId
    */
-  async cleanupConnection(tabId) {
+  async _cleanupConnection(tabId) {
-    const connection = this.activeConnections.get(tabId);
+    const info = this.activeConnections.get(tabId);
-    if (!connection) return;
+    if (!info) return;
-
+    this.activeConnections.delete(tabId);
-    // Remove listeners
-    if (connection.eventListener) {
-      chrome.debugger.onEvent.removeListener(connection.eventListener);
-    }
-    if (connection.detachListener) {
-      chrome.debugger.onDetach.removeListener(connection.detachListener);
-    }
 
     // Close WebSocket
-    if (connection.socket && connection.socket.readyState === WebSocket.OPEN) {
+    info.connection.close();
-      connection.socket.close();
-    }
 
     // Detach debugger
     try {
-      await chrome.debugger.detach(connection.debuggee);
+      await info.connection.detachDebugger();
     } catch (error) {
       // Ignore detach errors - might already be detached
+      debugLog('Error while detaching debugger:', error);
     }
-
-    this.activeConnections.delete(tabId);
   }
 
   /**
@@ -335,9 +199,8 @@ class TabShareExtension {
    * @param {number} tabId
    */
   async onTabRemoved(tabId) {
-    if (this.activeConnections.has(tabId)) {
+    if (this.activeConnections.has(tabId))
-      await this.cleanupConnection(tabId);
+      await this._cleanupConnection(tabId);
-    }
   }
 }
 

extension/connection.js

@@ -0,0 +1,169 @@
+/**
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// @ts-check
+
+function debugLog(...args) {
+  const enabled = true;
+  if (enabled) {
+    console.log('[Extension]', ...args);
+  }
+}
+
+export class Connection {
+  /**
+   * @param {number} tabId
+   * @param {WebSocket} ws
+   */
+  constructor(tabId, ws) {
+    /** @type {chrome.debugger.Debuggee} */
+    this._debuggee = { tabId };
+    this._rootSessionId = `pw-tab-${tabId}`;
+    this._ws = ws;
+    this._ws.onmessage = this._onMessage.bind(this);
+    // Store listeners for cleanup
+    this._eventListener = this._onDebuggerEvent.bind(this);
+    this._detachListener = this._onDebuggerDetach.bind(this);
+    chrome.debugger.onEvent.addListener(this._eventListener);
+    chrome.debugger.onDetach.addListener(this._detachListener);
+  }
+
+  close(message) {
+    chrome.debugger.onEvent.removeListener(this._eventListener);
+    chrome.debugger.onDetach.removeListener(this._detachListener);
+    this._ws.close(1000, message || 'Connection closed');
+  }
+
+  async detachDebugger() {
+    await chrome.debugger.detach(this._debuggee);
+  }
+
+  _onDebuggerEvent(source, method, params) {
+    if (source.tabId !== this._debuggee.tabId)
+      return;
+    // If the sessionId is not provided, use the root sessionId.
+    const event = {
+      sessionId: source.sessionId || this._rootSessionId,
+      method,
+      params,
+    };
+    debugLog('Forwarding CDP event:', event);
+    this._ws.send(JSON.stringify(event));
+  }
+
+  _onDebuggerDetach(source, reason) {
+    if (source.tabId !== this._debuggee.tabId)
+      return;
+    this._sendMessage({
+      method: 'PWExtension.detachedFromTab',
+      params: {
+        tabId: this._debuggee.tabId,
+        reason,
+      },
+    });
+  }
+
+  /**
+   * @param {MessageEvent} event
+   */
+  _onMessage(event) {
+    this._onMessageAsync(event).catch(e => debugLog('Error handling message:', e));
+  }
+
+  async _onMessageAsync(event) {
+    /** @type {import('../src/cdpRelay').ProtocolCommand} */
+    let message;
+    try {
+      message = JSON.parse(event.data);
+    } catch (error) {
+      debugLog('Error parsing message:', error);
+      this._sendError(-32700, `Error parsing message: ${error.message}`);
+      return;
+    }
+
+    debugLog('Received message:', message);
+
+    const sessionId = message.sessionId;
+    const response = {
+      id: message.id,
+      sessionId,
+    };
+    try {
+      if (message.method.startsWith('PWExtension.'))
+        response.result = await this._handleExtensionCommand(message);
+      else
+        response.result = await this._handleCDPCommand(message);
+    } catch (error) {
+      debugLog('Error handling message:', error);
+      response.error = {
+        code: -32000,
+        message: error.message,
+      };
+    }
+    debugLog('Sending response:', response);
+    this._sendMessage(response);
+  }
+
+  async _handleExtensionCommand(message) {
+    if (message.method === 'PWExtension.attachToTab') {
+      debugLog('Attaching debugger to tab:', this._debuggee);
+      await chrome.debugger.attach(this._debuggee, '1.3');
+      const result = /** @type {any} */ (await chrome.debugger.sendCommand(this._debuggee, 'Target.getTargetInfo'));
+      return {
+        sessionId: this._rootSessionId,
+        targetInfo: result.targetInfo,
+      };
+    }
+    if (message.method === 'PWExtension.detachFromTab') {
+      debugLog('Detaching debugger from tab:', this._debuggee);
+      await this.detachDebugger();
+      return;
+    }
+  }
+
+  async _handleCDPCommand(message) {
+    const sessionId = message.sessionId;
+    /** @type {chrome.debugger.DebuggerSession} */
+    const debuggerSession = { ...this._debuggee };
+    // Pass session id, unless it's the root session.
+    if (sessionId && sessionId !== this._rootSessionId)
+      debuggerSession.sessionId = sessionId;
+    // Forward CDP command to chrome.debugger
+    const result = await chrome.debugger.sendCommand(
+      debuggerSession,
+      message.method,
+      message.params
+    );
+    return result;
+  }
+
+  _sendError(code, message) {
+    this._sendMessage({
+      error: {
+        // @ts-ignore
+        code,
+        message
+      }
+    });
+  }
+
+  /**
+   * @param {import('../src/cdpRelay').ProtocolResponse} message
+   */
+  _sendMessage(message) {
+    this._ws.send(JSON.stringify(message));
+  }
+}

src/cdpRelay.ts

@@ -35,12 +35,20 @@ const debugLogger = debug('pw:mcp:relay');
 const CDP_PATH = '/cdp';
 const EXTENSION_PATH = '/extension';
 
+export type ProtocolCommand = {
+  id: number;
+  sessionId?: string;
+  method: string;
+  params?: any;
+};
+
 export class CDPRelayServer extends EventEmitter {
   private _wss: WebSocketServer;
   private _playwrightSocket: WebSocket | null = null;
-  private _extensionSocket: WebSocket | null = null;
+  private _extensionConnection: ExtensionConnection | null = null;
   private _connectionInfo: {
     targetInfo: any;
+    // Page sessionId that should be used by this connection.
     sessionId: string;
   } | undefined;
 
@@ -52,7 +60,7 @@ export class CDPRelayServer extends EventEmitter {
 
   stop(): void {
     this._playwrightSocket?.close();
-    this._extensionSocket?.close();
+    this._extensionConnection?.close();
   }
 
   private _onConnection(ws: WebSocket, request: http.IncomingMessage): void {
@@ -82,18 +90,20 @@ export class CDPRelayServer extends EventEmitter {
     this._playwrightSocket = ws;
     debugLogger('Playwright MCP connected');
 
-    ws.on('message', data => {
+    ws.on('message', async data => {
       try {
         const message = JSON.parse(data.toString());
-        this._handlePlaywrightMessage(message);
+        await this._handlePlaywrightMessage(message);
       } catch (error) {
         debugLogger('Error parsing Playwright message:', error);
       }
     });
 
     ws.on('close', () => {
-      if (this._playwrightSocket === ws)
+      if (this._playwrightSocket === ws) {
+        void this._detachDebugger();
         this._playwrightSocket = null;
+      }
 
       debugLogger('Playwright MCP disconnected');
     });
@@ -103,87 +113,72 @@ export class CDPRelayServer extends EventEmitter {
     });
   }
 
-  /**
+  private async _detachDebugger() {
-   * Handle Extension connection - forwards to chrome.debugger
+    this._connectionInfo = undefined;
-   */
+    await this._extensionConnection?.send('PWExtension.detachFromTab', {});
+  }
+
   private _handleExtensionConnection(ws: WebSocket): void {
-    if (this._extensionSocket?.readyState === WebSocket.OPEN) {
+    if (this._extensionConnection)
-      debugLogger('Closing previous extension connection');
+      this._extensionConnection.close('New connection established');
-      this._extensionSocket.close(1000, 'New connection established');
+    this._extensionConnection = new ExtensionConnection(ws);
+    this._extensionConnection.onclose = c => {
+      if (this._extensionConnection === c)
+        this._extensionConnection = null;
+    };
+    this._extensionConnection.onmessage = this._handleExtensionMessage.bind(this);
+  }
+
+  private _handleExtensionMessage(sessionId: string | undefined, method: string, params: any) {
+    if (!method.startsWith('PWExtension.')) {
+      this._sendToPlaywright({ sessionId, method, params });
+      return;
     }
 
-    this._extensionSocket = ws;
+    switch (method) {
-    debugLogger('Extension connected');
+      case 'PWExtension.detachedFromTab':
-
+        debugLogger('← Debugger detached from tab:', params);
-    ws.on('message', data => {
+        this._connectionInfo = undefined;
-      try {
+        this._extensionConnection?.close();
-        const message = JSON.parse(data.toString());
+        this._extensionConnection = null;
-        this._handleExtensionMessage(message);
+        break;
-      } catch (error) {
+    }
-        debugLogger('Error parsing extension message:', error);
-      }
-    });
-
-    ws.on('close', () => {
-      if (this._extensionSocket === ws)
-        this._extensionSocket = null;
-
-      debugLogger('Extension disconnected');
-    });
-
-    ws.on('error', error => {
-      debugLogger('Extension WebSocket error:', error);
-    });
   }
 
   /**
    * Handle messages from Playwright MCP
    */
-  private _handlePlaywrightMessage(message: any): void {
+  private async _handlePlaywrightMessage(message: ProtocolCommand): Promise<void> {
-    debugLogger('← Playwright:', message.method || `response(${message.id})`);
+    debugLogger('← Playwright:', `${message.method} (id=${message.id})`);
+    if (!this._extensionConnection) {
+      debugLogger('Extension not connected, sending error to Playwright');
+      this._sendToPlaywright({
+        id: message.id,
+        error: { message: 'Extension not connected' }
+      });
+      return;
+    }
 
     // Handle Browser domain methods locally
     if (message.method?.startsWith('Browser.')) {
-      this._handleBrowserDomainMethod(message);
+      await this._handleBrowserDomainMethod(message);
       return;
     }
 
     // Handle Target domain methods
     if (message.method?.startsWith('Target.')) {
-      this._handleTargetDomainMethod(message);
+      await this._handleTargetDomainMethod(message);
       return;
     }
 
     // Forward other commands to extension
-    if (message.method)
+    await this._forwardToExtension(message);
-      this._forwardToExtension(message);
-
-  }
-
-  /**
-   * Handle messages from Extension
-   */
-  private _handleExtensionMessage(message: any): void {
-    // Handle connection info from extension
-    if (message.type === 'connection_info') {
-      debugLogger('← Extension connected to tab:', message);
-      this._connectionInfo = {
-        targetInfo: message.targetInfo,
-        // Page sessionId that should be used by this connection.
-        sessionId: message.sessionId
-      };
-      return;
-    }
-
-    // CDP event from extension
-    debugLogger(`← Extension message: ${message.method ?? (message.id && `response(id=${message.id})`) ?? 'unknown'}`);
-    this._sendToPlaywright(message);
   }
 
   /**
    * Handle Browser domain methods locally
    */
-  private _handleBrowserDomainMethod(message: any): void {
+  private async _handleBrowserDomainMethod(message: any): Promise<void> {
     switch (message.method) {
       case 'Browser.getVersion':
         this._sendToPlaywright({
@@ -198,82 +193,75 @@ export class CDPRelayServer extends EventEmitter {
 
       case 'Browser.setDownloadBehavior':
         this._sendToPlaywright({
-          id: message.id,
+          id: message.id
-          result: {}
         });
         break;
 
       default:
         // Forward unknown Browser methods to extension
-        this._forwardToExtension(message);
+        await this._forwardToExtension(message);
     }
   }
 
   /**
    * Handle Target domain methods
    */
-  private _handleTargetDomainMethod(message: any): void {
+  private async _handleTargetDomainMethod(message: any): Promise<void> {
     switch (message.method) {
       case 'Target.setAutoAttach':
         // Simulate auto-attach behavior with real target info
-        if (this._connectionInfo && !message.sessionId) {
+        if (!message.sessionId) {
+          this._connectionInfo = await this._extensionConnection!.send('PWExtension.attachToTab');
           debugLogger('Simulating auto-attach for target:', JSON.stringify(message));
           this._sendToPlaywright({
             method: 'Target.attachedToTarget',
             params: {
-              sessionId: this._connectionInfo.sessionId,
+              sessionId: this._connectionInfo!.sessionId,
               targetInfo: {
-                ...this._connectionInfo.targetInfo,
+                ...this._connectionInfo!.targetInfo,
                 attached: true,
               },
               waitingForDebugger: false
             }
           });
           this._sendToPlaywright({
-            id: message.id,
+            id: message.id
-            result: {}
           });
         } else {
-          this._forwardToExtension(message);
+          await this._forwardToExtension(message);
         }
         break;
 
-      case 'Target.getTargets':
+      case 'Target.getTargetInfo':
-        const targetInfos = [];
+        debugLogger('Target.getTargetInfo', message);
-
-        if (this._connectionInfo) {
-          targetInfos.push({
-            ...this._connectionInfo.targetInfo,
-            attached: true,
-          });
-        }
-
         this._sendToPlaywright({
           id: message.id,
-          result: { targetInfos }
+          result: this._connectionInfo?.targetInfo
         });
         break;
 
       default:
-        this._forwardToExtension(message);
+        await this._forwardToExtension(message);
     }
   }
 
-  /**
+  private async _forwardToExtension(message: any): Promise<void> {
-   * Forward message to extension
+    try {
-   */
+      if (!this._extensionConnection)
-  private _forwardToExtension(message: any): void {
+        throw new Error('Extension not connected');
-    if (this._extensionSocket?.readyState === WebSocket.OPEN) {
+      const result = await this._extensionConnection.send(message.method, message.params, message.sessionId);
-      debugLogger('→ Extension:', message.method || `command(${message.id})`);
+      this._sendToPlaywright({
-      this._extensionSocket.send(JSON.stringify(message));
+        id: message.id,
-    } else {
+        sessionId: message.sessionId,
-      debugLogger('Extension not connected, cannot forward message');
+        result,
-      if (message.id) {
+      });
-        this._sendToPlaywright({
+    } catch (e) {
-          id: message.id,
+      debugLogger('Error in the extension:', e);
-          error: { message: 'Extension not connected' }
+      this._sendToPlaywright({
-        });
+        id: message.id,
-      }
+        sessionId: message.sessionId,
+        error: { message: (e as Error).message }
+      });
     }
   }
 
@@ -281,10 +269,8 @@ export class CDPRelayServer extends EventEmitter {
    * Forward message to Playwright
    */
   private _sendToPlaywright(message: any): void {
-    if (this._playwrightSocket?.readyState === WebSocket.OPEN) {
+    debugLogger('→ Playwright:', `${message.method ?? `response(id=${message.id})`}`);
-      debugLogger('→ Playwright:', JSON.stringify(message));
+    this._playwrightSocket?.send(JSON.stringify(message));
-      this._playwrightSocket.send(JSON.stringify(message));
-    }
   }
 }
 
@@ -292,7 +278,6 @@ export async function startCDPRelayServer(httpServer: http.Server) {
   const wsAddress = httpAddressToString(httpServer.address()).replace(/^http/, 'ws');
   const cdpRelayServer = new CDPRelayServer(httpServer);
   process.on('exit', () => cdpRelayServer.stop());
-  // eslint-disable-next-line no-console
   console.error(`CDP relay server started on ${wsAddress}${EXTENSION_PATH} - Connect to it using the browser extension.`);
   const cdpEndpoint = `${wsAddress}${CDP_PATH}`;
   return cdpEndpoint;
@@ -315,3 +300,84 @@ if (import.meta.url === `file://${process.argv[1]}`) {
     process.exit(0);
   });
 }
+
+class ExtensionConnection {
+  private readonly _ws: WebSocket;
+  private readonly _callbacks = new Map<number, { resolve: (o: any) => void, reject: (e: Error) => void }>();
+  private _lastId = 0;
+
+  onmessage?: (sessionId: string | undefined, method: string, params: any) => void;
+  onclose?: (self: ExtensionConnection) => void;
+
+  constructor(ws: WebSocket) {
+    this._ws = ws;
+    this._ws.on('message', this._onMessage.bind(this));
+    this._ws.on('close', this._onClose.bind(this));
+    this._ws.on('error', this._onError.bind(this));
+  }
+
+  async send(method: string, params?: any, sessionId?: string): Promise<any> {
+    if (this._ws.readyState !== WebSocket.OPEN)
+      throw new Error('WebSocket closed');
+    const id = ++this._lastId;
+    this._ws.send(JSON.stringify({ id, method, params, sessionId }));
+    return new Promise((resolve, reject) => {
+      this._callbacks.set(id, { resolve, reject });
+    });
+  }
+
+  close(message?: string) {
+    debugLogger('closing extension connection:', message);
+    this._ws.close(1000, message ?? 'Connection closed');
+    this.onclose?.(this);
+  }
+
+  private _onMessage(event: WebSocket.RawData) {
+    const eventData = event.toString();
+    let parsedJson;
+    try {
+      parsedJson = JSON.parse(eventData);
+    } catch (e: any) {
+      debugLogger(`<closing ws> Closing websocket due to malformed JSON. eventData=${eventData} e=${e?.message}`);
+      this._ws.close();
+      return;
+    }
+    try {
+      this._handleParsedMessage(parsedJson);
+    } catch (e: any) {
+      debugLogger(`<closing ws> Closing websocket due to failed onmessage callback. eventData=${eventData} e=${e?.message}`);
+      this._ws.close();
+    }
+  }
+
+  private _handleParsedMessage(object: any) {
+    if (object.id && this._callbacks.has(object.id)) {
+      const callback = this._callbacks.get(object.id)!;
+      this._callbacks.delete(object.id);
+      if (object.error)
+        callback.reject(new Error(object.error.message));
+      else
+        callback.resolve(object.result);
+    } else if (object.id) {
+      debugLogger('← Extension: unexpected response', object);
+    } else {
+      this.onmessage?.(object.sessionId, object.method, object.params);
+    }
+  }
+
+  private _onClose(event: WebSocket.CloseEvent) {
+    debugLogger(`<ws closed> code=${event.code} reason=${event.reason}`);
+    this._dispose();
+  }
+
+  private _onError(event: WebSocket.ErrorEvent) {
+    debugLogger(`<ws error> message=${event.message} type=${event.type} target=${event.target}`);
+    this._dispose();
+  }
+
+  private _dispose() {
+    for (const callback of this._callbacks.values())
+      callback.reject(new Error('WebSocket closed'));
+    this._callbacks.clear();
+  }
+}