[api] update root frame ID on init

.changeset/clean-olives-wash.md

@@ -0,0 +1,5 @@
+---
+"@browserbasehq/stagehand": patch
+---
+
+fix: tab handling on API

lib/StagehandContext.ts

@@ -27,6 +27,7 @@ export class StagehandContext {
           return async (): Promise<Page> => {
             const pwPage = await target.newPage();
             const stagehandPage = await this.createStagehandPage(pwPage);
+            await this.attachFrameNavigatedListener(pwPage);
             // Set as active page when created
             this.setActivePage(stagehandPage);
             return stagehandPage.page;
@@ -81,19 +82,8 @@ export class StagehandContext {
     stagehand: Stagehand,
   ): Promise<StagehandContext> {
     const instance = new StagehandContext(context, stagehand);
-
-    // Initialize existing pages
-    const existingPages = context.pages();
-    for (const page of existingPages) {
-      const stagehandPage = await instance.createStagehandPage(page);
-      await instance.attachFrameNavigatedListener(page);
-      // Set the first page as active
-      if (!instance.activeStagehandPage) {
-        instance.setActivePage(stagehandPage);
-      }
-    }
-
-    context.on("page", (pwPage) => {
+    context.on("page", async (pwPage) => {
+      await instance.handleNewPlaywrightPage(pwPage);
       instance
         .attachFrameNavigatedListener(pwPage)
         .catch((err) =>
@@ -114,6 +104,17 @@ export class StagehandContext {
         );
     });
 
+    // Initialize existing pages
+    const existingPages = context.pages();
+    for (const page of existingPages) {
+      const stagehandPage = await instance.createStagehandPage(page);
+      await instance.attachFrameNavigatedListener(page);
+      // Set the first page as active
+      if (!instance.activeStagehandPage) {
+        instance.setActivePage(stagehandPage);
+      }
+    }
+
     return instance;
   }
   public get frameIdLookup(): ReadonlyMap<string, StagehandPage> {
@@ -180,22 +181,19 @@ export class StagehandContext {
     await session.send("Page.enable");
 
     pwPage.once("close", () => {
-      this.unregisterFrameId(shPage.frameId);
+      if (shPage.frameId) this.unregisterFrameId(shPage.frameId);
     });
 
     session.on(
       "Page.frameNavigated",
       (evt: Protocol.Page.FrameNavigatedEvent): void => {
-        const { frame } = evt;
-
-        if (!frame.parentId) {
-          const oldId = shPage.frameId;
-          if (frame.id !== oldId) {
-            if (oldId) this.unregisterFrameId(oldId);
-            this.registerFrameId(frame.id, shPage);
-            shPage.updateRootFrameId(frame.id);
-          }
-        }
+        if (evt.frame.parentId) return;
+        if (evt.frame.id === shPage.frameId) return;
+
+        const oldId = shPage.frameId;
+        if (oldId) this.unregisterFrameId(oldId);
+        this.registerFrameId(evt.frame.id, shPage);
+        shPage.updateRootFrameId(evt.frame.id);
       },
     );
   }

lib/StagehandPage.ts

@@ -34,6 +34,13 @@ import { StagehandAPIError } from "@/types/stagehandApiErrors";
 import { scriptContent } from "@/lib/dom/build/scriptContent";
 import type { Protocol } from "devtools-protocol";
 
+async function getCurrentRootFrameId(session: CDPSession): Promise<string> {
+  const { frameTree } = (await session.send(
+    "Page.getFrameTree",
+  )) as Protocol.Page.GetFrameTreeResponse;
+  return frameTree.frame.id;
+}
+
 export class StagehandPage {
   private stagehand: Stagehand;
   private rawPage: PlaywrightPage;
@@ -367,7 +374,7 @@ ${scriptContent} \
               const result = this.api
                 ? await this.api.goto(url, {
                     ...options,
-                    frameId: this.frameId,
+                    frameId: this.rootFrameId,
                   })
                 : await rawGoto(url, options);
 
@@ -436,6 +443,13 @@ ${scriptContent} \
         },
       };
 
+      const session = await this.getCDPClient(this.rawPage);
+      await session.send("Page.enable");
+
+      const rootId = await getCurrentRootFrameId(session);
+      this.updateRootFrameId(rootId);
+      this.intContext.registerFrameId(rootId, this);
+
       this.intPage = new Proxy(page, handler) as unknown as Page;
       this.initialized = true;
       return this;
@@ -643,7 +657,7 @@ ${scriptContent} \
           if (this.api) {
             const result = await this.api.act({
               ...observeResult,
-              frameId: this.frameId,
+              frameId: this.rootFrameId,
             });
             await this._refreshPageFromAPI();
             this.stagehand.addToHistory("act", observeResult, result);
@@ -676,7 +690,7 @@ ${scriptContent} \
       const { action, modelName, modelClientOptions } = actionOrOptions;
 
       if (this.api) {
-        const opts = { ...actionOrOptions, frameId: this.frameId };
+        const opts = { ...actionOrOptions, frameId: this.rootFrameId };
         const result = await this.api.act(opts);
         await this._refreshPageFromAPI();
         this.stagehand.addToHistory("act", actionOrOptions, result);
@@ -738,7 +752,7 @@ ${scriptContent} \
       if (!instructionOrOptions) {
         let result: ExtractResult<T>;
         if (this.api) {
-          result = await this.api.extract<T>({ frameId: this.frameId });
+          result = await this.api.extract<T>({ frameId: this.rootFrameId });
         } else {
           result = await this.extractHandler.extract();
         }
@@ -775,7 +789,7 @@ ${scriptContent} \
       }
 
       if (this.api) {
-        const opts = { ...options, frameId: this.frameId };
+        const opts = { ...options, frameId: this.rootFrameId };
         const result = await this.api.extract<T>(opts);
         this.stagehand.addToHistory("extract", instructionOrOptions, result);
         return result;
@@ -883,7 +897,7 @@ ${scriptContent} \
       }
 
       if (this.api) {
-        const opts = { ...options, frameId: this.frameId };
+        const opts = { ...options, frameId: this.rootFrameId };
         const result = await this.api.observe(opts);
         this.stagehand.addToHistory("observe", instructionOrOptions, result);
         return result;