Bug 2065546 - [firefox-devtools-mcp] navigate should accept an option to wait for load

src/firefox/index.ts

@@ -11,7 +11,7 @@ import { remoteValueToNative } from '../utils/remote-value.js';
 import { ConsoleEvents, NetworkEvents, DebuggingEvents, DownloadEvents } from './events/index.js';
 import type { NetworkBodyResult } from './events/network.js';
 import { DomInteractions } from './dom.js';
-import { PageManagement } from './pages.js';
+import { PageManagement, type ReadinessState } from './pages.js';
 import { SnapshotManager, type Snapshot, type SnapshotOptions } from './snapshot/index.js';
 
 /**
@@ -232,11 +232,11 @@ export class FirefoxClient {
   // Pages / Navigation
   // ============================================================================
 
-  async navigate(url: string): Promise<void> {
+  async navigate(url: string, wait?: ReadinessState): Promise<void> {
     if (!this.pages) {
       throw new Error('Not connected');
     }
-    await this.pages.navigate(url);
+    await this.pages.navigate(url, wait);
   }
 
   async navigateBack(): Promise<void> {
@@ -302,11 +302,11 @@ export class FirefoxClient {
     return await this.pages.selectTab(index);
   }
 
-  async createNewPage(url: string): Promise<number> {
+  async createNewPage(url: string, wait?: ReadinessState): Promise<number> {
     if (!this.pages) {
       throw new Error('Not connected');
     }
-    return await this.pages.createNewPage(url);
+    return await this.pages.createNewPage(url, wait);
   }
 
   async closeTab(index: number): Promise<void> {

src/firefox/pages.ts

@@ -21,6 +21,20 @@ export function isCommonScheme(url: string): boolean {
 
 export type BiDiCommandFn = (method: string, params: Record<string, any>) => Promise<any>;
 
+/**
+ * WebDriver BiDi browsingContext.ReadinessState.
+ * - "none": return as soon as navigation starts
+ * - "interactive": wait for DOMContentLoaded
+ * - "complete": wait for the load event, including subresources
+ */
+export const READINESS_STATES = ['none', 'interactive', 'complete'] as const;
+
+export type ReadinessState = (typeof READINESS_STATES)[number];
+
+export function isReadinessState(value: unknown): value is ReadinessState {
+  return READINESS_STATES.includes(value as ReadinessState);
+}
+
 export class PageManagement {
   constructor(
     private driver: WebDriver,
@@ -31,16 +45,22 @@ export class PageManagement {
 
   /**
    * Navigate to URL using BiDi
+   *
+   * @param url - Target URL
+   * @param waitOverride - Explicit readiness state to wait for. When omitted,
+   *   common schemes wait for "interactive" and uncommon schemes do not wait.
    */
-  async navigate(url: string): Promise<void> {
+  async navigate(url: string, waitOverride?: ReadinessState): Promise<void> {
     const contextId = this.getCurrentContextId();
     if (!contextId) {
       throw new Error(`Cannot navigate: no browsing context ID`);
     }
 
     // Default wait time is "interactive" (DOMContentLoaded).
-    // All uncommon schemes use wait time "none"
-    const wait = isCommonScheme(url) ? 'interactive' : 'none';
+    // All uncommon schemes use wait time "none".
+    // An explicit override is honoured for every scheme: silently downgrading it
+    // would discard what the caller asked for with no way to tell.
+    const wait: ReadinessState = waitOverride ?? (isCommonScheme(url) ? 'interactive' : 'none');
 
     // Navigate using direct BiDi
     await this.sendBiDiCommand('browsingContext.navigate', {
@@ -186,13 +206,13 @@ export class PageManagement {
   /**
    * Create new page (tab)
    */
-  async createNewPage(url: string): Promise<number> {
+  async createNewPage(url: string, waitOverride?: ReadinessState): Promise<number> {
     await this.driver.switchTo().newWindow('tab');
     const handles = await this.driver.getAllWindowHandles();
     const newIdx = handles.length - 1;
     this.setCurrentContextId(handles[newIdx]!);
     this.cachedSelectedIdx = newIdx;
-    await this.navigate(url);
+    await this.navigate(url, waitOverride);
     return newIdx;
   }
 

src/tools/pages.ts

@@ -10,10 +10,46 @@ import {
 } from '../utils/response-helpers.js';
 import { saveOutput } from '../utils/save-output.js';
 import { defineModule } from './module.js';
+import { READINESS_STATES, isReadinessState, type ReadinessState } from '../firefox/pages.js';
 import type { McpToolResponse } from '../types/common.js';
 
 const DEFAULT_MAX_CONTENT_CHARS = 20_000;
 
+const WAIT_DESCRIPTION =
+  "When to return: 'none' (navigation started), 'interactive' (DOMContentLoaded), " +
+  "'complete' (load event fired, including subresources). Omit for the default: " +
+  "'interactive' for http/https/data/blob/file, 'none' for other schemes. Use " +
+  "'complete' when the page must be fully loaded, e.g. before stopping a performance recording.";
+
+const waitSchema = {
+  type: 'string',
+  enum: [...READINESS_STATES],
+  description: WAIT_DESCRIPTION,
+};
+
+/**
+ * Validate the optional `wait` argument.
+ *
+ * An unknown value is rejected rather than ignored: silently falling back to the
+ * default would leave the caller believing it waited for something it did not.
+ */
+function parseWait(value: unknown): ReadinessState | undefined {
+  if (value === undefined || value === null) {
+    return undefined;
+  }
+  if (!isReadinessState(value)) {
+    throw new Error(
+      `wait must be one of ${READINESS_STATES.join(', ')} (got ${JSON.stringify(value)})`
+    );
+  }
+  return value;
+}
+
+/** Echo the readiness state back only when the caller asked for one. */
+function waitSuffix(wait: ReadinessState | undefined): string {
+  return wait ? ` (waited for: ${wait})` : '';
+}
+
 // Tool definitions
 export const listPagesTool = {
   name: 'list_pages',
@@ -40,6 +76,7 @@ export const newPageTool = {
         type: 'string',
         description: 'Target URL',
       },
+      wait: waitSchema,
     },
     required: ['url'],
   },
@@ -58,6 +95,7 @@ export const navigatePageTool = {
         type: 'string',
         description: 'Target URL',
       },
+      wait: waitSchema,
     },
     required: ['url'],
   },
@@ -175,31 +213,35 @@ export async function handleListPages(_args: unknown): Promise<McpToolResponse>
 
 export async function handleNewPage(args: unknown): Promise<McpToolResponse> {
   try {
-    const { url } = args as { url: string };
+    const { url, wait } = args as { url: string; wait?: unknown };
 
     if (!url || typeof url !== 'string') {
       throw new Error('url parameter is required and must be a string');
     }
 
+    const waitFor = parseWait(wait);
+
     const { getFirefox } = await import('../index.js');
     const firefox = await getFirefox();
 
-    const newIdx = await firefox.createNewPage(url);
+    const newIdx = await firefox.createNewPage(url, waitFor);
 
-    return successResponse(`new page [${newIdx}] → ${url}`);
+    return successResponse(`new page [${newIdx}] → ${url}${waitSuffix(waitFor)}`);
   } catch (error) {
     return errorResponse(error as Error);
   }
 }
 
 export async function handleNavigatePage(args: unknown): Promise<McpToolResponse> {
   try {
-    const { url } = args as { url: string };
+    const { url, wait } = args as { url: string; wait?: unknown };
 
     if (!url || typeof url !== 'string') {
       throw new Error('url parameter is required and must be a string');
     }
 
+    const waitFor = parseWait(wait);
+
     const { getFirefox } = await import('../index.js');
     const firefox = await getFirefox();
 
@@ -213,9 +255,9 @@ export async function handleNavigatePage(args: unknown): Promise<McpToolResponse
       throw new Error('No page selected');
     }
 
-    await firefox.navigate(url);
+    await firefox.navigate(url, waitFor);
 
-    return successResponse(`[${selectedIdx}] → ${url}`);
+    return successResponse(`[${selectedIdx}] → ${url}${waitSuffix(waitFor)}`);
   } catch (error) {
     return errorResponse(error as Error);
   }

tests/firefox/pages.test.ts

@@ -7,7 +7,12 @@
  */
 
 import { describe, it, expect, vi } from 'vitest';
-import { isCommonScheme, PageManagement } from '@/firefox/pages.js';
+import {
+  isCommonScheme,
+  isReadinessState,
+  PageManagement,
+  READINESS_STATES,
+} from '@/firefox/pages.js';
 
 const HTTPS_URL = 'https://example.com/test.html';
 const HTTP_URL = 'http://example.com/test.html';
@@ -38,6 +43,22 @@ describe('isCommonScheme', () => {
   });
 });
 
+describe('isReadinessState', () => {
+  it('accepts the three BiDi readiness states', () => {
+    expect(READINESS_STATES).toEqual(['none', 'interactive', 'complete']);
+    for (const state of READINESS_STATES) {
+      expect(isReadinessState(state)).toBe(true);
+    }
+  });
+
+  it('rejects anything else, including the DOM event name "load"', () => {
+    expect(isReadinessState('load')).toBe(false);
+    expect(isReadinessState('COMPLETE')).toBe(false);
+    expect(isReadinessState(undefined)).toBe(false);
+    expect(isReadinessState(true)).toBe(false);
+  });
+});
+
 // -- PageManagement -----------------------------------------------------------
 
 describe('PageManagement', () => {
@@ -128,6 +149,57 @@ describe('PageManagement', () => {
         wait: 'none',
       });
     });
+
+    it('honours an explicit wait for common schemes', async () => {
+      const { pages, sendBiDiCommand } = createMocks();
+
+      await pages.navigate(HTTPS_URL, 'complete');
+      expect(sendBiDiCommand).toHaveBeenCalledWith('browsingContext.navigate', {
+        context: 'ctx-1',
+        url: HTTPS_URL,
+        wait: 'complete',
+      });
+    });
+
+    it('honours an explicit wait for uncommon schemes rather than downgrading it', async () => {
+      const { pages, sendBiDiCommand } = createMocks();
+
+      await pages.navigate(MOZ_EXT_URL, 'complete');
+      expect(sendBiDiCommand).toHaveBeenCalledWith('browsingContext.navigate', {
+        context: 'ctx-1',
+        url: MOZ_EXT_URL,
+        wait: 'complete',
+      });
+
+      await pages.navigate('about:blank', 'interactive');
+      expect(sendBiDiCommand).toHaveBeenCalledWith('browsingContext.navigate', {
+        context: 'ctx-1',
+        url: 'about:blank',
+        wait: 'interactive',
+      });
+    });
+
+    it('allows an explicit wait to opt out of waiting on a common scheme', async () => {
+      const { pages, sendBiDiCommand } = createMocks();
+
+      await pages.navigate(HTTPS_URL, 'none');
+      expect(sendBiDiCommand).toHaveBeenCalledWith('browsingContext.navigate', {
+        context: 'ctx-1',
+        url: HTTPS_URL,
+        wait: 'none',
+      });
+    });
+
+    it('keeps the scheme-based default when no wait is given', async () => {
+      const { pages, sendBiDiCommand } = createMocks();
+
+      await pages.navigate(HTTPS_URL, undefined);
+      expect(sendBiDiCommand).toHaveBeenCalledWith('browsingContext.navigate', {
+        context: 'ctx-1',
+        url: HTTPS_URL,
+        wait: 'interactive',
+      });
+    });
   });
 
   describe('createNewPage', () => {
@@ -170,5 +242,32 @@ describe('PageManagement', () => {
         wait: 'none',
       });
     });
+
+    it('forwards an explicit wait override to navigate', async () => {
+      const switchToMock = vi
+        .fn()
+        .mockReturnValue({ newWindow: vi.fn().mockResolvedValue(undefined) });
+      const getAllWindowHandlesMock = vi.fn().mockResolvedValue(['handle-1', 'handle-2']);
+
+      const driver = {
+        switchTo: switchToMock,
+        getAllWindowHandles: getAllWindowHandlesMock,
+      } as any;
+
+      const sendBiDiCommand = vi.fn().mockResolvedValue({});
+      const pages = new PageManagement(
+        driver,
+        vi.fn().mockReturnValue('handle-2'),
+        vi.fn(),
+        sendBiDiCommand
+      );
+
+      await pages.createNewPage(HTTPS_URL, 'complete');
+      expect(sendBiDiCommand).toHaveBeenCalledWith('browsingContext.navigate', {
+        context: 'handle-2',
+        url: HTTPS_URL,
+        wait: 'complete',
+      });
+    });
   });
 });