rea-agents 6.2.0 → 6.4.0

package/dist/browser/PlaywrightScenarioStorage.js

@@ -0,0 +1,507 @@
+import { randomUUID } from "node:crypto";
+import { z } from "zod";
+import { BrowserObservationError } from "../domain/browserObservationError.js";
+import { withPlaywrightExecutionBoundary } from "./PlaywrightExecutionBoundary.js";
+const breakpointResponse = z.object({ breakpointId: z.string() });
+const scriptIdentifierResponse = z.object({ identifier: z.string() });
+const storageKeyResponse = z.object({ storageKey: z.string() });
+const targetInfoResponse = z.object({
+    targetInfo: z.object({ targetId: z.string() }),
+});
+const STORAGE_EVENT_METHODS = new Set([
+    "Runtime.executionContextCreated",
+    "Runtime.executionContextDestroyed",
+    "Runtime.executionContextsCleared",
+    "Debugger.scriptParsed",
+    "Debugger.paused",
+    "Target.attachedToTarget",
+    "Target.detachedFromTarget",
+    "Target.receivedMessageFromTarget",
+]);
+const attachedTargetEvent = z.object({
+    sessionId: z.string(),
+    targetInfo: z.object({ type: z.string(), targetId: z.string() }),
+});
+const targetSessionEvent = z.object({ sessionId: z.string() });
+const nestedMessageEvent = z.object({
+    sessionId: z.string(),
+    message: z.string(),
+});
+const nestedProtocolMessage = z.object({
+    id: z.number().optional(),
+    method: z.string().optional(),
+    params: z.unknown().optional(),
+    result: z.unknown().optional(),
+    error: z.object({ message: z.string() }).optional(),
+});
+const executionContextEvent = z.object({
+    context: z.object({
+        id: z.number(),
+        origin: z.string(),
+        name: z.string(),
+        auxData: z.object({ frameId: z.string().optional() }).optional(),
+    }),
+});
+const scriptParsedEvent = z.object({
+    scriptId: z.string(),
+    executionContextId: z.number(),
+    url: z.string(),
+});
+const pausedEvent = z.object({
+    callFrames: z.array(z.object({
+        callFrameId: z.string(),
+        location: z.object({ scriptId: z.string() }),
+    })),
+});
+const evaluationResponse = z.object({
+    exceptionDetails: z.unknown().optional(),
+});
+const STORAGE_CLEANUP_TIMEOUT_MS = 1_000;
+/** Owns pre-script storage seeding and its CDP instrumentation for one scenario. */
+export class PlaywrightScenarioStorage {
+    cdp;
+    signal;
+    #seeds = new Map();
+    #initializedLocal = new Set();
+    #initializedSession = new Set();
+    #targets = new Map();
+    #replies = new Map();
+    #sourceUrl = `rea-storage-${randomUUID()}`;
+    #listeners = [];
+    #sequence = 0;
+    #generation = 0;
+    #closed = false;
+    #cleanupCompleted = false;
+    #failure;
+    #closePromise;
+    #detachPromise;
+    #rootInitialization;
+    #work = Promise.resolve();
+    #removeListeners = () => undefined;
+    constructor(cdp, signal) {
+        this.cdp = cdp;
+        this.signal = signal;
+    }
+    static async create(storagePage, scenario, secrets, options = {}) {
+        const { signal, retainCleanup } = options;
+        const { context, page } = storagePage;
+        if (scenario.storage.local_storage.length === 0 &&
+            scenario.storage.session_storage.length === 0)
+            return undefined;
+        const connection = context.newCDPSession(page);
+        let cdp;
+        try {
+            cdp = await withPlaywrightExecutionBoundary(() => connection, undefined, signal);
+        }
+        catch (cause) {
+            if (signal?.aborted === true) {
+                let detached = false;
+                const lateCleanup = async () => {
+                    let session;
+                    try {
+                        session = await connection;
+                    }
+                    catch {
+                        return;
+                    }
+                    if (!detached) {
+                        await session.detach();
+                        detached = true;
+                    }
+                };
+                retainCleanup?.(lateCleanup);
+                void withPlaywrightExecutionBoundary(lateCleanup, STORAGE_CLEANUP_TIMEOUT_MS).catch(() => undefined);
+            }
+            throw cause;
+        }
+        const owner = new PlaywrightScenarioStorage(cdp, signal);
+        owner.addSeeds(scenario.storage.local_storage, "local", secrets);
+        owner.addSeeds(scenario.storage.session_storage, "session", secrets);
+        return owner;
+    }
+    async initialize(page, context) {
+        const initialization = this.initializeRoot(page, context);
+        this.#rootInitialization = initialization;
+        await withPlaywrightExecutionBoundary(() => initialization, undefined, this.signal);
+    }
+    async initializeRoot(page, context) {
+        const rootInfo = targetInfoResponse.parse(await this.cdp.send("Target.getTargetInfo"));
+        const root = this.createTarget(undefined, rootInfo.targetInfo.targetId, (method, params) => this.cdp.send(method, params));
+        this.#targets.set("", root);
+        if (this.#closed)
+            return;
+        this.listen(page, context);
+        const ready = this.initializeTarget(root);
+        root.ready = ready;
+        await ready;
+    }
+    addSeeds(blocks, kind, secrets) {
+        for (const { origin, entries } of blocks) {
+            const seed = this.#seeds.get(origin) ?? { local: [], session: [] };
+            seed[kind].push(...entries.map(({ name, value }) => [
+                name,
+                secrets.value(value),
+            ]));
+            this.#seeds.set(origin, seed);
+        }
+    }
+    listen(page, context) {
+        const browser = context.browser();
+        const detached = () => this.targetDetached("");
+        page.on("close", detached);
+        context.on("close", detached);
+        browser?.on("disconnected", detached);
+        this.#removeListeners = () => {
+            page.off("close", detached);
+            context.off("close", detached);
+            browser?.off("disconnected", detached);
+        };
+        for (const method of [
+            "Runtime.executionContextCreated",
+            "Runtime.executionContextDestroyed",
+            "Runtime.executionContextsCleared",
+            "Debugger.scriptParsed",
+            "Debugger.paused",
+            "Target.attachedToTarget",
+            "Target.detachedFromTarget",
+            "Target.receivedMessageFromTarget",
+        ]) {
+            const listener = (params) => this.event("", method, params);
+            this.#listeners.push({ method, listener });
+            this.cdp.on(method, listener);
+        }
+    }
+    createTarget(parentId, topLevelTargetId, send) {
+        return {
+            parentId,
+            topLevelTargetId,
+            send,
+            contexts: new Map(),
+            scripts: new Map(),
+            ready: Promise.resolve(),
+            waitingForDebugger: false,
+        };
+    }
+    async initializeTarget(target) {
+        if (this.#closed)
+            return;
+        await target.send("Page.enable");
+        if (this.#closed)
+            return;
+        await target.send("Runtime.enable");
+        if (this.#closed)
+            return;
+        await target.send("Debugger.enable");
+        if (this.#closed)
+            return;
+        breakpointResponse.parse(await target.send("Debugger.setBreakpointByUrl", {
+            url: this.#sourceUrl,
+            lineNumber: 1,
+        }));
+        if (this.#closed)
+            return;
+        scriptIdentifierResponse.parse(await target.send("Page.addScriptToEvaluateOnNewDocument", {
+            worldName: this.#sourceUrl,
+            source: `(() => {\n  void 0;\n})()\n//# sourceURL=${this.#sourceUrl}`,
+        }));
+        if (this.#closed)
+            return;
+        await target.send("Target.setAutoAttach", {
+            autoAttach: true,
+            waitForDebuggerOnStart: true,
+            flatten: false,
+        });
+    }
+    sendObserved(target, method, params) {
+        return withPlaywrightExecutionBoundary(() => target.send(method, params), undefined, this.signal);
+    }
+    event(sessionId, method, params) {
+        if (this.#closed)
+            return;
+        const target = this.#targets.get(sessionId);
+        if (target === undefined)
+            return;
+        try {
+            if (method === "Runtime.executionContextCreated") {
+                const { context } = executionContextEvent.parse(params);
+                if (context.name === this.#sourceUrl)
+                    target.contexts.set(context.id, context);
+            }
+            else if (method === "Runtime.executionContextDestroyed") {
+                const { executionContextId } = z
+                    .object({ executionContextId: z.number() })
+                    .parse(params);
+                target.contexts.delete(executionContextId);
+            }
+            else if (method === "Runtime.executionContextsCleared") {
+                target.contexts.clear();
+                target.scripts.clear();
+            }
+            else if (method === "Debugger.scriptParsed") {
+                const script = scriptParsedEvent.parse(params);
+                if (script.url === this.#sourceUrl)
+                    target.scripts.set(script.scriptId, script.executionContextId);
+            }
+            else if (method === "Debugger.paused") {
+                const pause = pausedEvent.parse(params);
+                const frame = pause.callFrames[0];
+                // CDP resume is target-wide, so identify our private seed frame before
+                // doing any work that will resume the paused target.
+                if (frame === undefined || !target.scripts.has(frame.location.scriptId))
+                    return;
+                this.#generation += 1;
+                this.#work = this.#work
+                    .then(() => this.seed(target, pause))
+                    .catch((cause) => {
+                    if (sessionId === "" || this.#targets.get(sessionId) === target)
+                        this.fail(cause);
+                });
+            }
+            else if (method === "Target.attachedToTarget") {
+                this.#generation += 1;
+                this.attach(sessionId, attachedTargetEvent.parse(params));
+            }
+            else if (method === "Target.detachedFromTarget") {
+                this.#generation += 1;
+                this.targetDetached(targetSessionEvent.parse(params).sessionId);
+            }
+            else if (method === "Target.receivedMessageFromTarget") {
+                this.nestedMessage(nestedMessageEvent.parse(params));
+            }
+        }
+        catch (cause) {
+            this.fail(cause);
+        }
+    }
+    nestedMessage(event) {
+        const reply = nestedProtocolMessage.parse(JSON.parse(event.message));
+        if (reply.id !== undefined) {
+            const waiter = this.#replies.get(reply.id);
+            if (waiter?.sessionId !== event.sessionId)
+                return;
+            this.#replies.delete(reply.id);
+            if (reply.error !== undefined)
+                waiter.reject(new Error(reply.error.message));
+            else
+                waiter.resolve(reply.result);
+            return;
+        }
+        if (reply.method !== undefined && STORAGE_EVENT_METHODS.has(reply.method))
+            this.event(event.sessionId, reply.method, reply.params);
+    }
+    fail(cause) {
+        this.#generation += 1;
+        this.#failure ??= new BrowserObservationError("capture_browser_scenario", "protocol_error", { cause });
+    }
+    async seedPausedFrame(target, frame) {
+        const contextId = target.scripts.get(frame.location.scriptId);
+        const context = contextId === undefined ? undefined : target.contexts.get(contextId);
+        const seed = context === undefined ? undefined : this.#seeds.get(context.origin);
+        if (this.#closed ||
+            seed === undefined ||
+            context?.auxData?.frameId === undefined)
+            return;
+        const storageKey = storageKeyResponse.parse(await this.sendObserved(target, "Storage.getStorageKey", {
+            frameId: context.auxData.frameId,
+        })).storageKey;
+        const sessionKey = JSON.stringify([target.topLevelTargetId, storageKey]);
+        const local = this.#initializedLocal.has(storageKey) || seed.local.length === 0
+            ? []
+            : seed.local;
+        const session = this.#initializedSession.has(sessionKey) || seed.session.length === 0
+            ? []
+            : seed.session;
+        if (local.length === 0 && session.length === 0)
+            return;
+        const payload = JSON.stringify({ local, session }).replaceAll("<", "\\u003c");
+        const response = evaluationResponse.parse(await this.sendObserved(target, "Debugger.evaluateOnCallFrame", {
+            callFrameId: frame.callFrameId,
+            expression: `(() => { const seed = ${payload}; for (const [name, value] of seed.local) localStorage.setItem(name, value); for (const [name, value] of seed.session) sessionStorage.setItem(name, value); })()`,
+            returnByValue: true,
+        }));
+        if (response.exceptionDetails !== undefined)
+            throw new BrowserObservationError("capture_browser_scenario", "protocol_error", { detail: "The browser denied initial storage initialization." });
+        if (local.length > 0)
+            this.#initializedLocal.add(storageKey);
+        if (session.length > 0)
+            this.#initializedSession.add(sessionKey);
+    }
+    async seed(target, pause) {
+        const frame = pause.callFrames[0];
+        if (frame === undefined)
+            return;
+        let primaryFailure;
+        let hasPrimaryFailure = false;
+        try {
+            await this.seedPausedFrame(target, frame);
+        }
+        catch (cause) {
+            primaryFailure = cause;
+            hasPrimaryFailure = true;
+        }
+        let resumeFailure;
+        try {
+            await withPlaywrightExecutionBoundary(() => target.send("Debugger.resume"), STORAGE_CLEANUP_TIMEOUT_MS);
+        }
+        catch (cause) {
+            resumeFailure = cause;
+        }
+        if (hasPrimaryFailure && resumeFailure !== undefined)
+            throw new AggregateError([primaryFailure, resumeFailure], "Storage initialization and debugger resume both failed", { cause: primaryFailure });
+        if (this.signal?.aborted === true && resumeFailure === undefined)
+            return;
+        if (hasPrimaryFailure)
+            throw primaryFailure;
+        if (resumeFailure !== undefined)
+            throw resumeFailure;
+    }
+    attach(parentId, event) {
+        const parent = this.#targets.get(parentId);
+        if (parent === undefined)
+            return;
+        const { sessionId, targetInfo } = event;
+        const target = this.createTarget(parentId, event.targetInfo.type === "iframe"
+            ? parent.topLevelTargetId
+            : event.targetInfo.targetId, async (method, params = {}) => {
+            const id = ++this.#sequence;
+            const reply = new Promise((resolve, reject) => this.#replies.set(id, { sessionId, resolve, reject }));
+            try {
+                const sent = parent.send("Target.sendMessageToTarget", {
+                    sessionId,
+                    message: JSON.stringify({ id, method, params }),
+                });
+                const [, result] = await Promise.all([sent, reply]);
+                return result;
+            }
+            finally {
+                this.#replies.delete(id);
+            }
+        });
+        target.waitingForDebugger = true;
+        this.#targets.set(sessionId, target);
+        target.ready = (async () => {
+            try {
+                if (!this.#closed &&
+                    (targetInfo.type === "iframe" || targetInfo.type === "page"))
+                    await this.initializeTarget(target);
+            }
+            catch (cause) {
+                if (this.signal?.aborted !== true &&
+                    !this.#closed &&
+                    this.#targets.get(sessionId) === target)
+                    this.fail(cause);
+            }
+            if (this.#targets.get(sessionId) !== target)
+                return;
+            try {
+                await withPlaywrightExecutionBoundary(() => target.send("Runtime.runIfWaitingForDebugger"), STORAGE_CLEANUP_TIMEOUT_MS);
+                target.waitingForDebugger = false;
+            }
+            catch (cause) {
+                if (this.signal?.aborted !== true &&
+                    !this.#closed &&
+                    this.#targets.get(sessionId) === target)
+                    this.fail(cause);
+            }
+        })();
+    }
+    targetDetached(sessionId) {
+        for (const [id, target] of this.#targets)
+            if (target.parentId === sessionId)
+                this.targetDetached(id);
+        this.#targets.delete(sessionId);
+        for (const [id, reply] of this.#replies)
+            if (reply.sessionId === sessionId) {
+                this.#replies.delete(id);
+                reply.reject(new BrowserObservationError("capture_browser_scenario", "disconnected"));
+            }
+        this.#generation += 1;
+    }
+    async settle() {
+        for (;;) {
+            const generation = this.#generation;
+            const work = this.#work;
+            const targets = [...this.#targets.entries()];
+            await withPlaywrightExecutionBoundary(() => Promise.all(targets.map(([, target]) => target.ready)), undefined, this.signal);
+            await withPlaywrightExecutionBoundary(() => work, undefined, this.signal);
+            await Promise.resolve();
+            if (this.#failure !== undefined)
+                throw this.#failure;
+            if (generation === this.#generation &&
+                work === this.#work &&
+                targets.every(([id, target]) => this.#targets.get(id) === target))
+                return;
+        }
+    }
+    close(stopLoading = false) {
+        if (this.#cleanupCompleted)
+            return Promise.resolve();
+        if (this.#closePromise !== undefined)
+            return this.#closePromise;
+        const closing = this.closeResources(stopLoading).catch((cause) => {
+            throw new BrowserObservationError("capture_browser_scenario", "cleanup_failed", {
+                cause,
+                cleanup: {
+                    reason: cause instanceof Error ? cause.message : String(cause),
+                    resources: ["browser_transport", "browser_storage_bootstrap"],
+                },
+            });
+        });
+        this.#closePromise = closing;
+        void closing.catch(() => {
+            if (this.#closePromise === closing)
+                this.#closePromise = undefined;
+        });
+        return closing;
+    }
+    async closeResources(stopLoading) {
+        this.#closed = true;
+        let stopLoadingFailure;
+        if (stopLoading) {
+            try {
+                const root = this.#targets.get("");
+                if (root !== undefined)
+                    await withPlaywrightExecutionBoundary(() => root.send("Page.stopLoading"), STORAGE_CLEANUP_TIMEOUT_MS);
+            }
+            catch (cause) {
+                stopLoadingFailure = { cause };
+            }
+        }
+        const work = this.#work;
+        const rootInitialization = this.#rootInitialization;
+        const targets = [...this.#targets.values()];
+        try {
+            await withPlaywrightExecutionBoundary(() => Promise.all([
+                work.catch(() => undefined),
+                ...(rootInitialization === undefined
+                    ? []
+                    : [rootInitialization.catch(() => undefined)]),
+                ...targets.map((target) => target.ready.catch(() => undefined)),
+            ]), STORAGE_CLEANUP_TIMEOUT_MS);
+        }
+        catch {
+            // Detaching the owning CDP session below is the terminal release for
+            // scripts, breakpoints, debugger state, and attached child sessions.
+        }
+        let detach = this.#detachPromise;
+        if (detach === undefined) {
+            detach = Promise.resolve().then(() => this.cdp.detach());
+            this.#detachPromise = detach;
+            void detach.catch(() => {
+                if (this.#detachPromise === detach)
+                    this.#detachPromise = undefined;
+            });
+        }
+        await withPlaywrightExecutionBoundary(() => detach, STORAGE_CLEANUP_TIMEOUT_MS);
+        for (const { method, listener } of this.#listeners)
+            this.cdp.off(method, listener);
+        this.#removeListeners();
+        for (const reply of this.#replies.values())
+            reply.reject(new BrowserObservationError("capture_browser_scenario", "disconnected"));
+        this.#targets.clear();
+        this.#replies.clear();
+        this.#cleanupCompleted = true;
+        if (stopLoadingFailure !== undefined)
+            throw stopLoadingFailure.cause;
+    }
+}