feat: support timeouts in wait_for and navigations

docs/tool-reference.md

@@ -138,6 +138,7 @@
 
 **Parameters:**
 
+- **timeout** (integer) _(optional)_: Maximum wait time in milliseconds. If set to 0, the default timeout will be used.
 - **url** (string) **(required)**: URL to navigate the page to
 
 ---
@@ -149,6 +150,7 @@
 **Parameters:**
 
 - **navigate** (enum: "back", "forward") **(required)**: Whether to navigate back or navigate forward in the selected pages history
+- **timeout** (integer) _(optional)_: Maximum wait time in milliseconds. If set to 0, the default timeout will be used.
 
 ---
 
@@ -158,6 +160,7 @@
 
 **Parameters:**
 
+- **timeout** (integer) _(optional)_: Maximum wait time in milliseconds. If set to 0, the default timeout will be used.
 - **url** (string) **(required)**: URL to load in a new page.
 
 ---
@@ -179,6 +182,7 @@
 **Parameters:**
 
 - **text** (string) **(required)**: Text to appear on the page
+- **timeout** (integer) _(optional)_: Maximum wait time in milliseconds. If set to 0, the default timeout will be used.
 
 ---
 

src/McpResponse.ts

@@ -154,7 +154,7 @@ export class McpResponse implements Response {
       response.push(`## Network emulation`);
       response.push(`Emulating: ${networkConditions}`);
       response.push(
-        `Navigation timeout set to ${context.getNavigationTimeout()} ms`,
+        `Default navigation timeout set to ${context.getNavigationTimeout()} ms`,
       );
     }
 

src/tools/ToolDefinition.ts

@@ -5,7 +5,7 @@
  */
 
 import type {Dialog, ElementHandle, Page} from 'puppeteer-core';
-import type z from 'zod';
+import z from 'zod';
 
 import type {TraceResult} from '../trace-processing/parse.js';
 
@@ -85,3 +85,16 @@ export function defineTool<Schema extends z.ZodRawShape>(
 
 export const CLOSE_PAGE_ERROR =
   'The last open page cannot be closed. It is fine to keep it open.';
+
+export const timeoutSchema = {
+  timeout: z
+    .number()
+    .int()
+    .optional()
+    .describe(
+      `Maximum wait time in milliseconds. If set to 0, the default timeout will be used.`,
+    )
+    .transform(value => {
+      return value && value <= 0 ? undefined : value;
+    }),
+};

src/tools/pages.ts

@@ -9,7 +9,7 @@ import z from 'zod';
 import {logger} from '../logger.js';
 
 import {ToolCategories} from './categories.js';
-import {CLOSE_PAGE_ERROR, defineTool} from './ToolDefinition.js';
+import {CLOSE_PAGE_ERROR, defineTool, timeoutSchema} from './ToolDefinition.js';
 
 export const listPages = defineTool({
   name: 'list_pages',
@@ -83,12 +83,15 @@ export const newPage = defineTool({
   },
   schema: {
     url: z.string().describe('URL to load in a new page.'),
+    ...timeoutSchema,
   },
   handler: async (request, response, context) => {
     const page = await context.newPage();
 
     await context.waitForEventsAfterAction(async () => {
-      await page.goto(request.params.url);
+      await page.goto(request.params.url, {
+        timeout: request.params.timeout,
+      });
     });
 
     response.setIncludePages(true);
@@ -104,12 +107,15 @@ export const navigatePage = defineTool({
   },
   schema: {
     url: z.string().describe('URL to navigate the page to'),
+    ...timeoutSchema,
   },
   handler: async (request, response, context) => {
     const page = context.getSelectedPage();
 
     await context.waitForEventsAfterAction(async () => {
-      await page.goto(request.params.url);
+      await page.goto(request.params.url, {
+        timeout: request.params.timeout,
+      });
     });
 
     response.setIncludePages(true);
@@ -129,15 +135,18 @@ export const navigatePageHistory = defineTool({
       .describe(
         'Whether to navigate back or navigate forward in the selected pages history',
       ),
+    ...timeoutSchema,
   },
   handler: async (request, response, context) => {
     const page = context.getSelectedPage();
-
+    const options = {
+      timeout: request.params.timeout,
+    };
     try {
       if (request.params.navigate === 'back') {
-        await page.goBack();
+        await page.goBack(options);
       } else {
-        await page.goForward();
+        await page.goForward(options);
       }
     } catch {
       response.appendResponseLine(

src/tools/snapshot.ts

@@ -8,7 +8,7 @@ import {Locator} from 'puppeteer-core';
 import z from 'zod';
 
 import {ToolCategories} from './categories.js';
-import {defineTool} from './ToolDefinition.js';
+import {defineTool, timeoutSchema} from './ToolDefinition.js';
 
 export const takeSnapshot = defineTool({
   name: 'take_snapshot',
@@ -33,17 +33,24 @@ export const waitFor = defineTool({
   },
   schema: {
     text: z.string().describe('Text to appear on the page'),
+    ...timeoutSchema,
   },
   handler: async (request, response, context) => {
     const page = context.getSelectedPage();
     const frames = page.frames();
 
-    const locators = frames.flatMap(frame => [
-      frame.locator(`aria/${request.params.text}`),
-      frame.locator(`text/${request.params.text}`),
-    ]);
+    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.race(locators).wait();
+    await locator.wait();
 
     response.appendResponseLine(
       `Element with text "${request.params.text}" found.`,

tests/McpResponse.test.ts

@@ -105,7 +105,7 @@ uid=1_0 RootWebArea "My test page"
         `# test response
 ## Network emulation
 Emulating: Slow 3G
-Navigation timeout set to 100000 ms`,
+Default navigation timeout set to 100000 ms`,
       );
     });
   });