Work around overly strict check in puppeteer.

Author: Piotr Paulski

src/McpContext.ts

@@ -21,6 +21,7 @@ import type {
 
 import type {ListenerMap} from './PageCollector.js';
 import {NetworkCollector, PageCollector} from './PageCollector.js';
+import {Locator} from './third_party/puppeteer-core/index.js';
 import {listPages} from './tools/pages.js';
 import {takeSnapshot} from './tools/snapshot.js';
 import {CLOSE_PAGE_ERROR} from './tools/ToolDefinition.js';
@@ -91,9 +92,12 @@ export class McpContext implements Context {
   #nextSnapshotId = 1;
   #traceResults: TraceResult[] = [];
 
-  private constructor(browser: Browser, logger: Debugger) {
+  #locatorClass: typeof Locator;
+
+  private constructor(browser: Browser, logger: Debugger, locatorClass: typeof Locator) {
     this.browser = browser;
     this.logger = logger;
+    this.#locatorClass = locatorClass;
 
     this.#networkCollector = new NetworkCollector(this.browser);
 
@@ -116,8 +120,8 @@ export class McpContext implements Context {
     await this.#consoleCollector.init();
   }
 
-  static async from(browser: Browser, logger: Debugger) {
-    const context = new McpContext(browser, logger);
+  static async from(browser: Browser, logger: Debugger, locatorClass: typeof Locator = Locator) {
+    const context = new McpContext(browser, logger, locatorClass);
     await context.#init();
     return context;
   }
@@ -422,4 +426,22 @@ export class McpContext implements Context {
   getNetworkRequestStableId(request: HTTPRequest): number {
     return this.#networkCollector.getIdForResource(request);
   }
+
+  waitForTextOnPage({text, timeout}: {text: string, timeout?: number|undefined}): Promise<Element> {
+    const page = this.getSelectedPage();
+    const frames = page.frames();
+
+    const locator = this.#locatorClass.race(
+      frames.flatMap(frame => [
+        frame.locator(`aria/${text}`),
+        frame.locator(`text/${text}`),
+      ]),
+    );
+
+    if (timeout) {
+      locator.setTimeout(timeout);
+    }
+
+    return locator.wait();
+  }
 }

src/tools/ToolDefinition.ts

@@ -93,6 +93,7 @@ export type Context = Readonly<{
     filename: string,
   ): Promise<{filename: string}>;
   waitForEventsAfterAction(action: () => Promise<unknown>): Promise<void>;
+  waitForTextOnPage(params: {text: string, timeout?: number|undefined}): Promise<Element>;
 }>;
 
 export function defineTool<Schema extends zod.ZodRawShape>(

src/tools/snapshot.ts

@@ -44,21 +44,7 @@ export const waitFor = defineTool({
     ...timeoutSchema,
   },
   handler: async (request, response, context) => {
-    const page = context.getSelectedPage();
-    const frames = page.frames();
-
-    const locator = Locator.race(
-      frames.flatMap(frame => [
-        frame.locator(`aria/${request.params.text}`),
-        frame.locator(`text/${request.params.text}`),
-      ]),
-    );
-
-    if (request.params.timeout) {
-      locator.setTimeout(request.params.timeout);
-    }
-
-    await locator.wait();
+    await context.waitForTextOnPage(request.params);
 
     response.appendResponseLine(
       `Element with text "${request.params.text}" found.`,

tests/utils.ts

@@ -5,7 +5,7 @@
  */
 import logger from 'debug';
 import type {Browser} from 'puppeteer';
-import puppeteer from 'puppeteer';
+import puppeteer, {Locator} from 'puppeteer';
 import type {Frame, HTTPRequest, HTTPResponse} from 'puppeteer-core';
 
 import {McpContext} from '../src/McpContext.js';
@@ -36,7 +36,7 @@ export async function withBrowser(
     }),
   );
   const response = new McpResponse();
-  const context = await McpContext.from(browser, logger('test'));
+  const context = await McpContext.from(browser, logger('test'), Locator);
 
   await cb(response, context);
 }