chore: DOM elements as inputs for in-page tools

src/McpResponse.ts

@@ -19,6 +19,7 @@ import type {
   Page,
   ResourceType,
   TextContent,
+  JSONSchema7Definition,
 } from './third_party/index.js';
 import type {ToolGroup, ToolDefinition} from './tools/inPage.js';
 import {handleDialog} from './tools/pages.js';
@@ -41,6 +42,57 @@ interface TraceInsightData {
   insightName: InsightName;
 }
 
+export function replaceHtmlElementsWithUids(schema: JSONSchema7Definition) {
+  if (typeof schema === 'boolean') {
+    return;
+  }
+
+  let isHtmlElement = false;
+  for (const [key, value] of Object.entries(schema)) {
+    if (key === 'x-mcp-type' && value === 'HTMLElement') {
+      isHtmlElement = true;
+      break;
+    }
+  }
+
+  if (isHtmlElement) {
+    schema.properties = {uid: {type: 'string'}};
+    schema.required = ['uid'];
+  }
+
+  if (schema.properties) {
+    for (const key of Object.keys(schema.properties)) {
+      replaceHtmlElementsWithUids(schema.properties[key]);
+    }
+  }
+
+  if (schema.items) {
+    if (Array.isArray(schema.items)) {
+      for (const item of schema.items) {
+        replaceHtmlElementsWithUids(item);
+      }
+    } else {
+      replaceHtmlElementsWithUids(schema.items);
+    }
+  }
+
+  if (schema.anyOf) {
+    for (const s of schema.anyOf) {
+      replaceHtmlElementsWithUids(s);
+    }
+  }
+  if (schema.allOf) {
+    for (const s of schema.allOf) {
+      replaceHtmlElementsWithUids(s);
+    }
+  }
+  if (schema.oneOf) {
+    for (const s of schema.oneOf) {
+      replaceHtmlElementsWithUids(s);
+    }
+  }
+}
+
 async function getToolGroup(
   page: McpPage,
 ): Promise<ToolGroup<ToolDefinition> | undefined> {
@@ -91,6 +143,10 @@ async function getToolGroup(
       }, 0);
     });
   });
+
+  for (const tool of toolGroup?.tools ?? []) {
+    replaceHtmlElementsWithUids(tool.inputSchema);
+  }
   return toolGroup;
 }
 

src/third_party/index.ts

@@ -41,7 +41,7 @@ export {default as puppeteer} from 'puppeteer-core';
 export type * from 'puppeteer-core';
 export {PipeTransport} from 'puppeteer-core/internal/node/PipeTransport.js';
 export type {CdpPage} from 'puppeteer-core/internal/cdp/Page.js';
-export type {JSONSchema7} from 'json-schema';
+export type {JSONSchema7, JSONSchema7Definition} from 'json-schema';
 export {
   resolveDefaultUserDataDir,
   detectBrowserPlatform,

src/tools/inPage.ts

@@ -4,7 +4,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import {zod, ajv, type JSONSchema7} from '../third_party/index.js';
+import {
+  zod,
+  ajv,
+  type JSONSchema7,
+  type ElementHandle,
+} from '../third_party/index.js';
 
 import {ToolCategory} from './categories.js';
 import {definePageTool} from './ToolDefinition.js';
@@ -87,6 +92,22 @@ export const executeInPageTool = definePageTool({
       }
     }
 
+    // Creates array of ElementHandles from the UIDs in the params.
+    // We do not replace the uids with the ElementsHandles yet, because
+    // the `evaluate` function only turns them into DOM elements if they
+    // are passed as non-nested arguments.
+    const handles: ElementHandle[] = [];
+    for (const value of Object.values(params)) {
+      if (
+        value instanceof Object &&
+        'uid' in value &&
+        typeof value.uid === 'string' &&
+        Object.keys(value).length === 1
+      ) {
+        handles.push(await request.page.getElementByUid(value.uid));
+      }
+    }
+
     const toolGroup = request.page.getInPageTools();
     const tool = toolGroup?.tools.find(t => t.name === toolName);
     if (!tool) {
@@ -102,7 +123,19 @@ export const executeInPageTool = definePageTool({
     }
 
     const result = await request.page.pptrPage.evaluate(
-      async (name, args) => {
+      async (name, args, ...elements) => {
+        // Replace the UIDs with DOM elements.
+        for (const [key, value] of Object.entries(args)) {
+          if (
+            value instanceof Object &&
+            'uid' in value &&
+            typeof value.uid === 'string' &&
+            Object.keys(value).length === 1
+          ) {
+            args[key] = elements.shift();
+          }
+        }
+
         if (!window.__dtmcp?.executeTool) {
           throw new Error('No tools found on the page');
         }
@@ -114,6 +147,7 @@ export const executeInPageTool = definePageTool({
       },
       toolName,
       params,
+      ...handles,
     );
     response.appendResponseLine(JSON.stringify(result, null, 2));
   },