chore: optimize devtool tab connect (#279)

packages/visualizer/src/component/playground-component.tsx

@@ -157,11 +157,8 @@ const serverLaunchTip = (
 );
 
 // remember to destroy the agent when the tab is destroyed: agent.page.destroy()
-export const extensionAgentForTabId = (getTabId: () => number) => {
-  if (!getTabId()) {
-    return null;
-  }
-  const page = new ChromeExtensionProxyPage(getTabId);
+export const extensionAgentForTabId = (trackingActiveTab: () => boolean) => {
+  const page = new ChromeExtensionProxyPage(trackingActiveTab);
   return new ChromeExtensionProxyPageAgent(page);
 };
 
@@ -171,7 +168,9 @@ export function Playground({
   showContextPreview = true,
   dryMode = false,
 }: {
-  getAgent: () => StaticPageAgent | ChromeExtensionProxyPageAgent | null;
+  getAgent: (
+    trackingActiveTab: () => boolean,
+  ) => StaticPageAgent | ChromeExtensionProxyPageAgent | null;
   hideLogo?: boolean;
   showContextPreview?: boolean;
   dryMode?: boolean;
@@ -195,6 +194,13 @@ export function Playground({
   const [verticalMode, setVerticalMode] = useState(false);
 
   const { tabId } = useChromeTabInfo();
+  const tabActiveRef = useRef<boolean>(true);
+
+  useEffect(() => {
+    if (tabId) {
+      tabActiveRef.current = trackingActiveTab;
+    }
+  }, [trackingActiveTab]);
 
   // if the screen is narrow, we use vertical mode
   useEffect(() => {
@@ -226,7 +232,7 @@ export function Playground({
     if (uiContextPreview) return;
     if (!showContextPreview) return;
 
-    getAgent()
+    getAgent(() => tabActiveRef.current)
       ?.getUIContext()
       .then((context) => {
         setUiContextPreview(context);
@@ -289,7 +295,7 @@ export function Playground({
       error: null,
     };
 
-    const activeAgent = getAgent();
+    const activeAgent = getAgent(() => tabActiveRef.current);
     try {
       if (!activeAgent) {
         throw new Error('No agent found');

packages/visualizer/src/extension/popup.tsx

@@ -23,16 +23,8 @@ declare const __VERSION__: string;
 
 function PlaygroundPopup() {
   const extensionVersion = getExtensionVersion();
-  const { tabId } = useChromeTabInfo();
-  const tabIdRef = useRef<number | null>(null);
   const { popupTab, setPopupTab, trackingActiveTab } = useEnvConfig();
 
-  useEffect(() => {
-    if (tabId) {
-      tabIdRef.current = tabId;
-    }
-  }, [tabId]);
-
   const items = [
     {
       key: 'playground',
@@ -42,13 +34,8 @@ function PlaygroundPopup() {
         <div className="popup-playground-container">
           <Playground
             hideLogo
-            getAgent={() => {
-              return extensionAgentForTabId(() => {
-                if (trackingActiveTab) {
-                  return tabIdRef.current || 0;
-                }
-                return tabId || 0;
-              });
+            getAgent={(trackingActiveTab: () => boolean) => {
+              return extensionAgentForTabId(trackingActiveTab);
             }}
             showContextPreview={false}
           />

packages/web-integration/src/bridge-mode/page-browser-side.ts

@@ -21,7 +21,7 @@ export class ChromeExtensionPageBrowserSide extends ChromeExtensionProxyPage {
       type: 'log' | 'status',
     ) => void = () => {},
   ) {
-    super(() => 0);
+    super(() => true);
   }
 
   private async setupBridgeClient() {
@@ -44,7 +44,8 @@ export class ChromeExtensionPageBrowserSide extends ChromeExtensionProxyPage {
           return this.onLogMessage(args[0] as string, 'status');
         }
 
-        if (!this.getTabId || this.getTabId() === 0) {
+        const tabId = await this.getTabId();
+        if (!tabId || tabId === 0) {
           throw new Error('no tab is connected');
         }
 
@@ -93,31 +94,19 @@ export class ChromeExtensionPageBrowserSide extends ChromeExtensionProxyPage {
   }
 
   public async connectNewTabWithUrl(url: string) {
-    assert(url, 'url is required to create a new tab');
-    if (this.getTabId()) {
-      throw new Error('tab is already connected');
-    }
-
     const tab = await chrome.tabs.create({ url });
     const tabId = tab.id;
     assert(tabId, 'failed to get tabId after creating a new tab');
-    this.getTabId = () => tabId;
 
     // new tab
     this.onLogMessage(`Creating new tab: ${url}`, 'log');
   }
 
   public async connectCurrentTab() {
-    if (this.getTabId()) {
-      throw new Error(
-        `already connected with tab id ${this.getTabId()}, cannot reconnect`,
-      );
-    }
     const tabs = await chrome.tabs.query({ active: true, currentWindow: true });
     console.log('current tab', tabs);
     const tabId = tabs[0]?.id;
     assert(tabId, 'failed to get tabId');
-    this.getTabId = () => tabId;
 
     this.onLogMessage(`Connected to current tab: ${tabs[0]?.url}`, 'log');
   }
@@ -129,6 +118,5 @@ export class ChromeExtensionPageBrowserSide extends ChromeExtensionProxyPage {
       this.onDisconnect();
     }
     super.destroy();
-    this.getTabId = () => 0;
   }
 }

packages/web-integration/src/chrome-extension/page.ts

@@ -25,16 +25,28 @@ function sleep(ms: number) {
 export default class ChromeExtensionProxyPage implements AbstractPage {
   pageType = 'chrome-extension-proxy';
 
-  public getTabId: () => number;
+  public trackingActiveTab: () => boolean;
 
   private viewportSize?: Size;
 
-  private debuggerAttached = false;
+  private activeTabId: number | null = null;
 
   private attachingDebugger: Promise<void> | null = null;
 
-  constructor(getTabId: () => number) {
-    this.getTabId = getTabId;
+  constructor(trackingActiveTab: () => boolean) {
+    this.trackingActiveTab = trackingActiveTab;
+  }
+
+  public async getTabId() {
+    const trackingActiveTab = this.trackingActiveTab();
+    if (this.activeTabId && !trackingActiveTab) {
+      return this.activeTabId;
+    }
+    const tabId = await chrome.tabs
+      .query({ active: true, currentWindow: true })
+      .then((tabs) => tabs[0]?.id);
+    this.activeTabId = tabId || 0;
+    return this.activeTabId;
   }
 
   private async attachDebugger() {
@@ -54,7 +66,7 @@ export default class ChromeExtensionProxyPage implements AbstractPage {
       }
 
       try {
-        const currentTabId = this.getTabId();
+        const currentTabId = await this.getTabId();
         // check if debugger is already attached to the tab
         const targets = await chrome.debugger.getTargets();
         const target = targets.find(
@@ -63,16 +75,8 @@ export default class ChromeExtensionProxyPage implements AbstractPage {
         if (!target) {
           // await chrome.debugger.detach({ tabId: currentTabId });
           await chrome.debugger.attach({ tabId: currentTabId }, '1.3');
-          this.debuggerAttached = true;
           // Prevent AI logic from being influenced by changes in page width and height due to the debugger banner appearing on attach.
           await sleep(340);
-          // listen to the debugger detach event
-          chrome.debugger.onEvent.addListener((source, method, params) => {
-            console.log('debugger event', source, method, params);
-            if (method === 'Debugger.detached') {
-              this.debuggerAttached = false;
-            }
-          });
         }
       } finally {
         this.attachingDebugger = null;
@@ -113,7 +117,6 @@ export default class ChromeExtensionProxyPage implements AbstractPage {
           chrome.debugger.detach({ tabId: tab.tabId });
         }
       }
-      this.debuggerAttached = false;
     }
   }
 
@@ -122,7 +125,7 @@ export default class ChromeExtensionProxyPage implements AbstractPage {
     params: RequestType,
   ): Promise<ResponseType> {
     await this.attachDebugger();
-    const tabId = this.getTabId();
+    const tabId = await this.getTabId();
     this.enableWaterFlowAnimation(tabId);
     return (await chrome.debugger.sendCommand(
       { tabId },
@@ -222,7 +225,8 @@ export default class ChromeExtensionProxyPage implements AbstractPage {
   }
 
   async url() {
-    const url = await chrome.tabs.get(this.getTabId()).then((tab) => tab.url);
+    const tabId = await this.getTabId();
+    const url = await chrome.tabs.get(tabId).then((tab) => tab.url);
     return url || '';
   }
 
@@ -380,6 +384,7 @@ export default class ChromeExtensionProxyPage implements AbstractPage {
   };
 
   async destroy(): Promise<void> {
+    this.activeTabId = null;
     await this.detachDebugger();
   }
 }