From 46ce86f97e7da4a809462d78d4792178df5ed7b7 Mon Sep 17 00:00:00 2001 From: Yury Semikhatsky Date: Tue, 5 Aug 2025 09:47:39 -0700 Subject: [PATCH] chore(extension): terminate connection if nothing has been selected (#827) --- extension/src/background.ts | 91 ++++++++++++++++++++++++++------ extension/src/relayConnection.ts | 21 ++++++-- extension/src/ui/connect.tsx | 9 +++- 3 files changed, 99 insertions(+), 22 deletions(-) diff --git a/extension/src/background.ts b/extension/src/background.ts index 5fa700c..bfb672c 100644 --- a/extension/src/background.ts +++ b/extension/src/background.ts @@ -19,15 +19,19 @@ import { RelayConnection, debugLog } from './relayConnection.js'; type PageMessage = { type: 'connectToMCPRelay'; mcpRelayUrl: string; - tabId: number; - windowId: number; } | { type: 'getTabs'; +} | { + type: 'connectToTab'; + tabId: number; + windowId: number; + mcpRelayUrl: string; }; class TabShareExtension { private _activeConnection: RelayConnection | undefined; private _connectedTabId: number | null = null; + private _pendingTabSelection = new Map(); constructor() { chrome.tabs.onRemoved.addListener(this._onTabRemoved.bind(this)); @@ -39,22 +43,27 @@ class TabShareExtension { private _onMessage(message: PageMessage, sender: chrome.runtime.MessageSender, sendResponse: (response: any) => void) { switch (message.type) { case 'connectToMCPRelay': - this._connectTab(message.tabId, message.windowId, message.mcpRelayUrl!).then( + this._connectToRelay(message.mcpRelayUrl!, sender.tab!.id!).then( () => sendResponse({ success: true }), (error: any) => sendResponse({ success: false, error: error.message })); - return true; // Return true to indicate that the response will be sent asynchronously + return true; case 'getTabs': this._getTabs().then( tabs => sendResponse({ success: true, tabs, currentTabId: sender.tab?.id }), (error: any) => sendResponse({ success: false, error: error.message })); return true; + case 'connectToTab': + this._connectTab(message.tabId, message.windowId, message.mcpRelayUrl!).then( + () => sendResponse({ success: true }), + (error: any) => sendResponse({ success: false, error: error.message })); + return true; // Return true to indicate that the response will be sent asynchronously } return false; } - private async _connectTab(tabId: number, windowId: number, mcpRelayUrl: string): Promise { + private async _connectToRelay(mcpRelayUrl: string, selectorTabId: number): Promise { try { - debugLog(`Connecting tab ${tabId} to bridge at ${mcpRelayUrl}`); + debugLog(`Connecting to relay at ${mcpRelayUrl}`); const socket = new WebSocket(mcpRelayUrl); await new Promise((resolve, reject) => { socket.onopen = () => resolve(); @@ -62,17 +71,41 @@ class TabShareExtension { setTimeout(() => reject(new Error('Connection timeout')), 5000); }); - const connection = new RelayConnection(socket, tabId); - const connectionClosed = (m: string) => { - debugLog(m); - if (this._activeConnection === connection) { - this._activeConnection = undefined; - void this._setConnectedTabId(null); - } + const connection = new RelayConnection(socket); + connection.onclose = () => { + debugLog('Connection closed'); + this._pendingTabSelection.delete(selectorTabId); + // TODO: show error in the selector tab? + }; + this._pendingTabSelection.set(selectorTabId, { connection }); + debugLog(`Connected to MCP relay`); + } catch (error: any) { + debugLog(`Failed to connect to MCP relay:`, error.message); + throw error; + } + } + + private async _connectTab(tabId: number, windowId: number, mcpRelayUrl: string): Promise { + try { + debugLog(`Connecting tab ${tabId} to relay at ${mcpRelayUrl}`); + try { + this._activeConnection?.close('Another connection is requested'); + } catch (error: any) { + debugLog(`Error closing active connection:`, error); + } + await this._setConnectedTabId(null); + + this._activeConnection = this._pendingTabSelection.get(tabId)?.connection; + if (!this._activeConnection) + throw new Error('No active MCP relay connection'); + this._pendingTabSelection.delete(tabId); + + this._activeConnection.setTabId(tabId); + this._activeConnection.onclose = () => { + debugLog('MCP connection closed'); + this._activeConnection = undefined; + void this._setConnectedTabId(null); }; - socket.onclose = () => connectionClosed('WebSocket closed'); - socket.onerror = error => connectionClosed(`WebSocket error: ${error}`); - this._activeConnection = connection; await Promise.all([ this._setConnectedTabId(tabId), @@ -103,6 +136,12 @@ class TabShareExtension { } private async _onTabRemoved(tabId: number): Promise { + const pendingConnection = this._pendingTabSelection.get(tabId)?.connection; + if (pendingConnection) { + this._pendingTabSelection.delete(tabId); + pendingConnection.close('Browser tab closed'); + return; + } if (this._connectedTabId !== tabId) return; this._activeConnection?.close('Browser tab closed'); @@ -111,8 +150,26 @@ class TabShareExtension { } private async _onTabUpdated(tabId: number, changeInfo: chrome.tabs.TabChangeInfo, tab: chrome.tabs.Tab): Promise { - if (changeInfo.status === 'complete' && this._connectedTabId === tabId) + if (changeInfo.status === 'complete' && this._connectedTabId === tabId) { await this._setConnectedTabId(tabId); + return; + } + const pending = this._pendingTabSelection.get(tabId); + if (!pending) + return; + if (tab.active && pending.timerId) { + clearTimeout(pending.timerId); + pending.timerId = undefined; + return; + } + if (!tab.active && !pending.timerId) { + debugLog('Starting inactivity timer', tabId); + pending.timerId = window.setTimeout(() => { + this._pendingTabSelection.delete(tabId); + pending.connection.close('Tab is not active'); + }, 5000); + return; + } } private async _getTabs(): Promise { diff --git a/extension/src/relayConnection.ts b/extension/src/relayConnection.ts index a7fd4f7..15e835b 100644 --- a/extension/src/relayConnection.ts +++ b/extension/src/relayConnection.ts @@ -41,11 +41,17 @@ export class RelayConnection { private _ws: WebSocket; private _eventListener: (source: chrome.debugger.DebuggerSession, method: string, params: any) => void; private _detachListener: (source: chrome.debugger.Debuggee, reason: string) => void; + private _tabPromise: Promise; + private _tabPromiseResolve!: () => void; - constructor(ws: WebSocket, tabId: number) { - this._debuggee = { tabId }; + onclose?: () => void; + + constructor(ws: WebSocket) { + this._debuggee = { }; + this._tabPromise = new Promise(resolve => this._tabPromiseResolve = resolve); this._ws = ws; this._ws.onmessage = this._onMessage.bind(this); + this._ws.onclose = () => this.onclose?.(); // Store listeners for cleanup this._eventListener = this._onDebuggerEvent.bind(this); this._detachListener = this._onDebuggerDetach.bind(this); @@ -53,6 +59,12 @@ export class RelayConnection { chrome.debugger.onDetach.addListener(this._detachListener); } + // Either setTabId or close is called after creating the connection. + setTabId(tabId: number): void { + this._debuggee = { tabId }; + this._tabPromiseResolve(); + } + close(message: string): void { chrome.debugger.onEvent.removeListener(this._eventListener); chrome.debugger.onDetach.removeListener(this._detachListener); @@ -111,9 +123,8 @@ export class RelayConnection { } private async _handleCommand(message: ProtocolCommand): Promise { - if (!this._debuggee.tabId) - throw new Error('No tab is connected. Please go to the Playwright MCP extension and select the tab you want to connect to.'); if (message.method === 'attachToTab') { + await this._tabPromise; 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'); @@ -121,6 +132,8 @@ export class RelayConnection { targetInfo: result?.targetInfo, }; } + if (!this._debuggee.tabId) + throw new Error('No tab is connected. Please go to the Playwright MCP extension and select the tab you want to connect to.'); if (message.method === 'forwardCDPCommand') { const { sessionId, method, params } = message.params; debugLog('CDP command:', method, params); diff --git a/extension/src/ui/connect.tsx b/extension/src/ui/connect.tsx index 9d8f114..adcc79c 100644 --- a/extension/src/ui/connect.tsx +++ b/extension/src/ui/connect.tsx @@ -61,9 +61,16 @@ const ConnectApp: React.FC = () => { return; } + void connectToMCPRelay(relayUrl); void loadTabs(); }, []); + const connectToMCPRelay = useCallback(async (mcpRelayUrl: string) => { + const response = await chrome.runtime.sendMessage({ type: 'connectToMCPRelay', mcpRelayUrl }); + if (!response.success) + setStatus({ type: 'error', message: 'Failed to connect to MCP relay: ' + response.error }); + }, []); + const loadTabs = useCallback(async () => { const response = await chrome.runtime.sendMessage({ type: 'getTabs' }); if (response.success) { @@ -86,7 +93,7 @@ const ConnectApp: React.FC = () => { try { const response = await chrome.runtime.sendMessage({ - type: 'connectToMCPRelay', + type: 'connectToTab', mcpRelayUrl, tabId: selectedTab.id, windowId: selectedTab.windowId,