chore(extension): terminate connection if nothing has been selected (#827)

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<number, { connection: RelayConnection, timerId?: number }>();
 
   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<void> {
+  private async _connectToRelay(mcpRelayUrl: string, selectorTabId: number): Promise<void> {
     try {
-      debugLog(`Connecting tab ${tabId} to bridge at ${mcpRelayUrl}`);
+      debugLog(`Connecting to relay at ${mcpRelayUrl}`);
       const socket = new WebSocket(mcpRelayUrl);
       await new Promise<void>((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 connection = new RelayConnection(socket);
-      const connectionClosed = (m: string) => {
+      connection.onclose = () => {
-        debugLog(m);
+        debugLog('Connection closed');
-        if (this._activeConnection === connection) {
+        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<void> {
+    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<void> {
+    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<void> {
-    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<chrome.tabs.Tab[]> {

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<void>;
+  private _tabPromiseResolve!: () => void;
 
-  constructor(ws: WebSocket, tabId: number) {
+  onclose?: () => void;
-    this._debuggee = { tabId };
+
+  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<any> {
-    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);

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,