refactor: waitForResult helpers (#2041)

src/McpResponse.ts

@@ -40,6 +40,7 @@ import type {InsightName, TraceResult} from './trace-processing/parse.js';
 import {getInsightOutput, getTraceSummary} from './trace-processing/parse.js';
 import {paginate} from './utils/pagination.js';
 import type {PaginationOptions} from './utils/types.js';
+import type {WaitForEventsResult} from './WaitForHelper.js';
 
 interface TraceInsightData {
   trace: TraceResult;
@@ -206,6 +207,7 @@ export class McpResponse implements Response {
   #page?: McpPage;
   #redactNetworkHeaders = true;
   #error?: Error;
+  #attachedWaitForResult?: WaitForEventsResult;
 
   constructor(args: ParsedArguments) {
     this.#args = args;
@@ -388,6 +390,10 @@ export class McpResponse implements Response {
     this.#textResponseLines.push(value);
   }
 
+  attachWaitForResult(result: WaitForEventsResult): void {
+    this.#attachedWaitForResult = result;
+  }
+
   setHeapSnapshotAggregates(
     aggregates: Record<
       string,
@@ -739,6 +745,7 @@ export class McpResponse implements Response {
       extensionServiceWorkers?: object[];
       extensionPages?: object[];
       errorMessage?: string;
+      navigatedToUrl?: string;
     } = {};
 
     const response = [];
@@ -747,6 +754,16 @@ export class McpResponse implements Response {
       response.push(...this.#textResponseLines);
     }
 
+    if (this.#attachedWaitForResult) {
+      if (this.#attachedWaitForResult.navigatedToUrl) {
+        response.push(
+          `Page navigated to ${this.#attachedWaitForResult.navigatedToUrl}.`,
+        );
+        structuredContent.navigatedToUrl =
+          this.#attachedWaitForResult.navigatedToUrl;
+      }
+    }
+
     const networkConditions = this.#page?.networkConditions;
     if (networkConditions) {
       const timeout = this.#page!.pptrPage.getDefaultNavigationTimeout();

src/WaitForHelper.ts

@@ -16,6 +16,9 @@ export class WaitForHelper {
   #expectNavigationIn: number;
   #navigationTimeout: number;
 
+  #dialogOpened = false;
+  #initialUrl: string;
+
   constructor(
     page: Page,
     cpuTimeoutMultiplier: number,
@@ -26,6 +29,7 @@ export class WaitForHelper {
     this.#expectNavigationIn = 100 * cpuTimeoutMultiplier;
     this.#navigationTimeout = 3000 * networkTimeoutMultiplier;
     this.#page = page as unknown as CdpPage;
+    this.#initialUrl = page.url();
   }
 
   /**
@@ -128,10 +132,12 @@ export class WaitForHelper {
     action: () => Promise<unknown>,
     options?: {timeout?: number; handleDialog?: 'accept' | 'dismiss' | string},
   ): Promise<WaitForEventsResult> {
-    let dialogOpened = false;
+    if (this.#abortController.signal.aborted) {
+      throw new Error("Can't re-use a WaitForHelper");
+    }
     if (options?.handleDialog) {
       const dialogHandler = (dialog: Pick<Dialog, 'accept' | 'dismiss'>) => {
-        dialogOpened = true;
+        this.#dialogOpened = true;
         if (options.handleDialog === 'dismiss') {
           void dialog.dismiss();
         } else if (options.handleDialog === 'accept') {
@@ -146,7 +152,6 @@ export class WaitForHelper {
       });
     }
 
-    const urlBeforeAction = this.#page.url();
     const navigationFinished = this.waitForNavigationStarted()
       .then(navigationStated => {
         if (navigationStated) {
@@ -170,8 +175,8 @@ export class WaitForHelper {
     try {
       await navigationFinished;
 
-      if (dialogOpened) {
-        return {};
+      if (this.#dialogOpened) {
+        return this.#getResult();
       }
 
       // Wait for stable dom after navigation so we execute in
@@ -183,9 +188,13 @@ export class WaitForHelper {
       this.#abortController.abort();
     }
 
+    return this.#getResult();
+  }
+
+  #getResult(): WaitForEventsResult {
     const urlAfterAction = this.#page.url();
     return {
-      ...(urlAfterAction !== urlBeforeAction
+      ...(urlAfterAction !== this.#initialUrl
         ? {navigatedToUrl: urlAfterAction}
         : {}),
     };
@@ -200,15 +209,6 @@ export interface WaitForEventsResult {
   navigatedToUrl?: string;
 }
 
-export function appendWaitForResult(
-  response: {appendResponseLine(value: string): void},
-  result: WaitForEventsResult,
-): void {
-  if (result.navigatedToUrl) {
-    response.appendResponseLine(`Page navigated to ${result.navigatedToUrl}.`);
-  }
-}
-
 export function getNetworkMultiplierFromString(
   condition: string | null,
 ): number {

src/tools/ToolDefinition.ts

@@ -154,6 +154,7 @@ export interface Response {
   attachLighthouseResult(result: LighthouseData): void;
   setListThirdPartyDeveloperTools(): void;
   setListWebMcpTools(): void;
+  attachWaitForResult(result: WaitForEventsResult): void;
 }
 
 export type SupportedExtensions =

src/tools/input.ts

@@ -10,10 +10,7 @@ import {zod} from '../third_party/index.js';
 import type {ElementHandle, KeyInput} from '../third_party/index.js';
 import type {TextSnapshotNode} from '../types.js';
 import {parseKey} from '../utils/keyboard.js';
-import {
-  appendWaitForResult,
-  type WaitForEventsResult,
-} from '../WaitForHelper.js';
+import type {WaitForEventsResult} from '../WaitForHelper.js';
 
 import {ToolCategory} from './categories.js';
 import type {ContextPage} from './ToolDefinition.js';
@@ -130,7 +127,7 @@ export const click = definePageTool({
           ? `Successfully double clicked on the element`
           : `Successfully clicked on the element`,
       );
-      appendWaitForResult(response, result);
+      response.attachWaitForResult(result);
       if (request.params.includeSnapshot) {
         response.includeSnapshot();
       }
@@ -169,7 +166,7 @@ export const clickAt = definePageTool({
         ? `Successfully double clicked at the coordinates`
         : `Successfully clicked at the coordinates`,
     );
-    appendWaitForResult(response, result);
+    response.attachWaitForResult(result);
     if (request.params.includeSnapshot) {
       response.includeSnapshot();
     }
@@ -200,7 +197,7 @@ export const hover = definePageTool({
         await handle.asLocator().hover();
       });
       response.appendResponseLine(`Successfully hovered over the element`);
-      appendWaitForResult(response, result);
+      response.attachWaitForResult(result);
       if (request.params.includeSnapshot) {
         response.includeSnapshot();
       }
@@ -330,7 +327,7 @@ export const fill = definePageTool({
       );
     });
     response.appendResponseLine(`Successfully filled out the element`);
-    appendWaitForResult(response, result);
+    response.attachWaitForResult(result);
     if (request.params.includeSnapshot) {
       response.includeSnapshot();
     }
@@ -362,7 +359,7 @@ export const typeText = definePageTool({
     response.appendResponseLine(
       `Typed text "${request.params.text}${request.params.submitKey ? ` + ${request.params.submitKey}` : ''}"`,
     );
-    appendWaitForResult(response, result);
+    response.attachWaitForResult(result);
   },
 });
 
@@ -391,7 +388,7 @@ export const drag = definePageTool({
         await toHandle.drop(fromHandle);
       });
       response.appendResponseLine(`Successfully dragged an element`);
-      appendWaitForResult(response, result);
+      response.attachWaitForResult(result);
       if (request.params.includeSnapshot) {
         response.includeSnapshot();
       }
@@ -440,7 +437,7 @@ export const fillForm = definePageTool({
       });
     }
     response.appendResponseLine(`Successfully filled out the form`);
-    appendWaitForResult(response, lastResult);
+    response.attachWaitForResult(lastResult);
     if (request.params.includeSnapshot) {
       response.includeSnapshot();
     }
@@ -533,7 +530,7 @@ export const pressKey = definePageTool({
     response.appendResponseLine(
       `Successfully pressed key: ${request.params.key}`,
     );
-    appendWaitForResult(response, result);
+    response.attachWaitForResult(result);
     if (request.params.includeSnapshot) {
       response.includeSnapshot();
     }

src/tools/script.ts

@@ -7,7 +7,6 @@
 import {zod} from '../third_party/index.js';
 import type {Frame, JSHandle, Page, WebWorker} from '../third_party/index.js';
 import type {ExtensionServiceWorker} from '../types.js';
-import {appendWaitForResult} from '../WaitForHelper.js';
 
 import {ToolCategory} from './categories.js';
 import type {Context, Response} from './ToolDefinition.js';
@@ -94,7 +93,7 @@ Example with arguments: \`(el) => {
             },
             {handleDialog: dialogAction ?? 'accept'},
           );
-        appendWaitForResult(response, result);
+        response.attachWaitForResult(result);
         return;
       }
 
@@ -120,7 +119,7 @@ Example with arguments: \`(el) => {
           },
           {handleDialog: dialogAction ?? 'accept'},
         );
-        appendWaitForResult(response, result);
+        response.attachWaitForResult(result);
       } finally {
         void Promise.allSettled(args.map(arg => arg.dispose()));
       }

tests/tools/input.test.ts

@@ -25,7 +25,7 @@ import {
 } from '../../src/tools/input.js';
 import {parseKey} from '../../src/utils/keyboard.js';
 import {serverHooks} from '../server.js';
-import {html, withMcpContext} from '../utils.js';
+import {html, withMcpContext, getTextContent} from '../utils.js';
 
 describe('input', () => {
   const server = serverHooks();
@@ -153,12 +153,12 @@ describe('input', () => {
           response,
           context,
         );
+        const result = await response.handle('click', context);
+        const textContent = getTextContent(result.content[0]);
         const expectedUrl = server.getRoute('/after-click');
         assert.ok(
-          response.responseLines.some(
-            line => line === `Page navigated to ${expectedUrl}.`,
-          ),
-          `Expected response to mention navigation to ${expectedUrl}, got: ${response.responseLines.join(' | ')}`,
+          textContent.includes(`Page navigated to ${expectedUrl}.`),
+          `Expected response to mention navigation to ${expectedUrl}, got: ${textContent}`,
         );
       });
     });
@@ -182,11 +182,11 @@ describe('input', () => {
           response,
           context,
         );
+        const result = await response.handle('click', context);
+        const textContent = getTextContent(result.content[0]);
         assert.ok(
-          !response.responseLines.some(line =>
-            line.startsWith('Page navigated to '),
-          ),
-          `Did not expect a navigation line, got: ${response.responseLines.join(' | ')}`,
+          !textContent.includes('Page navigated to '),
+          `Did not expect a navigation line, got: ${textContent}`,
         );
       });
     });