chore: evaluate script on service workers

src/McpContext.ts

@@ -41,6 +41,7 @@ import type {
   GeolocationOptions,
   TextSnapshot,
   TextSnapshotNode,
+  ExtensionServiceWorker,
 } from './types.js';
 import {
   ExtensionRegistry,
@@ -55,12 +56,6 @@ export type {
   TextSnapshotNode,
 } from './types.js';
 
-export interface ExtensionServiceWorker {
-  url: string;
-  target: Target;
-  id: string;
-}
-
 interface McpContextOptions {
   // Whether the DevTools windows are exposed as pages for debugging of DevTools.
   experimentalDevToolsDebugging: boolean;

src/tools/ToolDefinition.ts

@@ -5,6 +5,7 @@
  */
 
 import type {ParsedArguments} from '../cli.js';
+import type {McpPage} from '../McpPage.js';
 import {zod} from '../third_party/index.js';
 import type {
   Dialog,
@@ -14,7 +15,11 @@ import type {
   Viewport,
 } from '../third_party/index.js';
 import type {InsightName, TraceResult} from '../trace-processing/parse.js';
-import type {TextSnapshotNode, GeolocationOptions} from '../types.js';
+import type {
+  TextSnapshotNode,
+  GeolocationOptions,
+  ExtensionServiceWorker,
+} from '../types.js';
 import type {InstalledExtension} from '../utils/ExtensionRegistry.js';
 import type {PaginationOptions} from '../utils/types.js';
 
@@ -187,6 +192,11 @@ export type Context = Readonly<{
   uninstallExtension(id: string): Promise<void>;
   listExtensions(): InstalledExtension[];
   getExtension(id: string): InstalledExtension | undefined;
+  getSelectedMcpPage(): McpPage;
+  getExtensionServiceWorkers(): ExtensionServiceWorker[];
+  getExtensionServiceWorkerId(
+    extensionServiceWorker: ExtensionServiceWorker,
+  ): string | undefined;
 }>;
 
 export type ContextPage = Readonly<{

src/tools/script.ts

@@ -5,22 +5,27 @@
  */
 
 import {zod} from '../third_party/index.js';
-import type {Frame, JSHandle, Page} from '../third_party/index.js';
+import type {Frame, JSHandle, Page, WebWorker} from '../third_party/index.js';
+import type {ExtensionServiceWorker} from '../types.js';
 
 import {ToolCategory} from './categories.js';
-import {definePageTool} from './ToolDefinition.js';
+import type {Context, Response} from './ToolDefinition.js';
+import {defineTool, pageIdSchema} from './ToolDefinition.js';
 
-export const evaluateScript = definePageTool({
-  name: 'evaluate_script',
-  description: `Evaluate a JavaScript function inside the currently selected page. Returns the response as JSON,
+export type Evaluatable = Page | Frame | WebWorker;
+
+export const evaluateScript = defineTool(cliArgs => {
+  return {
+    name: 'evaluate_script',
+    description: `Evaluate a JavaScript function inside the currently selected page. Returns the response as JSON,
 so returned values have to be JSON-serializable.`,
-  annotations: {
-    category: ToolCategory.DEBUGGING,
-    readOnlyHint: false,
-  },
-  schema: {
-    function: zod.string().describe(
-      `A JavaScript function declaration to be executed by the tool in the currently selected page.
+    annotations: {
+      category: ToolCategory.DEBUGGING,
+      readOnlyHint: false,
+    },
+    schema: {
+      function: zod.string().describe(
+        `A JavaScript function declaration to be executed by the tool in the currently selected page.
 Example without arguments: \`() => {
   return document.title
 }\` or \`async () => {
@@ -30,57 +35,143 @@ Example with arguments: \`(el) => {
   return el.innerText;
 }\`
 `,
-    ),
-    args: zod
-      .array(
-        zod.object({
-          uid: zod
-            .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 args: Array<JSHandle<unknown>> = [];
-    try {
-      const frames = new Set<Frame>();
-      for (const el of request.params.args ?? []) {
-        const handle = await request.page.getElementByUid(el.uid);
-        frames.add(handle.frame);
-        args.push(handle);
+      ),
+      args: zod
+        .array(
+          zod.object({
+            uid: zod
+              .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.`),
+      ...(cliArgs?.experimentalPageIdRouting ? pageIdSchema : {}),
+      ...(cliArgs?.categoryExtensions
+        ? {
+            serviceWorkerId: zod
+              .string()
+              .optional()
+              .describe(
+                `An optional service worker id to evaluate the script in.`,
+              ),
+          }
+        : {}),
+    },
+    handler: async (request, response, context) => {
+      const {
+        serviceWorkerId,
+        args: uidArgs,
+        function: fnString,
+        pageId,
+      } = request.params;
+
+      if (cliArgs?.categoryExtensions && serviceWorkerId) {
+        if (uidArgs && uidArgs.length > 0) {
+          throw new Error(
+            'args (element uids) cannot be used when evaluating in a service worker.',
+          );
+        }
+        if (pageId) {
+          throw new Error('specify either a pageId or a serviceWorkerId.');
+        }
+
+        const worker = await getWebWorker(context, serviceWorkerId);
+        await performEvaluation(worker, fnString, [], response, context);
+        return;
       }
-      let pageOrFrame: Page | Frame;
-      // We can't evaluate the element handle across frames
-      if (frames.size > 1) {
-        throw new Error(
-          "Elements from different frames can't be evaluated together.",
-        );
-      } else {
-        pageOrFrame = [...frames.values()][0] ?? request.page.pptrPage;
+
+      const mcpPage = cliArgs?.experimentalPageIdRouting
+        ? context.getPageById(request.params.pageId)
+        : context.getSelectedMcpPage();
+      const page: Page = mcpPage.pptrPage;
+
+      const args: Array<JSHandle<unknown>> = [];
+      try {
+        const frames = new Set<Frame>();
+        for (const el of uidArgs ?? []) {
+          const handle = await mcpPage.getElementByUid(el.uid);
+          frames.add(handle.frame);
+          args.push(handle);
+        }
+
+        const evaluatable = await getPageOrFrame(page, frames);
+
+        await performEvaluation(evaluatable, fnString, args, response, context);
+      } finally {
+        void Promise.allSettled(args.map(arg => arg.dispose()));
       }
-      const fn = await pageOrFrame.evaluateHandle(
-        `(${request.params.function})`,
+    },
+  };
+});
+
+const performEvaluation = async (
+  evaluatable: Evaluatable,
+  fnString: string,
+  args: Array<JSHandle<unknown>>,
+  response: Response,
+  context: Context,
+) => {
+  const fn = await evaluatable.evaluateHandle(`(${fnString})`);
+  try {
+    await context.waitForEventsAfterAction(async () => {
+      const result = await evaluatable.evaluate(
+        async (fn, ...args) => {
+          // @ts-expect-error no types for function fn
+          return JSON.stringify(await fn(...args));
+        },
+        fn,
+        ...args,
       );
-      args.unshift(fn);
-      await context.waitForEventsAfterAction(async () => {
-        const result = await pageOrFrame.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 {
-      void Promise.allSettled(args.map(arg => arg.dispose()));
+      response.appendResponseLine('Script ran on page and returned:');
+      response.appendResponseLine('```json');
+      response.appendResponseLine(`${result}`);
+      response.appendResponseLine('```');
+    });
+  } finally {
+    void fn.dispose();
+  }
+};
+
+const getPageOrFrame = async (
+  page: Page,
+  frames: Set<Frame>,
+): Promise<Page | Frame> => {
+  let pageOrFrame: Page | Frame;
+  // We can't evaluate the element handle across frames
+  if (frames.size > 1) {
+    throw new Error(
+      "Elements from different frames can't be evaluated together.",
+    );
+  } else {
+    pageOrFrame = [...frames.values()][0] ?? page;
+  }
+
+  return pageOrFrame;
+};
+
+const getWebWorker = async (
+  context: Context,
+  serviceWorkerId: string,
+): Promise<WebWorker> => {
+  const serviceWorkers = context.getExtensionServiceWorkers();
+
+  const serviceWorker = serviceWorkers.find(
+    (sw: ExtensionServiceWorker) =>
+      context.getExtensionServiceWorkerId(sw) === serviceWorkerId,
+  );
+
+  if (serviceWorker && serviceWorker.target) {
+    const worker = await serviceWorker.target.worker();
+
+    if (!worker) {
+      throw new Error('Service worker target not found.');
     }
-  },
-});
+
+    return worker;
+  } else {
+    throw new Error('Service worker not found.');
+  }
+};

src/tools/tools.ts

@@ -44,7 +44,7 @@ export const createTools = (args: ParsedArguments) => {
   const tools = [];
   for (const tool of rawTools) {
     if (typeof tool === 'function') {
-      tools.push(tool(args));
+      tools.push(tool(args) as unknown as ToolDefinition);
     } else {
       tools.push(tool as ToolDefinition);
     }

src/types.ts

@@ -4,7 +4,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type {SerializedAXNode, Viewport} from './third_party/index.js';
+import type {SerializedAXNode, Viewport, Target} from './third_party/index.js';
+
+export interface ExtensionServiceWorker {
+  url: string;
+  target: Target;
+  id: string;
+}
 
 export interface TextSnapshotNode extends SerializedAXNode {
   id: string;

tests/tools/extensions.test.ts

@@ -24,7 +24,7 @@ const EXTENSION_PATH = path.join(
   '../../../tests/tools/fixtures/extension',
 );
 
-function extractId(response: McpResponse) {
+export function extractId(response: McpResponse) {
   const responseLine = response.responseLines[0];
   assert.ok(responseLine, 'Response should not be empty');
   const match = responseLine.match(/Extension installed\. Id: (.+)/);