From 908323ac2fcb18080aa83ee93b296980a5c20619 Mon Sep 17 00:00:00 2001 From: miguel Date: Mon, 27 Jul 2026 15:35:19 -0700 Subject: [PATCH] Harden resident target bootstrap --- .../tests/resident-target-safety.test.ts | 310 ++++++++++++++++++ .../understudy-context-lifecycle.test.ts | 63 ++++ .../understudy/browserWebSocketTransport.ts | 3 +- packages/server/understudy/cdp.ts | 10 +- packages/server/understudy/context.ts | 154 +++++++-- packages/server/understudy/networkManager.ts | 45 ++- packages/server/understudy/page.ts | 69 +++- 7 files changed, 609 insertions(+), 45 deletions(-) create mode 100644 packages/server/tests/resident-target-safety.test.ts create mode 100644 packages/server/tests/understudy-context-lifecycle.test.ts diff --git a/packages/server/tests/resident-target-safety.test.ts b/packages/server/tests/resident-target-safety.test.ts new file mode 100644 index 000000000..dad03000e --- /dev/null +++ b/packages/server/tests/resident-target-safety.test.ts @@ -0,0 +1,310 @@ +import type { Protocol } from "devtools-protocol"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { ChromeTabTargetController } from "../understudy/chromeTabs.js"; +import { V3Context, isSupportedWebTarget } from "../understudy/context.js"; +import { CdpConnection, STAGEHAND_WEB_TARGET_FILTER } from "../understudy/cdp.js"; +import { Page } from "../understudy/page.js"; + +function target( + targetId: string, + type: string, + url: string, + attached = false, +): Protocol.Target.TargetInfo { + return { targetId, type, title: targetId, url, attached, canAccessOpener: false }; +} + +const chromeTabs: ChromeTabTargetController = { + activeTargetId: async () => undefined, + targetIdForTabId: async () => undefined, + tabIdForTargetId: async () => undefined, + activateTarget: async () => {}, +}; + +describe("resident target attachment safety", () => { + afterEach(() => vi.restoreAllMocks()); + + it("allowlists page and injectable iframe targets", () => { + expect(isSupportedWebTarget(target("page", "page", "about:blank"))).toBe(true); + expect(isSupportedWebTarget(target("internal", "page", "chrome://newtab"))).toBe(false); + expect( + isSupportedWebTarget(target("other-extension", "page", "chrome-extension://other")), + ).toBe(false); + expect(isSupportedWebTarget(target("iframe", "iframe", "https://example.com/frame"))).toBe( + true, + ); + expect(isSupportedWebTarget(target("extension", "iframe", "chrome-extension://other"))).toBe( + false, + ); + for (const type of [ + "browser", + "service_worker", + "worker", + "shared_worker", + "background_page", + ]) { + expect(isSupportedWebTarget(target(type, type, "https://example.com"))).toBe(false); + } + }); + + it("filters the initial sweep before manually attaching", async () => { + const attachToTarget = vi.fn(async () => ({})); + const context = new V3Context( + { + on: vi.fn(), + enableAutoAttach: vi.fn(async () => {}), + getTargets: vi.fn(async () => [ + target("page", "page", "about:blank"), + target("iframe", "iframe", "https://example.com/frame"), + target("worker", "worker", "https://example.com/worker.js"), + target("stagehand-worker", "service_worker", "chrome-extension://stagehand/sw.js"), + target("background", "background_page", "chrome-extension://other/background.html"), + ]), + attachToTarget, + } as never, + {} as never, + chromeTabs, + ); + vi.spyOn(context, "waitForInitialTopLevelTargets").mockResolvedValue(); + + await context.bootstrap(); + + expect(attachToTarget.mock.calls).toStrictEqual([["page"], ["iframe"]]); + }); + + it("attaches a previously ignored page after it navigates to an injectable URL", async () => { + let targetInfoChanged: ((event: Protocol.Target.TargetInfoChangedEvent) => void) | undefined; + const attachToTarget = vi.fn(async () => ({})); + const context = new V3Context( + { + on: vi.fn((event: string, handler: (event: never) => void) => { + if (event === "Target.targetInfoChanged") targetInfoChanged = handler; + }), + enableAutoAttach: vi.fn(async () => {}), + getTargets: vi.fn(async () => [target("new-tab", "page", "chrome://newtab")]), + attachToTarget, + } as never, + {} as never, + chromeTabs, + ); + vi.spyOn(context, "waitForInitialTopLevelTargets").mockResolvedValue(); + + await context.bootstrap(); + expect(attachToTarget).not.toHaveBeenCalled(); + + targetInfoChanged?.({ targetInfo: target("new-tab", "page", "https://example.com") }); + await vi.waitFor(() => expect(attachToTarget).toHaveBeenCalledWith("new-tab")); + }); + + it("coalesces repeated target-info attachment attempts", async () => { + let targetInfoChanged: ((event: Protocol.Target.TargetInfoChangedEvent) => void) | undefined; + let resolveAttachment: (() => void) | undefined; + const attachToTarget = vi.fn( + async () => + await new Promise((resolve) => { + resolveAttachment = resolve; + }), + ); + const context = new V3Context( + { + on: vi.fn((event: string, handler: (event: never) => void) => { + if (event === "Target.targetInfoChanged") targetInfoChanged = handler; + }), + enableAutoAttach: vi.fn(async () => {}), + getTargets: vi.fn(async () => []), + attachToTarget, + } as never, + {} as never, + chromeTabs, + ); + vi.spyOn(context, "waitForInitialTopLevelTargets").mockResolvedValue(); + await context.bootstrap(); + + const event = { targetInfo: target("page", "page", "https://example.com") }; + targetInfoChanged?.(event); + targetInfoChanged?.(event); + + await vi.waitFor(() => expect(attachToTarget).toHaveBeenCalledTimes(1)); + resolveAttachment?.(); + }); + + it("uses the restricted filter for future root auto-attachment", async () => { + const send = vi.fn(async () => ({})); + const connection = new CdpConnection( + { + connected: true, + send: vi.fn(), + close: vi.fn(async () => {}), + onMessage: vi.fn(), + onClose: vi.fn(), + onError: vi.fn(), + }, + { debug: () => {}, error: () => {} }, + ); + vi.spyOn(connection, "send").mockImplementation(send); + + await connection.enableAutoAttach(); + + expect(send).toHaveBeenNthCalledWith(1, "Target.setAutoAttach", { + autoAttach: true, + flatten: true, + waitForDebuggerOnStart: true, + filter: STAGEHAND_WEB_TARGET_FILTER, + }); + }); + + it("resumes and detaches an ignored target that arrives paused", async () => { + const send = vi.fn(async () => ({})); + const close = vi.fn(async () => {}); + const context = new V3Context( + { + getSession: vi.fn(() => ({ send, close })), + } as never, + {} as never, + chromeTabs, + ); + + await context.onAttachedToTarget( + target("worker", "worker", "https://example.com/worker.js"), + "session-worker", + ); + + expect(send).toHaveBeenCalledWith("Runtime.runIfWaitingForDebugger"); + expect(close).toHaveBeenCalledOnce(); + }); + + it("does not create a blank page during resident V3Context bootstrap", async () => { + const connection = { + connected: true, + onTransportClosed: vi.fn(), + close: vi.fn(async () => {}), + } as unknown as CdpConnection; + vi.spyOn(CdpConnection, "connect").mockResolvedValue(connection); + vi.spyOn(V3Context.prototype, "bootstrap").mockResolvedValue(); + const ensureFirstTopLevelPage = vi + .spyOn(V3Context.prototype, "ensureFirstTopLevelPage") + .mockResolvedValue(); + + await V3Context.create("ws://browser-proxy.test", { + websocketFactory: vi.fn(), + blankPageUrl: "chrome-extension://stagehand/blank.html", + fallbackLocatorScriptSource: "", + chromeTabs, + logger: {} as never, + ensureInitialPage: false, + }); + + expect(ensureFirstTopLevelPage).not.toHaveBeenCalled(); + }); + + it("defers page and network instrumentation until Stagehand initialization", async () => { + const send = vi.fn(async () => ({})); + const session = { + id: "page-session", + send, + on: vi.fn(), + off: vi.fn(), + }; + const page = new Page( + {} as never, + session as never, + "page-target", + "main-frame", + {} as never, + null, + true, + true, + ); + + expect(send).not.toHaveBeenCalled(); + await page.prepareForInitialization(); + + expect(send).toHaveBeenCalledWith("Page.enable"); + expect(send).toHaveBeenCalledWith("Runtime.enable"); + expect(send).toHaveBeenCalledWith("Network.enable"); + }); + + it("retries deferred page instrumentation after a transient CDP failure", async () => { + let pageEnableAttempts = 0; + const send = vi.fn(async (method: string) => { + if (method === "Page.enable" && pageEnableAttempts++ === 0) { + throw new Error("transient Page.enable failure"); + } + return {}; + }); + const page = new Page( + {} as never, + { id: "page-session", send, on: vi.fn(), off: vi.fn() } as never, + "page-target", + "main-frame", + {} as never, + null, + true, + true, + ); + + await expect(page.prepareForInitialization()).rejects.toThrow("transient Page.enable failure"); + expect(page.isInstrumentationReady()).toBe(false); + + await page.prepareForInitialization(); + expect(page.isInstrumentationReady()).toBe(true); + expect(pageEnableAttempts).toBe(2); + }); + + it("defers Page.enable for OOPIF sessions adopted before initialization", async () => { + const mainSession = { id: "main", send: vi.fn(), on: vi.fn(), off: vi.fn() }; + const childSend = vi.fn(async (method: string) => { + if (method === "Page.getFrameTree") { + return { frameTree: { frame: { id: "child-frame", url: "https://example.com" } } }; + } + return {}; + }); + const childSession = { id: "child", send: childSend, on: vi.fn(), off: vi.fn() }; + const page = new Page( + {} as never, + mainSession as never, + "page-target", + "main-frame", + {} as never, + null, + true, + true, + ); + + page.adoptOopifSession(childSession as never, "child-frame"); + await vi.waitFor(() => expect(childSend).toHaveBeenCalledWith("Page.getFrameTree")); + + expect(childSend).not.toHaveBeenCalledWith("Page.enable"); + }); + + it("tolerates a page closing during deferred context instrumentation", async () => { + let rejectPreparation: ((error: Error) => void) | undefined; + const page = { + isInstrumentationReady: () => false, + prepareForInitialization: () => + new Promise((_resolve, reject) => { + rejectPreparation = reject; + }), + }; + const context = new V3Context( + {} as never, + {} as never, + chromeTabs, + undefined, + undefined, + undefined, + undefined, + undefined, + true, + ); + context.pagesByTarget.set("closing-page", page as never); + context.typeByTarget.set("closing-page", "page"); + + const preparation = context.prepareForInitialization(); + await vi.waitFor(() => expect(rejectPreparation).toBeDefined()); + context.pagesByTarget.delete("closing-page"); + rejectPreparation?.(new Error("target closed")); + + await expect(preparation).resolves.toBeUndefined(); + }); +}); diff --git a/packages/server/tests/understudy-context-lifecycle.test.ts b/packages/server/tests/understudy-context-lifecycle.test.ts new file mode 100644 index 000000000..6c8bd6484 --- /dev/null +++ b/packages/server/tests/understudy-context-lifecycle.test.ts @@ -0,0 +1,63 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { CdpConnection } from "../understudy/cdp.js"; +import { V3Context } from "../understudy/context.js"; + +function contextOptions(onConnected: () => void, onDisconnected: () => void) { + return { + websocketFactory: vi.fn(), + blankPageUrl: "chrome-extension://stagehand/blank.html", + fallbackLocatorScriptSource: "", + chromeTabs: {} as never, + logger: { debug: vi.fn(), error: vi.fn() } as never, + onConnected, + onDisconnected, + }; +} + +describe("V3Context connection lifecycle", () => { + afterEach(() => vi.restoreAllMocks()); + + it("reports a successful connection and deduplicates disconnect notification", async () => { + let transportClosed: (() => void) | undefined; + const connection = { + onTransportClosed: vi.fn((handler: () => void) => { + transportClosed = handler; + }), + close: vi.fn(async () => {}), + }; + vi.spyOn(CdpConnection, "connect").mockResolvedValue(connection as never); + vi.spyOn(V3Context.prototype, "bootstrap").mockResolvedValue(); + vi.spyOn(V3Context.prototype, "ensureFirstTopLevelPage").mockResolvedValue({} as never); + const onConnected = vi.fn(); + const onDisconnected = vi.fn(); + + await V3Context.create("ws://browser.example", contextOptions(onConnected, onDisconnected)); + transportClosed?.(); + transportClosed?.(); + + expect(onConnected).toHaveBeenCalledOnce(); + expect(onDisconnected).toHaveBeenCalledOnce(); + }); + + it("closes and reports disconnection when bootstrap fails", async () => { + let transportClosed: (() => void) | undefined; + const connection = { + onTransportClosed: vi.fn((handler: () => void) => { + transportClosed = handler; + }), + close: vi.fn(async () => transportClosed?.()), + }; + vi.spyOn(CdpConnection, "connect").mockResolvedValue(connection as never); + vi.spyOn(V3Context.prototype, "bootstrap").mockRejectedValue(new Error("bootstrap failed")); + const onConnected = vi.fn(); + const onDisconnected = vi.fn(); + + await expect( + V3Context.create("ws://browser.example", contextOptions(onConnected, onDisconnected)), + ).rejects.toThrow("bootstrap failed"); + + expect(onConnected).toHaveBeenCalledOnce(); + expect(onDisconnected).toHaveBeenCalledOnce(); + expect(connection.close).toHaveBeenCalledOnce(); + }); +}); diff --git a/packages/server/understudy/browserWebSocketTransport.ts b/packages/server/understudy/browserWebSocketTransport.ts index 91e94c90e..7fbc27a71 100644 --- a/packages/server/understudy/browserWebSocketTransport.ts +++ b/packages/server/understudy/browserWebSocketTransport.ts @@ -39,8 +39,9 @@ class BrowserWebSocketTransport implements CdpWebSocketTransport { }); } - onMessage(handler: (data: string) => void): void { + onMessage(handler: (data: string) => void): () => void { this.messageHandlers.add(handler); + return () => this.messageHandlers.delete(handler); } onClose(handler: (event: CdpWebSocketCloseEvent) => void): void { diff --git a/packages/server/understudy/cdp.ts b/packages/server/understudy/cdp.ts index 92d112c0e..20d767a55 100644 --- a/packages/server/understudy/cdp.ts +++ b/packages/server/understudy/cdp.ts @@ -3,6 +3,13 @@ import type { Protocol } from "devtools-protocol"; import { z } from "zod/v4"; import type { StagehandLogger } from "../logger.js"; +/** CDP uses the first matching entry, so allowed web types must precede the catch-all exclusion. */ +export const STAGEHAND_WEB_TARGET_FILTER: Protocol.Target.TargetFilter = [ + { type: "page" }, + { type: "iframe" }, + { exclude: true }, +]; + /** * CDP transport & session multiplexer * @@ -28,7 +35,7 @@ export interface CdpWebSocketTransport { readonly connected: boolean; send(payload: string): void; close(): Promise; - onMessage(handler: (data: string) => void): void; + onMessage(handler: (data: string) => void): () => void; onClose(handler: (event: CdpWebSocketCloseEvent) => void): void; onError(handler: (error: Error) => void): void; } @@ -161,6 +168,7 @@ export class CdpConnection implements CDPSessionLike { autoAttach: true, flatten: true, waitForDebuggerOnStart: true, + filter: STAGEHAND_WEB_TARGET_FILTER, }); await this.send("Target.setDiscoverTargets", { discover: true }); } diff --git a/packages/server/understudy/context.ts b/packages/server/understudy/context.ts index e29a74588..6cdcdcbb6 100644 --- a/packages/server/understudy/context.ts +++ b/packages/server/understudy/context.ts @@ -1,7 +1,7 @@ // lib/v3/understudy/context.ts import type { Protocol } from "devtools-protocol"; import type { StagehandLogger } from "../logger.js"; -import { CdpConnection } from "./cdp.js"; +import { CdpConnection, STAGEHAND_WEB_TARGET_FILTER } from "./cdp.js"; import type { CDPSessionLike, CdpWebSocketFactory } from "./cdp.js"; import { Page } from "./page.js"; import { executionContexts } from "./executionContextRegistry.js"; @@ -60,13 +60,14 @@ function hasInjectableDOM(url: string | undefined): boolean { } function isNonWebTarget(info: Protocol.Target.TargetInfo): boolean { - // Top-level pages should always be tracked — the initial URL may be a - // non-web scheme (e.g. chrome://newtab/) but the user can navigate to - // web content, and the target identity stays the same. - if (info.type === "page") return false; + if (info.type === "page") return !hasInjectableDOM(info.url); return info.type !== "iframe" || !hasInjectableDOM(info.url); } +export function isSupportedWebTarget(info: Protocol.Target.TargetInfo): boolean { + return !isNonWebTarget(info); +} + function isTopLevelPage(info: Protocol.Target.TargetInfo): boolean { const ti = info as unknown as { subtype?: string }; return info.type === "page" && ti.subtype !== "iframe"; @@ -97,6 +98,7 @@ export class V3Context { readonly localBrowserLaunchOptions: LocalBrowserLaunchOptions | null = null, readonly blankPageUrl: string = "about:blank", readonly fallbackLocatorScriptSource: string | null = null, + private deferPageInstrumentation = false, ) {} readonly _targetSessionListeners = new Set(); @@ -114,6 +116,8 @@ export class V3Context { createdAtByTarget = new Map(); typeByTarget = new Map(); pendingCreatedTargetUrl = new Map(); + private readonly pendingTargetAttachments = new Set(); + private pageInstrumentationTask?: Promise; pageCreationFailures = new Map(); // Popup close attempts can race targetCreated, targetInfoChanged, and attached. // In-flight promises let attach wait for a close result before deciding whether @@ -166,30 +170,51 @@ export class V3Context { fallbackLocatorScriptSource: string; chromeTabs: ChromeTabTargetController; logger: StagehandLogger; + onConnected?(): void; + onDisconnected?(): void; + ensureInitialPage?: boolean; + deferPageInstrumentation?: boolean; }, ): Promise { const connectTask = async () => { const conn = await CdpConnection.connect(wsUrl, opts.websocketFactory, opts.logger); - const ctx = new V3Context( - conn, - opts.logger, - opts.chromeTabs, - opts?.env ?? "LOCAL", - opts?.apiClient ?? null, - opts?.localBrowserLaunchOptions ?? null, - opts.blankPageUrl, - opts.fallbackLocatorScriptSource, - ); - await ctx.bootstrap(); - // Allow connectTimeoutMs to also govern how long we wait for the first - // top-level page to appear. On slow machines the browser may need more - // time after the CDP socket is open before the initial page registers. - const firstPageTimeoutMs = Math.max( - opts?.localBrowserLaunchOptions?.connectTimeoutMs ?? 0, - DEFAULT_FIRST_TOP_LEVEL_PAGE_TIMEOUT_MS, - ); - await ctx.ensureFirstTopLevelPage(firstPageTimeoutMs); - return ctx; + if (opts.onDisconnected) { + let disconnected = false; + conn.onTransportClosed(() => { + if (disconnected) return; + disconnected = true; + opts.onDisconnected?.(); + }); + } + try { + opts.onConnected?.(); + const ctx = new V3Context( + conn, + opts.logger, + opts.chromeTabs, + opts?.env ?? "LOCAL", + opts?.apiClient ?? null, + opts?.localBrowserLaunchOptions ?? null, + opts.blankPageUrl, + opts.fallbackLocatorScriptSource, + opts.deferPageInstrumentation ?? false, + ); + await ctx.bootstrap(); + if (opts.ensureInitialPage ?? true) { + // Allow connectTimeoutMs to also govern how long we wait for the first + // top-level page to appear. On slow machines the browser may need more + // time after the CDP socket is open before the initial page registers. + const firstPageTimeoutMs = Math.max( + opts?.localBrowserLaunchOptions?.connectTimeoutMs ?? 0, + DEFAULT_FIRST_TOP_LEVEL_PAGE_TIMEOUT_MS, + ); + await ctx.ensureFirstTopLevelPage(firstPageTimeoutMs); + } + return ctx; + } catch (error) { + await conn.close(); + throw error; + } }; return await connectTask(); @@ -204,6 +229,35 @@ export class V3Context { return false; } + /** Complete page-domain setup only when the client begins Stagehand initialization. */ + async prepareForInitialization(): Promise { + if (!this.deferPageInstrumentation) return; + if (this.pageInstrumentationTask) return await this.pageInstrumentationTask; + + const task = (async () => { + while (true) { + const pending = this.pages().filter((page) => !page.isInstrumentationReady()); + if (pending.length === 0) break; + await Promise.all( + pending.map(async (page) => { + try { + await page.prepareForInitialization(); + } catch (error) { + if (this.pages().includes(page)) throw error; + } + }), + ); + } + this.deferPageInstrumentation = false; + })(); + this.pageInstrumentationTask = task; + try { + await task; + } finally { + if (this.pageInstrumentationTask === task) this.pageInstrumentationTask = undefined; + } + } + async ensureFirstTopLevelPage(timeout: number): Promise { if (this.hasTopLevelPage()) return; @@ -603,6 +657,7 @@ export class V3Context { this.createdAtByTarget.clear(); this.typeByTarget.clear(); this.pendingCreatedTargetUrl.clear(); + this.pendingTargetAttachments.clear(); this.pageCreationFailures.clear(); this.domainPolicyClosingTargets.clear(); this.domainPolicyClosePromises.clear(); @@ -637,7 +692,7 @@ export class V3Context { } }); this.conn.on("Target.targetInfoChanged", (evt) => { - void this.closePopupIfBlockedByDomainPolicy(evt.targetInfo, "targetInfoChanged"); + void this.onTargetInfoChanged(evt.targetInfo); }); // Only enable auto-attach after listeners are ready so replayed targets are captured. @@ -646,17 +701,35 @@ export class V3Context { const targets = await this.conn.getTargets(); for (const t of targets) { if (t.attached) continue; // auto-attach already handled this target - try { - await this.conn.attachToTarget(t.targetId); - } catch { - // ignore attach race - } + if (!isSupportedWebTarget(t)) continue; + await this.attachToTarget(t.targetId); } - const topLevelTargetIds = targets.filter((t) => isTopLevelPage(t)).map((t) => t.targetId); + const topLevelTargetIds = targets + .filter((target) => isTopLevelPage(target) && isSupportedWebTarget(target)) + .map((target) => target.targetId); await this.waitForInitialTopLevelTargets(topLevelTargetIds); } + private async onTargetInfoChanged(info: Protocol.Target.TargetInfo): Promise { + if (await this.closePopupIfBlockedByDomainPolicy(info, "targetInfoChanged")) return; + if (info.attached || !isSupportedWebTarget(info) || this.pagesByTarget.has(info.targetId)) + return; + await this.attachToTarget(info.targetId); + } + + private async attachToTarget(targetId: TargetId): Promise { + if (this.pendingTargetAttachments.has(targetId) || this.pagesByTarget.has(targetId)) return; + this.pendingTargetAttachments.add(targetId); + try { + await this.conn.attachToTarget(targetId); + } catch { + // Target discovery and auto-attachment can win this race. + } finally { + this.pendingTargetAttachments.delete(targetId); + } + } + /** * Handle a newly attached target (top-level or potential OOPIF): * - Enable Page domain and lifecycle events. @@ -674,10 +747,14 @@ export class V3Context { // They still need to be resumed so we don't leave them paused by // waitForDebuggerOnStart. Trying to initialize these targets can throw or // corrupt their internal state (e.g. Chrome's PDF viewer). - if (isNonWebTarget(info)) { + const isStagehandBlankPage = + info.type === "page" && + (info.url === this.blankPageUrl || this.pendingCreatedTargetUrl.has(info.targetId)); + if (isNonWebTarget(info) && !isStagehandBlankPage) { const session = this.conn.getSession(sessionId); if (session) { await session.send("Runtime.runIfWaitingForDebugger").catch(() => {}); + await session.close().catch(() => {}); } return; } @@ -759,12 +836,14 @@ export class V3Context { // - register init scripts. // Commands are sent in-order on the same session before resume. const corePreResumeOps = [ - queuePreResume("Page.enable"), - queuePreResume("Runtime.enable"), + ...(this.deferPageInstrumentation + ? [] + : [queuePreResume("Page.enable"), queuePreResume("Runtime.enable")]), queuePreResume("Target.setAutoAttach", { autoAttach: true, waitForDebuggerOnStart: true, flatten: true, + filter: STAGEHAND_WEB_TARGET_FILTER, }), ]; const headerPreResumeOps: Array<{ @@ -882,7 +961,9 @@ export class V3Context { try { // Best-effort lifecycle events; do not block top-level page registration // on this optional signal stream. - void session.send("Page.setLifecycleEventsEnabled", { enabled: true }).catch(() => {}); + if (!this.deferPageInstrumentation) { + void session.send("Page.setLifecycleEventsEnabled", { enabled: true }).catch(() => {}); + } // Top-level handling if (isTopLevelPage(info)) { @@ -900,6 +981,7 @@ export class V3Context { this.apiClient, this.localBrowserLaunchOptions, this.env === "BROWSERBASE", + this.deferPageInstrumentation, ); } catch (error) { createError = error; diff --git a/packages/server/understudy/networkManager.ts b/packages/server/understudy/networkManager.ts index 67dd2deca..6672a17e8 100644 --- a/packages/server/understudy/networkManager.ts +++ b/packages/server/understudy/networkManager.ts @@ -23,6 +23,8 @@ import { * Aggregates network information for all CDP sessions owned by a Page. */ export class NetworkManager { + private enabled: boolean; + private enableTask?: Promise; readonly sessions = new Map< string, { @@ -37,6 +39,10 @@ export class NetworkManager { readonly documentRequestsByFrame = new Map(); + constructor(enabled = true) { + this.enabled = enabled; + } + /** * Begin tracking network traffic for a CDP session (top-level or OOPIF). * Safe to call multiple times; duplicate registrations are ignored. @@ -139,8 +145,9 @@ export class NetworkManager { session.on("Network.responseReceived", onResponse); session.on("Page.frameStoppedLoading", onFrameStopped); - void session.send("Network.enable").catch(() => {}); - void session.send("Page.enable").catch(() => {}); + if (this.enabled) { + void session.send("Network.enable").catch(() => {}); + } this.sessions.set(sid, { session, @@ -155,6 +162,40 @@ export class NetworkManager { }); } + /** Enable event-producing domains after resident bootstrap has been initialized. */ + public async enable(): Promise { + if (this.enabled) return; + if (this.enableTask) return await this.enableTask; + + const task = (async () => { + const enabledSessions = new Set(); + while (true) { + const pending = [...this.sessions.values()] + .map(({ session }) => session) + .filter((session) => !enabledSessions.has(session)); + if (pending.length === 0) break; + await Promise.all( + pending.map(async (session) => { + try { + await session.send("Network.enable"); + enabledSessions.add(session); + } catch (error) { + if (this.sessions.get(this.sessionKey(session))?.session !== session) return; + throw error; + } + }), + ); + } + this.enabled = true; + })(); + this.enableTask = task; + try { + await task; + } finally { + if (this.enableTask === task) this.enableTask = undefined; + } + } + /** * Stop tracking a session and discard any inflight bookkeeping owned by it. */ diff --git a/packages/server/understudy/page.ts b/packages/server/understudy/page.ts index dd4f8a055..2bff2439c 100644 --- a/packages/server/understudy/page.ts +++ b/packages/server/understudy/page.ts @@ -91,6 +91,9 @@ export class Page { /** Document-start scripts installed across every session this page owns. */ readonly initScripts: string[] = []; extraHTTPHeaders: Record = {}; + private instrumentationReady: boolean; + private instrumentationTask?: Promise; + private readonly instrumentedSessions = new WeakSet(); constructor( readonly conn: CdpConnection, @@ -100,10 +103,12 @@ export class Page { public readonly logger: StagehandLogger, apiClient?: StagehandAPIClient | null, browserIsRemote = false, + deferInstrumentation = false, ) { this.pageId = _targetId; this.apiClient = apiClient ?? null; this.browserIsRemote = browserIsRemote; + this.instrumentationReady = !deferInstrumentation; // own the main session if (mainSession.id) this.sessions.set(mainSession.id, mainSession); @@ -120,10 +125,50 @@ export class Page { this.logger, ); - this.networkManager = new NetworkManager(); + this.networkManager = new NetworkManager(!deferInstrumentation); this.networkManager.trackSession(this.mainSession); } + /** Enable domains intentionally deferred while the resident runtime was only becoming ready. */ + public async prepareForInitialization(): Promise { + if (this.instrumentationReady) return; + if (this.instrumentationTask) return await this.instrumentationTask; + + const task = (async () => { + await this.networkManager.enable(); + while (true) { + const pending = [...this.sessions.values()].filter( + (session) => !this.instrumentedSessions.has(session), + ); + if (pending.length === 0) break; + await Promise.all( + pending.map(async (session) => { + try { + await session.send("Page.enable"); + await session.send("Runtime.enable"); + await session.send("Page.setLifecycleEventsEnabled", { enabled: true }); + this.instrumentedSessions.add(session); + } catch (error) { + if (session.id && !this.sessions.has(session.id)) return; + throw error; + } + }), + ); + } + this.instrumentationReady = true; + })(); + this.instrumentationTask = task; + try { + await task; + } finally { + if (this.instrumentationTask === task) this.instrumentationTask = undefined; + } + } + + public isInstrumentationReady(): boolean { + return this.instrumentationReady; + } + // Send a single init script to a specific CDP session. async installInitScriptOnSession(session: CDPSessionLike, source: string): Promise { await session.send("Page.addScriptToEvaluateOnNewDocument", { @@ -200,19 +245,31 @@ export class Page { apiClient?: StagehandAPIClient | null, localBrowserLaunchOptions?: LocalBrowserLaunchOptions | null, browserIsRemote = false, + deferInstrumentation = false, ): Promise { // Context already issues Page.enable + lifecycle enable before resume. // Re-issue here only as best-effort and do not block page registration on // their acknowledgements; some remote CDP backends can delay these replies // long after the target is otherwise ready. - void session.send("Page.enable").catch(() => {}); - void session.send("Page.setLifecycleEventsEnabled", { enabled: true }).catch(() => {}); + if (!deferInstrumentation) { + void session.send("Page.enable").catch(() => {}); + void session.send("Page.setLifecycleEventsEnabled", { enabled: true }).catch(() => {}); + } const { frameTree } = await session.send<{ frameTree: Protocol.Page.FrameTree; }>("Page.getFrameTree"); const mainFrameId = frameTree.frame.id; - const page = new Page(conn, session, targetId, mainFrameId, logger, apiClient, browserIsRemote); + const page = new Page( + conn, + session, + targetId, + mainFrameId, + logger, + apiClient, + browserIsRemote, + deferInstrumentation, + ); // Seed current URL from initial frame tree try { page._currentUrl = String(frameTree?.frame?.url ?? page._currentUrl); @@ -334,7 +391,9 @@ export class Page { // One-shot seed the child's subtree ownership from its current tree void (async () => { try { - await childSession.send("Page.enable").catch(() => {}); + if (this.instrumentationReady) { + await childSession.send("Page.enable").catch(() => {}); + } let { frameTree } = await childSession.send("Page.getFrameTree");