use CDP listener instead of playwright

Author: seanmcguire12

lib/StagehandContext.ts

@@ -1,11 +1,13 @@
 import type {
   BrowserContext as PlaywrightContext,
+  CDPSession,
   Page as PlaywrightPage,
 } from "playwright";
 import { Stagehand } from "./index";
 import { StagehandPage } from "./StagehandPage";
 import { Page } from "../types/page";
 import { EnhancedContext } from "../types/context";
+import { Protocol } from "devtools-protocol";
 
 export class StagehandContext {
   private readonly stagehand: Stagehand;
@@ -84,20 +86,32 @@ export class StagehandContext {
     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) => {
-      instance.handleNewPlaywrightPage(pwPage).catch((err) =>
-        stagehand.logger({
-          category: "context",
-          message: `Failed to initialise new page: ${err}`,
-          level: 0,
-        }),
-      );
+      instance
+        .attachFrameNavigatedListener(pwPage)
+        .catch((err) =>
+          stagehand.logger({
+            category: "cdp",
+            message: `Failed to attach frameNavigated listener: ${err}`,
+            level: 0,
+          }),
+        )
+        .finally(() =>
+          instance.handleNewPlaywrightPage(pwPage).catch((err) =>
+            stagehand.logger({
+              category: "context",
+              message: `Failed to initialise new page: ${err}`,
+              level: 0,
+            }),
+          ),
+        );
     });
 
     return instance;
@@ -156,4 +170,33 @@ export class StagehandContext {
     }
     this.setActivePage(stagehandPage);
   }
+
+  private async attachFrameNavigatedListener(
+    pwPage: PlaywrightPage,
+  ): Promise<void> {
+    const shPage = this.pageMap.get(pwPage);
+    if (!shPage) return;
+    const session: CDPSession = await this.intContext.newCDPSession(pwPage);
+    await session.send("Page.enable");
+
+    pwPage.once("close", () => {
+      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);
+          }
+        }
+      },
+    );
+  }
 }

lib/StagehandPage.ts

@@ -62,6 +62,10 @@ export class StagehandPage {
     return this.rootFrameId;
   }
 
+  public updateRootFrameId(newId: string): void {
+    this.rootFrameId = newId;
+  }
+
   constructor(
     page: PlaywrightPage,
     stagehand: Stagehand,
@@ -434,31 +438,6 @@ ${scriptContent} \
 
       this.intPage = new Proxy(page, handler) as unknown as Page;
       this.initialized = true;
-      const { frameTree } =
-        await this.sendCDP<Protocol.Page.GetFrameTreeResponse>(
-          "Page.getFrameTree",
-        );
-      this.rootFrameId = frameTree.frame.id;
-      this.intContext.registerFrameId(this.rootFrameId, this);
-
-      this.page.once("close", () => {
-        this.intContext.unregisterFrameId(this.rootFrameId);
-      });
-
-      this.page.on("framenavigated", async (frame) => {
-        if (frame.parentFrame() === null) {
-          const { frameTree: ft } =
-            await this.sendCDP<Protocol.Page.GetFrameTreeResponse>(
-              "Page.getFrameTree",
-            );
-          const newId = ft.frame.id;
-          if (newId !== this.rootFrameId) {
-            this.intContext.unregisterFrameId(this.rootFrameId);
-            this.rootFrameId = newId;
-            this.intContext.registerFrameId(this.rootFrameId, this);
-          }
-        }
-      });
       return this;
     } catch (err: unknown) {
       if (err instanceof StagehandError || err instanceof StagehandAPIError) {