refactor: move waitForEventsAfterAction to McpPage (#1780)

Splits out from #1244 per review feedback.

'waitForEventsAfterAction' previously lived in 'McpContext' and always
used the selected page's CPU/network throttling settings. With pageId
routing, a tool can target a different page than the selected one,
meaning wrong throttling multipliers were applied.

Moving the method to 'McpPage' fixes this: each tool now calls
'page.waitForEventsAfterAction(...)' and gets the correct page's
emulation settings.

'getNetworkMultiplierFromString' is extracted to 'WaitForHelper.ts' to
avoid a circular import (McpContext → McpPage already exists).

Unblocks #1777.

Author: bcfmtolgahan

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,25 @@ export class McpPage implements ContextPage {
     return this.emulationSettings.colorScheme ?? null;
   }
 
+  // Public for testability: tests spy on this method to verify throttle multipliers.
+  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 type {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/inPage.ts

@@ -71,7 +71,7 @@ export const executeInPageTool = definePageTool({
       .describe('The JSON-stringified parameters to pass to the tool'),
   },
   handler: async (request, response, context) => {
-    const page = context.getSelectedMcpPage();
+    const page = request.page;
     const toolName = request.params.toolName;
     let params: Record<string, unknown> = {};
     if (request.params.params) {

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`);
@@ -217,7 +217,6 @@ async function fillFormElement(
   }
 }
 
-// here
 export const fill = definePageTool({
   name: 'fill',
   description: `Type text into a input, text area or select an option from a <select> element.`,
@@ -236,7 +235,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 +261,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 +289,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);
@@ -334,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,
@@ -415,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,
@@ -166,7 +166,7 @@ export const navigatePage = definePageTool({
       ),
     ...timeoutSchema,
   },
-  handler: async (request, response, context) => {
+  handler: async (request, response) => {
     const page = request.page;
     const options = {
       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

@@ -77,7 +77,11 @@ Example with arguments: \`(el) => {
         }
 
         const worker = await getWebWorker(context, serviceWorkerId);
-        await performEvaluation(worker, fnString, [], response, context);
+        await context
+          .getSelectedMcpPage()
+          .waitForEventsAfterAction(async () => {
+            await performEvaluation(worker, fnString, [], response);
+          });
         return;
       }
 
@@ -97,7 +101,9 @@ Example with arguments: \`(el) => {
 
         const evaluatable = await getPageOrFrame(page, frames);
 
-        await performEvaluation(evaluatable, fnString, args, response, context);
+        await mcpPage.waitForEventsAfterAction(async () => {
+          await performEvaluation(evaluatable, fnString, args, response);
+        });
       } finally {
         void Promise.allSettled(args.map(arg => arg.dispose()));
       }
@@ -110,24 +116,21 @@ const performEvaluation = async (
   fnString: string,
   args: Array<JSHandle<unknown>>,
   response: Response,
-  context: Context,
 ) => {
   const fn = await evaluatable.evaluateHandle(`(${fnString})`);
   try {
-    await context.waitForEventsAfterAction(async () => {
-      const result = await evaluatable.evaluate(
-        async (fn, ...args) => {
-          // @ts-expect-error no types for function fn
-          return JSON.stringify(await fn(...args));
-        },
-        fn,
-        ...args,
-      );
-      response.appendResponseLine('Script ran on page and returned:');
-      response.appendResponseLine('```json');
-      response.appendResponseLine(`${result}`);
-      response.appendResponseLine('```');
-    });
+    const result = await evaluatable.evaluate(
+      async (fn, ...args) => {
+        // @ts-expect-error no types for function fn
+        return JSON.stringify(await fn(...args));
+      },
+      fn,
+      ...args,
+    );
+    response.appendResponseLine('Script ran on page and returned:');
+    response.appendResponseLine('```json');
+    response.appendResponseLine(`${result}`);
+    response.appendResponseLine('```');
   } finally {
     void fn.dispose();
   }

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);
     });
   });