refactor: move waitForEventsAfterAction to McpPage

src/McpContext.ts

@@ -48,7 +48,7 @@ import {
   type InstalledExtension,
 } from './utils/ExtensionRegistry.js';
 import {saveTemporaryFile} from './utils/files.js';
-import {WaitForHelper} from './WaitForHelper.js';
+import {getNetworkMultiplierFromString} from './WaitForHelper.js';
 
 interface McpContextOptions {
   // Whether the DevTools windows are exposed as pages for debugging of DevTools.
@@ -62,23 +62,6 @@ interface McpContextOptions {
 const DEFAULT_TIMEOUT = 5_000;
 const NAVIGATION_TIMEOUT = 10_000;
 
-function getNetworkMultiplierFromString(condition: string | null): number {
-  const puppeteerCondition =
-    condition as keyof typeof PredefinedNetworkConditions;
-
-  switch (puppeteerCondition) {
-    case 'Fast 4G':
-      return 1;
-    case 'Slow 4G':
-      return 2.5;
-    case 'Fast 3G':
-      return 5;
-    case 'Slow 3G':
-      return 10;
-  }
-  return 1;
-}
-
 export class McpContext implements Context {
   browser: Browser;
   logger: Debugger;
@@ -851,31 +834,6 @@ export class McpContext implements Context {
     return this.#traceResults;
   }
 
-  getWaitForHelper(
-    page: Page,
-    cpuMultiplier: number,
-    networkMultiplier: number,
-  ) {
-    return new WaitForHelper(page, cpuMultiplier, networkMultiplier);
-  }
-
-  waitForEventsAfterAction(
-    action: () => Promise<unknown>,
-    options?: {timeout?: number},
-  ): Promise<void> {
-    const page = this.#getSelectedMcpPage();
-    const cpuMultiplier = page.cpuThrottlingRate;
-    const networkMultiplier = getNetworkMultiplierFromString(
-      page.networkConditions,
-    );
-    const waitForHelper = this.getWaitForHelper(
-      page.pptrPage,
-      cpuMultiplier,
-      networkMultiplier,
-    );
-    return waitForHelper.waitForEventsAfterAction(action, options);
-  }
-
   getNetworkRequestStableId(request: HTTPRequest): number {
     return this.#networkCollector.getIdForResource(request);
   }

src/McpPage.ts

@@ -18,6 +18,10 @@ import type {
   TextSnapshot,
   TextSnapshotNode,
 } from './types.js';
+import {
+  getNetworkMultiplierFromString,
+  WaitForHelper,
+} from './WaitForHelper.js';
 
 /**
  * Per-page state wrapper. Consolidates dialog, snapshot, emulation,
@@ -91,6 +95,24 @@ export class McpPage implements ContextPage {
     return this.emulationSettings.colorScheme ?? null;
   }
 
+  createWaitForHelper(
+    cpuMultiplier: number,
+    networkMultiplier: number,
+  ): WaitForHelper {
+    return new WaitForHelper(this.pptrPage, cpuMultiplier, networkMultiplier);
+  }
+
+  waitForEventsAfterAction(
+    action: () => Promise<unknown>,
+    options?: {timeout?: number},
+  ): Promise<void> {
+    const helper = this.createWaitForHelper(
+      this.cpuThrottlingRate,
+      getNetworkMultiplierFromString(this.networkConditions),
+    );
+    return helper.waitForEventsAfterAction(action, options);
+  }
+
   dispose(): void {
     this.pptrPage.off('dialog', this.#dialogHandler);
   }

src/WaitForHelper.ts

@@ -6,6 +6,7 @@
 
 import {logger} from './logger.js';
 import type {Page, Protocol, CdpPage} from './third_party/index.js';
+import {PredefinedNetworkConditions} from './third_party/index.js';
 
 export class WaitForHelper {
   #abortController = new AbortController();
@@ -160,3 +161,22 @@ export class WaitForHelper {
     }
   }
 }
+
+export function getNetworkMultiplierFromString(
+  condition: string | null,
+): number {
+  const puppeteerCondition =
+    condition as keyof typeof PredefinedNetworkConditions;
+
+  switch (puppeteerCondition) {
+    case 'Fast 4G':
+      return 1;
+    case 'Slow 4G':
+      return 2.5;
+    case 'Fast 3G':
+      return 5;
+    case 'Slow 3G':
+      return 10;
+  }
+  return 1;
+}

src/tools/ToolDefinition.ts

@@ -172,10 +172,6 @@ export type Context = Readonly<{
     data: Uint8Array<ArrayBufferLike>,
     filename: string,
   ): Promise<{filename: string}>;
-  waitForEventsAfterAction(
-    action: () => Promise<unknown>,
-    options?: {timeout?: number},
-  ): Promise<void>;
   waitForTextOnPage(
     text: string[],
     timeout?: number,
@@ -213,6 +209,10 @@ export type ContextPage = Readonly<{
 
   getDialog(): Dialog | undefined;
   clearDialog(): void;
+  waitForEventsAfterAction(
+    action: () => Promise<unknown>,
+    options?: {timeout?: number},
+  ): Promise<void>;
 }>;
 
 export function defineTool<Schema extends zod.ZodRawShape>(

src/tools/input.ts

@@ -58,11 +58,11 @@ export const click = definePageTool({
     dblClick: dblClickSchema,
     includeSnapshot: includeSnapshotSchema,
   },
-  handler: async (request, response, context) => {
+  handler: async (request, response) => {
     const uid = request.params.uid;
     const handle = await request.page.getElementByUid(uid);
     try {
-      await context.waitForEventsAfterAction(async () => {
+      await request.page.waitForEventsAfterAction(async () => {
         await handle.asLocator().click({
           count: request.params.dblClick ? 2 : 1,
         });
@@ -97,9 +97,9 @@ export const clickAt = definePageTool({
     dblClick: dblClickSchema,
     includeSnapshot: includeSnapshotSchema,
   },
-  handler: async (request, response, context) => {
+  handler: async (request, response) => {
     const page = request.page;
-    await context.waitForEventsAfterAction(async () => {
+    await page.waitForEventsAfterAction(async () => {
       await page.pptrPage.mouse.click(request.params.x, request.params.y, {
         clickCount: request.params.dblClick ? 2 : 1,
       });
@@ -130,11 +130,11 @@ export const hover = definePageTool({
       ),
     includeSnapshot: includeSnapshotSchema,
   },
-  handler: async (request, response, context) => {
+  handler: async (request, response) => {
     const uid = request.params.uid;
     const handle = await request.page.getElementByUid(uid);
     try {
-      await context.waitForEventsAfterAction(async () => {
+      await request.page.waitForEventsAfterAction(async () => {
         await handle.asLocator().hover();
       });
       response.appendResponseLine(`Successfully hovered over the element`);
@@ -236,7 +236,7 @@ export const fill = definePageTool({
   },
   handler: async (request, response, context) => {
     const page = request.page;
-    await context.waitForEventsAfterAction(async () => {
+    await page.waitForEventsAfterAction(async () => {
       await fillFormElement(
         request.params.uid,
         request.params.value,
@@ -262,9 +262,9 @@ export const typeText = definePageTool({
     text: zod.string().describe('The text to type'),
     submitKey: submitKeySchema,
   },
-  handler: async (request, response, context) => {
+  handler: async (request, response) => {
     const page = request.page;
-    await context.waitForEventsAfterAction(async () => {
+    await page.waitForEventsAfterAction(async () => {
       await page.pptrPage.keyboard.type(request.params.text);
       if (request.params.submitKey) {
         await page.pptrPage.keyboard.press(
@@ -290,13 +290,13 @@ export const drag = definePageTool({
     to_uid: zod.string().describe('The uid of the element to drop into'),
     includeSnapshot: includeSnapshotSchema,
   },
-  handler: async (request, response, context) => {
+  handler: async (request, response) => {
     const fromHandle = await request.page.getElementByUid(
       request.params.from_uid,
     );
     const toHandle = await request.page.getElementByUid(request.params.to_uid);
     try {
-      await context.waitForEventsAfterAction(async () => {
+      await request.page.waitForEventsAfterAction(async () => {
         await fromHandle.drag(toHandle);
         await new Promise(resolve => setTimeout(resolve, 50));
         await toHandle.drop(fromHandle);
@@ -333,7 +333,7 @@ export const fillForm = definePageTool({
   handler: async (request, response, context) => {
     const page = request.page;
     for (const element of request.params.elements) {
-      await context.waitForEventsAfterAction(async () => {
+      await page.waitForEventsAfterAction(async () => {
         await fillFormElement(
           element.uid,
           element.value,
@@ -414,12 +414,12 @@ export const pressKey = definePageTool({
       ),
     includeSnapshot: includeSnapshotSchema,
   },
-  handler: async (request, response, context) => {
+  handler: async (request, response) => {
     const page = request.page;
     const tokens = parseKey(request.params.key);
     const [key, ...modifiers] = tokens;
 
-    await context.waitForEventsAfterAction(async () => {
+    await page.waitForEventsAfterAction(async () => {
       for (const modifier of modifiers) {
         await page.pptrPage.keyboard.down(modifier);
       }

src/tools/pages.ts

@@ -119,7 +119,7 @@ export const newPage = defineTool({
       request.params.isolatedContext,
     );
 
-    await context.waitForEventsAfterAction(
+    await page.waitForEventsAfterAction(
       async () => {
         await page.pptrPage.goto(request.params.url, {
           timeout: request.params.timeout,
@@ -206,7 +206,7 @@ export const navigatePage = definePageTool({
     page.pptrPage.on('dialog', dialogHandler);
 
     try {
-      await context.waitForEventsAfterAction(
+      await page.waitForEventsAfterAction(
         async () => {
           switch (request.params.type) {
             case 'url':

src/tools/script.ts

@@ -9,7 +9,7 @@ import type {Frame, JSHandle, Page, WebWorker} from '../third_party/index.js';
 import type {ExtensionServiceWorker} from '../types.js';
 
 import {ToolCategory} from './categories.js';
-import type {Context, Response} from './ToolDefinition.js';
+import type {Context, ContextPage, Response} from './ToolDefinition.js';
 import {defineTool, pageIdSchema} from './ToolDefinition.js';
 
 export type Evaluatable = Page | Frame | WebWorker;
@@ -77,7 +77,7 @@ Example with arguments: \`(el) => {
         }
 
         const worker = await getWebWorker(context, serviceWorkerId);
-        await performEvaluation(worker, fnString, [], response, context);
+        await performEvaluation(worker, fnString, [], response, context.getSelectedMcpPage());
         return;
       }
 
@@ -97,7 +97,7 @@ Example with arguments: \`(el) => {
 
         const evaluatable = await getPageOrFrame(page, frames);
 
-        await performEvaluation(evaluatable, fnString, args, response, context);
+        await performEvaluation(evaluatable, fnString, args, response, mcpPage);
       } finally {
         void Promise.allSettled(args.map(arg => arg.dispose()));
       }
@@ -110,11 +110,11 @@ const performEvaluation = async (
   fnString: string,
   args: Array<JSHandle<unknown>>,
   response: Response,
-  context: Context,
+  page: ContextPage,
 ) => {
   const fn = await evaluatable.evaluateHandle(`(${fnString})`);
   try {
-    await context.waitForEventsAfterAction(async () => {
+    await page.waitForEventsAfterAction(async () => {
       const result = await evaluatable.evaluate(
         async (fn, ...args) => {
           // @ts-expect-error no types for function fn

tests/McpContext.test.ts

@@ -75,13 +75,13 @@ describe('McpContext', () => {
         cpuThrottlingRate: 2,
         networkConditions: 'Slow 3G',
       });
-      const stub = sinon.spy(context, 'getWaitForHelper');
+      const stub = sinon.spy(page, 'createWaitForHelper');
 
-      await context.waitForEventsAfterAction(async () => {
+      await page.waitForEventsAfterAction(async () => {
         // trigger the waiting only
       });
 
-      sinon.assert.calledWithExactly(stub, page.pptrPage, 2, 10);
+      sinon.assert.calledWithExactly(stub, 2, 10);
     });
   });