Bug 2089768 - Add a wait_for tool

docs/tools.md

@@ -14,7 +14,7 @@ the Mozilla-internal build; the public package drops both.
 | Module                    | Tools | slim | basic | developer | mozilla | all |
 | ------------------------- | ----- | ---- | ----- | --------- | ------- | --- |
 | `pages`                   | 6     | yes  | yes   | yes       | yes     | yes |
-| `snapshot`                | 3     | yes  | yes   | yes       | yes     | yes |
+| `snapshot`                | 4     | yes  | yes   | yes       | yes     | yes |
 | `input`                   | 8     | yes  | yes   | yes       | yes     | yes |
 | `network`                 | 3     | -    | -     | yes       | yes     | yes |
 | `console`                 | 2     | -    | -     | yes       | yes     | yes |
@@ -167,6 +167,19 @@ Parameters:
 
 - `uid` (string, required) - UID from snapshot
 
+### `wait_for`
+
+_Read-only._
+
+Wait until text or an element matching a CSS selector is visible, or with gone until it disappears. Also searches open shadow roots and same-origin iframes.
+
+Parameters:
+
+- `text` (string, optional) - Visible text to wait for
+- `selector` (string, optional) - CSS selector to wait for
+- `gone` (boolean, optional) - Wait until it disappears (default: false)
+- `timeout` (number, optional) - Timeout in ms (default: 5000, max: 30000)
+
 ### `clear_snapshot`
 
 Clear snapshot UIDs. Usually not needed.

src/firefox/dom.ts

@@ -4,6 +4,7 @@
 
 import { Key } from 'selenium-webdriver';
 import type { BrowsingContext, Input, Script } from 'webdriver-bidi-protocol';
+import { nativeToLocalValue } from '../utils/local-value.js';
 import type { BiDiFacade } from './bidi';
 
 /**
@@ -153,6 +154,41 @@ function comboKeyActions({ modifiers, key }: KeyCombo): Input.KeySourceAction[]
   return actions;
 }
 
+export interface WaitForOptions {
+  /** CSS selector of an element that has to be visible */
+  selector?: string | undefined;
+  /** Text that has to be visible */
+  text?: string | undefined;
+  /** Wait until it is no longer there instead */
+  gone?: boolean | undefined;
+  timeout?: number | undefined;
+}
+
+// Looks through the document, its open shadow roots and same-origin frames, like a snapshot does.
+const PRESENCE_SCRIPT = `(selector, text) => {
+  const roots = [document];
+  for (let i = 0; i < roots.length; i++) {
+    for (const el of roots[i].querySelectorAll('*')) {
+      if (el.shadowRoot) {
+        roots.push(el.shadowRoot);
+      }
+      try {
+        if (el.contentDocument) {
+          roots.push(el.contentDocument);
+        }
+      } catch {}
+    }
+  }
+  if (selector) {
+    return roots.some((root) =>
+      [...root.querySelectorAll(selector)].some((el) => el.checkVisibility({ visibilityProperty: true }))
+    );
+  }
+  const textOf = (root) =>
+    root.body ? root.body.innerText : [...root.children].map((el) => el.innerText).join('\\n');
+  return roots.some((root) => textOf(root).includes(text));
+}`;
+
 export class DomInteractions {
   constructor(
     private bidi: BiDiFacade,
@@ -235,6 +271,45 @@ export class DomInteractions {
     });
   }
 
+  /**
+   * Wait until an element or text shows up, or with gone until it disappears. A page that
+   * navigates in the meantime is searched again once its new document is there.
+   * @returns Milliseconds waited
+   */
+  async waitFor(
+    context: BrowsingContext.BrowsingContext,
+    options: WaitForOptions
+  ): Promise<number> {
+    const { selector, text, gone = false, timeout = 5000 } = options;
+    const started = Date.now();
+    for (;;) {
+      let present: boolean | undefined;
+      try {
+        present = await this.bidi.callFunction<boolean>(
+          context,
+          PRESENCE_SCRIPT,
+          [nativeToLocalValue(selector), nativeToLocalValue(text)],
+          'wait_for'
+        );
+      } catch (error) {
+        if (/not a valid selector/.test((error as Error).message)) {
+          throw error;
+        }
+      }
+      const waited = Date.now() - started;
+      if (present !== undefined && present !== gone) {
+        return waited;
+      }
+      if (waited >= timeout) {
+        const target = selector ? `selector ${selector}` : `text "${text}"`;
+        throw new Error(
+          `${target} ${gone ? 'was still there' : 'did not appear'} after ${timeout} ms`
+        );
+      }
+      await new Promise((r) => setTimeout(r, 100));
+    }
+  }
+
   // ============================================================================
   // UID-based input methods
   // ============================================================================

src/firefox/index.ts

@@ -10,7 +10,7 @@ import { listExtensions, type ListExtensionsFilters } from './webextension.js';
 import { logDebug } from '../utils/logger.js';
 import { ConsoleEvents, NetworkEvents, DebuggingEvents, DownloadEvents } from './events/index.js';
 import type { NetworkBodyResult } from './events/network.js';
-import { DomInteractions } from './dom.js';
+import { DomInteractions, type WaitForOptions } from './dom.js';
 import { PageManagement } from './pages.js';
 import { CacheManagement, type CacheBehavior } from './cache.js';
 import { SnapshotManager, type Snapshot, type SnapshotOptions } from './snapshot/index.js';
@@ -174,6 +174,10 @@ export class FirefoxClient {
     return await this.getDom().uploadFileByUid(this.getContext(), uid, filePath);
   }
 
+  async waitFor(options: WaitForOptions): Promise<number> {
+    return await this.getDom().waitFor(this.getContext(), options);
+  }
+
   async pressKey(key: string, uid?: string): Promise<void> {
     return await this.getDom().pressKey(this.getContext(), key, uid);
   }

src/tools/snapshot.ts

@@ -83,6 +83,36 @@ export const resolveUidToSelectorTool = {
   },
 } as const satisfies ToolDefinition;
 
+export const waitForTool = {
+  name: 'wait_for',
+  description:
+    'Wait until text or an element matching a CSS selector is visible, or with gone until it disappears. Also searches open shadow roots and same-origin iframes.',
+  annotations: {
+    readOnlyHint: true,
+  },
+  inputSchema: {
+    type: 'object',
+    properties: {
+      text: {
+        type: 'string',
+        description: 'Visible text to wait for',
+      },
+      selector: {
+        type: 'string',
+        description: 'CSS selector to wait for',
+      },
+      gone: {
+        type: 'boolean',
+        description: 'Wait until it disappears (default: false)',
+      },
+      timeout: {
+        type: 'number',
+        description: 'Timeout in ms (default: 5000, max: 30000)',
+      },
+    },
+  },
+} as const satisfies ToolDefinition;
+
 export const clearSnapshotTool = {
   name: 'clear_snapshot',
   description: 'Clear snapshot UIDs. Usually not needed.',
@@ -230,6 +260,29 @@ export const handleResolveUidToSelector = defineToolHandler<typeof resolveUidToS
   }
 );
 
+const MAX_WAIT_TIMEOUT = 30000;
+
+export const handleWaitFor = defineToolHandler<typeof waitForTool>(
+  async function handleWaitFor(args): Promise<McpToolResponse> {
+    const { text, selector, gone = false, timeout = 5000 } = args ?? {};
+    if ((text === undefined) === (selector === undefined)) {
+      throw new Error('wait_for needs either text or selector, not both');
+    }
+
+    const { getFirefox } = await import('../index.js');
+    const firefox = await getFirefox();
+
+    const waited = await firefox.waitFor({
+      text,
+      selector,
+      gone,
+      timeout: Math.min(Math.max(0, timeout), MAX_WAIT_TIMEOUT),
+    });
+    const target = selector ?? `"${text}"`;
+    return successResponse(`${target} ${gone ? 'disappeared' : 'appeared'} after ${waited} ms`);
+  }
+);
+
 export const handleClearSnapshot = defineToolHandler<typeof clearSnapshotTool>(
   async function handleClearSnapshot(_args): Promise<McpToolResponse> {
     const { getFirefox } = await import('../index.js');
@@ -247,6 +300,7 @@ export const module = defineModule({
   tools: [
     [takeSnapshotTool, handleTakeSnapshot],
     [resolveUidToSelectorTool, handleResolveUidToSelector],
+    [waitForTool, handleWaitFor],
     [clearSnapshotTool, handleClearSnapshot],
   ],
 });

tests/integration/firefox/wait-for.test.ts

@@ -0,0 +1,68 @@
+/**
+ * Integration tests for waiting on page content
+ */
+
+import { describe, it, expect, beforeAll, afterAll } from 'vitest';
+import { createTestFirefox, closeFirefox } from '../../helpers/firefox.js';
+import type { FirefoxClient } from '@/firefox/index.js';
+
+const PAGE =
+  'data:text/html,<div id=slot></div><p id=old>Loading</p><x-box></x-box><script>' +
+  'setTimeout(() => { slot.innerHTML = "<button id=late>Late</button>"; old.remove(); }, 500);' +
+  'customElements.define("x-box", class extends HTMLElement { connectedCallback() {' +
+  ' const root = this.attachShadow({ mode: "open" });' +
+  ' setTimeout(() => { root.innerHTML = "<span>Deep text</span>"; }, 500); } });</script>';
+
+describe('waitFor Integration Tests', () => {
+  let firefox: FirefoxClient;
+
+  beforeAll(async () => {
+    firefox = await createTestFirefox();
+  });
+
+  afterAll(async () => {
+    await closeFirefox(firefox);
+  });
+
+  it('should wait until an element matching the selector appears', async () => {
+    await firefox.navigate(PAGE);
+
+    await firefox.waitFor({ selector: '#late' });
+
+    expect(await firefox.evaluate('!!document.getElementById("late")')).toBe(true);
+  }, 10000);
+
+  it('should wait until text appears inside an open shadow root', async () => {
+    await firefox.navigate(PAGE);
+
+    await firefox.waitFor({ text: 'Deep text' });
+
+    expect(await firefox.evaluate('document.querySelector("x-box").shadowRoot.textContent')).toBe(
+      'Deep text'
+    );
+  }, 10000);
+
+  it('should wait until the text is gone', async () => {
+    await firefox.navigate(PAGE);
+
+    await firefox.waitFor({ text: 'Loading', gone: true });
+
+    expect(await firefox.evaluate('!!document.getElementById("old")')).toBe(false);
+  }, 10000);
+
+  it('should fail after the timeout and name what it waited for', async () => {
+    await firefox.navigate(PAGE);
+
+    await expect(firefox.waitFor({ selector: '#never', timeout: 300 })).rejects.toThrow(/#never/);
+  }, 10000);
+
+  it('should reject an invalid selector without waiting for the timeout', async () => {
+    await firefox.navigate(PAGE);
+    const started = Date.now();
+
+    await expect(firefox.waitFor({ selector: 'div[', timeout: 5000 })).rejects.toThrow(
+      /not a valid selector/
+    );
+    expect(Date.now() - started).toBeLessThan(2000);
+  }, 10000);
+});

tests/integration/mcp/snapshot.test.ts

@@ -0,0 +1,30 @@
+/**
+ * Integration tests for the snapshot module, called through the MCP server
+ */
+
+import { expect } from 'vitest';
+import { mcpTest } from '../../helpers/mcp.js';
+
+mcpTest.override('mcpServerArgs', ['--tools', 'pages', 'snapshot']);
+
+mcpTest.describe('wait_for (MCP)', () => {
+  mcpTest.beforeEach(async ({ mcp }) => {
+    await mcp.callTool('navigate_page', { url: 'data:text/html,<p>Ready</p>' });
+  });
+
+  mcpTest('reports how long it waited', async ({ mcp }) => {
+    const result = await mcp.callTool('wait_for', { text: 'Ready' });
+
+    expect(result.isError).toBe(false);
+    expect(result.text).toMatch(/"Ready" appeared after \d+ ms/);
+  });
+
+  mcpTest('asks for either text or a selector', async ({ mcp }) => {
+    const neither = await mcp.callTool('wait_for', {});
+    const both = await mcp.callTool('wait_for', { text: 'Ready', selector: 'p' });
+
+    expect(neither.isError).toBe(true);
+    expect(both.isError).toBe(true);
+    expect(neither.text).toContain('either text or selector');
+  });
+});