diff --git a/.changeset/tender-moons-laugh.md b/.changeset/tender-moons-laugh.md new file mode 100644 index 000000000..66af857b2 --- /dev/null +++ b/.changeset/tender-moons-laugh.md @@ -0,0 +1,5 @@ +--- +"@browserbasehq/stagehand": patch +--- + +add support for context.addInitScript() diff --git a/packages/core/lib/v3/tests/context-addInitScript.spec.ts b/packages/core/lib/v3/tests/context-addInitScript.spec.ts new file mode 100644 index 000000000..939ef98b7 --- /dev/null +++ b/packages/core/lib/v3/tests/context-addInitScript.spec.ts @@ -0,0 +1,116 @@ +import { test, expect } from "@playwright/test"; +import { V3 } from "../v3"; +import { v3TestConfig } from "./v3.config"; +import { V3Context } from "../understudy/context"; + +const toDataUrl = (html: string): string => + `data:text/html,${encodeURIComponent(html)}`; + +test.describe("context.addInitScript", () => { + let v3: V3; + let ctx: V3Context; + + test.beforeEach(async () => { + v3 = new V3(v3TestConfig); + await v3.init(); + ctx = v3.context; + }); + + test.afterEach(async () => { + await v3?.close?.().catch(() => {}); + }); + + test("runs before inline document scripts on navigation", async () => { + const page = await ctx.awaitActivePage(); + + await ctx.addInitScript(() => { + (window as unknown as { __fromContextInit?: string }).__fromContextInit = + "injected-value"; + }); + + const html = ` + + + + + `; + + await page.goto(toDataUrl(html), { waitUntil: "load" }); + + const observed = await page.evaluate(() => { + return document.body.dataset.initWitness; + }); + expect(observed).toBe("injected-value"); + }); + + test("re-applies the script on every navigation for the same page", async () => { + const page = await ctx.awaitActivePage(); + + await ctx.addInitScript(() => { + function markVisit() { + const root = document.documentElement; + if (!root) return; + const current = Number(window.name || "0"); + const next = current + 1; + window.name = String(next); + root.dataset.visitCount = String(next); + } + if (document.readyState === "loading") { + document.addEventListener("DOMContentLoaded", markVisit, { + once: true, + }); + } else { + markVisit(); + } + }); + + await page.goto(toDataUrl("first"), { + waitUntil: "load", + }); + const first = await page.evaluate(() => { + return Number(document.documentElement.dataset.visitCount ?? "0"); + }); + expect(first).toBe(1); + + await page.goto(toDataUrl("second"), { + waitUntil: "load", + }); + const second = await page.evaluate(() => { + return Number(document.documentElement.dataset.visitCount ?? "0"); + }); + expect(second).toBe(2); + }); + + test("applies script (with args) to newly created pages", async () => { + const payload = { greeting: "hi", nested: { count: 2 } }; + + await ctx.addInitScript((arg) => { + function setPayload(): void { + const root = document.documentElement; + if (!root) return; + root.dataset.initPayload = JSON.stringify(arg); + } + if (document.readyState === "loading") { + document.addEventListener("DOMContentLoaded", setPayload, { + once: true, + }); + } else { + setPayload(); + } + }, payload); + + const newPage = await ctx.newPage(); + await newPage.goto(toDataUrl("child"), { + waitUntil: "load", + }); + + const observed = await newPage.evaluate(() => { + const raw = document.documentElement.dataset.initPayload; + return raw ? JSON.parse(raw) : undefined; + }); + expect(observed).toEqual(payload); + }); +}); diff --git a/packages/core/lib/v3/understudy/context.ts b/packages/core/lib/v3/understudy/context.ts index 2fbf19afb..d5d8ae171 100644 --- a/packages/core/lib/v3/understudy/context.ts +++ b/packages/core/lib/v3/understudy/context.ts @@ -1,4 +1,5 @@ // lib/v3/understudy/context.ts +import { promises as fs } from "fs"; import type { Protocol } from "devtools-protocol"; import { v3Logger } from "../logger"; import { CdpConnection, CDPSessionLike } from "./cdp"; @@ -7,13 +8,70 @@ import { installV3PiercerIntoSession } from "./piercer"; import { executionContexts } from "./executionContextRegistry"; import type { StagehandAPIClient } from "../api"; import { LocalBrowserLaunchOptions } from "../types/public"; -import { TimeoutError, PageNotFoundError } from "../types/public/sdkErrors"; +import { + StagehandInvalidArgumentError, + TimeoutError, + PageNotFoundError, +} from "../types/public/sdkErrors"; type TargetId = string; type SessionId = string; type TargetType = "page" | "iframe" | string; +type InitScriptSource = + | string + | { path?: string; content?: string } + | ((arg: Arg) => unknown); + +async function normalizeInitScriptSource( + script: InitScriptSource, + arg?: Arg, +): Promise { + if (typeof script === "function") { + const argString = Object.is(arg, undefined) + ? "undefined" + : JSON.stringify(arg); + return `(${script.toString()})(${argString})`; + } + + if (!Object.is(arg, undefined)) { + throw new StagehandInvalidArgumentError( + "context.addInitScript: 'arg' is only supported when passing a function.", + ); + } + + if (typeof script === "string") { + return script; + } + + if (!script || typeof script !== "object") { + throw new StagehandInvalidArgumentError( + "context.addInitScript: provide a string, function, or an object with path/content.", + ); + } + + if (typeof script.content === "string") { + return script.content; + } + + if (typeof script.path === "string" && script.path.trim()) { + const raw = await fs.readFile(script.path, "utf8"); + return appendSourceURL(raw, script.path); + } + + throw new StagehandInvalidArgumentError( + "context.addInitScript: provide a string, function, or an object with path/content.", + ); +} + +// Chrome surfaces injected scripts using a //# sourceURL tag; mirroring Playwright keeps +// stack traces and console errors pointing back to the preload file when path is used. +function appendSourceURL(source: string, filePath: string): string { + const sanitized = filePath.replace(/\n/g, ""); + return `${source}\n//# sourceURL=${sanitized}`; +} + function isTopLevelPage(info: Protocol.Target.TargetInfo): boolean { const ti = info as unknown as { subtype?: string }; return info.type === "page" && ti.subtype !== "iframe"; @@ -55,6 +113,7 @@ export class V3Context { private typeByTarget = new Map(); private _pageOrder: TargetId[] = []; private pendingCreatedTargetUrl = new Map(); + private readonly initScripts: string[] = []; /** * Create a Context for a given CDP websocket URL and bootstrap target wiring. @@ -206,6 +265,16 @@ export class V3Context { void this.conn.send("Target.activateTarget", { targetId }).catch(() => {}); } + public async addInitScript( + script: InitScriptSource, + arg?: Arg, + ): Promise { + const source = await normalizeInitScriptSource(script, arg); + this.initScripts.push(source); + const pages = this.pages(); + await Promise.all(pages.map((page) => page.registerInitScript(source))); + } + /** * Return top-level `Page`s (oldest → newest). OOPIF targets are not included. */ @@ -220,6 +289,12 @@ export class V3Context { return rows.map((r) => r.page); } + private async applyInitScriptsToPage(page: Page): Promise { + for (const source of this.initScripts) { + await page.registerInitScript(source); + } + } + /** * Resolve an owning `Page` by the **top-level main frame id**. * Note: child (OOPIF) roots are intentionally not present in this mapping. @@ -407,6 +482,7 @@ export class V3Context { page.seedCurrentUrl(pendingSeedUrl ?? info.url ?? ""); this._pushActive(info.targetId); this.installFrameEventBridges(sessionId, page); + await this.applyInitScriptsToPage(page); return; } diff --git a/packages/core/lib/v3/understudy/page.ts b/packages/core/lib/v3/understudy/page.ts index 579552237..72ed0b27b 100644 --- a/packages/core/lib/v3/understudy/page.ts +++ b/packages/core/lib/v3/understudy/page.ts @@ -94,6 +94,8 @@ export class Page { string, (evt: Protocol.Runtime.ConsoleAPICalledEvent) => void >(); + /** Document-start scripts installed across every session this page owns. */ + private readonly initScripts: string[] = []; private constructor( private readonly conn: CdpConnection, @@ -125,6 +127,37 @@ export class Page { this.networkManager.trackSession(this.mainSession); } + // Send a single init script to a specific CDP session. + private async installInitScriptOnSession( + session: CDPSessionLike, + source: string, + ): Promise { + await session.send("Page.addScriptToEvaluateOnNewDocument", { source }); + } + + // Replay every previously registered init script onto a newly adopted session. + private async applyInitScriptsToSession( + session: CDPSessionLike, + ): Promise { + for (const source of this.initScripts) { + await this.installInitScriptOnSession(session, source); + } + } + + // Register a new init script and fan it out to all active sessions for this page. + public async registerInitScript(source: string): Promise { + if (this.initScripts.includes(source)) return; + this.initScripts.push(source); + + const installs: Array> = []; + installs.push(this.installInitScriptOnSession(this.mainSession, source)); + for (const session of this.sessions.values()) { + if (session === this.mainSession) continue; + installs.push(this.installInitScriptOnSession(session, source)); + } + await Promise.all(installs); + } + // --- Optional visual cursor overlay management --- private cursorEnabled = false; private async ensureCursorScript(): Promise { @@ -371,6 +404,8 @@ export class Page { this.networkManager.trackSession(childSession); + void this.applyInitScriptsToSession(childSession).catch(() => {}); + if (this.consoleListeners.size > 0) { this.installConsoleTap(childSession); }