diff --git a/src/common/DOMWorld.ts b/src/common/DOMWorld.ts
index 54b5dc206418c..c9adf8ff68269 100644
--- a/src/common/DOMWorld.ts
+++ b/src/common/DOMWorld.ts
@@ -645,6 +645,7 @@ export class DOMWorld {
     const waitTaskOptions: WaitTaskOptions = {
       domWorld: this,
       predicateBody: helper.makePredicateString(predicate, queryOne),
+      predicateAcceptsContextElement: true,
       title,
       polling,
       timeout,
@@ -691,6 +692,7 @@ export class DOMWorld {
     const waitTaskOptions: WaitTaskOptions = {
       domWorld: this,
       predicateBody: helper.makePredicateString(predicate),
+      predicateAcceptsContextElement: true,
       title,
       polling,
       timeout,
@@ -717,6 +719,7 @@ export class DOMWorld {
     const waitTaskOptions: WaitTaskOptions = {
       domWorld: this,
       predicateBody: pageFunction,
+      predicateAcceptsContextElement: false,
       title: 'function',
       polling,
       timeout,
@@ -737,6 +740,7 @@ export class DOMWorld {
 export interface WaitTaskOptions {
   domWorld: DOMWorld;
   predicateBody: Function | string;
+  predicateAcceptsContextElement: boolean;
   title: string;
   polling: string | number;
   timeout: number;
@@ -753,6 +757,7 @@ export class WaitTask {
   _polling: string | number;
   _timeout: number;
   _predicateBody: string;
+  _predicateAcceptsContextElement: boolean;
   _args: SerializableOrJSHandle[];
   _binding: PageBinding;
   _runCount = 0;
@@ -786,6 +791,8 @@ export class WaitTask {
     this._timeout = options.timeout;
     this._root = options.root;
     this._predicateBody = getPredicateBody(options.predicateBody);
+    this._predicateAcceptsContextElement =
+      options.predicateAcceptsContextElement;
     this._args = options.args;
     this._binding = options.binding;
     this._runCount = 0;
@@ -835,6 +842,7 @@ export class WaitTask {
         success = await this._root.evaluateHandle(
           waitForPredicatePageFunction,
           this._predicateBody,
+          this._predicateAcceptsContextElement,
           this._polling,
           this._timeout,
           ...this._args
@@ -844,6 +852,7 @@ export class WaitTask {
           waitForPredicatePageFunction,
           null,
           this._predicateBody,
+          this._predicateAcceptsContextElement,
           this._polling,
           this._timeout,
           ...this._args
@@ -911,6 +920,7 @@ export class WaitTask {
 async function waitForPredicatePageFunction(
   root: Element | Document | null,
   predicateBody: string,
+  predicateAcceptsContextElement: boolean,
   polling: string,
   timeout: number,
   ...args: unknown[]
@@ -927,7 +937,9 @@ async function waitForPredicatePageFunction(
    * @returns {!Promise<*>}
    */
   async function pollMutation(): Promise<unknown> {
-    const success = await predicate(root, ...args);
+    const success = predicateAcceptsContextElement
+      ? await predicate(root, ...args)
+      : await predicate(...args);
     if (success) return Promise.resolve(success);
 
     let fulfill;
@@ -937,7 +949,9 @@ async function waitForPredicatePageFunction(
         observer.disconnect();
         fulfill();
       }
-      const success = await predicate(root, ...args);
+      const success = predicateAcceptsContextElement
+        ? await predicate(root, ...args)
+        : await predicate(...args);
       if (success) {
         observer.disconnect();
         fulfill(success);
@@ -962,7 +976,9 @@ async function waitForPredicatePageFunction(
         fulfill();
         return;
       }
-      const success = await predicate(root, ...args);
+      const success = predicateAcceptsContextElement
+        ? await predicate(root, ...args)
+        : await predicate(...args);
       if (success) fulfill(success);
       else requestAnimationFrame(onRaf);
     }
@@ -979,7 +995,9 @@ async function waitForPredicatePageFunction(
         fulfill();
         return;
       }
-      const success = await predicate(root, ...args);
+      const success = predicateAcceptsContextElement
+        ? await predicate(root, ...args)
+        : await predicate(...args);
       if (success) fulfill(success);
       else setTimeout(onTimeout, pollInterval);
     }
diff --git a/test/waittask.spec.ts b/test/waittask.spec.ts
index bce546adbb4e1..7b7f409e8cea1 100644
--- a/test/waittask.spec.ts
+++ b/test/waittask.spec.ts
@@ -293,7 +293,11 @@ describe('waittask specs', function () {
       const div = await page.$('div');
       let resolved = false;
       const waitForFunction = page
-        .waitForFunction((element) => !element.parentElement, {}, div)
+        .waitForFunction(
+          (element) => element.localName === 'div' && !element.parentElement,
+          {},
+          div
+        )
         .then(() => (resolved = true));
       expect(resolved).toBe(false);
       await page.evaluate((element: HTMLElement) => element.remove(), div);