refactor: move WaitForHelper to class and use it from inside context

src/McpContext.ts

@@ -20,6 +20,7 @@ import os from 'node:os';
 import path from 'node:path';
 import {listPages} from './tools/pages.js';
 import {TraceResult} from './trace-processing/parse.js';
+import {WaitForHelper} from './WaitForHelper.js';
 
 export interface TextSnapshotNode extends SerializedAXNode {
   id: string;
@@ -32,6 +33,9 @@ export interface TextSnapshot {
   snapshotId: string;
 }
 
+const DEFAULT_TIMEOUT = 5_000;
+const NAVIGATION_TIMEOUT = 10_000;
+
 export class McpContext implements Context {
   browser: Browser;
   logger: Debugger;
@@ -68,10 +72,10 @@ export class McpContext implements Context {
     this.#consoleCollector = new PageCollector(
       this.browser,
       (page, collect) => {
-        page.on('console', (event: ConsoleMessage) => {
+        page.on('console', event => {
           collect(event);
         });
-        page.on('pageerror', (event: Error) => {
+        page.on('pageerror', event => {
           collect(event);
         });
       },
@@ -195,10 +199,10 @@ export class McpContext implements Context {
     newPage.on('dialog', this.#dialogHandler);
 
     // For waiters 5sec timeout should be sufficient.
-    newPage.setDefaultTimeout(5_000);
+    newPage.setDefaultTimeout(DEFAULT_TIMEOUT);
     // 10sec should be enough for the load event to be emitted during
     // navigations.
-    newPage.setDefaultNavigationTimeout(10_000);
+    newPage.setDefaultNavigationTimeout(NAVIGATION_TIMEOUT);
   }
 
   async getElementByUid(uid: string): Promise<ElementHandle<Element>> {
@@ -302,4 +306,10 @@ export class McpContext implements Context {
   recordedTraces(): TraceResult[] {
     return this.#traceResults;
   }
+
+  waitForEventsAfterAction(action: () => Promise<unknown>): Promise<void> {
+    const page = this.getSelectedPage();
+    const waitForHelper = new WaitForHelper(page);
+    return waitForHelper.waitForEventsAfterAction(action);
+  }
 }

src/McpResponse.ts

@@ -21,7 +21,7 @@ export class McpResponse implements Response {
   #attachedNetworkRequestUrl?: string;
   #includeConsoleData: boolean = false;
   #textResponseLines: string[] = [];
-  #formatedConsoleData?: string[];
+  #formattedConsoleData?: string[];
   #images: ImageContentData[] = [];
 
   setIncludePages(value: boolean): void {
@@ -97,7 +97,7 @@ export class McpResponse implements Response {
         formattedConsoleMessages = await Promise.all(
           consoleMessages.map(message => formatConsoleEvent(message)),
         );
-        this.#formatedConsoleData = formattedConsoleMessages;
+        this.#formattedConsoleData = formattedConsoleMessages;
       }
     }
 
@@ -167,9 +167,9 @@ Call browser_handle_dialog to handle it before continuing.`);
       }
     }
 
-    if (this.#includeConsoleData && this.#formatedConsoleData) {
+    if (this.#includeConsoleData && this.#formattedConsoleData) {
       response.push('## Console messages');
-      response.push(...this.#formatedConsoleData);
+      response.push(...this.#formattedConsoleData);
     }
 
     const text: TextContent = {

src/WaitForHelper.ts

@@ -0,0 +1,145 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import type {CdpPage} from 'puppeteer-core/internal/cdp/Page.js';
+import {Page, Protocol} from 'puppeteer-core';
+import {logger} from './logger.js';
+
+export class WaitForHelper {
+  #abortController = new AbortController();
+  #genericTimeout: number;
+  #stableDomFor: number;
+  #expectNavigationIn: number;
+  #page: CdpPage;
+
+  constructor(page: Page) {
+    this.#genericTimeout = 3000;
+    this.#stableDomFor = 100;
+    this.#expectNavigationIn = 100;
+    this.#page = page as unknown as CdpPage;
+  }
+
+  /**
+   * A wrapper that executes a action and waits for
+   * a potential navigation, after which it waits
+   * for the DOM to be stable before returning.
+   */
+  async waitForStableDom(): Promise<void> {
+    const stableDomObserver = await this.#page.evaluateHandle(timeout => {
+      let timeoutId: ReturnType<typeof setTimeout>;
+      function callback() {
+        clearTimeout(timeoutId);
+        timeoutId = setTimeout(() => {
+          domObserver.resolver.resolve();
+          domObserver.observer.disconnect();
+        }, timeout);
+      }
+      const domObserver = {
+        resolver: Promise.withResolvers<void>(),
+        observer: new MutationObserver(callback),
+      };
+      // It's possible that the DOM is not gonna change so we
+      // need to start the timeout initially.
+      callback();
+
+      domObserver.observer.observe(document.body, {
+        childList: true,
+        subtree: true,
+        attributes: true,
+      });
+
+      return domObserver;
+    }, this.#stableDomFor);
+
+    this.#abortController.signal.addEventListener('abort', async () => {
+      try {
+        await stableDomObserver.evaluate(observer => {
+          observer.observer.disconnect();
+          observer.resolver.resolve();
+        });
+        await stableDomObserver.dispose();
+      } catch {
+        // Ignored cleanup errors
+      }
+    });
+
+    return Promise.race([
+      stableDomObserver.evaluate(async observer => {
+        return await observer.resolver.promise;
+      }),
+      this.timeout(this.#genericTimeout).then(() => {
+        throw new Error('Timeout');
+      }),
+    ]);
+  }
+
+  async waitForNavigationStarted() {
+    // Currently Puppeteer does not have API
+    // For when a navigation is about to start
+    const navigationStartedPromise = new Promise<boolean>(resolve => {
+      const listener = (event: Protocol.Page.FrameStartedNavigatingEvent) => {
+        if (
+          [
+            'historySameDocument',
+            'historyDifferentDocument',
+            'sameDocument',
+          ].includes(event.navigationType)
+        ) {
+          resolve(false);
+          return;
+        }
+
+        resolve(true);
+      };
+
+      this.#page._client().on('Page.frameStartedNavigating', listener);
+      this.#abortController.signal.addEventListener('abort', () => {
+        resolve(false);
+        this.#page._client().off('Page.frameStartedNavigating', listener);
+      });
+    });
+
+    return await Promise.race([
+      navigationStartedPromise,
+      this.timeout(this.#expectNavigationIn).then(() => false),
+    ]);
+  }
+
+  timeout(time: number): Promise<void> {
+    return new Promise<void>(res => {
+      const id = setTimeout(res, time);
+      this.#abortController.signal.addEventListener('abort', () => {
+        res();
+        clearTimeout(id);
+      });
+    });
+  }
+
+  async waitForEventsAfterAction(
+    action: () => Promise<unknown>,
+  ): Promise<void> {
+    const navigationStartedPromise = this.waitForNavigationStarted();
+
+    await action();
+
+    try {
+      const navigationStated = await navigationStartedPromise;
+      if (navigationStated) {
+        await this.#page.waitForNavigation({
+          timeout: this.#genericTimeout,
+          signal: this.#abortController.signal,
+        });
+      }
+
+      // Wait for stable dom after navigation so we execute in
+      // the correct context
+      await this.waitForStableDom();
+    } catch (error) {
+      logger(error);
+    } finally {
+      this.#abortController.abort();
+    }
+  }
+}

src/tools/ToolDefinition.ts

@@ -70,6 +70,7 @@ export type Context = Readonly<{
     data: Uint8Array<ArrayBufferLike>,
     mimeType: 'image/png' | 'image/jpeg',
   ): Promise<{filename: string}>;
+  waitForEventsAfterAction(action: () => Promise<unknown>): Promise<void>;
 }>;
 
 export function defineTool<Schema extends Zod.ZodRawShape>(

src/tools/input.ts

@@ -8,7 +8,6 @@ import z from 'zod';
 import {defineTool} from './ToolDefinition.js';
 import {ElementHandle} from 'puppeteer-core';
 import {ToolCategories} from './categories.js';
-import {waitForEventsAfterAction} from '../waitForHelpers.js';
 
 export const click = defineTool({
   name: 'click',
@@ -32,7 +31,7 @@ export const click = defineTool({
     const uid = request.params.uid;
     const handle = await context.getElementByUid(uid);
     try {
-      await waitForEventsAfterAction(handle.frame.page(), async () => {
+      await context.waitForEventsAfterAction(async () => {
         await handle.asLocator().click({
           count: request.params.dblClick ? 2 : 1,
         });
@@ -67,7 +66,7 @@ export const hover = defineTool({
     const uid = request.params.uid;
     const handle = await context.getElementByUid(uid);
     try {
-      await waitForEventsAfterAction(handle.frame.page(), async () => {
+      await context.waitForEventsAfterAction(async () => {
         await handle.asLocator().hover();
       });
       response.appendResponseLine(`Successfully hovered over the element`);
@@ -96,7 +95,7 @@ export const fill = defineTool({
   handler: async (request, response, context) => {
     const handle = await context.getElementByUid(request.params.uid);
     try {
-      await waitForEventsAfterAction(handle.frame.page(), async () => {
+      await context.waitForEventsAfterAction(async () => {
         await handle.asLocator().fill(request.params.value);
       });
       response.appendResponseLine(`Successfully filled out the element`);
@@ -122,7 +121,7 @@ export const drag = defineTool({
     const fromHandle = await context.getElementByUid(request.params.from_uid);
     const toHandle = await context.getElementByUid(request.params.to_uid);
     try {
-      await waitForEventsAfterAction(fromHandle.frame.page(), async () => {
+      await context.waitForEventsAfterAction(async () => {
         await fromHandle.drag(toHandle);
         await new Promise(resolve => setTimeout(resolve, 50));
         await toHandle.drop(fromHandle);
@@ -157,7 +156,7 @@ export const fillForm = defineTool({
     for (const element of request.params.elements) {
       const handle = await context.getElementByUid(element.uid);
       try {
-        await waitForEventsAfterAction(handle.frame.page(), async () => {
+        await context.waitForEventsAfterAction(async () => {
           await handle.asLocator().fill(element.value);
         });
       } finally {

src/tools/pages.ts

@@ -7,7 +7,6 @@
 import z from 'zod';
 import {defineTool} from './ToolDefinition.js';
 import {ToolCategories} from './categories.js';
-import {waitForEventsAfterAction} from '../waitForHelpers.js';
 
 export const listPages = defineTool({
   name: 'list_pages',
@@ -79,7 +78,7 @@ export const newPage = defineTool({
   handler: async (request, response, context) => {
     const page = await context.newPage();
 
-    await waitForEventsAfterAction(page, async () => {
+    await context.waitForEventsAfterAction(async () => {
       await page.goto(request.params.url);
     });
 
@@ -100,7 +99,7 @@ export const navigatePage = defineTool({
   handler: async (request, response, context) => {
     const page = context.getSelectedPage();
 
-    await waitForEventsAfterAction(page, async () => {
+    await context.waitForEventsAfterAction(async () => {
       await page.goto(request.params.url);
     });
 

src/tools/script.ts

@@ -6,8 +6,7 @@
 import z from 'zod';
 import {defineTool} from './ToolDefinition.js';
 import {ToolCategories} from './categories.js';
-import {waitForEventsAfterAction} from '../waitForHelpers.js';
-import {JSHandle} from 'puppeteer-core';
+import type {JSHandle} from 'puppeteer-core';
 
 export const evaluateScript = defineTool({
   name: 'evaluate_script',
@@ -51,7 +50,7 @@ Example with arguments: \`(el) => {
       for (const el of request.params.args ?? []) {
         args.push(await context.getElementByUid(el.uid));
       }
-      await waitForEventsAfterAction(page, async () => {
+      await context.waitForEventsAfterAction(async () => {
         const result = await page.evaluate(
           async (fn, ...args) => {
             // @ts-expect-error no types.