[api] update root frame ID on init (#920)

# why
- tab handling for `act`, `extract` and `observe` was failing because we
sometimes sent `frameId: undefined` after a new page had been opened
- the `rootFrameId` only became available after `Page.frameNavigated`
fired, so the time between the event firing and a stagehand API call
caused undefined `frameId`s
# what changed
- Added `getCurrentRootFrameId` helper which calls `Page.getFrameTree`
to fetch the current root frame ID on page creation
- `attachFrameNavigatedListener` is now called right after a page is
wrapped with `createStagehandPage`
- `handleNewPlaywrightPage` now runs before attaching the listener
# test plan
- tested in local dev 
- run all evals

Author: seanmcguire12

.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

@@ -33,6 +33,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;
@@ -366,7 +373,7 @@ ${scriptContent} \
               const result = this.api
                 ? await this.api.goto(url, {
                     ...options,
-                    frameId: this.frameId,
+                    frameId: this.rootFrameId,
                   })
                 : await rawGoto(url, options);
 
@@ -435,6 +442,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;
@@ -639,7 +653,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);
@@ -672,7 +686,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);
@@ -734,7 +748,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();
         }
@@ -767,7 +781,7 @@ ${scriptContent} \
       } = options;
 
       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;
@@ -871,7 +885,7 @@ ${scriptContent} \
       } = options;
 
       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;