fix: flakiness in wait for helper (#2541)

Start 100ms timer only after the action is done to avoid the race with
the action code itself.

src/McpPage.ts

@@ -418,6 +418,7 @@ export class McpPage implements ContextPage {
     options?: {
       timeout?: number;
       waitForStableDom?: boolean;
+      expectNavigationIn?: number;
       handleDialog?:
         DialogAction | Partial<Record<Protocol.Page.DialogType, DialogAction>>;
     },

src/WaitForHelper.ts

@@ -99,38 +99,6 @@ export class WaitForHelper {
     ]);
   }
 
-  async waitForNavigationStarted() {
-    // Currently Puppeteer does not have API
-    // For when a navigation is about to start
-    const navigationStartedPromise = new Promise<boolean>(resolve => {
-      const listener = (event: Protocol.Page.FrameStartedNavigatingEvent) => {
-        if (
-          [
-            'historySameDocument',
-            'historyDifferentDocument',
-            'sameDocument',
-          ].includes(event.navigationType)
-        ) {
-          resolve(false);
-          return;
-        }
-
-        resolve(true);
-      };
-
-      this.#page._client().on('Page.frameStartedNavigating', listener);
-      this.#abortController.signal.addEventListener('abort', () => {
-        resolve(false);
-        this.#page._client().off('Page.frameStartedNavigating', listener);
-      });
-    });
-
-    return await Promise.race([
-      navigationStartedPromise,
-      this.timeout(this.#expectNavigationIn).then(() => false),
-    ]);
-  }
-
   timeout(time: number): Promise<void> {
     return new Promise<void>(res => {
       const id = setTimeout(res, time);
@@ -146,6 +114,7 @@ export class WaitForHelper {
     options?: {
       timeout?: number;
       waitForStableDom?: boolean;
+      expectNavigationIn?: number;
       handleDialog?:
         DialogAction | Partial<Record<Protocol.Page.DialogType, DialogAction>>;
     },
@@ -186,17 +155,67 @@ export class WaitForHelper {
       this.#page.off('dialog', dialogHandler);
     });
 
-    const navigationFinished = this.waitForNavigationStarted()
-      .then(navigationStated => {
-        if (navigationStated) {
-          return this.#page.waitForNavigation({
-            timeout: options?.timeout ?? this.#navigationTimeout,
-            signal: this.#abortController.signal,
-          });
-        }
+    // A scoped AbortController used to clean up navigation probe listeners.
+    // When aborted (either after navigation detection finishes or if this.#abortController
+    // aborts), it removes the CDP Page.frameStartedNavigating listener and automatically
+    // detaches the abort listener from this.#abortController.signal.
+    const navigationAbortController = new AbortController();
+    const navigationStartedResolvers = Promise.withResolvers<boolean>();
+
+    const navigationListener = (
+      event: Protocol.Page.FrameStartedNavigatingEvent,
+    ) => {
+      if (event.frameId !== this.#page.mainFrame()._id) {
         return;
+      }
+      if (
+        event.navigationType === 'sameDocument' ||
+        event.navigationType === 'historySameDocument'
+      ) {
+        return;
+      }
+
+      navigationStartedResolvers.resolve(true);
+    };
+
+    this.#page._client().on('Page.frameStartedNavigating', navigationListener);
+    navigationAbortController.signal.addEventListener('abort', () => {
+      this.#page
+        ._client()
+        .off('Page.frameStartedNavigating', navigationListener);
+    });
+    this.#abortController.signal.addEventListener(
+      'abort',
+      () => {
+        navigationStartedResolvers.resolve(false);
+        navigationAbortController.abort();
+      },
+      {signal: navigationAbortController.signal},
+    );
+
+    // Puppeteer's waitForNavigation must be started before the action runs so that
+    // it captures the pre-action loader ID. If started after the action triggers navigation,
+    // it risks recording the new loader ID and hanging until timeout.
+    // If no navigation occurs, this.#abortController will cancel it in the finally block.
+    const navigationFinished = this.#page
+      .waitForNavigation({
+        timeout: options?.timeout ?? this.#navigationTimeout,
+        signal: this.#abortController.signal,
+        ignoreSameDocumentNavigation: true,
       })
-      .catch(error => logger?.(error));
+      .then(result => {
+        navigationStartedResolvers.resolve(true);
+        return result;
+      })
+      .catch(error => {
+        if (
+          this.#abortController.signal.aborted ||
+          (error instanceof Error && error.name === 'AbortError')
+        ) {
+          return;
+        }
+        logger?.(error);
+      });
 
     try {
       await action();
@@ -206,8 +225,23 @@ export class WaitForHelper {
       throw error;
     }
 
+    const expectNavigationIn =
+      options?.expectNavigationIn ?? this.#expectNavigationIn;
+    const navigationStarted = await Promise.race([
+      navigationStartedResolvers.promise,
+      this.timeout(expectNavigationIn).then(() => {
+        navigationStartedResolvers.resolve(false);
+        return false;
+      }),
+    ]);
+    navigationAbortController.abort();
+
     try {
-      await navigationFinished;
+      // Only await navigation if one was actually initiated; otherwise, the
+      // pending waitForNavigation promise will be cancelled when this.#abortController aborts.
+      if (navigationStarted) {
+        await navigationFinished;
+      }
 
       if (this.#dialogDetected) {
         return this.#getResult();

tests/WaitForHelper.test.ts

@@ -7,9 +7,12 @@
 import assert from 'node:assert';
 import {describe, it} from 'node:test';
 
+import {serverHooks} from './server.js';
 import {html, withMcpContext} from './utils.js';
 
 describe('WaitForHelper', () => {
+  const server = serverHooks();
+
   it('does not stall when an action opens a dialog without handleDialog', async () => {
     await withMcpContext(async (response, context) => {
       const mcpPage = context.getSelectedMcpPage();
@@ -19,28 +22,136 @@ describe('WaitForHelper', () => {
       // The dialog leaves the renderer paused; without the fix,
       // waitForStableDom's setup evaluation would hang until protocolTimeout
       // (~180s) while the tool mutex is held, freezing the session.
-      const result = await Promise.race([
-        mcpPage.waitForEventsAfterAction(async () => {
-          await mcpPage.pptrPage.evaluate(() => {
-            setTimeout(() => confirm('blocked?'), 0);
-          });
-        }),
-        // Comfortably above WaitForHelper.#stableDomTimeout (3s): the call
-        // should return well within this once the dialog is detected.
-        new Promise<'stalled'>(resolve =>
-          setTimeout(() => resolve('stalled'), 5_000),
-        ),
-      ]);
+      try {
+        const result = await Promise.race([
+          mcpPage.waitForEventsAfterAction(async () => {
+            await mcpPage.pptrPage.evaluate(() => {
+              setTimeout(() => confirm('blocked?'), 0);
+            });
+          }),
+          // Comfortably above WaitForHelper.#stableDomTimeout (3s): the call
+          // should return well within this once the dialog is detected.
+          new Promise<'stalled'>(resolve =>
+            setTimeout(() => resolve('stalled'), 5_000),
+          ),
+        ]);
 
-      assert(
-        result !== 'stalled',
-        'stalled because a dialog was shown; would time out with ProtocolError',
+        assert(
+          result !== 'stalled',
+          'stalled because a dialog was shown; would time out with ProtocolError',
+        );
+        // The dialog was detected but not handled (no handleDialog was passed).
+        assert.strictEqual(result.dialogHandled, false);
+        // The dialog is still open and recorded, so the next blockedByDialog tool
+        // correctly refuses to run.
+        assert.throws(() => mcpPage.throwIfDialogOpen());
+      } finally {
+        await mcpPage.getDialog()?.dismiss();
+      }
+    });
+  });
+
+  it('awaits navigation when action takes longer than expectNavigationIn', async () => {
+    await withMcpContext(async (response, context) => {
+      server.addHtmlRoute('/nav-target', html`<main>navigated</main>`);
+      const url = server.getRoute('/nav-target');
+      const mcpPage = context.getSelectedMcpPage();
+
+      const result = await mcpPage.waitForEventsAfterAction(
+        async () => {
+          // Simulate an action that takes longer than expectNavigationIn (300ms)
+          // before triggering the navigation.
+          await new Promise(resolve => setTimeout(resolve, 600));
+          await mcpPage.pptrPage.evaluate(targetUrl => {
+            location.href = targetUrl;
+          }, url);
+        },
+        {waitForStableDom: false, expectNavigationIn: 300},
       );
-      // The dialog was detected but not handled (no handleDialog was passed).
-      assert.strictEqual(result.dialogHandled, false);
-      // The dialog is still open and recorded, so the next blockedByDialog tool
-      // correctly refuses to run.
-      assert.throws(() => mcpPage.throwIfDialogOpen());
+
+      assert.strictEqual(result.navigatedToUrl, url);
+    });
+  });
+
+  it('does not hang when an iframe navigates', async () => {
+    await withMcpContext(async (response, context) => {
+      server.addHtmlRoute('/iframe-src', html`<p>iframe</p>`);
+      server.addHtmlRoute('/iframe-target', html`<p>iframe navigated</p>`);
+      const iframeSrc = server.getRoute('/iframe-src');
+      const iframeTarget = server.getRoute('/iframe-target');
+      const mcpPage = context.getSelectedMcpPage();
+      await mcpPage.pptrPage.setContent(
+        html`<iframe
+          id="subframe"
+          src="${iframeSrc}"
+        ></iframe>`,
+      );
+
+      const startTime = Date.now();
+      const result = await mcpPage.waitForEventsAfterAction(
+        async () => {
+          await mcpPage.pptrPage.evaluate(targetUrl => {
+            const frame = document.querySelector('iframe');
+            if (!frame) {
+              throw new Error('iframe not found');
+            }
+            frame.src = targetUrl;
+          }, iframeTarget);
+        },
+        {waitForStableDom: false, expectNavigationIn: 50, timeout: 2000},
+      );
+
+      const elapsed = Date.now() - startTime;
+      assert(
+        elapsed < 1500,
+        `Took ${elapsed}ms; should not hang waiting for iframe`,
+      );
+      assert.strictEqual(result.navigatedToUrl, undefined);
+    });
+  });
+
+  it('awaits navigation when preceded by same-document navigation', async () => {
+    await withMcpContext(async (response, context) => {
+      server.addHtmlRoute('/nav-start', html`<main>start</main>`);
+      server.addHtmlRoute('/nav-target-2', html`<main>navigated 2</main>`);
+      const startUrl = server.getRoute('/nav-start');
+      const targetUrl = server.getRoute('/nav-target-2');
+      const mcpPage = context.getSelectedMcpPage();
+      await mcpPage.pptrPage.goto(startUrl);
+
+      const result = await mcpPage.waitForEventsAfterAction(
+        async () => {
+          await mcpPage.pptrPage.evaluate(url => {
+            history.pushState({}, '', '/intermediate-state');
+            location.href = url;
+          }, targetUrl);
+        },
+        {waitForStableDom: false, expectNavigationIn: 1000},
+      );
+
+      assert.strictEqual(result.navigatedToUrl, targetUrl);
+    });
+  });
+
+  it('captures navigatedToUrl for same-document navigation alone', async () => {
+    await withMcpContext(async (response, context) => {
+      server.addHtmlRoute('/nav-start-push', html`<main>start push</main>`);
+      server.addHtmlRoute('/same-doc-target', html`<main>target</main>`);
+      const startUrl = server.getRoute('/nav-start-push');
+      const targetUrl = server.getRoute('/same-doc-target');
+      const mcpPage = context.getSelectedMcpPage();
+      await mcpPage.pptrPage.goto(startUrl);
+
+      const result = await mcpPage.waitForEventsAfterAction(
+        async () => {
+          await mcpPage.pptrPage.evaluate(url => {
+            history.pushState({}, '', url);
+          }, targetUrl);
+        },
+        {waitForStableDom: false},
+      );
+
+      assert.strictEqual(result.navigatedToUrl, targetUrl);
     });
   });
 });