fix(domworld): fix missing binding for waittasks (#6562)

src/common/AriaQueryHandler.ts

@@ -18,7 +18,7 @@ import { InternalQueryHandler } from './QueryHandler.js';
 import { ElementHandle, JSHandle } from './JSHandle.js';
 import { Protocol } from 'devtools-protocol';
 import { CDPSession } from './Connection.js';
-import { DOMWorld, WaitForSelectorOptions } from './DOMWorld.js';
+import { DOMWorld, PageBinding, WaitForSelectorOptions } from './DOMWorld.js';
 
 async function queryAXTree(
   client: CDPSession,
@@ -87,12 +87,20 @@ const waitFor = async (
   domWorld: DOMWorld,
   selector: string,
   options: WaitForSelectorOptions
-) => {
+): Promise<ElementHandle<Element>> => {
-  await addHandlerToWorld(domWorld);
+  const binding: PageBinding = {
+    name: 'ariaQuerySelector',
+    pptrFunction: async (selector: string) => {
+      const document = await domWorld._document();
+      const element = await queryOne(document, selector);
+      return element;
+    },
+  };
   return domWorld.waitForSelectorInPage(
     (_: Element, selector: string) => globalThis.ariaQuerySelector(selector),
     selector,
-    options
+    options,
+    binding
   );
 };
 
@@ -121,14 +129,6 @@ const queryAllArray = async (
   return jsHandle;
 };
 
-async function addHandlerToWorld(domWorld: DOMWorld) {
-  await domWorld.addBinding('ariaQuerySelector', async (selector: string) => {
-    const document = await domWorld._document();
-    const element = await queryOne(document, selector);
-    return element;
-  });
-}
-
 /**
  * @internal
  */

src/common/DOMWorld.ts

@@ -59,6 +59,14 @@ export interface WaitForSelectorOptions {
   timeout?: number;
 }
 
+/**
+ * @internal
+ */
+export interface PageBinding {
+  name: string;
+  pptrFunction: Function;
+}
+
 /**
  * @internal
  */
@@ -73,14 +81,19 @@ export class DOMWorld {
 
   private _detached = false;
   /**
-   * internal
+   * @internal
    */
   _waitTasks = new Set<WaitTask>();
 
-  // Contains mapping from functions that should be bound to Puppeteer functions.
+  /**
-  private _boundFunctions = new Map<string, Function>();
+   * @internal
+   * Contains mapping from functions that should be bound to Puppeteer functions.
+   */
+  _boundFunctions = new Map<string, Function>();
   // Set of bindings that have been registered in the current context.
   private _ctxBindings = new Set<string>();
+  private static bindingIdentifier = (name: string, contextId: number) =>
+    `${name}_${contextId}`;
 
   constructor(
     frameManager: FrameManager,
@@ -104,10 +117,6 @@ export class DOMWorld {
     if (context) {
       this._contextResolveCallback.call(null, context);
       this._contextResolveCallback = null;
-      this._ctxBindings.clear();
-      for (const name of this._boundFunctions.keys()) {
-        await this.addBindingToContext(name);
-      }
       for (const waitTask of this._waitTasks) waitTask.rerun();
     } else {
       this._documentPromise = null;
@@ -482,19 +491,27 @@ export class DOMWorld {
   /**
    * @internal
    */
-  async addBindingToContext(name: string) {
+  async addBindingToContext(
+    context: ExecutionContext,
+    name: string
+  ): Promise<void> {
     // Previous operation added the binding so we are done.
-    if (this._ctxBindings.has(name)) return;
+    if (
+      this._ctxBindings.has(
+        DOMWorld.bindingIdentifier(name, context._contextId)
+      )
+    ) {
+      return;
+    }
     // Wait for other operation to finish
     if (this._settingUpBinding) {
       await this._settingUpBinding;
-      return this.addBindingToContext(name);
+      return this.addBindingToContext(context, name);
     }
 
     const bind = async (name: string) => {
       const expression = helper.pageBindingInitString('internal', name);
       try {
-        const context = await this.executionContext();
         await context._client.send('Runtime.addBinding', {
           name,
           executionContextId: context._contextId,
@@ -511,14 +528,15 @@ export class DOMWorld {
           'Cannot find context with specified id'
         );
         if (ctxDestroyed || ctxNotFound) {
-          // Retry adding the binding in the next context
+          return;
-          await bind(name);
         } else {
           debugError(error);
           return;
         }
       }
-      this._ctxBindings.add(name);
+      this._ctxBindings.add(
+        DOMWorld.bindingIdentifier(name, context._contextId)
+      );
     };
 
     this._settingUpBinding = bind(name);
@@ -526,18 +544,12 @@ export class DOMWorld {
     this._settingUpBinding = null;
   }
 
-  /**
-   * @internal
-   */
-  async addBinding(name: string, puppeteerFunction: Function): Promise<void> {
-    this._boundFunctions.set(name, puppeteerFunction);
-    await this.addBindingToContext(name);
-  }
-
   private async _onBindingCalled(
     event: Protocol.Runtime.BindingCalledEvent
   ): Promise<void> {
     let payload: { type: string; name: string; seq: number; args: unknown[] };
+    if (!this._hasContext()) return;
+    const context = await this.executionContext();
     try {
       payload = JSON.parse(event.payload);
     } catch {
@@ -546,9 +558,13 @@ export class DOMWorld {
       return;
     }
     const { type, name, seq, args } = payload;
-    if (type !== 'internal' || !this._ctxBindings.has(name)) return;
+    if (
-    if (!this._hasContext()) return;
+      type !== 'internal' ||
-    const context = await this.executionContext();
+      !this._ctxBindings.has(
+        DOMWorld.bindingIdentifier(name, context._contextId)
+      )
+    )
+      return;
     if (context._contextId !== event.executionContextId) return;
     try {
       const result = await this._boundFunctions.get(name)(...args);
@@ -574,7 +590,8 @@ export class DOMWorld {
   async waitForSelectorInPage(
     queryOne: Function,
     selector: string,
-    options: WaitForSelectorOptions
+    options: WaitForSelectorOptions,
+    binding?: PageBinding
   ): Promise<ElementHandle | null> {
     const {
       visible: waitForVisible = false,
@@ -595,16 +612,16 @@ export class DOMWorld {
         : document.querySelector(selector);
       return checkWaitForOptions(node, waitForVisible, waitForHidden);
     }
-    const waitTask = new WaitTask(
+    const waitTaskOptions: WaitTaskOptions = {
-      this,
+      domWorld: this,
-      helper.makePredicateString(predicate, queryOne),
+      predicateBody: helper.makePredicateString(predicate, queryOne),
       title,
       polling,
       timeout,
-      selector,
+      args: [selector, waitForVisible, waitForHidden],
-      waitForVisible,
+      binding,
-      waitForHidden
+    };
-    );
+    const waitTask = new WaitTask(waitTaskOptions);
     const jsHandle = await waitTask.promise;
     const elementHandle = jsHandle.asElement();
     if (!elementHandle) {
@@ -639,16 +656,15 @@ export class DOMWorld {
       ).singleNodeValue;
       return checkWaitForOptions(node, waitForVisible, waitForHidden);
     }
-    const waitTask = new WaitTask(
+    const waitTaskOptions: WaitTaskOptions = {
-      this,
+      domWorld: this,
-      helper.makePredicateString(predicate),
+      predicateBody: helper.makePredicateString(predicate),
       title,
       polling,
       timeout,
-      xpath,
+      args: [xpath, waitForVisible, waitForHidden],
-      waitForVisible,
+    };
-      waitForHidden
+    const waitTask = new WaitTask(waitTaskOptions);
-    );
     const jsHandle = await waitTask.promise;
     const elementHandle = jsHandle.asElement();
     if (!elementHandle) {
@@ -667,14 +683,16 @@ export class DOMWorld {
       polling = 'raf',
       timeout = this._timeoutSettings.timeout(),
     } = options;
-    return new WaitTask(
+    const waitTaskOptions: WaitTaskOptions = {
-      this,
+      domWorld: this,
-      pageFunction,
+      predicateBody: pageFunction,
-      'function',
+      title: 'function',
       polling,
       timeout,
-      ...args
+      args,
-    ).promise;
+    };
+    const waitTask = new WaitTask(waitTaskOptions);
+    return waitTask.promise;
   }
 
   async title(): Promise<string> {
@@ -682,6 +700,19 @@ export class DOMWorld {
   }
 }
 
+/**
+ * @internal
+ */
+export interface WaitTaskOptions {
+  domWorld: DOMWorld;
+  predicateBody: Function | string;
+  title: string;
+  polling: string | number;
+  timeout: number;
+  binding?: PageBinding;
+  args: SerializableOrJSHandle[];
+}
+
 /**
  * @internal
  */
@@ -691,6 +722,7 @@ export class WaitTask {
   _timeout: number;
   _predicateBody: string;
   _args: SerializableOrJSHandle[];
+  _binding: PageBinding;
   _runCount = 0;
   promise: Promise<JSHandle>;
   _resolve: (x: JSHandle) => void;
@@ -698,48 +730,51 @@ export class WaitTask {
   _timeoutTimer?: NodeJS.Timeout;
   _terminated = false;
 
-  constructor(
+  constructor(options: WaitTaskOptions) {
-    domWorld: DOMWorld,
+    if (helper.isString(options.polling))
-    predicateBody: Function | string,
-    title: string,
-    polling: string | number,
-    timeout: number,
-    ...args: SerializableOrJSHandle[]
-  ) {
-    if (helper.isString(polling))
       assert(
-        polling === 'raf' || polling === 'mutation',
+        options.polling === 'raf' || options.polling === 'mutation',
-        'Unknown polling option: ' + polling
+        'Unknown polling option: ' + options.polling
       );
-    else if (helper.isNumber(polling))
+    else if (helper.isNumber(options.polling))
-      assert(polling > 0, 'Cannot poll with non-positive interval: ' + polling);
+      assert(
-    else throw new Error('Unknown polling options: ' + polling);
+        options.polling > 0,
+        'Cannot poll with non-positive interval: ' + options.polling
+      );
+    else throw new Error('Unknown polling options: ' + options.polling);
 
     function getPredicateBody(predicateBody: Function | string) {
       if (helper.isString(predicateBody)) return `return (${predicateBody});`;
       return `return (${predicateBody})(...args);`;
     }
 
-    this._domWorld = domWorld;
+    this._domWorld = options.domWorld;
-    this._polling = polling;
+    this._polling = options.polling;
-    this._timeout = timeout;
+    this._timeout = options.timeout;
-    this._predicateBody = getPredicateBody(predicateBody);
+    this._predicateBody = getPredicateBody(options.predicateBody);
-    this._args = args;
+    this._args = options.args;
+    this._binding = options.binding;
     this._runCount = 0;
-    domWorld._waitTasks.add(this);
+    this._domWorld._waitTasks.add(this);
+    if (this._binding) {
+      this._domWorld._boundFunctions.set(
+        this._binding.name,
+        this._binding.pptrFunction
+      );
+    }
     this.promise = new Promise<JSHandle>((resolve, reject) => {
       this._resolve = resolve;
       this._reject = reject;
     });
     // Since page navigation requires us to re-install the pageScript, we should track
     // timeout on our end.
-    if (timeout) {
+    if (options.timeout) {
       const timeoutError = new TimeoutError(
-        `waiting for ${title} failed: timeout ${timeout}ms exceeded`
+        `waiting for ${options.title} failed: timeout ${options.timeout}ms exceeded`
       );
       this._timeoutTimer = setTimeout(
         () => this.terminate(timeoutError),
-        timeout
+        options.timeout
       );
     }
     this.rerun();
@@ -753,11 +788,16 @@ export class WaitTask {
 
   async rerun(): Promise<void> {
     const runCount = ++this._runCount;
-    /** @type {?JSHandle} */
+    let success: JSHandle = null;
-    let success = null;
+    let error: Error = null;
-    let error = null;
+    const context = await this._domWorld.executionContext();
+    if (this._terminated || runCount !== this._runCount) return;
+    if (this._binding) {
+      await this._domWorld.addBindingToContext(context, this._binding.name);
+    }
+    if (this._terminated || runCount !== this._runCount) return;
     try {
-      success = await (await this._domWorld.executionContext()).evaluateHandle(
+      success = await context.evaluateHandle(
         waitForPredicatePageFunction,
         this._predicateBody,
         this._polling,
@@ -783,10 +823,13 @@ export class WaitTask {
       await success.dispose();
       return;
     }
-    // When frame is detached the task should have been terminated by the DOMWorld.
-    // This can fail if we were adding this task while the frame was detached,
-    // so we terminate here instead.
     if (error) {
+      if (error.message.includes('TypeError: binding is not a function')) {
+        return this.rerun();
+      }
+      // When frame is detached the task should have been terminated by the DOMWorld.
+      // This can fail if we were adding this task while the frame was detached,
+      // so we terminate here instead.
       if (
         error.message.includes(
           'Execution context is not available in detached frame'