Skip to content

Commit 2b70347

Browse files
seanmcguire12seanmcguire12
authored
[feat]: add context.addInitScript() (#1300)
# why - adds a public facing API to inject scripts into new pages before existing JS is run # what changed - added `context.addInitScripts()` which calls `Page.addScriptToEvaluateOnNewDocument` whenever we attach to a page target # test plan - added unit tests which check that the injected scripts are run before existing scripts
1 parent dbcb892 commit 2b70347

File tree

4 files changed

+233
-1
lines changed

4 files changed

+233
-1
lines changed

‎.changeset/tender-moons-laugh.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"@browserbasehq/stagehand": patch
3+
---
4+
5+
add support for context.addInitScript()

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

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
import { test, expect } from "@playwright/test";
2+
import { V3 } from "../v3";
3+
import { v3TestConfig } from "./v3.config";
4+
import { V3Context } from "../understudy/context";
5+
6+
const toDataUrl = (html: string): string =>
7+
`data:text/html,${encodeURIComponent(html)}`;
8+
9+
test.describe("context.addInitScript", () => {
10+
let v3: V3;
11+
let ctx: V3Context;
12+
13+
test.beforeEach(async () => {
14+
v3 = new V3(v3TestConfig);
15+
await v3.init();
16+
ctx = v3.context;
17+
});
18+
19+
test.afterEach(async () => {
20+
await v3?.close?.().catch(() => {});
21+
});
22+
23+
test("runs before inline document scripts on navigation", async () => {
24+
const page = await ctx.awaitActivePage();
25+
26+
await ctx.addInitScript(() => {
27+
(window as unknown as { __fromContextInit?: string }).__fromContextInit =
28+
"injected-value";
29+
});
30+
31+
const html = `<!DOCTYPE html>
32+
<html>
33+
<body>
34+
<script>
35+
var value = (window && window.__fromContextInit) || 'missing';
36+
document.body.dataset.initWitness = value;
37+
</script>
38+
</body>
39+
</html>`;
40+
41+
await page.goto(toDataUrl(html), { waitUntil: "load" });
42+
43+
const observed = await page.evaluate(() => {
44+
return document.body.dataset.initWitness;
45+
});
46+
expect(observed).toBe("injected-value");
47+
});
48+
49+
test("re-applies the script on every navigation for the same page", async () => {
50+
const page = await ctx.awaitActivePage();
51+
52+
await ctx.addInitScript(() => {
53+
function markVisit() {
54+
const root = document.documentElement;
55+
if (!root) return;
56+
const current = Number(window.name || "0");
57+
const next = current + 1;
58+
window.name = String(next);
59+
root.dataset.visitCount = String(next);
60+
}
61+
if (document.readyState === "loading") {
62+
document.addEventListener("DOMContentLoaded", markVisit, {
63+
once: true,
64+
});
65+
} else {
66+
markVisit();
67+
}
68+
});
69+
70+
await page.goto(toDataUrl("<html><body>first</body></html>"), {
71+
waitUntil: "load",
72+
});
73+
const first = await page.evaluate(() => {
74+
return Number(document.documentElement.dataset.visitCount ?? "0");
75+
});
76+
expect(first).toBe(1);
77+
78+
await page.goto(toDataUrl("<html><body>second</body></html>"), {
79+
waitUntil: "load",
80+
});
81+
const second = await page.evaluate(() => {
82+
return Number(document.documentElement.dataset.visitCount ?? "0");
83+
});
84+
expect(second).toBe(2);
85+
});
86+
87+
test("applies script (with args) to newly created pages", async () => {
88+
const payload = { greeting: "hi", nested: { count: 2 } };
89+
90+
await ctx.addInitScript((arg) => {
91+
function setPayload(): void {
92+
const root = document.documentElement;
93+
if (!root) return;
94+
root.dataset.initPayload = JSON.stringify(arg);
95+
}
96+
if (document.readyState === "loading") {
97+
document.addEventListener("DOMContentLoaded", setPayload, {
98+
once: true,
99+
});
100+
} else {
101+
setPayload();
102+
}
103+
}, payload);
104+
105+
const newPage = await ctx.newPage();
106+
await newPage.goto(toDataUrl("<html><body>child</body></html>"), {
107+
waitUntil: "load",
108+
});
109+
110+
const observed = await newPage.evaluate(() => {
111+
const raw = document.documentElement.dataset.initPayload;
112+
return raw ? JSON.parse(raw) : undefined;
113+
});
114+
expect(observed).toEqual(payload);
115+
});
116+
});

‎packages/core/lib/v3/understudy/context.ts‎

Lines changed: 77 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
// lib/v3/understudy/context.ts
2+
import { promises as fs } from "fs";
23
import type { Protocol } from "devtools-protocol";
34
import { v3Logger } from "../logger";
45
import { CdpConnection, CDPSessionLike } from "./cdp";
@@ -7,13 +8,70 @@ import { installV3PiercerIntoSession } from "./piercer";
78
import { executionContexts } from "./executionContextRegistry";
89
import type { StagehandAPIClient } from "../api";
910
import { LocalBrowserLaunchOptions } from "../types/public";
10-
import { TimeoutError, PageNotFoundError } from "../types/public/sdkErrors";
11+
import {
12+
StagehandInvalidArgumentError,
13+
TimeoutError,
14+
PageNotFoundError,
15+
} from "../types/public/sdkErrors";
1116

1217
type TargetId = string;
1318
type SessionId = string;
1419

1520
type TargetType = "page" | "iframe" | string;
1621

22+
type InitScriptSource<Arg> =
23+
| string
24+
| { path?: string; content?: string }
25+
| ((arg: Arg) => unknown);
26+
27+
async function normalizeInitScriptSource<Arg>(
28+
script: InitScriptSource<Arg>,
29+
arg?: Arg,
30+
): Promise<string> {
31+
if (typeof script === "function") {
32+
const argString = Object.is(arg, undefined)
33+
? "undefined"
34+
: JSON.stringify(arg);
35+
return `(${script.toString()})(${argString})`;
36+
}
37+
38+
if (!Object.is(arg, undefined)) {
39+
throw new StagehandInvalidArgumentError(
40+
"context.addInitScript: 'arg' is only supported when passing a function.",
41+
);
42+
}
43+
44+
if (typeof script === "string") {
45+
return script;
46+
}
47+
48+
if (!script || typeof script !== "object") {
49+
throw new StagehandInvalidArgumentError(
50+
"context.addInitScript: provide a string, function, or an object with path/content.",
51+
);
52+
}
53+
54+
if (typeof script.content === "string") {
55+
return script.content;
56+
}
57+
58+
if (typeof script.path === "string" && script.path.trim()) {
59+
const raw = await fs.readFile(script.path, "utf8");
60+
return appendSourceURL(raw, script.path);
61+
}
62+
63+
throw new StagehandInvalidArgumentError(
64+
"context.addInitScript: provide a string, function, or an object with path/content.",
65+
);
66+
}
67+
68+
// Chrome surfaces injected scripts using a //# sourceURL tag; mirroring Playwright keeps
69+
// stack traces and console errors pointing back to the preload file when path is used.
70+
function appendSourceURL(source: string, filePath: string): string {
71+
const sanitized = filePath.replace(/\n/g, "");
72+
return `${source}\n//# sourceURL=${sanitized}`;
73+
}
74+
1775
function isTopLevelPage(info: Protocol.Target.TargetInfo): boolean {
1876
const ti = info as unknown as { subtype?: string };
1977
return info.type === "page" && ti.subtype !== "iframe";
@@ -55,6 +113,7 @@ export class V3Context {
55113
private typeByTarget = new Map<TargetId, TargetType>();
56114
private _pageOrder: TargetId[] = [];
57115
private pendingCreatedTargetUrl = new Map<TargetId, string>();
116+
private readonly initScripts: string[] = [];
58117

59118
/**
60119
* Create a Context for a given CDP websocket URL and bootstrap target wiring.
@@ -206,6 +265,16 @@ export class V3Context {
206265
void this.conn.send("Target.activateTarget", { targetId }).catch(() => {});
207266
}
208267

268+
public async addInitScript<Arg>(
269+
script: InitScriptSource<Arg>,
270+
arg?: Arg,
271+
): Promise<void> {
272+
const source = await normalizeInitScriptSource(script, arg);
273+
this.initScripts.push(source);
274+
const pages = this.pages();
275+
await Promise.all(pages.map((page) => page.registerInitScript(source)));
276+
}
277+
209278
/**
210279
* Return top-level `Page`s (oldest → newest). OOPIF targets are not included.
211280
*/
@@ -220,6 +289,12 @@ export class V3Context {
220289
return rows.map((r) => r.page);
221290
}
222291

292+
private async applyInitScriptsToPage(page: Page): Promise<void> {
293+
for (const source of this.initScripts) {
294+
await page.registerInitScript(source);
295+
}
296+
}
297+
223298
/**
224299
* Resolve an owning `Page` by the **top-level main frame id**.
225300
* Note: child (OOPIF) roots are intentionally not present in this mapping.
@@ -407,6 +482,7 @@ export class V3Context {
407482
page.seedCurrentUrl(pendingSeedUrl ?? info.url ?? "");
408483
this._pushActive(info.targetId);
409484
this.installFrameEventBridges(sessionId, page);
485+
await this.applyInitScriptsToPage(page);
410486

411487
return;
412488
}

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,8 @@ export class Page {
9494
string,
9595
(evt: Protocol.Runtime.ConsoleAPICalledEvent) => void
9696
>();
97+
/** Document-start scripts installed across every session this page owns. */
98+
private readonly initScripts: string[] = [];
9799

98100
private constructor(
99101
private readonly conn: CdpConnection,
@@ -125,6 +127,37 @@ export class Page {
125127
this.networkManager.trackSession(this.mainSession);
126128
}
127129

130+
// Send a single init script to a specific CDP session.
131+
private async installInitScriptOnSession(
132+
session: CDPSessionLike,
133+
source: string,
134+
): Promise<void> {
135+
await session.send("Page.addScriptToEvaluateOnNewDocument", { source });
136+
}
137+
138+
// Replay every previously registered init script onto a newly adopted session.
139+
private async applyInitScriptsToSession(
140+
session: CDPSessionLike,
141+
): Promise<void> {
142+
for (const source of this.initScripts) {
143+
await this.installInitScriptOnSession(session, source);
144+
}
145+
}
146+
147+
// Register a new init script and fan it out to all active sessions for this page.
148+
public async registerInitScript(source: string): Promise<void> {
149+
if (this.initScripts.includes(source)) return;
150+
this.initScripts.push(source);
151+
152+
const installs: Array<Promise<void>> = [];
153+
installs.push(this.installInitScriptOnSession(this.mainSession, source));
154+
for (const session of this.sessions.values()) {
155+
if (session === this.mainSession) continue;
156+
installs.push(this.installInitScriptOnSession(session, source));
157+
}
158+
await Promise.all(installs);
159+
}
160+
128161
// --- Optional visual cursor overlay management ---
129162
private cursorEnabled = false;
130163
private async ensureCursorScript(): Promise<void> {
@@ -371,6 +404,8 @@ export class Page {
371404

372405
this.networkManager.trackSession(childSession);
373406

407+
void this.applyInitScriptsToSession(childSession).catch(() => {});
408+
374409
if (this.consoleListeners.size > 0) {
375410
this.installConsoleTap(childSession);
376411
}

0 commit comments

Comments
 (0)