chore: extract page snapshot, prep for multipage (#120)

src/tools/common.ts

@@ -20,7 +20,7 @@ import path from 'path';
 import { z } from 'zod';
 import { zodToJsonSchema } from 'zod-to-json-schema';
 
-import { captureAriaSnapshot, runAndWait, sanitizeForFilePath } from './utils';
+import { sanitizeForFilePath } from './utils';
 
 import type { ToolFactory, Tool } from './tool';
 
@@ -28,7 +28,7 @@ const navigateSchema = z.object({
   url: z.string().describe('The URL to navigate to'),
 });
 
-export const navigate: ToolFactory = snapshot => ({
+export const navigate: ToolFactory = captureSnapshot => ({
   schema: {
     name: 'browser_navigate',
     description: 'Navigate to a URL',
@@ -36,18 +36,15 @@ export const navigate: ToolFactory = snapshot => ({
   },
   handle: async (context, params) => {
     const validatedParams = navigateSchema.parse(params);
-    const page = await context.createPage();
-    await page.goto(validatedParams.url, { waitUntil: 'domcontentloaded' });
-    // Cap load event to 5 seconds, the page is operational at this point.
-    await page.waitForLoadState('load', { timeout: 5000 }).catch(() => {});
-    if (snapshot)
-      return captureAriaSnapshot(context);
-    return {
-      content: [{
-        type: 'text',
-        text: `Navigated to ${validatedParams.url}`,
-      }],
-    };
+    await context.createPage();
+    return await context.run(async page => {
+      await page.goto(validatedParams.url, { waitUntil: 'domcontentloaded' });
+      // Cap load event to 5 seconds, the page is operational at this point.
+      await page.waitForLoadState('load', { timeout: 5000 }).catch(() => {});
+    }, {
+      status: `Navigated to ${validatedParams.url}`,
+      captureSnapshot,
+    });
   },
 });
 
@@ -60,7 +57,12 @@ export const goBack: ToolFactory = snapshot => ({
     inputSchema: zodToJsonSchema(goBackSchema),
   },
   handle: async context => {
-    return await runAndWait(context, 'Navigated back', async page => page.goBack(), snapshot);
+    return await context.runAndWait(async page => {
+      await page.goBack();
+    }, {
+      status: 'Navigated back',
+      captureSnapshot: snapshot,
+    });
   },
 });
 
@@ -73,7 +75,12 @@ export const goForward: ToolFactory = snapshot => ({
     inputSchema: zodToJsonSchema(goForwardSchema),
   },
   handle: async context => {
-    return await runAndWait(context, 'Navigated forward', async page => page.goForward(), snapshot);
+    return await context.runAndWait(async page => {
+      await page.goForward();
+    }, {
+      status: 'Navigated forward',
+      captureSnapshot: snapshot,
+    });
   },
 });
 
@@ -103,7 +110,7 @@ const pressKeySchema = z.object({
   key: z.string().describe('Name of the key to press or a character to generate, such as `ArrowLeft` or `a`'),
 });
 
-export const pressKey: Tool = {
+export const pressKey: (captureSnapshot: boolean) => Tool = captureSnapshot => ({
   schema: {
     name: 'browser_press_key',
     description: 'Press a key on the keyboard',
@@ -111,11 +118,14 @@ export const pressKey: Tool = {
   },
   handle: async (context, params) => {
     const validatedParams = pressKeySchema.parse(params);
-    return await runAndWait(context, `Pressed key ${validatedParams.key}`, async page => {
+    return await context.runAndWait(async page => {
       await page.keyboard.press(validatedParams.key);
+    }, {
+      status: `Pressed key ${validatedParams.key}`,
+      captureSnapshot,
     });
   },
-};
+});
 
 const pdfSchema = z.object({});
 
@@ -161,7 +171,7 @@ const chooseFileSchema = z.object({
   paths: z.array(z.string()).describe('The absolute paths to the files to upload. Can be a single file or multiple files.'),
 });
 
-export const chooseFile: ToolFactory = snapshot => ({
+export const chooseFile: ToolFactory = captureSnapshot => ({
   schema: {
     name: 'browser_choose_file',
     description: 'Choose one or multiple files to upload',
@@ -169,9 +179,13 @@ export const chooseFile: ToolFactory = snapshot => ({
   },
   handle: async (context, params) => {
     const validatedParams = chooseFileSchema.parse(params);
-    return await runAndWait(context, `Chose files ${validatedParams.paths.join(', ')}`, async () => {
+    return await context.runAndWait(async () => {
       await context.submitFileChooser(validatedParams.paths);
-    }, snapshot);
+    }, {
+      status: `Chose files ${validatedParams.paths.join(', ')}`,
+      captureSnapshot,
+      noClearFileChooser: true,
+    });
   },
 });