feat: improve tools with awaiting common events (#10)

Author: Lightning00Blade

src/tools/input.ts

@@ -8,7 +8,7 @@ import z from 'zod';
 import {defineTool} from './ToolDefinition.js';
 import {ElementHandle} from 'puppeteer-core';
 import {ToolCategories} from './categories.js';
-import {performAction} from '../performAction.js';
+import {waitForEventsAfterAction} from '../waitForHelpers.js';
 
 export const click = defineTool({
   name: 'click',
@@ -32,7 +32,7 @@ export const click = defineTool({
     const uid = request.params.uid;
     const handle = await context.getElementByUid(uid);
     try {
-      await performAction(handle.frame.page(), async () => {
+      await waitForEventsAfterAction(handle.frame.page(), async () => {
         await handle.asLocator().click({
           count: request.params.dblClick ? 2 : 1,
         });
@@ -67,7 +67,9 @@ export const hover = defineTool({
     const uid = request.params.uid;
     const handle = await context.getElementByUid(uid);
     try {
-      await handle.asLocator().hover();
+      await waitForEventsAfterAction(handle.frame.page(), async () => {
+        await handle.asLocator().hover();
+      });
       response.appendResponseLine(`Successfully hovered over the element`);
       response.setIncludeSnapshot(true);
     } finally {
@@ -92,10 +94,11 @@ export const fill = defineTool({
     value: z.string().describe('The value to fill in'),
   },
   handler: async (request, response, context) => {
-    const uid = request.params.uid;
-    const handle = await context.getElementByUid(uid);
+    const handle = await context.getElementByUid(request.params.uid);
     try {
-      await handle.asLocator().fill(request.params.value);
+      await waitForEventsAfterAction(handle.frame.page(), async () => {
+        await handle.asLocator().fill(request.params.value);
+      });
       response.appendResponseLine(`Successfully filled out the element`);
       response.setIncludeSnapshot(true);
     } finally {
@@ -119,9 +122,11 @@ export const drag = defineTool({
     const fromHandle = await context.getElementByUid(request.params.from_uid);
     const toHandle = await context.getElementByUid(request.params.to_uid);
     try {
-      await fromHandle.drag(toHandle);
-      await new Promise(resolve => setTimeout(resolve, 50));
-      await toHandle.drop(fromHandle);
+      await waitForEventsAfterAction(fromHandle.frame.page(), async () => {
+        await fromHandle.drag(toHandle);
+        await new Promise(resolve => setTimeout(resolve, 50));
+        await toHandle.drop(fromHandle);
+      });
       response.appendResponseLine(`Successfully dragged an element`);
       response.setIncludeSnapshot(true);
     } finally {
@@ -152,7 +157,9 @@ export const fillForm = defineTool({
     for (const element of request.params.elements) {
       const handle = await context.getElementByUid(element.uid);
       try {
-        await handle.asLocator().fill(element.value);
+        await waitForEventsAfterAction(handle.frame.page(), async () => {
+          await handle.asLocator().fill(element.value);
+        });
       } finally {
         handle.dispose();
       }

src/tools/pages.ts

@@ -7,6 +7,7 @@
 import z from 'zod';
 import {defineTool} from './ToolDefinition.js';
 import {ToolCategories} from './categories.js';
+import {waitForEventsAfterAction} from '../waitForHelpers.js';
 
 export const listPages = defineTool({
   name: 'list_pages',
@@ -77,7 +78,11 @@ export const newPage = defineTool({
   },
   handler: async (request, response, context) => {
     const page = await context.newPage();
-    await page.goto(request.params.url);
+
+    await waitForEventsAfterAction(page, async () => {
+      await page.goto(request.params.url);
+    });
+
     response.setIncludePages(true);
   },
 });
@@ -95,7 +100,9 @@ export const navigatePage = defineTool({
   handler: async (request, response, context) => {
     const page = context.getSelectedPage();
 
-    await page.goto(request.params.url);
+    await waitForEventsAfterAction(page, async () => {
+      await page.goto(request.params.url);
+    });
 
     response.setIncludePages(true);
   },

src/tools/script.ts

@@ -6,6 +6,7 @@
 import z from 'zod';
 import {defineTool} from './ToolDefinition.js';
 import {ToolCategories} from './categories.js';
+import {waitForEventsAfterAction} from '../waitForHelpers.js';
 
 export const evaluateScript = defineTool({
   name: 'evaluate_script',
@@ -26,13 +27,14 @@ export const evaluateScript = defineTool({
 
     const script = `(async () => {
       return JSON.stringify(await (${request.params.function})());
-      })()`;
+    })()`;
 
-    const result = await page.evaluate(script);
-
-    response.appendResponseLine('Script ran on page and returned:');
-    response.appendResponseLine('```json');
-    response.appendResponseLine(`${result}`);
-    response.appendResponseLine('```');
+    await waitForEventsAfterAction(page, async () => {
+      const result = await page.evaluate(script);
+      response.appendResponseLine('Script ran on page and returned:');
+      response.appendResponseLine('```json');
+      response.appendResponseLine(`${result}`);
+      response.appendResponseLine('```');
+    });
   },
 });

src/performAction.ts renamed to src/waitForHelpers.ts

@@ -49,9 +49,14 @@ async function waitForStableDom(
     }
   });
 
-  return stableDomObserver.evaluate(async observer => {
-    return await observer.resolver.promise;
-  });
+  return Promise.race([
+    stableDomObserver.evaluate(async observer => {
+      return await observer.resolver.promise;
+    }),
+    timeout(3000, signal).then(() => {
+      throw new Error('Timeout');
+    }),
+  ]);
 }
 
 async function waitForNavigationStarted(page: CdpPage, signal: AbortSignal) {
@@ -101,7 +106,10 @@ function timeout(time: number, signal?: AbortSignal): Promise<void> {
  * a potential navigation, after which it waits
  * for the DOM to be stable before returning.
  */
-export async function performAction(page: Page, callback: () => Promise<void>) {
+export async function waitForEventsAfterAction(
+  page: Page,
+  callback: () => Promise<void>,
+) {
   const controller = new AbortController();
 
   const navigationStartedPromise = waitForNavigationStarted(
@@ -122,12 +130,7 @@ export async function performAction(page: Page, callback: () => Promise<void>) {
 
     // Wait for stable dom after navigation so we execute in
     // the correct context
-    await Promise.race([
-      waitForStableDom(page, controller.signal),
-      timeout(3000, controller.signal).then(() => {
-        throw new Error('Timeout');
-      }),
-    ]);
+    await waitForStableDom(page, controller.signal);
   } catch (error) {
     logger(error);
   } finally {