chore(extension): wrap CDP protocol (#604)

extension/src/background.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { Connection, debugLog } from './connection.js';
+import { RelayConnection, debugLog } from './relayConnection.js';
 
 /**
  * Simple Chrome Extension that pumps CDP messages between chrome.debugger and WebSocket
@@ -29,7 +29,7 @@ type PopupMessage = {
 type SendResponse = (response: any) => void;
 
 class TabShareExtension {
-  private activeConnections: Map<number, Connection>;
+  private activeConnections: Map<number, RelayConnection>;
 
   constructor() {
     this.activeConnections = new Map(); // tabId -> connection
@@ -75,7 +75,7 @@ class TabShareExtension {
     let activeTabId: number | null = null;
 
     if (isConnected) {
-      const [tabId] = this.activeConnections.entries().next().value as [number, Connection];
+      const [tabId] = this.activeConnections.entries().next().value as [number, RelayConnection];
       activeTabId = tabId;
 
       // Get tab info
@@ -123,14 +123,14 @@ class TabShareExtension {
       // Store connection
       this.activeConnections.set(tabId, info);
 
-      void this._updateUI(tabId, { text: '●', color: '#4CAF50', title: 'Disconnect from Playwright MCP' });
+      await this._updateUI(tabId, { text: '●', color: '#4CAF50', title: 'Disconnect from Playwright MCP' });
       debugLog(`Tab ${tabId} connected successfully`);
     } catch (error: any) {
       debugLog(`Failed to connect tab ${tabId}:`, error.message);
       await this._cleanupConnection(tabId);
 
       // Show error to user
-      void this._updateUI(tabId, { text: '!', color: '#F44336', title: `Connection failed: ${error.message}` });
+      await this._updateUI(tabId, { text: '!', color: '#F44336', title: `Connection failed: ${error.message}` });
 
       throw error;
     }
@@ -143,8 +143,8 @@ class TabShareExtension {
     await chrome.action.setTitle({ tabId, title });
   }
 
-  private _createConnection(tabId: number, socket: WebSocket): Connection {
-    const connection = new Connection(tabId, socket);
+  private _createConnection(tabId: number, socket: WebSocket): RelayConnection {
+    const connection = new RelayConnection(tabId, socket);
     socket.onclose = () => {
       debugLog(`WebSocket closed for tab ${tabId}`);
       void this.disconnectTab(tabId);

extension/src/relayConnection.ts

@@ -22,14 +22,21 @@ export function debugLog(...args: unknown[]): void {
   }
 }
 
-export type ProtocolCommand = {
+type ProtocolCommand = {
   id: number;
-  sessionId?: string;
   method: string;
   params?: any;
 };
 
-export class Connection {
+type ProtocolResponse = {
+  id?: number;
+  method?: string;
+  params?: any;
+  result?: any;
+  error?: string;
+};
+
+export class RelayConnection {
   private _debuggee: chrome.debugger.Debuggee;
   private _rootSessionId: string;
   private _ws: WebSocket;
@@ -61,21 +68,23 @@ export class Connection {
   private _onDebuggerEvent(source: chrome.debugger.DebuggerSession, method: string, params: any): void {
     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));
+    debugLog('Forwarding CDP event:', method, params);
+    const sessionId = source.sessionId || this._rootSessionId;
+    this._sendMessage({
+      method: 'forwardCDPEvent',
+      params: {
+        sessionId,
+        method,
+        params,
+      },
+    });
   }
 
   private _onDebuggerDetach(source: chrome.debugger.Debuggee, reason: string): void {
     if (source.tabId !== this._debuggee.tabId)
       return;
     this._sendMessage({
-      method: 'PWExtension.detachedFromTab',
+      method: 'detachedFromTab',
       params: {
         tabId: this._debuggee.tabId,
         reason,
@@ -99,29 +108,21 @@ export class Connection {
 
     debugLog('Received message:', message);
 
-    const sessionId = message.sessionId;
-    const response: { id: any; sessionId: any; result?: any; error?: { code: number; message: string } } = {
+    const response: ProtocolResponse = {
       id: message.id,
-      sessionId,
     };
     try {
-      if (message.method.startsWith('PWExtension.'))
-        response.result = await this._handleExtensionCommand(message);
-      else
-        response.result = await this._handleCDPCommand(message);
+      response.result = await this._handleCommand(message);
     } catch (error: any) {
-      debugLog('Error handling message:', error);
-      response.error = {
-        code: -32000,
-        message: error.message,
-      };
+      debugLog('Error handling command:', error);
+      response.error = error.message;
     }
     debugLog('Sending response:', response);
     this._sendMessage(response);
   }
 
-  private async _handleExtensionCommand(message: ProtocolCommand): Promise<any> {
-    if (message.method === 'PWExtension.attachToTab') {
+  private async _handleCommand(message: ProtocolCommand): Promise<any> {
+    if (message.method === 'attachToTab') {
       debugLog('Attaching debugger to tab:', this._debuggee);
       await chrome.debugger.attach(this._debuggee, '1.3');
       const result: any = await chrome.debugger.sendCommand(this._debuggee, 'Target.getTargetInfo');
@@ -130,26 +131,24 @@ export class Connection {
         targetInfo: result?.targetInfo,
       };
     }
-    if (message.method === 'PWExtension.detachFromTab') {
+    if (message.method === 'detachFromTab') {
       debugLog('Detaching debugger from tab:', this._debuggee);
-      await this.detachDebugger();
-      return;
+      return await this.detachDebugger();
+    }
+    if (message.method === 'forwardCDPCommand') {
+      const { sessionId, method, params } = message.params;
+      debugLog('CDP command:', method, params);
+      const debuggerSession: chrome.debugger.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
+      return await chrome.debugger.sendCommand(
+          debuggerSession,
+          method,
+          params
+      );
     }
-  }
-
-  private async _handleCDPCommand(message: ProtocolCommand): Promise<any> {
-    const sessionId = message.sessionId;
-    const debuggerSession: chrome.debugger.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;
   }
 
   private _sendError(code: number, message: string): void {
@@ -161,7 +160,7 @@ export class Connection {
     });
   }
 
-  private _sendMessage(message: object): void {
+  private _sendMessage(message: any): void {
     this._ws.send(JSON.stringify(message));
   }
 }

extension/tsconfig.json

@@ -7,7 +7,7 @@
     "module": "ESNext",
     "rootDir": "src",
     "outDir": "./lib",
-    "resolveJsonModule": true
+    "resolveJsonModule": true,
   },
   "include": [
     "src",

package.json

@@ -28,7 +28,7 @@
     "wtest": "playwright test --project=webkit",
     "etest": "playwright test --project=chromium-extension",
     "run-server": "node lib/browserServer.js",
-    "clean": "rm -rf lib",
+    "clean": "rm -rf lib && rm -rf extension/lib",
     "npm-publish": "npm run clean && npm run build && npm run test && npm publish"
   },
   "exports": {

src/cdpRelay.ts

@@ -26,7 +26,6 @@
 
 import { WebSocket, WebSocketServer } from 'ws';
 import http from 'node:http';
-import { EventEmitter } from 'node:events';
 import debug from 'debug';
 import { httpAddressToString } from './transport.js';
 
@@ -35,14 +34,23 @@ const debugLogger = debug('pw:mcp:relay');
 const CDP_PATH = '/cdp';
 const EXTENSION_PATH = '/extension';
 
-export type ProtocolCommand = {
+type CDPCommand = {
   id: number;
   sessionId?: string;
   method: string;
   params?: any;
 };
 
-export class CDPRelayServer extends EventEmitter {
+type CDPResponse = {
+  id?: number;
+  sessionId?: string;
+  method?: string;
+  params?: any;
+  result?: any;
+  error?: { code?: number; message: string };
+};
+
+export class CDPRelayServer {
   private _wss: WebSocketServer;
   private _playwrightSocket: WebSocket | null = null;
   private _extensionConnection: ExtensionConnection | null = null;
@@ -53,7 +61,6 @@ export class CDPRelayServer extends EventEmitter {
   } | undefined;
 
   constructor(server: http.Server) {
-    super();
     this._wss = new WebSocketServer({ server });
     this._wss.on('connection', this._onConnection.bind(this));
   }
@@ -65,9 +72,7 @@ export class CDPRelayServer extends EventEmitter {
 
   private _onConnection(ws: WebSocket, request: http.IncomingMessage): void {
     const url = new URL(`http://localhost${request.url}`);
-
     debugLogger(`New connection to ${url.pathname}`);
-
     if (url.pathname === CDP_PATH) {
       this._handlePlaywrightConnection(ws);
     } else if (url.pathname === EXTENSION_PATH) {
@@ -86,10 +91,8 @@ export class CDPRelayServer extends EventEmitter {
       debugLogger('Closing previous Playwright connection');
       this._playwrightSocket.close(1000, 'New connection established');
     }
-
     this._playwrightSocket = ws;
     debugLogger('Playwright MCP connected');
-
     ws.on('message', async data => {
       try {
         const message = JSON.parse(data.toString());
@@ -98,16 +101,13 @@ export class CDPRelayServer extends EventEmitter {
         debugLogger('Error parsing Playwright message:', error);
       }
     });
-
     ws.on('close', () => {
       if (this._playwrightSocket === ws) {
         void this._detachDebugger();
         this._playwrightSocket = null;
       }
-
       debugLogger('Playwright MCP disconnected');
     });
-
     ws.on('error', error => {
       debugLogger('Playwright WebSocket error:', error);
     });
@@ -115,7 +115,7 @@ export class CDPRelayServer extends EventEmitter {
 
   private async _detachDebugger() {
     this._connectionInfo = undefined;
-    await this._extensionConnection?.send('PWExtension.detachFromTab', {});
+    await this._extensionConnection?.send('detachFromTab', {});
   }
 
   private _handleExtensionConnection(ws: WebSocket): void {
@@ -129,14 +129,16 @@ export class CDPRelayServer extends EventEmitter {
     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;
-    }
-
+  private _handleExtensionMessage(method: string, params: any) {
     switch (method) {
-      case 'PWExtension.detachedFromTab':
+      case 'forwardCDPEvent':
+        this._sendToPlaywright({
+          sessionId: params.sessionId,
+          method: params.method,
+          params: params.params
+        });
+        break;
+      case 'detachedFromTab':
         debugLogger('← Debugger detached from tab:', params);
         this._connectionInfo = undefined;
         this._extensionConnection?.close();
@@ -145,10 +147,7 @@ export class CDPRelayServer extends EventEmitter {
     }
   }
 
-  /**
-   * Handle messages from Playwright MCP
-   */
-  private async _handlePlaywrightMessage(message: ProtocolCommand): Promise<void> {
+  private async _handlePlaywrightMessage(message: CDPCommand): Promise<void> {
     debugLogger('← Playwright:', `${message.method} (id=${message.id})`);
     if (!this._extensionConnection) {
       debugLogger('Extension not connected, sending error to Playwright');
@@ -158,29 +157,14 @@ export class CDPRelayServer extends EventEmitter {
       });
       return;
     }
-
-    // Handle Browser domain methods locally
-    if (message.method?.startsWith('Browser.')) {
-      await this._handleBrowserDomainMethod(message);
+    if (await this._interceptCDPCommand(message))
       return;
-    }
-
-    // Handle Target domain methods
-    if (message.method?.startsWith('Target.')) {
-      await this._handleTargetDomainMethod(message);
-      return;
-    }
-
-    // Forward other commands to extension
     await this._forwardToExtension(message);
   }
 
-  /**
-   * Handle Browser domain methods locally
-   */
-  private async _handleBrowserDomainMethod(message: any): Promise<void> {
+  private async _interceptCDPCommand(message: CDPCommand): Promise<boolean> {
     switch (message.method) {
-      case 'Browser.getVersion':
+      case 'Browser.getVersion': {
         this._sendToPlaywright({
           id: message.id,
           result: {
@@ -189,30 +173,19 @@ export class CDPRelayServer extends EventEmitter {
             userAgent: 'CDP-Bridge-Server/1.0.0',
           }
         });
-        break;
-
-      case 'Browser.setDownloadBehavior':
+        return true;
+      }
+      case 'Browser.setDownloadBehavior': {
         this._sendToPlaywright({
           id: message.id
         });
-        break;
-
-      default:
-        // Forward unknown Browser methods to extension
-        await this._forwardToExtension(message);
-    }
-  }
-
-  /**
-   * Handle Target domain methods
-   */
-  private async _handleTargetDomainMethod(message: any): Promise<void> {
-    switch (message.method) {
-      case 'Target.setAutoAttach':
+        return true;
+      }
+      case 'Target.setAutoAttach': {
         // Simulate auto-attach behavior with real target info
         if (!message.sessionId) {
-          this._connectionInfo = await this._extensionConnection!.send('PWExtension.attachToTab');
-          debugLogger('Simulating auto-attach for target:', JSON.stringify(message));
+          this._connectionInfo = await this._extensionConnection!.send('attachToTab');
+          debugLogger('Simulating auto-attach for target:', message);
           this._sendToPlaywright({
             method: 'Target.attachedToTarget',
             params: {
@@ -230,31 +203,27 @@ export class CDPRelayServer extends EventEmitter {
         } else {
           await this._forwardToExtension(message);
         }
-        break;
-
-      case 'Target.getTargetInfo':
+        return true;
+      }
+      case 'Target.getTargetInfo': {
         debugLogger('Target.getTargetInfo', message);
         this._sendToPlaywright({
           id: message.id,
           result: this._connectionInfo?.targetInfo
         });
-        break;
-
-      default:
-        await this._forwardToExtension(message);
+        return true;
+      }
     }
+    return false;
   }
 
-  private async _forwardToExtension(message: any): Promise<void> {
+  private async _forwardToExtension(message: CDPCommand): Promise<void> {
     try {
       if (!this._extensionConnection)
         throw new Error('Extension not connected');
-      const result = await this._extensionConnection.send(message.method, message.params, message.sessionId);
-      this._sendToPlaywright({
-        id: message.id,
-        sessionId: message.sessionId,
-        result,
-      });
+      const { id, sessionId, method, params } = message;
+      const result = await this._extensionConnection.send('forwardCDPCommand', { sessionId, method, params });
+      this._sendToPlaywright({ id, sessionId, result });
     } catch (e) {
       debugLogger('Error in the extension:', e);
       this._sendToPlaywright({
@@ -265,10 +234,7 @@ export class CDPRelayServer extends EventEmitter {
     }
   }
 
-  /**
-   * Forward message to Playwright
-   */
-  private _sendToPlaywright(message: any): void {
+  private _sendToPlaywright(message: CDPResponse): void {
     debugLogger('→ Playwright:', `${message.method ?? `response(id=${message.id})`}`);
     this._playwrightSocket?.send(JSON.stringify(message));
   }
@@ -306,7 +272,7 @@ class ExtensionConnection {
   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;
+  onmessage?: (method: string, params: any) => void;
   onclose?: (self: ExtensionConnection) => void;
 
   constructor(ws: WebSocket) {
@@ -361,7 +327,7 @@ class ExtensionConnection {
     } else if (object.id) {
       debugLogger('← Extension: unexpected response', object);
     } else {
-      this.onmessage?.(object.sessionId, object.method, object.params);
+      this.onmessage?.(object.method, object.params);
     }
   }