Add handling

Author: Lightning00Blade

src/McpPage.ts

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

src/WaitForHelper.ts

@@ -126,11 +126,17 @@ export class WaitForHelper {
 
   async waitForEventsAfterAction(
     action: () => Promise<unknown>,
-    options?: {timeout?: number; handleDialog?: boolean},
+    options?: {timeout?: number; handleDialog?: 'accept' | 'dismiss' | string},
   ): Promise<void> {
     if (options?.handleDialog) {
-      const dialogHandler = (dialog: Pick<Dialog, 'accept'>) => {
-        void dialog.accept();
+      const dialogHandler = (dialog: Pick<Dialog, 'accept' | 'dismiss'>) => {
+        if (options.handleDialog === 'dismiss') {
+          void dialog.dismiss();
+        } else if (options.handleDialog === 'accept') {
+          void dialog.accept();
+        } else {
+          void dialog.accept(options.handleDialog);
+        }
       };
       this.#page.on('dialog', dialogHandler);
       this.#abortController.signal.addEventListener('abort', () => {

src/tools/ToolDefinition.ts

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

src/tools/script.ts

@@ -46,6 +46,12 @@ Example with arguments: \`(el) => {
         )
         .optional()
         .describe(`An optional list of arguments to pass to the function.`),
+      dialogAction: zod
+        .string()
+        .optional()
+        .describe(
+          'Handle dialogs while execution. "accept", "dismiss", or a string for window.prompt. Defaults to accept.',
+        ),
       ...(cliArgs?.experimentalPageIdRouting ? pageIdSchema : {}),
       ...(cliArgs?.categoryExtensions
         ? {
@@ -64,6 +70,7 @@ Example with arguments: \`(el) => {
         args: uidArgs,
         function: fnString,
         pageId,
+        dialogAction,
       } = request.params;
 
       if (cliArgs?.categoryExtensions && serviceWorkerId) {
@@ -81,7 +88,7 @@ Example with arguments: \`(el) => {
           async () => {
             await performEvaluation(worker, fnString, [], response);
           },
-          {handleDialog: true},
+          {handleDialog: dialogAction ?? 'accept'},
         );
         return;
       }
@@ -106,7 +113,7 @@ Example with arguments: \`(el) => {
           async () => {
             await performEvaluation(evaluatable, fnString, args, response);
           },
-          {handleDialog: true},
+          {handleDialog: dialogAction ?? 'accept'},
         );
       } finally {
         void Promise.allSettled(args.map(arg => arg.dispose()));

tests/tools/script.test.ts

@@ -121,6 +121,52 @@ describe('script', () => {
       });
     });
 
+    it('work for scripts that trigger dialogs and dismiss them', async () => {
+      await withMcpContext(async (response, context) => {
+        const page = context.getSelectedPptrPage();
+
+        await page.setContent(html`<button id="test">test</button>`);
+
+        await evaluateScript().handler(
+          {
+            params: {
+              function: String(() => {
+                return confirm('hello');
+              }),
+              dialogAction: 'dismiss',
+            },
+          },
+          response,
+          context,
+        );
+        const lineEvaluation = response.responseLines.at(2)!;
+        assert.strictEqual(JSON.parse(lineEvaluation), false);
+      });
+    });
+
+    it('work for scripts that trigger prompts and fill them', async () => {
+      await withMcpContext(async (response, context) => {
+        const page = context.getSelectedPptrPage();
+
+        await page.setContent(html`<button id="test">test</button>`);
+
+        await evaluateScript().handler(
+          {
+            params: {
+              function: String(() => {
+                return prompt('Enter your name:');
+              }),
+              dialogAction: 'John Doe',
+            },
+          },
+          response,
+          context,
+        );
+        const lineEvaluation = response.responseLines.at(2)!;
+        assert.strictEqual(JSON.parse(lineEvaluation), 'John Doe');
+      });
+    });
+
     it('work for async functions', async () => {
       await withMcpContext(async (response, context) => {
         const page = context.getSelectedPptrPage();