feat: support script evaluate arguments

Author: OrKoN

src/tools/script.ts

@@ -7,34 +7,65 @@ import z from 'zod';
 import {defineTool} from './ToolDefinition.js';
 import {ToolCategories} from './categories.js';
 import {waitForEventsAfterAction} from '../waitForHelpers.js';
+import {JSHandle} from 'puppeteer-core';
 
 export const evaluateScript = defineTool({
   name: 'evaluate_script',
-  description: `Evaluate a JavaScript function inside the currently selected page. Returns the response as JSON.`,
+  description: `Evaluate a JavaScript function inside the currently selected page. Returns the response as JSON
+so returned values have to JSON-serializable.`,
   annotations: {
     category: ToolCategories.DEBUGGING,
     readOnlyHint: false,
   },
   schema: {
-    function: z
-      .string()
-      .describe(
-        'A JavaScript function to run in the currently selected page. Example: `() => {return document.title}` or `async () => {return await fetch("example.com")}`',
-      ),
+    function: z.string().describe(
+      `A JavaScript function to run in the currently selected page.
+Example without arguments: \`() => {
+  return document.title
+}\` or \`async () => {
+  return await fetch("example.com")
+}\`.
+Example with arguments: \`(el) => {
+  return el.innerText;
+}\`
+`,
+    ),
+    args: z
+      .array(
+        z.object({
+          uid: z
+            .string()
+            .describe(
+              'The uid of an element on the page from the page content snapshot',
+            ),
+        }),
+      )
+      .optional()
+      .describe(`An optional list of arguments to pass to the function.`),
   },
   handler: async (request, response, context) => {
     const page = context.getSelectedPage();
-
-    const script = `(async () => {
-      return JSON.stringify(await (${request.params.function})());
-    })()`;
-
-    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('```');
-    });
+    const fn = await page.evaluateHandle(`(${request.params.function})`);
+    const args: JSHandle<unknown>[] = [fn];
+    try {
+      for (const el of request.params.args ?? []) {
+        args.push(await context.getElementByUid(el.uid));
+      }
+      await waitForEventsAfterAction(page, async () => {
+        const result = await page.evaluate(
+          async (fn, ...args) => {
+            // @ts-expect-error no types.
+            return JSON.stringify(await fn(...args));
+          },
+          ...args,
+        );
+        response.appendResponseLine('Script ran on page and returned:');
+        response.appendResponseLine('```json');
+        response.appendResponseLine(`${result}`);
+        response.appendResponseLine('```');
+      });
+    } finally {
+      Promise.allSettled(args.map(arg => arg.dispose())).catch(() => {});
+    }
   },
 });

tests/tools/script.test.ts

@@ -103,5 +103,55 @@ describe('script', () => {
         assert.strictEqual(JSON.parse(lineEvaluation), 'Works');
       });
     });
+
+    it('work with one argument', async () => {
+      await withBrowser(async (response, context) => {
+        const page = context.getSelectedPage();
+
+        await page.setContent(html`<button id="test">test</button>`);
+
+        await context.createTextSnapshot();
+
+        await evaluateScript.handler(
+          {
+            params: {
+              function: String(async (el: Element) => {
+                return el.id;
+              }),
+              args: [{uid: '1_1'}],
+            },
+          },
+          response,
+          context,
+        );
+        const lineEvaluation = response.responseLines.at(2)!;
+        assert.strictEqual(JSON.parse(lineEvaluation), 'test');
+      });
+    });
+
+    it('work with multiple args', async () => {
+      await withBrowser(async (response, context) => {
+        const page = context.getSelectedPage();
+
+        await page.setContent(html`<button id="test">test</button>`);
+
+        await context.createTextSnapshot();
+
+        await evaluateScript.handler(
+          {
+            params: {
+              function: String((container: Element, child: Element) => {
+                return container.contains(child);
+              }),
+              args: [{uid: '1_0'}, {uid: '1_1'}],
+            },
+          },
+          response,
+          context,
+        );
+        const lineEvaluation = response.responseLines.at(2)!;
+        assert.strictEqual(JSON.parse(lineEvaluation), true);
+      });
+    });
   });
 });