fix: move dialogHandler inside conditional block

Author: Lightning00Blade

src/McpPage.ts

@@ -112,7 +112,7 @@ export class McpPage implements ContextPage {
 
   waitForEventsAfterAction(
     action: () => Promise<unknown>,
-    options?: {timeout?: number; dialog?: 'accept' | 'dismiss'},
+    options?: {timeout?: number; handleDialog?: boolean},
   ): Promise<void> {
     const helper = this.createWaitForHelper(
       this.cpuThrottlingRate,

src/WaitForHelper.ts

@@ -5,7 +5,7 @@
  */
 
 import {logger} from './logger.js';
-import type {Page, Protocol, CdpPage} from './third_party/index.js';
+import type {Page, Protocol, CdpPage, Dialog} from './third_party/index.js';
 import type {PredefinedNetworkConditions} from './third_party/index.js';
 
 export class WaitForHelper {
@@ -126,17 +126,16 @@ export class WaitForHelper {
 
   async waitForEventsAfterAction(
     action: () => Promise<unknown>,
-    options?: {timeout?: number; dialog?: 'accept' | 'dismiss'},
+    options?: {timeout?: number; handleDialog?: boolean},
   ): Promise<void> {
-    const dialogHandler = (dialog: any) => {
-      if (options?.dialog === 'dismiss') {
-        void dialog.dismiss();
-      } else {
+    if (options?.handleDialog) {
+      const dialogHandler = (dialog: Pick<Dialog, 'accept'>) => {
         void dialog.accept();
-      }
-    };
-    if (options?.dialog) {
+      };
       this.#page.on('dialog', dialogHandler);
+      this.#abortController.signal.addEventListener('abort', () => {
+        this.#page.off('dialog', dialogHandler);
+      });
     }
 
     const navigationFinished = this.waitForNavigationStarted()
@@ -169,9 +168,6 @@ export class WaitForHelper {
       logger(error);
     } finally {
       this.#abortController.abort();
-      if (options?.dialog) {
-        this.#page.off('dialog', dialogHandler);
-      }
     }
   }
 }

src/tools/ToolDefinition.ts

@@ -210,7 +210,7 @@ export type ContextPage = Readonly<{
   clearDialog(): void;
   waitForEventsAfterAction(
     action: () => Promise<unknown>,
-    options?: {timeout?: number; dialog?: 'accept' | 'dismiss'},
+    options?: {timeout?: number; handleDialog?: boolean},
   ): Promise<void>;
   getInPageTools(): ToolGroup<InPageToolDefinition> | undefined;
 }>;

src/tools/script.ts

@@ -77,11 +77,12 @@ Example with arguments: \`(el) => {
         }
 
         const worker = await getWebWorker(context, serviceWorkerId);
-        await context
-          .getSelectedMcpPage()
-          .waitForEventsAfterAction(async () => {
+        await context.getSelectedMcpPage().waitForEventsAfterAction(
+          async () => {
             await performEvaluation(worker, fnString, [], response);
-          }, {dialog: 'accept'});
+          },
+          {handleDialog: true},
+        );
         return;
       }
 
@@ -101,9 +102,12 @@ Example with arguments: \`(el) => {
 
         const evaluatable = await getPageOrFrame(page, frames);
 
-        await mcpPage.waitForEventsAfterAction(async () => {
-          await performEvaluation(evaluatable, fnString, args, response);
-        }, {dialog: 'accept'});
+        await mcpPage.waitForEventsAfterAction(
+          async () => {
+            await performEvaluation(evaluatable, fnString, args, response);
+          },
+          {handleDialog: true},
+        );
       } finally {
         void Promise.allSettled(args.map(arg => arg.dispose()));
       }