adding support for multiple tabs

extension-playwright/src/background.ts

@@ -21,7 +21,7 @@ const RELAY_URL = 'ws://localhost:9988/extension';
 
 class SimplifiedExtension {
   private _connection: RelayConnection | undefined;
-  private _connectedTabId: number | null = null;
+  private _connectedTabs: Map<number, string> = new Map();
 
   constructor() {
     debugLog(`Using relay URL: ${RELAY_URL}`);
@@ -36,31 +36,22 @@ class SimplifiedExtension {
       return;
     }
 
-    // Toggle: if connected to this tab, disconnect; otherwise connect
-    if (this._connectedTabId === tab.id) {
-      await this._disconnect();
+    if (this._connectedTabs.has(tab.id)) {
+      await this._disconnectTab(tab.id);
     } else {
-      await this._connect(tab.id);
+      await this._connectTab(tab.id);
     }
   }
 
-  private async _waitForServer(): Promise<WebSocket> {
+  private async _waitForServer(): Promise<void> {
     const httpUrl = 'http://localhost:9988';
 
     while (true) {
       try {
         debugLog('Checking if relay server is available...');
         await fetch(httpUrl, { method: 'HEAD' });
-        debugLog('Server is available, connecting WebSocket...');
-
-        const socket = new WebSocket(RELAY_URL);
-        await new Promise<void>((resolve, reject) => {
-          socket.onopen = () => resolve();
-          socket.onerror = (e) => reject(e);
-          setTimeout(() => reject(new Error('Connection timeout')), 2000);
-        });
-        debugLog('Connected to relay server');
-        return socket;
+        debugLog('Server is available');
+        return;
       } catch (error: any) {
         debugLog(`Server not available, retrying in 1 second...`);
         await new Promise(resolve => setTimeout(resolve, 1000));
@@ -68,55 +59,103 @@ class SimplifiedExtension {
     }
   }
 
-  private async _connect(tabId: number): Promise<void> {
+  private async _connectTab(tabId: number): Promise<void> {
     try {
-      debugLog(`Connecting to tab ${tabId}`);
-
-      // Disconnect from any existing connection
-      if (this._connection) {
-        this._connection.close('Switching to new tab');
-        this._connection = undefined;
-      }
-
-      // Update icon to show connecting state
+      debugLog(`=== Starting connection to tab ${tabId} ===`);
       await this._updateIcon(tabId, 'connecting');
 
-      await this._waitForServer()
-      // Connect to WebSocket relay
-      const socket = new WebSocket(RELAY_URL);
-      await new Promise<void>((resolve, reject) => {
-        socket.onopen = () => resolve();
-        socket.onerror = () => reject(new Error('WebSocket connection failed'));
-        setTimeout(() => reject(new Error('Connection timeout')), 5000);
-      });
+      if (!this._connection) {
+        debugLog('No existing connection, creating new relay connection');
+        debugLog('Waiting for server at http://localhost:9988...');
+        await this._waitForServer();
+        
+        debugLog('Server is ready, creating WebSocket connection to:', RELAY_URL);
+        const socket = new WebSocket(RELAY_URL);
+        debugLog('WebSocket created, initial readyState:', socket.readyState, '(0=CONNECTING, 1=OPEN, 2=CLOSING, 3=CLOSED)');
+        
+        await new Promise<void>((resolve, reject) => {
+          let timeoutFired = false;
+          const timeout = setTimeout(() => {
+            timeoutFired = true;
+            debugLog('=== WebSocket connection TIMEOUT after 5 seconds ===');
+            debugLog('Final WebSocket readyState:', socket.readyState);
+            debugLog('WebSocket URL:', socket.url);
+            debugLog('Socket protocol:', socket.protocol);
+            reject(new Error('Connection timeout'));
+          }, 5000);
+          
+          socket.onopen = () => {
+            if (timeoutFired) {
+              debugLog('WebSocket opened but timeout already fired!');
+              return;
+            }
+            debugLog('WebSocket onopen fired! readyState:', socket.readyState);
+            clearTimeout(timeout);
+            resolve();
+          };
+          
+          socket.onerror = (error) => {
+            debugLog('WebSocket onerror during connection:', error);
+            debugLog('Error type:', error.type);
+            debugLog('Current readyState:', socket.readyState);
+            if (!timeoutFired) {
+              clearTimeout(timeout);
+              reject(new Error('WebSocket connection failed'));
+            }
+          };
+          
+          socket.onclose = (event) => {
+            debugLog('WebSocket onclose during connection setup:', {
+              code: event.code,
+              reason: event.reason,
+              wasClean: event.wasClean,
+              readyState: socket.readyState
+            });
+            if (!timeoutFired) {
+              clearTimeout(timeout);
+              reject(new Error(`WebSocket closed: ${event.reason || event.code}`));
+            }
+          };
+          
+          debugLog('Event handlers set, waiting for connection...');
+        });
 
-      // Create relay connection
-      this._connection = new RelayConnection(socket);
-      this._connection.onclose = () => {
-        debugLog('Connection closed');
-        this._connection = undefined;
-        void this._setConnectedTabId(null);
-      };
+        debugLog('WebSocket connected successfully, creating RelayConnection instance');
+        this._connection = new RelayConnection(socket);
+        this._connection.onclose = () => {
+          debugLog('=== Relay connection onclose callback triggered ===');
+          debugLog('Connected tabs before clearing:', Array.from(this._connectedTabs.keys()));
+          this._connection = undefined;
+          for (const tabId of this._connectedTabs.keys()) {
+            debugLog('Updating icon to disconnected for tab:', tabId);
+            void this._updateIcon(tabId, 'disconnected');
+          }
+          this._connectedTabs.clear();
+          debugLog('All tabs cleared');
+        };
+      } else {
+        debugLog('Reusing existing connection');
+      }
 
-      // Set the tab ID (this will attach debugger)
-      this._connection.setTabId(tabId);
-
-      // Update state
-      await this._setConnectedTabId(tabId);
-
-      debugLog(`Successfully connected to tab ${tabId}`);
+      debugLog('Calling attachTab for tab:', tabId);
+      const targetInfo = await this._connection.attachTab(tabId);
+      debugLog('attachTab completed, storing in connectedTabs map');
+      this._connectedTabs.set(tabId, targetInfo.targetId);
+      
+      await this._updateIcon(tabId, 'connected');
+      debugLog(`=== Successfully connected to tab ${tabId} ===`);
     } catch (error: any) {
-      debugLog(`Failed to connect: ${error.message}`);
+      debugLog(`=== Failed to connect to tab ${tabId} ===`);
+      debugLog('Error details:', error);
+      debugLog('Error stack:', error.stack);
       await this._updateIcon(tabId, 'disconnected');
 
-      // Show error notification
       chrome.action.setBadgeText({ tabId, text: '!' });
       chrome.action.setBadgeBackgroundColor({ tabId, color: '#f44336' });
       chrome.action.setTitle({ tabId, title: `Error: ${error.message}` });
 
-      // Clear error after 3 seconds
       setTimeout(() => {
-        if (this._connectedTabId !== tabId) {
+        if (!this._connectedTabs.has(tabId)) {
           chrome.action.setBadgeText({ tabId, text: '' });
           chrome.action.setTitle({ tabId, title: 'Click to attach debugger' });
         }
@@ -124,33 +163,26 @@ class SimplifiedExtension {
     }
   }
 
-  private async _disconnect(): Promise<void> {
-    debugLog('Disconnecting');
-
-    const tabId = this._connectedTabId;
-
-    this._connection?.close('User disconnected');
-    this._connection = undefined;
-
-    await this._setConnectedTabId(null);
-
-    if (tabId) {
-      await this._updateIcon(tabId, 'disconnected');
+  private async _disconnectTab(tabId: number): Promise<void> {
+    debugLog(`=== Disconnecting tab ${tabId} ===`);
+    
+    if (!this._connectedTabs.has(tabId)) {
+      debugLog('Tab not in connectedTabs map, ignoring disconnect');
+      return;
     }
-  }
-
-  private async _setConnectedTabId(tabId: number | null): Promise<void> {
-    const oldTabId = this._connectedTabId;
-    this._connectedTabId = tabId;
-
-    // Clear old tab icon
-    if (oldTabId && oldTabId !== tabId) {
-      await this._updateIcon(oldTabId, 'disconnected');
-    }
-
-    // Set new tab icon
-    if (tabId) {
-      await this._updateIcon(tabId, 'connected');
+    
+    debugLog('Calling detachTab on connection');
+    this._connection?.detachTab(tabId);
+    this._connectedTabs.delete(tabId);
+    debugLog('Tab removed from connectedTabs map');
+    
+    await this._updateIcon(tabId, 'disconnected');
+    
+    debugLog('Connected tabs remaining:', this._connectedTabs.size);
+    if (this._connectedTabs.size === 0 && this._connection) {
+      debugLog('No tabs remaining, closing relay connection');
+      this._connection.close('All tabs disconnected');
+      this._connection = undefined;
     }
   }
 
@@ -208,23 +240,17 @@ class SimplifiedExtension {
   }
 
   private async _onTabRemoved(tabId: number): Promise<void> {
-    if (this._connectedTabId !== tabId) {
-      return;
-    }
-
-    debugLog(`Connected tab ${tabId} was closed`);
-    this._connection?.close('Browser tab closed');
-    this._connection = undefined;
-    this._connectedTabId = null;
+    debugLog('Tab removed event for tab:', tabId, 'is connected:', this._connectedTabs.has(tabId));
+    if (!this._connectedTabs.has(tabId)) return;
+    
+    debugLog(`Connected tab ${tabId} was closed, disconnecting`);
+    await this._disconnectTab(tabId);
   }
 
   private async _onTabActivated(activeInfo: chrome.tabs.TabActiveInfo): Promise<void> {
-    // Update icon for the newly active tab
-    if (this._connectedTabId === activeInfo.tabId) {
-      await this._updateIcon(activeInfo.tabId, 'connected');
-    } else {
-      await this._updateIcon(activeInfo.tabId, 'disconnected');
-    }
+    const isConnected = this._connectedTabs.has(activeInfo.tabId);
+    debugLog('Tab activated:', activeInfo.tabId, 'is connected:', isConnected);
+    await this._updateIcon(activeInfo.tabId, isConnected ? 'connected' : 'disconnected');
   }
 }
 

extension-playwright/src/relayConnection.ts

@@ -25,62 +25,156 @@ export function debugLog(...args: unknown[]): void {
   }
 }
 
+interface AttachedTab {
+  debuggee: chrome.debugger.Debuggee;
+  targetId: string;
+  sessionId: string;
+  targetInfo: Protocol.Target.TargetInfo;
+}
+
 export class RelayConnection {
-  private _debuggee: chrome.debugger.Debuggee;
+  private _attachedTabs: Map<number, AttachedTab> = new Map();
+  private _nextSessionId: number = 1;
   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;
   private _closed = false;
 
   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._ws.onclose = (event) => {
+      debugLog('WebSocket onclose event:', {
+        code: event.code,
+        reason: event.reason,
+        wasClean: event.wasClean
+      });
+      this._onClose();
+    };
+    this._ws.onerror = (event) => {
+      debugLog('WebSocket onerror event:', event);
+    };
     this._eventListener = this._onDebuggerEvent.bind(this);
     this._detachListener = this._onDebuggerDetach.bind(this);
     chrome.debugger.onEvent.addListener(this._eventListener);
     chrome.debugger.onDetach.addListener(this._detachListener);
+    debugLog('RelayConnection created, WebSocket readyState:', this._ws.readyState);
   }
 
-  // Either setTabId or close is called after creating the connection.
-  setTabId(tabId: number): void {
-    this._debuggee = { tabId };
-    this._tabPromiseResolve();
+  async attachTab(tabId: number): Promise<Protocol.Target.TargetInfo> {
+    const debuggee = { tabId };
+    
+    debugLog('Attaching debugger to tab:', tabId, 'WebSocket state:', this._ws.readyState);
+    
+    try {
+      await chrome.debugger.attach(debuggee, '1.3');
+      debugLog('Debugger attached successfully to tab:', tabId);
+    } catch (error: any) {
+      debugLog('ERROR attaching debugger to tab:', tabId, error);
+      throw error;
+    }
+    
+    debugLog('Sending Target.getTargetInfo command for tab:', tabId);
+    const result = await chrome.debugger.sendCommand(
+      debuggee, 
+      'Target.getTargetInfo'
+    ) as Protocol.Target.GetTargetInfoResponse;
+    
+    debugLog('Received targetInfo for tab:', tabId, result.targetInfo);
+    
+    const targetInfo = result.targetInfo;
+    const sessionId = `pw-tab-${this._nextSessionId++}`;
+    
+    this._attachedTabs.set(tabId, {
+      debuggee,
+      targetId: targetInfo.targetId,
+      sessionId,
+      targetInfo
+    });
+    
+    debugLog('Sending Target.attachedToTarget event, WebSocket state:', this._ws.readyState);
+    this._sendMessage({
+      method: 'forwardCDPEvent',
+      params: {
+        method: 'Target.attachedToTarget',
+        params: {
+          sessionId,
+          targetInfo: {
+            ...targetInfo,
+            attached: true
+          },
+          waitingForDebugger: false
+        }
+      }
+    });
+    
+    debugLog('Tab attached successfully:', tabId, 'sessionId:', sessionId, 'targetId:', targetInfo.targetId);
+    return targetInfo;
+  }
+
+  detachTab(tabId: number): void {
+    const tab = this._attachedTabs.get(tabId);
+    if (!tab) return;
+    
+    debugLog('Detaching tab:', tabId, 'sessionId:', tab.sessionId);
+    
+    this._sendMessage({
+      method: 'forwardCDPEvent',
+      params: {
+        method: 'Target.detachedFromTarget',
+        params: {
+          sessionId: tab.sessionId,
+          targetId: tab.targetId
+        }
+      }
+    });
+    
+    chrome.debugger.detach(tab.debuggee).catch(() => {});
+    this._attachedTabs.delete(tabId);
   }
 
   close(message: string): void {
+    debugLog('Closing RelayConnection, reason:', message, 'current state:', this._ws.readyState);
     this._ws.close(1000, message);
-    // ws.onclose is called asynchronously, so we call it here to avoid forwarding
-    // CDP events to the closed connection.
     this._onClose();
   }
 
   private _onClose() {
-    if (this._closed)
+    if (this._closed) {
+      debugLog('_onClose called but already closed');
       return;
+    }
+    
+    debugLog('Connection closing, attached tabs count:', this._attachedTabs.size);
     this._closed = true;
+    
     chrome.debugger.onEvent.removeListener(this._eventListener);
     chrome.debugger.onDetach.removeListener(this._detachListener);
-    chrome.debugger.detach(this._debuggee).catch(() => {});
+    
+    for (const [tabId, tab] of this._attachedTabs) {
+      debugLog('Detaching debugger from tab:', tabId);
+      chrome.debugger.detach(tab.debuggee).catch((err) => {
+        debugLog('Error detaching debugger from tab:', tabId, err);
+      });
+    }
+    this._attachedTabs.clear();
+    
+    debugLog('Connection closed, calling onclose callback');
     this.onclose?.();
   }
 
   private _onDebuggerEvent(source: chrome.debugger.DebuggerSession, method: string, params: any): void {
-    if (source.tabId !== this._debuggee.tabId)
-      return;
-    debugLog('Forwarding CDP event:', method, params);
-    const sessionId = source.sessionId;
+    const tab = this._attachedTabs.get(source.tabId!);
+    if (!tab) return;
+    
+    debugLog('Forwarding CDP event:', method, 'from tab:', source.tabId);
+    
     this._sendMessage({
       method: 'forwardCDPEvent',
       params: {
-        sessionId,
+        sessionId: source.sessionId || tab.sessionId,
         method,
         params,
       },
@@ -88,10 +182,16 @@ export class RelayConnection {
   }
 
   private _onDebuggerDetach(source: chrome.debugger.Debuggee, reason: string): void {
-    if (source.tabId !== this._debuggee.tabId)
+    const tabId = source.tabId;
+    debugLog('_onDebuggerDetach called for tab:', tabId, 'reason:', reason, 'isAttached:', tabId ? this._attachedTabs.has(tabId) : false);
+    
+    if (!tabId || !this._attachedTabs.has(tabId)) {
+      debugLog('Ignoring debugger detach event for untracked tab:', tabId);
       return;
-    this.close(`Debugger detached: ${reason}`);
-    this._debuggee = { };
+    }
+    
+    debugLog(`Debugger detached from tab ${tabId}: ${reason}`);
+    this.detachTab(tabId);
   }
 
   private _onMessage(event: MessageEvent): void {
@@ -125,28 +225,52 @@ export class RelayConnection {
 
   private async _handleCommand(message: ExtensionCommandMessage): Promise<any> {
     if (message.method === 'attachToTab') {
-      await this._tabPromise;
-      debugLog('Attaching debugger to tab:', this._debuggee);
-      await chrome.debugger.attach(this._debuggee, '1.3');
-      const result = await chrome.debugger.sendCommand(this._debuggee, 'Target.getTargetInfo') as Protocol.Target.GetTargetInfoResponse;
-      return {
-        targetInfo: result?.targetInfo,
-      };
+      return {};
     }
-    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 forwardParams = message.params as { sessionId?: string; method: string; params?: any };
-      const { sessionId, method, params } = forwardParams;
-      debugLog('CDP command:', method, params);
+      const { sessionId, method, params } = message.params;
+      
+      if (method === 'Target.closeTarget' && params?.targetId) {
+        for (const [tabId, tab] of this._attachedTabs) {
+          if (tab.targetId === params.targetId) {
+            await chrome.tabs.remove(tabId);
+            return { success: true };
+          }
+        }
+        throw new Error(`Target not found: ${params.targetId}`);
+      }
+      
+      let targetTab: AttachedTab | undefined;
+      
+      for (const [tabId, tab] of this._attachedTabs) {
+        if (tab.sessionId === sessionId) {
+          targetTab = tab;
+          break;
+        }
+      }
+      
+      if (!targetTab) {
+        if (method === 'Browser.getVersion' || method === 'Target.getTargets') {
+          targetTab = this._attachedTabs.values().next().value;
+        }
+        
+        if (!targetTab) {
+          throw new Error(`No tab found for sessionId: ${sessionId}`);
+        }
+      }
+      
+      debugLog('CDP command:', method, 'for tab:', targetTab.debuggee.tabId);
+      
       const debuggerSession: chrome.debugger.DebuggerSession = {
-        ...this._debuggee,
-        sessionId,
+        ...targetTab.debuggee,
+        sessionId: sessionId !== targetTab.sessionId ? sessionId : undefined,
       };
+      
       return await chrome.debugger.sendCommand(
-          debuggerSession,
-          method,
-          params
+        debuggerSession,
+        method,
+        params
       );
     }
   }
@@ -161,7 +285,15 @@ export class RelayConnection {
   }
 
   private _sendMessage(message: any): void {
-    if (this._ws.readyState === WebSocket.OPEN)
-      this._ws.send(JSON.stringify(message));
+    if (this._ws.readyState === WebSocket.OPEN) {
+      try {
+        this._ws.send(JSON.stringify(message));
+        debugLog('Message sent successfully, type:', message.method || 'response');
+      } catch (error: any) {
+        debugLog('ERROR sending message:', error, 'message type:', message.method || 'response');
+      }
+    } else {
+      debugLog('Cannot send message, WebSocket not open. State:', this._ws.readyState, 'message type:', message.method || 'response');
+    }
   }
 }