rewrite extension. use zustand for state

extension-playwright/src/background.ts

@@ -1,151 +1,285 @@
 
 import { RelayConnection, debugLog } from './relayConnection';
+import { create } from 'zustand';
 
 // Relay URL - fixed port for MCP bridge
 const RELAY_URL = 'ws://localhost:9988/extension';
 
-class SimplifiedExtension {
+type ConnectionState = 'disconnected' | 'reconnecting' | 'connected';
-  private _connection: RelayConnection | undefined;
-  private _connectedTabs: Map<number, string> = new Map();
 
-  constructor() {
+interface ExtensionState {
-    debugLog(`Using relay URL: ${RELAY_URL}`);
+  connection: RelayConnection | undefined;
-    chrome.tabs.onRemoved.addListener(this._onTabRemoved);
+  connectedTabs: Map<number, string>;
-    chrome.tabs.onActivated.addListener(this._onTabActivated);
+  connectionState: ConnectionState;
-    chrome.action.onClicked.addListener(this._onActionClicked);
+  currentTabId: number | undefined;
+}
+
+const useExtensionStore = create<ExtensionState>(() => ({
+  connection: undefined,
+  connectedTabs: new Map(),
+  connectionState: 'disconnected',
+  currentTabId: undefined,
+}));
+
+async function updateIcon(tabId: number, state: 'connected' | 'disconnected' | 'connecting'): Promise<void> {
+  try {
+    switch (state) {
+      case 'connected':
+        await chrome.action.setIcon({
+          tabId,
+          path: {
+            '16': '/icons/icon-green-16.png',
+            '32': '/icons/icon-green-32.png',
+            '48': '/icons/icon-green-48.png',
+            '128': '/icons/icon-green-128.png'
+          }
+        });
+        await chrome.action.setBadgeText({ tabId, text: '' });
+        await chrome.action.setTitle({ tabId, title: 'Connected - Click to disconnect' });
+        break;
+
+      case 'connecting':
+        await chrome.action.setIcon({
+          tabId,
+          path: {
+            '16': '/icons/icon-gray-16.png',
+            '32': '/icons/icon-gray-32.png',
+            '48': '/icons/icon-gray-48.png',
+            '128': '/icons/icon-gray-128.png'
+          }
+        });
+        await chrome.action.setBadgeText({ tabId, text: '...' });
+        await chrome.action.setBadgeBackgroundColor({ tabId, color: '#FF9800' });
+        await chrome.action.setTitle({ tabId, title: 'Connecting...' });
+        break;
+
+      case 'disconnected':
+      default:
+        await chrome.action.setIcon({
+          tabId,
+          path: {
+            '16': '/icons/icon-gray-16.png',
+            '32': '/icons/icon-gray-32.png',
+            '48': '/icons/icon-gray-48.png',
+            '128': '/icons/icon-gray-128.png'
+          }
+        });
+        await chrome.action.setBadgeText({ tabId, text: '' });
+        await chrome.action.setTitle({ tabId, title: 'Click to attach debugger' });
+        break;
+    }
+  } catch (error: any) {
+    debugLog(`Error updating icon: ${error.message}`);
+  }
+}
+
+useExtensionStore.subscribe((state, prevState) => {
+  const prevTabs = new Set(prevState.connectedTabs.keys());
+  const currentTabs = new Set(state.connectedTabs.keys());
+  const prevConnectionState = prevState.connectionState;
+  const currentConnectionState = state.connectionState;
+  const prevCurrentTabId = prevState.currentTabId;
+  const currentTabId = state.currentTabId;
+
+  const connectionStateChanged = prevConnectionState !== currentConnectionState;
+  const currentTabChanged = prevCurrentTabId !== currentTabId;
+
+  const tabsToUpdate = new Set<number>();
+
+  if (connectionStateChanged) {
+    debugLog('Connection state changed:', prevConnectionState, '->', currentConnectionState);
+    for (const tabId of currentTabs) {
+      tabsToUpdate.add(tabId);
+    }
   }
 
-  private _onActionClicked = async (tab: chrome.tabs.Tab): Promise<void> => {
+  const addedTabs = [...currentTabs].filter(id => !prevTabs.has(id));
-    if (!tab.id) {
+  const removedTabs = [...prevTabs].filter(id => !currentTabs.has(id));
-      debugLog('No tab ID available');
+
+  for (const tabId of addedTabs) {
+    tabsToUpdate.add(tabId);
+  }
+
+  for (const tabId of removedTabs) {
+    tabsToUpdate.add(tabId);
+  }
+
+  if (currentTabChanged && currentTabId !== undefined) {
+    tabsToUpdate.add(currentTabId);
+  }
+
+  for (const tabId of tabsToUpdate) {
+    const isTracked = currentTabs.has(tabId);
+    let iconState: 'connected' | 'disconnected' | 'connecting';
+
+    if (currentConnectionState === 'disconnected') {
+      iconState = 'disconnected';
+    } else    if (!isTracked) {
+      iconState = 'disconnected';
+    } else if (currentConnectionState === 'connected') {
+      iconState = 'connected';
+    } else {
+      iconState = 'connecting';
+    }
+
+    void updateIcon(tabId, iconState);
+  }
+});
+
+async function ensureConnection(): Promise<void> {
+    const { connection } = useExtensionStore.getState();
+    if (connection) {
+      debugLog('Connection already exists, reusing');
       return;
     }
 
-    if (this._connectedTabs.has(tab.id)) {
+    useExtensionStore.setState({ connectionState: 'reconnecting' });
-      await this._disconnectTab(tab.id);
+    debugLog('No existing connection, creating new relay connection');
-    } else {
+    debugLog('Waiting for server at http://localhost:9988...');
-      await this._connectTab(tab.id);
-    }
-  };
 
-  private async _connectTab(tabId: number): Promise<void> {
+    while (useExtensionStore.getState().connectionState !== 'disconnected') {
+      try {
+        await fetch('http://localhost:9988', { method: 'HEAD' });
+        debugLog('Server is available');
+        break;
+      } catch (error: any) {
+        debugLog('Server not available, retrying in 1 second...');
+        await new Promise(resolve => setTimeout(resolve, 1000));
+      }
+    }
+
+    if (useExtensionStore.getState().connectionState === 'disconnected') {
+      debugLog('Connection cancelled by user');
+      return;
+    }
+
+    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...');
+    });
+
+    debugLog('WebSocket connected successfully, creating RelayConnection instance');
+    const newConnection = new RelayConnection({
+      ws: socket,
+      onClose: () => {
+        debugLog('=== Relay connection onClose callback triggered ===');
+        const { connectedTabs } = useExtensionStore.getState();
+        debugLog('Connected tabs before potential reconnection:', Array.from(connectedTabs.keys()));
+        useExtensionStore.setState({ connection: undefined, connectionState: 'disconnected' });
+
+        if (connectedTabs.size > 0) {
+          debugLog('Tabs still connected, triggering reconnection');
+          void reconnect();
+        } else {
+          debugLog('No tabs to reconnect');
+        }
+      },
+      onTabDetached: (tabId) => {
+        debugLog('=== Manual tab detachment detected for tab:', tabId, '===');
+        debugLog('User closed debugger via Chrome automation bar');
+
+        useExtensionStore.setState((state) => {
+          const newTabs = new Map(state.connectedTabs);
+          newTabs.delete(tabId);
+          return { connectionState: 'disconnected', connectedTabs: newTabs };
+        });
+        debugLog('Removed tab from _connectedTabs map');
+      }
+    });
+
+    useExtensionStore.setState({ connection: newConnection });
+    debugLog('Connection established, WebSocket open (caller should set connectionState)');
+  }
+
+async function connectTab(tabId: number): Promise<void> {
     try {
       debugLog(`=== Starting connection to tab ${tabId} ===`);
-      await this._updateIcon(tabId, 'connecting');
 
-      if (!this._connection) {
+      useExtensionStore.setState((state) => {
-        debugLog('No existing connection, creating new relay connection');
+        const newTabs = new Map(state.connectedTabs);
-        debugLog('Waiting for server at http://localhost:9988...');
+        newTabs.set(tabId, '');
+        return { connectedTabs: newTabs };
+      });
 
-        // Wait for server to be available
+      await ensureConnection();
-        while (true) {
-          try {
-            await fetch('http://localhost:9988', { method: 'HEAD' });
-            debugLog('Server is available');
-            break;
-          } catch (error: any) {
-            debugLog('Server not available, retrying in 1 second...');
-            await new Promise(resolve => setTimeout(resolve, 1000));
-          }
-        }
-
-        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...');
-        });
-
-        debugLog('WebSocket connected successfully, creating RelayConnection instance');
-        this._connection = new RelayConnection({
-          ws: socket,
-          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');
-          },
-          onTabDetached: (tabId) => {
-            debugLog('=== Manual tab detachment detected for tab:', tabId, '===');
-            debugLog('User closed debugger via Chrome automation bar');
-            this._connectedTabs.delete(tabId);
-            debugLog('Removed tab from _connectedTabs map');
-            void this._updateIcon(tabId, 'disconnected');
-            debugLog('Updated icon to disconnected state');
-          }
-        });
-      } else {
-        debugLog('Reusing existing connection');
-      }
 
       debugLog('Calling attachTab for tab:', tabId);
-      const targetInfo = await this._connection.attachTab(tabId);
+      const { connection } = useExtensionStore.getState();
-      debugLog('attachTab completed, storing in connectedTabs map');
+      if (!connection) return
-      this._connectedTabs.set(tabId, targetInfo.targetId);
+      const targetInfo = await connection.attachTab(tabId);
+      debugLog('attachTab completed, updating targetId in connectedTabs map');
+      useExtensionStore.setState((state) => {
+        const newTabs = new Map(state.connectedTabs);
+
+        newTabs.set(tabId, targetInfo?.targetId);
+        return { connectedTabs: newTabs, connectionState: 'connected' };
+      });
 
-      await this._updateIcon(tabId, 'connected');
       debugLog(`=== Successfully connected to tab ${tabId} ===`);
     } catch (error: any) {
       debugLog(`=== Failed to connect to tab ${tabId} ===`);
       debugLog('Error details:', error);
       debugLog('Error stack:', error.stack);
-      await this._updateIcon(tabId, 'disconnected');
+
+      useExtensionStore.setState((state) => {
+        const newTabs = new Map(state.connectedTabs);
+        newTabs.delete(tabId);
+        return { connectedTabs: newTabs };
+      });
 
       chrome.action.setBadgeText({ tabId, text: '!' });
       chrome.action.setBadgeBackgroundColor({ tabId, color: '#f44336' });
       chrome.action.setTitle({ tabId, title: `Error: ${error.message}` });
 
       setTimeout(() => {
-        if (!this._connectedTabs.has(tabId)) {
+        const { connectedTabs } = useExtensionStore.getState();
+        if (!connectedTabs.has(tabId)) {
           chrome.action.setBadgeText({ tabId, text: '' });
           chrome.action.setTitle({ tabId, title: 'Click to attach debugger' });
         }
@@ -153,96 +287,148 @@ class SimplifiedExtension {
     }
   }
 
-  private async _disconnectTab(tabId: number): Promise<void> {
+async function disconnectTab(tabId: number): Promise<void> {
     debugLog(`=== Disconnecting tab ${tabId} ===`);
 
-    if (!this._connectedTabs.has(tabId)) {
+    const { connectedTabs, connection } = useExtensionStore.getState();
+    if (!connectedTabs.has(tabId)) {
       debugLog('Tab not in connectedTabs map, ignoring disconnect');
       return;
     }
 
     debugLog('Calling detachTab on connection');
-    this._connection?.detachTab(tabId);
+    connection?.detachTab(tabId);
-    this._connectedTabs.delete(tabId);
+    useExtensionStore.setState((state) => {
+      const newTabs = new Map(state.connectedTabs);
+      newTabs.delete(tabId);
+      return { connectedTabs: newTabs };
+    });
     debugLog('Tab removed from connectedTabs map');
 
-    await this._updateIcon(tabId, 'disconnected');
+    const { connectedTabs: updatedTabs, connection: updatedConnection } = useExtensionStore.getState();
-
+    debugLog('Connected tabs remaining:', updatedTabs.size);
-    debugLog('Connected tabs remaining:', this._connectedTabs.size);
+    if (updatedTabs.size === 0 && updatedConnection) {
-    if (this._connectedTabs.size === 0 && this._connection) {
       debugLog('No tabs remaining, closing relay connection');
-      this._connection.close('All tabs disconnected');
+      updatedConnection.close('All tabs disconnected');
-      this._connection = undefined;
+      useExtensionStore.setState({ connection: undefined, connectionState: 'disconnected' });
     }
   }
 
-  private async _updateIcon(tabId: number, state: 'connected' | 'disconnected' | 'connecting'): Promise<void> {
+async function reconnect(): Promise<void> {
+    debugLog('=== Starting reconnection ===');
+    const { connectedTabs } = useExtensionStore.getState();
+    debugLog('Tabs to reconnect:', Array.from(connectedTabs.keys()));
+
     try {
-      switch (state) {
+      await ensureConnection();
-        case 'connected':
-          await chrome.action.setIcon({
-            tabId,
-            path: {
-              '16': '/icons/icon-green-16.png',
-              '32': '/icons/icon-green-32.png',
-              '48': '/icons/icon-green-48.png',
-              '128': '/icons/icon-green-128.png'
-            }
-          });
-          await chrome.action.setBadgeText({ tabId, text: '' });
-          await chrome.action.setTitle({ tabId, title: 'Connected - Click to disconnect' });
-          break;
 
-        case 'connecting':
+      const tabsToReconnect = Array.from(connectedTabs.keys());
-          await chrome.action.setIcon({
+      debugLog('Re-attaching', tabsToReconnect.length, 'tabs');
-            tabId,
-            path: {
-              '16': '/icons/icon-gray-16.png',
-              '32': '/icons/icon-gray-32.png',
-              '48': '/icons/icon-gray-48.png',
-              '128': '/icons/icon-gray-128.png'
-            }
-          });
-          await chrome.action.setBadgeText({ tabId, text: '...' });
-          await chrome.action.setBadgeBackgroundColor({ tabId, color: '#FF9800' });
-          await chrome.action.setTitle({ tabId, title: 'Connecting...' });
-          break;
 
-        case 'disconnected':
+      for (const tabId of tabsToReconnect) {
-        default:
+        const { connectedTabs: currentTabs } = useExtensionStore.getState();
-          await chrome.action.setIcon({
+        if (!currentTabs.has(tabId)) {
-            tabId,
+          debugLog('Tab', tabId, 'was manually disconnected during reconnection, skipping');
-            path: {
+          continue;
-              '16': '/icons/icon-gray-16.png',
+        }
-              '32': '/icons/icon-gray-32.png',
+
-              '48': '/icons/icon-gray-48.png',
+        try {
-              '128': '/icons/icon-gray-128.png'
+          debugLog('Checking if tab', tabId, 'still exists');
-            }
+          await chrome.tabs.get(tabId);
+
+          debugLog('Re-attaching tab:', tabId);
+          const { connection } = useExtensionStore.getState();
+          if (!connection) return;
+          const targetInfo = await connection.attachTab(tabId);
+          useExtensionStore.setState((state) => {
+            const newTabs = new Map(state.connectedTabs);
+            newTabs.set(tabId, targetInfo.targetId);
+            return { connectedTabs: newTabs };
           });
-          await chrome.action.setBadgeText({ tabId, text: '' });
+          debugLog('Successfully re-attached tab:', tabId);
-          await chrome.action.setTitle({ tabId, title: 'Click to attach debugger' });
+        } catch (error: any) {
-          break;
+          debugLog('Failed to re-attach tab:', tabId, error.message);
+          useExtensionStore.setState((state) => {
+            const newTabs = new Map(state.connectedTabs);
+            newTabs.delete(tabId);
+            return { connectedTabs: newTabs };
+          });
+        }
+      }
+
+      const { connectedTabs: finalTabs } = useExtensionStore.getState();
+      debugLog('=== Reconnection complete ===');
+      debugLog('Successfully reconnected tabs:', finalTabs.size);
+
+      if (finalTabs.size > 0) {
+        useExtensionStore.setState({ connectionState: 'connected' });
+        debugLog('Set connectionState to connected');
+      } else {
+        debugLog('No tabs successfully reconnected, staying in reconnecting state');
+        useExtensionStore.setState({ connectionState: 'disconnected' });
       }
     } catch (error: any) {
-      // Ignore errors (tab may be closed)
+      debugLog('=== Reconnection failed ===', error);
-      debugLog(`Error updating icon: ${error.message}`);
+
+      const { connectedTabs: failedTabs } = useExtensionStore.getState();
+      for (const tabId of failedTabs.keys()) {
+        chrome.action.setBadgeText({ tabId, text: '!' });
+        chrome.action.setBadgeBackgroundColor({ tabId, color: '#f44336' });
+        chrome.action.setTitle({ tabId, title: 'Reconnection failed - Click to retry' });
+      }
+
+      useExtensionStore.setState({ connectedTabs: new Map(), connectionState: 'disconnected' });
     }
   }
 
-  private _onTabRemoved = async (tabId: number): Promise<void> => {
+async function onTabRemoved(tabId: number): Promise<void> {
-    debugLog('Tab removed event for tab:', tabId, 'is connected:', this._connectedTabs.has(tabId));
+  const { connectedTabs } = useExtensionStore.getState();
-    if (!this._connectedTabs.has(tabId)) return;
+  debugLog('Tab removed event for tab:', tabId, 'is connected:', connectedTabs.has(tabId));
+  if (!connectedTabs.has(tabId)) return;
 
-    debugLog(`Connected tab ${tabId} was closed, disconnecting`);
+  debugLog(`Connected tab ${tabId} was closed, disconnecting`);
-    await this._disconnectTab(tabId);
+  await disconnectTab(tabId);
-  };
-
-  private _onTabActivated = async (activeInfo: chrome.tabs.TabActiveInfo): Promise<void> => {
-    const isConnected = this._connectedTabs.has(activeInfo.tabId);
-    debugLog('Tab activated:', activeInfo.tabId, 'is connected:', isConnected);
-    await this._updateIcon(activeInfo.tabId, isConnected ? 'connected' : 'disconnected');
-  };
 }
 
-// @ts-ignore
+async function onTabActivated(activeInfo: chrome.tabs.TabActiveInfo): Promise<void> {
-globalThis.state = new SimplifiedExtension();
+  debugLog('Tab activated:', activeInfo.tabId);
+  useExtensionStore.setState({ currentTabId: activeInfo.tabId });
+}
+
+async function onActionClicked(tab: chrome.tabs.Tab): Promise<void> {
+  if (!tab.id) {
+    debugLog('No tab ID available');
+    return;
+  }
+
+  const { connectedTabs, connectionState, connection } = useExtensionStore.getState();
+
+  if (connectionState === 'reconnecting') {
+    debugLog('User clicked during reconnection, canceling reconnection and disconnecting all tabs');
+
+    const tabsToDisconnect = Array.from(connectedTabs.keys());
+
+    for (const tabId of tabsToDisconnect) {
+      connection?.detachTab(tabId);
+    }
+
+    useExtensionStore.setState({ connectionState: 'disconnected', connectedTabs: new Map() });
+
+    if (connection) {
+      connection.close('User cancelled reconnection');
+    }
+
+    return;
+  }
+
+  if (connectedTabs.has(tab.id)) {
+    await disconnectTab(tab.id);
+  } else {
+    await connectTab(tab.id);
+  }
+}
+
+debugLog(`Using relay URL: ${RELAY_URL}`);
+chrome.tabs.onRemoved.addListener(onTabRemoved);
+chrome.tabs.onActivated.addListener(onTabActivated);
+chrome.action.onClicked.addListener(onActionClicked);