fix: atomically get Puppeteer utilities (#9597)

packages/puppeteer-core/src/common/ExecutionContext.ts

@@ -15,6 +15,8 @@
  */
 
 import {Protocol} from 'devtools-protocol';
+import {source as injectedSource} from '../generated/injected.js';
+import type PuppeteerUtil from '../injected/injected.js';
 import {CDPSession} from './Connection.js';
 import {IsolatedWorld} from './IsolatedWorld.js';
 import {JSHandle} from './JSHandle.js';
@@ -87,6 +89,20 @@ export class ExecutionContext {
     this._contextName = contextPayload.name;
   }
 
+  #puppeteerUtil?: Promise<JSHandle<PuppeteerUtil>>;
+  get puppeteerUtil(): Promise<JSHandle<PuppeteerUtil>> {
+    if (!this.#puppeteerUtil) {
+      this.#puppeteerUtil = this.evaluateHandle(
+        `(() => {
+            const module = {};
+            ${injectedSource}
+            return module.exports.default;
+          })()`
+      ) as Promise<JSHandle<PuppeteerUtil>>;
+    }
+    return this.#puppeteerUtil;
+  }
+
   /**
    * Evaluates the given function.
    *
@@ -304,7 +320,7 @@ export class ExecutionContext {
       arg: unknown
     ): Promise<Protocol.Runtime.CallArgument> {
       if (arg instanceof LazyArg) {
-        arg = await arg.get();
+        arg = await arg.get(this);
       }
       if (typeof arg === 'bigint') {
         // eslint-disable-line valid-typeof

packages/puppeteer-core/src/common/Frame.ts

@@ -34,6 +34,7 @@ import {Page} from '../api/Page.js';
 import {getQueryHandlerAndSelector} from './QueryHandler.js';
 import {EvaluateFunc, HandleFor, NodeFor} from './types.js';
 import {importFS} from './util.js';
+import {LazyArg} from './LazyArg.js';
 
 /**
  * @public
@@ -839,7 +840,9 @@ export class Frame {
           await promise;
           return script;
         },
-        await this.worlds[PUPPETEER_WORLD].puppeteerUtil,
+        LazyArg.create(context => {
+          return context.puppeteerUtil;
+        }),
         {...options, type, content}
       )
     );
@@ -923,7 +926,9 @@ export class Frame {
           await promise;
           return element;
         },
-        await this.worlds[PUPPETEER_WORLD].puppeteerUtil,
+        LazyArg.create(context => {
+          return context.puppeteerUtil;
+        }),
         options
       )
     );

packages/puppeteer-core/src/common/IsolatedWorld.ts

@@ -15,7 +15,6 @@
  */
 
 import {Protocol} from 'devtools-protocol';
-import {source as injectedSource} from '../generated/injected.js';
 import {assert} from '../util/assert.js';
 import {createDeferredPromise} from '../util/DeferredPromise.js';
 import {isErrorLike} from '../util/ErrorLike.js';
@@ -24,15 +23,14 @@ import {ExecutionContext} from './ExecutionContext.js';
 import {Frame} from './Frame.js';
 import {FrameManager} from './FrameManager.js';
 import {MouseButton} from './Input.js';
+import {MAIN_WORLD, PUPPETEER_WORLD} from './IsolatedWorlds.js';
 import {JSHandle} from './JSHandle.js';
 import {LifecycleWatcher, PuppeteerLifeCycleEvent} from './LifecycleWatcher.js';
 import {TimeoutSettings} from './TimeoutSettings.js';
 import {EvaluateFunc, HandleFor, InnerLazyParams, NodeFor} from './types.js';
 import {createJSHandle, debugError, pageBindingInitString} from './util.js';
 import {TaskManager, WaitTask} from './WaitTask.js';
-import {MAIN_WORLD, PUPPETEER_WORLD} from './IsolatedWorlds.js';
 
-import type PuppeteerUtil from '../injected/injected.js';
 import type {ElementHandle} from './ElementHandle.js';
 import {LazyArg} from './LazyArg.js';
 
@@ -96,21 +94,6 @@ export class IsolatedWorld {
   // Contains mapping from functions that should be bound to Puppeteer functions.
   #boundFunctions = new Map<string, Function>();
   #taskManager = new TaskManager();
-  #puppeteerUtil?: Promise<JSHandle<PuppeteerUtil> | undefined>;
-
-  get puppeteerUtil(): Promise<JSHandle<PuppeteerUtil>> {
-    /**
-     * This is supposed to mimic what happens when evaluating Puppeteer utilities
-     * break due to navigation.
-     */
-    return (async () => {
-      const util = await this.#puppeteerUtil;
-      if (util) {
-        return util;
-      }
-      throw new Error('Execution context was destroyed!');
-    })();
-  }
 
   get taskManager(): TaskManager {
     return this.#taskManager;
@@ -149,35 +132,13 @@ export class IsolatedWorld {
 
   clearContext(): void {
     this.#document = undefined;
-    this.#puppeteerUtil = createDeferredPromise();
     this.#context = createDeferredPromise();
   }
 
   setContext(context: ExecutionContext): void {
-    this.#injectPuppeteerUtil(context);
     this.#ctxBindings.clear();
     this.#context.resolve(context);
-  }
+    this.#taskManager.rerunAll();
-
-  async #injectPuppeteerUtil(context: ExecutionContext): Promise<void> {
-    try {
-      this.#puppeteerUtil = (async () => {
-        try {
-          return (await context.evaluateHandle(
-            `(() => {
-            const module = {};
-            ${injectedSource}
-            return module.exports.default;
-          })()`
-          )) as JSHandle<PuppeteerUtil>;
-        } catch {
-          return undefined;
-        }
-      })();
-      this.#taskManager.rerunAll();
-    } catch (error: unknown) {
-      debugError(error);
-    }
   }
 
   hasContext(): boolean {
@@ -531,13 +492,8 @@ export class IsolatedWorld {
           root,
           timeout,
         },
-        new LazyArg(async () => {
+        LazyArg.create(context => {
-          try {
+          return context.puppeteerUtil;
-            // In case CDP fails.
-            return await this.puppeteerUtil;
-          } catch {
-            return undefined;
-          }
         }),
         queryOne.toString(),
         selector,

packages/puppeteer-core/src/common/LazyArg.ts

@@ -14,16 +14,26 @@
  * limitations under the License.
  */
 
+import {ExecutionContext} from './ExecutionContext.js';
+
 /**
  * @internal
  */
 export class LazyArg<T> {
-  #get: () => Promise<T>;
+  static create = <T>(
-  constructor(get: () => Promise<T>) {
+    get: (context: ExecutionContext) => Promise<T> | T
+  ): T => {
+    // We don't want to introduce LazyArg to the type system, otherwise we would
+    // have to make it public.
+    return new LazyArg(get) as unknown as T;
+  };
+
+  #get: (context: ExecutionContext) => Promise<T> | T;
+  private constructor(get: (context: ExecutionContext) => Promise<T> | T) {
     this.#get = get;
   }
 
-  get(): Promise<T> {
+  async get(context: ExecutionContext): Promise<T> {
-    return this.#get();
+    return this.#get(context);
   }
 }

packages/puppeteer-core/src/common/QueryHandler.ts

@@ -21,6 +21,7 @@ import {ElementHandle} from './ElementHandle.js';
 import {Frame} from './Frame.js';
 import {WaitForSelectorOptions} from './IsolatedWorld.js';
 import {MAIN_WORLD, PUPPETEER_WORLD} from './IsolatedWorlds.js';
+import {LazyArg} from './LazyArg.js';
 
 /**
  * @public
@@ -105,7 +106,9 @@ function createPuppeteerQueryHandler(
       const jsHandle = await element.evaluateHandle(
         queryOne,
         selector,
-        await world.puppeteerUtil
+        LazyArg.create(context => {
+          return context.puppeteerUtil;
+        })
       );
       const elementHandle = jsHandle.asElement();
       if (elementHandle) {
@@ -153,7 +156,9 @@ function createPuppeteerQueryHandler(
       const jsHandle = await element.evaluateHandle(
         queryAll,
         selector,
-        await element.executionContext()._world!.puppeteerUtil
+        LazyArg.create(context => {
+          return context.puppeteerUtil;
+        })
       );
       const properties = await jsHandle.getProperties();
       await jsHandle.dispose();

packages/puppeteer-core/src/common/WaitTask.ts

@@ -20,6 +20,7 @@ import {ElementHandle} from './ElementHandle.js';
 import {TimeoutError} from './Errors.js';
 import {IsolatedWorld} from './IsolatedWorld.js';
 import {JSHandle} from './JSHandle.js';
+import {LazyArg} from './LazyArg.js';
 import {HandleFor} from './types.js';
 
 /**
@@ -114,7 +115,9 @@ export class WaitTask<T = unknown> {
                 return fun(...args) as Promise<T>;
               });
             },
-            await this.#world.puppeteerUtil,
+            LazyArg.create(context => {
+              return context.puppeteerUtil;
+            }),
             this.#fn,
             ...this.#args
           );
@@ -127,7 +130,9 @@ export class WaitTask<T = unknown> {
                 return fun(...args) as Promise<T>;
               }, root || document);
             },
-            await this.#world.puppeteerUtil,
+            LazyArg.create(context => {
+              return context.puppeteerUtil;
+            }),
             this.#root,
             this.#fn,
             ...this.#args
@@ -141,7 +146,9 @@ export class WaitTask<T = unknown> {
                 return fun(...args) as Promise<T>;
               }, ms);
             },
-            await this.#world.puppeteerUtil,
+            LazyArg.create(context => {
+              return context.puppeteerUtil;
+            }),
             this.#polling,
             this.#fn,
             ...this.#args

test/TestExpectations.json

@@ -3013,7 +3013,7 @@
   },
   {
     "testIdPattern": "[waittask.spec] waittask specs Frame.waitForSelector should work with removed MutationObserver",
-    "platforms": ["darwin", "win32"],
+    "platforms": ["darwin", "linux", "win32"],
     "parameters": ["firefox"],
     "expectations": ["SKIP"]
   },

test/src/injected.spec.ts

@@ -16,6 +16,7 @@
 
 import expect from 'expect';
 import {PUPPETEER_WORLD} from 'puppeteer-core/internal/common/IsolatedWorlds.js';
+import {LazyArg} from 'puppeteer-core/internal/common/LazyArg.js';
 import {
   getTestState,
   setupTestBrowserHooks,
@@ -30,9 +31,14 @@ describe('PuppeteerUtil tests', function () {
     const {page} = getTestState();
 
     const world = page.mainFrame().worlds[PUPPETEER_WORLD];
-    const value = await world.evaluate(PuppeteerUtil => {
+    const value = await world.evaluate(
-      return typeof PuppeteerUtil === 'object';
+      PuppeteerUtil => {
-    }, world.puppeteerUtil);
+        return typeof PuppeteerUtil === 'object';
+      },
+      LazyArg.create(context => {
+        return context.puppeteerUtil;
+      })
+    );
     expect(value).toBeTruthy();
   });
 
@@ -45,7 +51,9 @@ describe('PuppeteerUtil tests', function () {
         ({createFunction}, fnString) => {
           return createFunction(fnString)(4);
         },
-        await world.puppeteerUtil,
+        LazyArg.create(context => {
+          return context.puppeteerUtil;
+        }),
         (() => {
           return 4;
         }).toString()