From 5df6c2431b4e8e97d79cf84a00bfe0849aea1321 Mon Sep 17 00:00:00 2001 From: Yury Semikhatsky Date: Thu, 26 Jun 2025 11:12:23 -0700 Subject: [PATCH] chore(extension): support reconnect, implement relay-extension protocol (#602) --- extension/background.js | 195 +++++----------------------- extension/connection.js | 169 +++++++++++++++++++++++++ src/cdpRelay.ts | 272 +++++++++++++++++++++++++--------------- 3 files changed, 367 insertions(+), 269 deletions(-) create mode 100644 extension/connection.js diff --git a/extension/background.js b/extension/background.js index 0142809..2e2d2df 100644 --- a/extension/background.js +++ b/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 = () => { - 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.onopen = () => resolve(undefined); socket.onerror = reject; setTimeout(() => reject(new Error('Connection timeout')), 5000); }); - // Set up message handling - this.setupMessageHandling(connection); - + const info = this._createConnection(tabId, socket); // Store connection - this.activeConnections.set(tabId, connection); - - // Update UI - chrome.action.setBadgeText({ tabId, text: '●' }); - chrome.action.setBadgeBackgroundColor({ tabId, color: '#4CAF50' }); - chrome.action.setTitle({ tabId, title: 'Disconnect from Playwright MCP' }); + this.activeConnections.set(tabId, info); + 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: '!' }); - chrome.action.setBadgeBackgroundColor({ tabId, color: '#F44336' }); - chrome.action.setTitle({ tabId, title: `Connection failed: ${error.message}` }); + this._updateUI(tabId, { text: '!', color: '#F44336', title: `Connection failed: ${error.message}` }); throw error; // Re-throw for popup to handle } } - /** - * Set up bidirectional message handling between debugger and WebSocket - * @param {Object} connection - */ - setupMessageHandling(connection) { - const { debuggee, socket, tabId, sessionId: rootSessionId } = connection; + _updateUI(tabId, { text, color, title }) { + chrome.action.setBadgeText({ tabId, text }); + if (color) + chrome.action.setBadgeBackgroundColor({ tabId, color }); + chrome.action.setTitle({ tabId, title }); + } - // WebSocket -> chrome.debugger - socket.onmessage = async (event) => { - 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 + _createConnection(tabId, socket) { + const connection = new Connection(tabId, socket); 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); - - // Update UI - chrome.action.setBadgeText({ tabId, text: '' }); - chrome.action.setTitle({ tabId, title: 'Share tab with Playwright MCP' }); - + await this._cleanupConnection(tabId); + this._updateUI(tabId, { text: '', color: null, 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) { - const connection = this.activeConnections.get(tabId); - if (!connection) return; - - // Remove listeners - if (connection.eventListener) { - chrome.debugger.onEvent.removeListener(connection.eventListener); - } - if (connection.detachListener) { - chrome.debugger.onDetach.removeListener(connection.detachListener); - } + async _cleanupConnection(tabId) { + const info = this.activeConnections.get(tabId); + if (!info) return; + this.activeConnections.delete(tabId); // Close WebSocket - if (connection.socket && connection.socket.readyState === WebSocket.OPEN) { - connection.socket.close(); - } + info.connection.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)) { - await this.cleanupConnection(tabId); - } + if (this.activeConnections.has(tabId)) + await this._cleanupConnection(tabId); } } diff --git a/extension/connection.js b/extension/connection.js new file mode 100644 index 0000000..887663f --- /dev/null +++ b/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)); + } +} diff --git a/src/cdpRelay.ts b/src/cdpRelay.ts index 0c8c3cb..f84e7e3 100644 --- a/src/cdpRelay.ts +++ b/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 { }); } - /** - * Handle Extension connection - forwards to chrome.debugger - */ + private async _detachDebugger() { + this._connectionInfo = undefined; + await this._extensionConnection?.send('PWExtension.detachFromTab', {}); + } + private _handleExtensionConnection(ws: WebSocket): void { - if (this._extensionSocket?.readyState === WebSocket.OPEN) { - debugLogger('Closing previous extension connection'); - this._extensionSocket.close(1000, 'New connection established'); + if (this._extensionConnection) + this._extensionConnection.close('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; - debugLogger('Extension connected'); - - ws.on('message', data => { - try { - const message = JSON.parse(data.toString()); - this._handleExtensionMessage(message); - } 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); - }); + switch (method) { + case 'PWExtension.detachedFromTab': + debugLogger('← Debugger detached from tab:', params); + this._connectionInfo = undefined; + this._extensionConnection?.close(); + this._extensionConnection = null; + break; + } } /** * Handle messages from Playwright MCP */ - private _handlePlaywrightMessage(message: any): void { - debugLogger('← Playwright:', message.method || `response(${message.id})`); + private async _handlePlaywrightMessage(message: ProtocolCommand): Promise { + 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) - 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); + await this._forwardToExtension(message); } /** * Handle Browser domain methods locally */ - private _handleBrowserDomainMethod(message: any): void { + private async _handleBrowserDomainMethod(message: any): Promise { 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, - result: {} + id: message.id }); 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 { 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, - result: {} + id: message.id }); } else { - this._forwardToExtension(message); + await this._forwardToExtension(message); } break; - case 'Target.getTargets': - const targetInfos = []; - - if (this._connectionInfo) { - targetInfos.push({ - ...this._connectionInfo.targetInfo, - attached: true, - }); - } - + case 'Target.getTargetInfo': + debugLogger('Target.getTargetInfo', message); this._sendToPlaywright({ id: message.id, - result: { targetInfos } + result: this._connectionInfo?.targetInfo }); break; default: - this._forwardToExtension(message); + await this._forwardToExtension(message); } } - /** - * Forward message to extension - */ - private _forwardToExtension(message: any): void { - if (this._extensionSocket?.readyState === WebSocket.OPEN) { - debugLogger('→ Extension:', message.method || `command(${message.id})`); - this._extensionSocket.send(JSON.stringify(message)); - } else { - debugLogger('Extension not connected, cannot forward message'); - if (message.id) { - this._sendToPlaywright({ - id: message.id, - error: { message: 'Extension not connected' } - }); - } + private async _forwardToExtension(message: any): Promise { + 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, + }); + } catch (e) { + debugLogger('Error in the extension:', e); + 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:', JSON.stringify(message)); - this._playwrightSocket.send(JSON.stringify(message)); - } + debugLogger('→ Playwright:', `${message.method ?? `response(id=${message.id})`}`); + 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 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 { + 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 websocket due to malformed JSON. eventData=${eventData} e=${e?.message}`); + this._ws.close(); + return; + } + try { + this._handleParsedMessage(parsedJson); + } catch (e: any) { + debugLogger(` 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(` code=${event.code} reason=${event.reason}`); + this._dispose(); + } + + private _onError(event: WebSocket.ErrorEvent) { + debugLogger(` 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(); + } +}