Bug 2018268 - fix: migrate script evaluation tools to WebDriver BiDi

src/tools/privileged-context.ts

@@ -4,6 +4,8 @@
  */
 
 import { successResponse, errorResponse } from '../utils/response-helpers.js';
+import { validateFunction } from '../utils/js-validation.js';
+import { remoteValueToNative } from '../utils/remote-value.js';
 import type { McpToolResponse } from '../types/common.js';
 
 export const listPrivilegedContextsTool = {
@@ -35,28 +37,19 @@ export const selectPrivilegedContextTool = {
 export const evaluatePrivilegedScriptTool = {
   name: 'evaluate_privileged_script',
   description:
-    'Evaluate JavaScript in the current privileged context. Requires MOZ_REMOTE_ALLOW_SYSTEM_ACCESS=1 env var. Returns the result of the expression. IMPORTANT: Only provide expressions, not statements. Do not use const, let, or var declarations as they will cause syntax errors. For complex logic, wrap in an IIFE: (function() { const x = 1; return x; })()',
+    'Execute JS function in the current privileged context. Requires MOZ_REMOTE_ALLOW_SYSTEM_ACCESS=1 env var. Use select_privileged_context first to target a chrome context.',
   inputSchema: {
     type: 'object',
     properties: {
-      expression: {
+      function: {
         type: 'string',
-        description: 'JavaScript expression to evaluate in the privileged context',
+        description: 'JS function string, e.g. () => Services.prefs.getBoolPref("foo")',
       },
     },
-    required: ['expression'],
+    required: ['function'],
   },
 };
 
-/**
- * Detects if the input looks like a JavaScript statement rather than an expression.
- * Statements like const/let/var declarations cannot be used with return().
- */
-export function isLikelyStatement(input: string): boolean {
-  const trimmed = input.trim();
-  return /^(const|let|var)\s/.test(trimmed);
-}
-
 function formatContextList(contexts: any[]): string {
   if (contexts.length === 0) {
     return 'No privileged contexts found';
@@ -132,47 +125,45 @@ export async function handleSelectPrivilegedContext(args: unknown): Promise<McpT
   }
 }
 
+const EvaluateResultType = {
+  Exception: 'exception',
+  Success: 'success',
+};
+
 export async function handleEvaluatePrivilegedScript(args: unknown): Promise<McpToolResponse> {
   try {
-    const { expression } = args as { expression: string };
-
-    if (!expression || typeof expression !== 'string') {
-      throw new Error('expression parameter is required and must be a string');
-    }
+    const { function: fnString } = args as { function: string };
 
-    if (isLikelyStatement(expression)) {
-      return errorResponse(
-        new Error(
-          `Cannot evaluate statement: "${expression.substring(0, 50)}${expression.length > 50 ? '...' : ''}". ` +
-            'This tool expects an expression, not a statement (const/let/var declarations are statements). ' +
-            'To use statements, wrap them in an IIFE: (function() { const x = 1; return x; })()'
-        )
-      );
-    }
+    validateFunction(fnString);
 
     const { getFirefox } = await import('../index.js');
     const firefox = await getFirefox();
 
-    const driver = firefox.getDriver();
+    const result = await firefox.sendBiDiCommand('script.callFunction', {
+      functionDeclaration: fnString,
+      awaitPromise: true,
+      arguments: [],
+      target: { context: firefox.getCurrentContextId() },
+    });
 
-    try {
-      const result = await driver.executeScript(`return (${expression});`);
-      const resultText =
-        typeof result === 'string'
-          ? result
-          : result === null
-            ? 'null'
-            : result === undefined
-              ? 'undefined'
-              : JSON.stringify(result, null, 2);
-
-      return successResponse(`Result:\n${resultText}`);
-    } catch (executeError) {
+    if (result.type === EvaluateResultType.Success) {
+      let output = 'Script ran in chrome context and returned:\n';
+      output += '```json\n';
+      output += JSON.stringify(remoteValueToNative(result.result), null, 2);
+      output += '\n```';
+      return successResponse(output);
+    } else if (result.type === EvaluateResultType.Exception) {
+      const exceptionDetails = result.exceptionDetails;
       return errorResponse(
         new Error(
-          `Script execution failed: ${executeError instanceof Error ? executeError.message : String(executeError)}`
+          `Script execution failed: ${exceptionDetails.text}\n\n` +
+            '```json\n' +
+            JSON.stringify(remoteValueToNative(exceptionDetails.exception), null, 2) +
+            '\n```'
         )
       );
+    } else {
+      return errorResponse(`Unexpected script.callFunction result type: ${result.type}`);
     }
   } catch (error) {
     return errorResponse(error as Error);

src/tools/script.ts

@@ -1,8 +1,10 @@
 /**
- * JavaScript evaluation tool (currently disabled - see docs/future-features.md)
+ * JavaScript evaluation tool
  */
 
 import { successResponse, errorResponse } from '../utils/response-helpers.js';
+import { remoteValueToNative } from '../utils/remote-value.js';
+import { validateFunction } from '../utils/js-validation.js';
 import type { McpToolResponse } from '../types/common.js';
 
 export const evaluateScriptTool = {
@@ -39,43 +41,14 @@ export const evaluateScriptTool = {
 };
 
 // Constants
-const MAX_FUNCTION_SIZE = 16 * 1024; // 16 KB
 const DEFAULT_TIMEOUT = 5000; // 5 seconds
+const TIMEOUT = Symbol('Timeout');
 
-/**
- * Validate function string format
- */
-function validateFunction(fnString: string): void {
-  if (!fnString || typeof fnString !== 'string') {
-    throw new Error('function parameter is required and must be a string');
-  }
-
-  if (fnString.length > MAX_FUNCTION_SIZE) {
-    throw new Error(
-      `Function too large (${fnString.length} bytes, max ${MAX_FUNCTION_SIZE} bytes). ` +
-        'This tool is not designed for massive scripts.'
-    );
-  }
-
-  // Check if it looks like a function or arrow function
-  const trimmed = fnString.trim();
-  const isFunctionLike =
-    trimmed.startsWith('function') ||
-    trimmed.startsWith('async function') ||
-    trimmed.startsWith('(') ||
-    trimmed.startsWith('async (');
-
-  if (!isFunctionLike) {
-    throw new Error(
-      `Invalid function format. Expected a function or arrow function, got: "${trimmed.substring(0, 50)}...".\n\n` +
-        'Valid examples:\n' +
-        '  () => document.title\n' +
-        '  async () => { return await fetch("/api") }\n' +
-        '  (el) => el.innerText\n' +
-        '  function() { return window.location.href }'
-    );
-  }
-}
+// Types from the WebDriver BiDi specification.
+const EvaluateResultType = {
+  Exception: 'exception',
+  Success: 'success',
+};
 
 export async function handleEvaluateScript(args: unknown): Promise<McpToolResponse> {
   try {
@@ -94,21 +67,16 @@ export async function handleEvaluateScript(args: unknown): Promise<McpToolRespon
 
     const { getFirefox } = await import('../index.js');
     const firefox = await getFirefox();
-    const driver = firefox.getDriver();
-
-    if (!driver) {
-      throw new Error('WebDriver not available');
-    }
 
     const scriptTimeout = timeout ?? DEFAULT_TIMEOUT;
 
-    // Prepare arguments: resolve UIDs to WebElements if provided
+    // Prepare arguments: resolve UIDs to references shared ids if provided
     const resolvedArgs: unknown[] = [];
     if (fnArgs && fnArgs.length > 0) {
       for (const arg of fnArgs) {
         try {
           const element = await firefox.resolveUidToElement(arg.uid);
-          resolvedArgs.push(element);
+          resolvedArgs.push({ sharedId: await element.getId() });
         } catch (error) {
           const errorMsg = (error as Error).message;
 
@@ -130,42 +98,51 @@ export async function handleEvaluateScript(args: unknown): Promise<McpToolRespon
       }
     }
 
-    // Unified execution path: use executeScript with optional args
-    const evalCode = `
-      const fn = ${fnString};
-      const args = Array.from(arguments);
-      const result = fn(...args);
-      return result instanceof Promise ? result : Promise.resolve(result);
-    `;
-
-    // Set script timeout
-    await driver.manage().setTimeouts({ script: scriptTimeout });
-
     // Execute with resolved args (empty array if no args)
-    const result = await driver.executeScript(evalCode, ...resolvedArgs);
-
-    // Format output
-    let output = 'Script ran on page and returned:\n';
-    output += '```json\n';
-    output += JSON.stringify(result, null, 2);
-    output += '\n```';
-
-    return successResponse(output);
-  } catch (error) {
-    const errorMsg = (error as Error).message;
-
-    // Enhance timeout errors
-    if (errorMsg.includes('timeout') || errorMsg.includes('Timeout')) {
-      const timeoutValue = (args as { timeout?: number })?.timeout ?? DEFAULT_TIMEOUT;
+    const callFunctionPromise = firefox.sendBiDiCommand('script.callFunction', {
+      functionDeclaration: fnString,
+      awaitPromise: true,
+      arguments: resolvedArgs,
+      target: { context: firefox.getCurrentContextId() },
+    });
+
+    // Race against timeout as WebDriver BiDi callFunction has no built-in
+    // timeout feature.
+    const result = await Promise.race([
+      new Promise((r) => setTimeout(() => r(TIMEOUT), scriptTimeout)),
+      callFunctionPromise,
+    ]);
+
+    if (result === TIMEOUT) {
       return errorResponse(
         new Error(
-          `Script execution timed out (exceeded ${timeoutValue}ms).\n\n` +
+          `Script execution timed out (exceeded ${scriptTimeout}ms).\n\n` +
             'The function may contain an infinite loop or be waiting for a slow operation.\n' +
             'Try simplifying the script or increasing the timeout parameter.'
         )
       );
+    } else if (result.type === EvaluateResultType.Success) {
+      // Format output
+      let output = 'Script ran on page and returned:\n';
+      output += '```json\n';
+      output += JSON.stringify(remoteValueToNative(result.result), null, 2);
+      output += '\n```';
+
+      return successResponse(output);
+    } else if (result.type === EvaluateResultType.Exception) {
+      const exceptionDetails = result.exceptionDetails;
+      return errorResponse(
+        new Error(
+          `Script execution failed: ${exceptionDetails.text}\n\n` +
+            '```json\n' +
+            JSON.stringify(remoteValueToNative(exceptionDetails.exception), null, 2) +
+            '\n```'
+        )
+      );
+    } else {
+      return errorResponse(`Unexpected script.callFunction result type: ${result.type}`);
     }
-
+  } catch (error) {
     return errorResponse(error as Error);
   }
 }

src/utils/js-validation.ts

@@ -0,0 +1,37 @@
+const MAX_FUNCTION_SIZE = 16 * 1024; // 16 KB
+
+/**
+ * Validates that a string is a function or arrow function declaration suitable
+ * for use with WebDriver BiDi script.callFunction.
+ * Throws with a descriptive error if validation fails.
+ */
+export function validateFunction(fnString: string): void {
+  if (!fnString || typeof fnString !== 'string') {
+    throw new Error('function parameter is required and must be a string');
+  }
+
+  if (fnString.length > MAX_FUNCTION_SIZE) {
+    throw new Error(
+      `Function too large (${fnString.length} bytes, max ${MAX_FUNCTION_SIZE} bytes). ` +
+        'This tool is not designed for massive scripts.'
+    );
+  }
+
+  const trimmed = fnString.trim();
+  const isFunctionLike =
+    trimmed.startsWith('function') ||
+    trimmed.startsWith('async function') ||
+    trimmed.startsWith('(') ||
+    trimmed.startsWith('async (');
+
+  if (!isFunctionLike) {
+    throw new Error(
+      `Invalid function format. Expected a function or arrow function, got: "${trimmed.substring(0, 50)}...".\n\n` +
+        'Valid examples:\n' +
+        '  () => document.title\n' +
+        '  async () => { return await fetch("/api") }\n' +
+        '  (el) => el.innerText\n' +
+        '  function() { return window.location.href }'
+    );
+  }
+}

src/utils/remote-value.ts

@@ -0,0 +1,63 @@
+/**
+ * Converts a WebDriver BiDi RemoteValue to a native JavaScript value.
+ * Special number values (NaN, Infinity, -0) are returned as strings since
+ * they cannot be represented in JSON.
+ */
+export function remoteValueToNative(rv: unknown): unknown {
+  if (!rv || typeof rv !== 'object') {
+    return rv;
+  }
+
+  const { type, value } = rv as { type: string; value?: unknown };
+
+  switch (type) {
+    case 'undefined':
+      return undefined;
+    case 'null':
+      return null;
+    case 'string':
+    case 'boolean':
+      return value;
+    case 'number':
+      if (value === 'NaN') {
+        return 'NaN';
+      }
+      if (value === 'Infinity') {
+        return 'Infinity';
+      }
+      if (value === '-Infinity') {
+        return '-Infinity';
+      }
+      if (value === '-0') {
+        return '-0';
+      }
+      return value;
+    case 'bigint':
+      return `${value as string}n`;
+    case 'array':
+      return (value as unknown[]).map(remoteValueToNative);
+    case 'object':
+      return Object.fromEntries(
+        (value as [string, unknown][]).map(([k, v]) => [k, remoteValueToNative(v)])
+      );
+    case 'map':
+      return Object.fromEntries(
+        (value as [unknown, unknown][]).map(([k, v]) => [
+          typeof k === 'object'
+            ? JSON.stringify(remoteValueToNative(k))
+            : String(k as string | number | boolean),
+          remoteValueToNative(v),
+        ])
+      );
+    case 'set':
+      return (value as unknown[]).map(remoteValueToNative);
+    case 'regexp': {
+      const { pattern, flags } = value as { pattern: string; flags?: string };
+      return `/${pattern}/${flags ?? ''}`;
+    }
+    case 'date':
+      return value;
+    default:
+      return `[${type}]`;
+  }
+}