[fix]: make page.evaluate() use main world

.changeset/seven-mice-draw.md

@@ -0,0 +1,5 @@
+---
+"@browserbasehq/stagehand": patch
+---
+
+fix: page.evaluate() now works with scripts injected via context.addInitScript()

packages/core/lib/v3/tests/context-addInitScript.spec.ts

@@ -113,4 +113,25 @@ test.describe("context.addInitScript", () => {
     });
     expect(observed).toEqual(payload);
   });
+
+  test("context.addInitScript installs a function callable from page.evaluate", async () => {
+    const page = await ctx.awaitActivePage();
+
+    await ctx.addInitScript(() => {
+      // installed before any navigation
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-expect-error
+      window.sayHelloFromStagehand = () => "hello from stagehand";
+    });
+
+    await page.goto("https://example.com", { waitUntil: "domcontentloaded" });
+
+    const result = await page.evaluate(() => {
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-expect-error
+      return window.sayHelloFromStagehand();
+    });
+
+    expect(result).toBe("hello from stagehand");
+  });
 });

packages/core/lib/v3/understudy/frame.ts

@@ -3,6 +3,7 @@ import { Protocol } from "devtools-protocol";
 import type { CDPSessionLike } from "./cdp";
 import { Locator } from "./locator";
 import { StagehandEvalError } from "../types/public/sdkErrors";
+import { executionContexts } from "./executionContextRegistry";
 
 interface FrameManager {
   session: CDPSessionLike;
@@ -116,7 +117,7 @@ export class Frame implements FrameManager {
   }
 
   /**
-   * Evaluate a function or expression in this frame's isolated world.
+   * Evaluate a function or expression in this frame's main world.
    * - If a string is provided, treated as a JS expression.
    * - If a function is provided, it is stringified and invoked with the optional argument.
    */
@@ -125,7 +126,7 @@ export class Frame implements FrameManager {
     arg?: Arg,
   ): Promise<R> {
     await this.session.send("Runtime.enable").catch(() => {});
-    const contextId = await this.getExecutionContextId();
+    const contextId = await this.getMainWorldExecutionContextId();
 
     const isString = typeof pageFunctionOrExpression === "string";
     let expression: string;
@@ -293,16 +294,8 @@ export class Frame implements FrameManager {
     return new Locator(this, selector, options);
   }
 
-  /** Create/get an isolated world for this frame and return its executionContextId */
-  private async getExecutionContextId(): Promise<number> {
-    await this.session.send("Page.enable");
-    await this.session.send("Runtime.enable");
-    const { executionContextId } = await this.session.send<{
-      executionContextId: number;
-    }>("Page.createIsolatedWorld", {
-      frameId: this.frameId,
-      worldName: "v3-world",
-    });
-    return executionContextId;
+  /** Resolve the main-world execution context id for this frame. */
+  private async getMainWorldExecutionContextId(): Promise<number> {
+    return executionContexts.waitForMainWorld(this.session, this.frameId, 1000);
   }
 }

packages/core/lib/v3/understudy/page.ts

@@ -8,6 +8,7 @@ import { FrameLocator } from "./frameLocator";
 import { deepLocatorFromPage } from "./deepLocator";
 import { resolveXpathForLocation } from "./a11y/snapshot";
 import { FrameRegistry } from "./frameRegistry";
+import { executionContexts } from "./executionContextRegistry";
 import { LoadState } from "../types/public/page";
 import { NetworkManager } from "./networkManager";
 import { LifecycleWatcher } from "./lifecycleWatcher";
@@ -132,7 +133,9 @@ export class Page {
     session: CDPSessionLike,
     source: string,
   ): Promise<void> {
-    await session.send("Page.addScriptToEvaluateOnNewDocument", { source });
+    await session.send("Page.addScriptToEvaluateOnNewDocument", {
+      source: source,
+    });
   }
 
   // Replay every previously registered init script onto a newly adopted session.
@@ -975,7 +978,7 @@ export class Page {
   async title(): Promise<string> {
     try {
       await this.mainSession.send("Runtime.enable").catch(() => {});
-      const ctxId = await this.createIsolatedWorldForCurrentMain();
+      const ctxId = await this.mainWorldExecutionContextId();
       const { result } =
         await this.mainSession.send<Protocol.Runtime.EvaluateResponse>(
           "Runtime.evaluate",
@@ -1157,7 +1160,7 @@ export class Page {
   }
 
   /**
-   * Evaluate a function or expression in the current main frame's isolated world.
+   * Evaluate a function or expression in the current main frame's main world.
    * - If a string is provided, it is treated as a JS expression.
    * - If a function is provided, it is stringified and invoked with the optional argument.
    * - The return value should be JSON-serializable. Non-serializable objects will
@@ -1168,7 +1171,7 @@ export class Page {
     arg?: Arg,
   ): Promise<R> {
     await this.mainSession.send("Runtime.enable").catch(() => {});
-    const ctxId = await this.createIsolatedWorldForCurrentMain();
+    const ctxId = await this.mainWorldExecutionContextId();
 
     const isString = typeof pageFunctionOrExpression === "string";
     let expression: string;
@@ -1979,18 +1982,13 @@ export class Page {
 
   // ---- Page-level lifecycle waiter that follows main frame id swaps ----
 
-  /**
-   * Create an isolated world for the **current** main frame and return its context id.
-   */
-  private async createIsolatedWorldForCurrentMain(): Promise<number> {
-    await this.mainSession.send("Runtime.enable").catch(() => {});
-    const { executionContextId } = await this.mainSession.send<{
-      executionContextId: number;
-    }>("Page.createIsolatedWorld", {
-      frameId: this.mainFrameId(),
-      worldName: "v3-world",
-    });
-    return executionContextId;
+  /** Resolve the main-world execution context for the current main frame. */
+  private async mainWorldExecutionContextId(): Promise<number> {
+    return executionContexts.waitForMainWorld(
+      this.mainSession,
+      this.mainFrameId(),
+      1000,
+    );
   }
 
   /**
@@ -2009,7 +2007,7 @@ export class Page {
 
     // Fast path: check the *current* main frame's readyState.
     try {
-      const ctxId = await this.createIsolatedWorldForCurrentMain();
+      const ctxId = await this.mainWorldExecutionContextId();
       const { result } =
         await this.mainSession.send<Protocol.Runtime.EvaluateResponse>(
           "Runtime.evaluate",