From 457b8af2f7db1741d0097417c70cf9c668424514 Mon Sep 17 00:00:00 2001
From: Julian Descottes <jdescottes@mozilla.com>
Date: Fri, 12 Jun 2026 10:13:28 +0200
Subject: [PATCH 1/4] Bug 2018268 - refactor: move js validation helpers to
 dedicated module

---
 src/tools/privileged-context.ts        | 10 +--------
 tests/tools/privileged-context.test.ts | 31 --------------------------
 2 files changed, 1 insertion(+), 40 deletions(-)

diff --git a/src/tools/privileged-context.ts b/src/tools/privileged-context.ts
index 643a235..30b59c5 100644
--- a/src/tools/privileged-context.ts
+++ b/src/tools/privileged-context.ts
@@ -4,6 +4,7 @@
  */
 
 import { successResponse, errorResponse } from '../utils/response-helpers.js';
+import { isLikelyStatement } from '../utils/js-validation.js';
 import type { McpToolResponse } from '../types/common.js';
 
 export const listPrivilegedContextsTool = {
@@ -48,15 +49,6 @@ export const evaluatePrivilegedScriptTool = {
   },
 };
 
-/**
- * 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';
diff --git a/tests/tools/privileged-context.test.ts b/tests/tools/privileged-context.test.ts
index 5dfa37c..e8948e7 100644
--- a/tests/tools/privileged-context.test.ts
+++ b/tests/tools/privileged-context.test.ts
@@ -6,7 +6,6 @@ import { describe, it, expect, vi, beforeEach } from 'vitest';
 import {
   evaluatePrivilegedScriptTool,
   handleEvaluatePrivilegedScript,
-  isLikelyStatement,
 } from '../../src/tools/privileged-context.js';
 
 // Mock the index module (used by handler tests)
@@ -28,36 +27,6 @@ describe('Privileged Context Tool Definitions', () => {
   });
 });
 
-describe('isLikelyStatement', () => {
-  it('should detect const declarations', () => {
-    expect(isLikelyStatement('const x = 1')).toBe(true);
-  });
-
-  it('should detect let declarations', () => {
-    expect(isLikelyStatement('let x = 1')).toBe(true);
-  });
-
-  it('should detect var declarations', () => {
-    expect(isLikelyStatement('var x = 1')).toBe(true);
-  });
-
-  it('should allow function calls', () => {
-    expect(isLikelyStatement('Services.prefs.getBoolPref("foo")')).toBe(false);
-  });
-
-  it('should allow simple expressions', () => {
-    expect(isLikelyStatement('1 + 2')).toBe(false);
-  });
-
-  it('should allow property access', () => {
-    expect(isLikelyStatement('document.title')).toBe(false);
-  });
-
-  it('should handle leading whitespace', () => {
-    expect(isLikelyStatement('  const x = 1')).toBe(true);
-  });
-});
-
 describe('Privileged Context Tool Handlers', () => {
   const mockExecuteScript = vi.fn();
   const mockSetContext = vi.fn();

From 0a06d4265746eee98447a68962baef21097da69f Mon Sep 17 00:00:00 2001
From: Julian Descottes <jdescottes@mozilla.com>
Date: Fri, 12 Jun 2026 11:29:07 +0200
Subject: [PATCH 2/4] Bug 2018268 - refactor: update script module to use
 WebDriver BiDi for script evaluation

---
 src/tools/script.ts               |  84 ++++++++++-------
 src/utils/js-validation.ts        |  18 ++++
 src/utils/remote-value.ts         |  63 +++++++++++++
 tests/utils/js-validation.test.ts |  62 ++++++++++++
 tests/utils/remote-value.test.ts  | 152 ++++++++++++++++++++++++++++++
 5 files changed, 343 insertions(+), 36 deletions(-)
 create mode 100644 src/utils/js-validation.ts
 create mode 100644 src/utils/remote-value.ts
 create mode 100644 tests/utils/js-validation.test.ts
 create mode 100644 tests/utils/remote-value.test.ts

diff --git a/src/tools/script.ts b/src/tools/script.ts
index 5237afb..5baf1b2 100644
--- a/src/tools/script.ts
+++ b/src/tools/script.ts
@@ -1,8 +1,9 @@
 /**
- * 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 type { McpToolResponse } from '../types/common.js';
 
 export const evaluateScriptTool = {
@@ -41,6 +42,13 @@ export const evaluateScriptTool = {
 // Constants
 const MAX_FUNCTION_SIZE = 16 * 1024; // 16 KB
 const DEFAULT_TIMEOUT = 5000; // 5 seconds
+const TIMEOUT = Symbol('Timeout');
+
+// Types from the WebDriver BiDi specification.
+const EvaluateResultType = {
+  Exception: 'exception',
+  Success: 'success',
+};
 
 /**
  * Validate function string format
@@ -94,21 +102,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 +133,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);
   }
 }
diff --git a/src/utils/js-validation.ts b/src/utils/js-validation.ts
new file mode 100644
index 0000000..a7bb526
--- /dev/null
+++ b/src/utils/js-validation.ts
@@ -0,0 +1,18 @@
+/**
+ * Detects if the input looks like a JavaScript statement rather than an expression.
+ * Statements like const/let/var declarations cannot be used in expression contexts.
+ */
+export function isLikelyStatement(input: string): boolean {
+  const trimmed = input.trim();
+  return /^(const|let|var)\s/.test(trimmed);
+}
+
+/**
+ * Detects if the input looks like a function body (starts with "return").
+ * Function bodies are valid in WebDriver Classic executeScript but not in BiDi script.evaluate,
+ * which expects an expression.
+ */
+export function isLikelyFunctionBody(input: string): boolean {
+  const trimmed = input.trim();
+  return /^return(\s|;|$)/.test(trimmed);
+}
diff --git a/src/utils/remote-value.ts b/src/utils/remote-value.ts
new file mode 100644
index 0000000..ed3aa73
--- /dev/null
+++ b/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}]`;
+  }
+}
diff --git a/tests/utils/js-validation.test.ts b/tests/utils/js-validation.test.ts
new file mode 100644
index 0000000..6e021d5
--- /dev/null
+++ b/tests/utils/js-validation.test.ts
@@ -0,0 +1,62 @@
+import { describe, it, expect } from 'vitest';
+import { isLikelyFunctionBody, isLikelyStatement } from '../../src/utils/js-validation.js';
+
+describe('isLikelyStatement', () => {
+  it('should detect const declarations', () => {
+    expect(isLikelyStatement('const x = 1')).toBe(true);
+  });
+
+  it('should detect let declarations', () => {
+    expect(isLikelyStatement('let x = 1')).toBe(true);
+  });
+
+  it('should detect var declarations', () => {
+    expect(isLikelyStatement('var x = 1')).toBe(true);
+  });
+
+  it('should allow function calls', () => {
+    expect(isLikelyStatement('Services.prefs.getBoolPref("foo")')).toBe(false);
+  });
+
+  it('should allow simple expressions', () => {
+    expect(isLikelyStatement('1 + 2')).toBe(false);
+  });
+
+  it('should allow property access', () => {
+    expect(isLikelyStatement('document.title')).toBe(false);
+  });
+
+  it('should handle leading whitespace', () => {
+    expect(isLikelyStatement('  const x = 1')).toBe(true);
+  });
+});
+
+describe('isLikelyFunctionBody', () => {
+  it('should detect return with a value', () => {
+    expect(isLikelyFunctionBody('return document.title')).toBe(true);
+  });
+
+  it('should detect return followed by semicolon', () => {
+    expect(isLikelyFunctionBody('return;')).toBe(true);
+  });
+
+  it('should detect bare return', () => {
+    expect(isLikelyFunctionBody('return')).toBe(true);
+  });
+
+  it('should handle leading whitespace', () => {
+    expect(isLikelyFunctionBody('  return x')).toBe(true);
+  });
+
+  it('should allow expressions', () => {
+    expect(isLikelyFunctionBody('document.title')).toBe(false);
+  });
+
+  it('should allow IIFEs', () => {
+    expect(isLikelyFunctionBody('(function() { return 1; })()')).toBe(false);
+  });
+
+  it('should not match identifiers starting with return', () => {
+    expect(isLikelyFunctionBody('returnValue')).toBe(false);
+  });
+});
diff --git a/tests/utils/remote-value.test.ts b/tests/utils/remote-value.test.ts
new file mode 100644
index 0000000..d972669
--- /dev/null
+++ b/tests/utils/remote-value.test.ts
@@ -0,0 +1,152 @@
+import { describe, it, expect } from 'vitest';
+import { remoteValueToNative } from '../../src/utils/remote-value.js';
+
+describe('remoteValueToNative', () => {
+  describe('primitives', () => {
+    it('converts undefined', () => {
+      expect(remoteValueToNative({ type: 'undefined' })).toBeUndefined();
+    });
+
+    it('converts null', () => {
+      expect(remoteValueToNative({ type: 'null' })).toBeNull();
+    });
+
+    it('converts string', () => {
+      expect(remoteValueToNative({ type: 'string', value: 'hello' })).toBe('hello');
+    });
+
+    it('converts boolean true', () => {
+      expect(remoteValueToNative({ type: 'boolean', value: true })).toBe(true);
+    });
+
+    it('converts boolean false', () => {
+      expect(remoteValueToNative({ type: 'boolean', value: false })).toBe(false);
+    });
+
+    it('converts number', () => {
+      expect(remoteValueToNative({ type: 'number', value: 42 })).toBe(42);
+    });
+
+    it('converts NaN to string', () => {
+      expect(remoteValueToNative({ type: 'number', value: 'NaN' })).toBe('NaN');
+    });
+
+    it('converts Infinity to string', () => {
+      expect(remoteValueToNative({ type: 'number', value: 'Infinity' })).toBe('Infinity');
+    });
+
+    it('converts -Infinity to string', () => {
+      expect(remoteValueToNative({ type: 'number', value: '-Infinity' })).toBe('-Infinity');
+    });
+
+    it('converts -0 to string', () => {
+      expect(remoteValueToNative({ type: 'number', value: '-0' })).toBe('-0');
+    });
+
+    it('converts bigint', () => {
+      expect(remoteValueToNative({ type: 'bigint', value: '9007199254740993' })).toBe(
+        '9007199254740993n'
+      );
+    });
+  });
+
+  describe('collections', () => {
+    it('converts array', () => {
+      expect(
+        remoteValueToNative({
+          type: 'array',
+          value: [
+            { type: 'number', value: 1 },
+            { type: 'string', value: 'two' },
+          ],
+        })
+      ).toEqual([1, 'two']);
+    });
+
+    it('converts nested array', () => {
+      expect(
+        remoteValueToNative({
+          type: 'array',
+          value: [{ type: 'array', value: [{ type: 'number', value: 1 }] }],
+        })
+      ).toEqual([[1]]);
+    });
+
+    it('converts object', () => {
+      expect(
+        remoteValueToNative({
+          type: 'object',
+          value: [
+            ['title', { type: 'string', value: 'Google' }],
+            ['count', { type: 'number', value: 641 }],
+          ],
+        })
+      ).toEqual({ title: 'Google', count: 641 });
+    });
+
+    it('converts nested object', () => {
+      expect(
+        remoteValueToNative({
+          type: 'object',
+          value: [['nested', { type: 'object', value: [['x', { type: 'number', value: 1 }]] }]],
+        })
+      ).toEqual({ nested: { x: 1 } });
+    });
+
+    it('converts map', () => {
+      expect(
+        remoteValueToNative({
+          type: 'map',
+          value: [
+            ['key1', { type: 'string', value: 'value1' }],
+            ['key2', { type: 'number', value: 2 }],
+          ],
+        })
+      ).toEqual({ key1: 'value1', key2: 2 });
+    });
+
+    it('converts set', () => {
+      expect(
+        remoteValueToNative({
+          type: 'set',
+          value: [
+            { type: 'number', value: 1 },
+            { type: 'number', value: 2 },
+          ],
+        })
+      ).toEqual([1, 2]);
+    });
+  });
+
+  describe('special types', () => {
+    it('converts regexp', () => {
+      expect(remoteValueToNative({ type: 'regexp', value: { pattern: 'foo', flags: 'gi' } })).toBe(
+        '/foo/gi'
+      );
+    });
+
+    it('converts regexp without flags', () => {
+      expect(remoteValueToNative({ type: 'regexp', value: { pattern: 'bar' } })).toBe('/bar/');
+    });
+
+    it('converts date', () => {
+      expect(remoteValueToNative({ type: 'date', value: '2024-01-01T00:00:00.000Z' })).toBe(
+        '2024-01-01T00:00:00.000Z'
+      );
+    });
+  });
+
+  describe('fallback', () => {
+    it('returns [type] for node', () => {
+      expect(remoteValueToNative({ type: 'node' })).toBe('[node]');
+    });
+
+    it('returns [type] for function', () => {
+      expect(remoteValueToNative({ type: 'function' })).toBe('[function]');
+    });
+
+    it('returns [type] for promise', () => {
+      expect(remoteValueToNative({ type: 'promise' })).toBe('[promise]');
+    });
+  });
+});

From 9fda9dbbc1a2921e2d27cc3c599b600473992a09 Mon Sep 17 00:00:00 2001
From: Julian Descottes <jdescottes@mozilla.com>
Date: Fri, 12 Jun 2026 11:53:52 +0200
Subject: [PATCH 3/4] Bug 2018268 - refactor: update privileged-context module
 to use WebDriver BiDi for script evaluation

---
 src/tools/privileged-context.ts        | 67 ++++++++++++-----------
 src/tools/script.ts                    | 37 +------------
 src/utils/js-validation.ts             | 38 +++++++++++++
 tests/tools/privileged-context.test.ts | 74 ++++++++++++--------------
 tests/utils/js-validation.test.ts      | 36 ++++++++++++-
 5 files changed, 140 insertions(+), 112 deletions(-)

diff --git a/src/tools/privileged-context.ts b/src/tools/privileged-context.ts
index 30b59c5..644e0de 100644
--- a/src/tools/privileged-context.ts
+++ b/src/tools/privileged-context.ts
@@ -4,7 +4,8 @@
  */
 
 import { successResponse, errorResponse } from '../utils/response-helpers.js';
-import { isLikelyStatement } from '../utils/js-validation.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 = {
@@ -36,16 +37,16 @@ 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'],
   },
 };
 
@@ -124,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);
diff --git a/src/tools/script.ts b/src/tools/script.ts
index 5baf1b2..d60e91f 100644
--- a/src/tools/script.ts
+++ b/src/tools/script.ts
@@ -4,6 +4,7 @@
 
 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 = {
@@ -40,7 +41,6 @@ export const evaluateScriptTool = {
 };
 
 // Constants
-const MAX_FUNCTION_SIZE = 16 * 1024; // 16 KB
 const DEFAULT_TIMEOUT = 5000; // 5 seconds
 const TIMEOUT = Symbol('Timeout');
 
@@ -50,41 +50,6 @@ const EvaluateResultType = {
   Success: 'success',
 };
 
-/**
- * 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 }'
-    );
-  }
-}
-
 export async function handleEvaluateScript(args: unknown): Promise<McpToolResponse> {
   try {
     const {
diff --git a/src/utils/js-validation.ts b/src/utils/js-validation.ts
index a7bb526..0f89c4e 100644
--- a/src/utils/js-validation.ts
+++ b/src/utils/js-validation.ts
@@ -1,3 +1,41 @@
+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 }'
+    );
+  }
+}
+
 /**
  * Detects if the input looks like a JavaScript statement rather than an expression.
  * Statements like const/let/var declarations cannot be used in expression contexts.
diff --git a/tests/tools/privileged-context.test.ts b/tests/tools/privileged-context.test.ts
index e8948e7..b913bf1 100644
--- a/tests/tools/privileged-context.test.ts
+++ b/tests/tools/privileged-context.test.ts
@@ -1,7 +1,3 @@
-/**
- * Tests for statement detection and rejection in privileged context tools
- */
-
 import { describe, it, expect, vi, beforeEach } from 'vitest';
 import {
   evaluatePrivilegedScriptTool,
@@ -21,76 +17,72 @@ describe('Privileged Context Tool Definitions', () => {
       expect(evaluatePrivilegedScriptTool.name).toBe('evaluate_privileged_script');
     });
 
-    it('should mention expression in description', () => {
-      expect(evaluatePrivilegedScriptTool.description).toContain('expression');
+    it('should require function parameter', () => {
+      expect(evaluatePrivilegedScriptTool.inputSchema.required).toContain('function');
     });
   });
 });
 
 describe('Privileged Context Tool Handlers', () => {
-  const mockExecuteScript = vi.fn();
-  const mockSetContext = vi.fn();
-  const mockSwitchToWindow = vi.fn();
-
   beforeEach(() => {
     vi.clearAllMocks();
   });
 
   describe('handleEvaluatePrivilegedScript', () => {
-    it('should reject const statements with error', async () => {
-      const result = await handleEvaluatePrivilegedScript({ expression: 'const x = 1' });
+    it('should return error when function parameter is missing', async () => {
+      const result = await handleEvaluatePrivilegedScript({});
 
       expect(result.isError).toBe(true);
+      expect(result.content[0].text).toContain('function parameter is required');
     });
 
-    it('should mention "statement" in error message', async () => {
-      const result = await handleEvaluatePrivilegedScript({ expression: 'const x = 1' });
-
-      expect(result.content[0]).toHaveProperty('text', expect.stringMatching(/statement/i));
-    });
-
-    it('should suggest IIFE workaround in error message', async () => {
-      const result = await handleEvaluatePrivilegedScript({ expression: 'const x = 1' });
-
-      expect(result.content[0].text).toContain('function()');
-    });
-
-    it('should return error when expression parameter is missing', async () => {
-      const result = await handleEvaluatePrivilegedScript({});
+    it('should reject plain expressions (not function strings)', async () => {
+      const result = await handleEvaluatePrivilegedScript({ function: 'document.title' });
 
       expect(result.isError).toBe(true);
-      expect(result.content[0].text).toContain('expression parameter is required');
+      expect(result.content[0].text).toContain('Invalid function format');
     });
 
-    it('should execute valid expressions successfully', async () => {
+    it('should execute valid function successfully', async () => {
       const mockFirefox = {
-        getDriver: vi.fn().mockReturnValue({
-          switchTo: () => ({ window: mockSwitchToWindow }),
-          setContext: mockSetContext,
-          executeScript: mockExecuteScript.mockResolvedValue('test-result'),
+        getCurrentContextId: vi.fn().mockReturnValue('context-1'),
+        sendBiDiCommand: vi.fn().mockResolvedValue({
+          type: 'success',
+          result: { type: 'string', value: 'test-result' },
         }),
       };
 
       mockGetFirefox.mockResolvedValue(mockFirefox);
 
-      const result = await handleEvaluatePrivilegedScript({ expression: 'document.title' });
+      const result = await handleEvaluatePrivilegedScript({
+        function: '() => Services.prefs.getBoolPref("foo")',
+      });
 
       expect(result.isError).toBeUndefined();
+      expect(result.content[0].text).toContain('chrome context');
       expect(result.content[0].text).toContain('test-result');
     });
 
-    it('should reject let statements', async () => {
-      const result = await handleEvaluatePrivilegedScript({ expression: 'let y = 2' });
+    it('should surface BiDi exception details', async () => {
+      const mockFirefox = {
+        getCurrentContextId: vi.fn().mockReturnValue('context-1'),
+        sendBiDiCommand: vi.fn().mockResolvedValue({
+          type: 'exception',
+          exceptionDetails: {
+            text: 'ReferenceError: Services is not defined',
+            exception: { type: 'object', value: [] },
+          },
+        }),
+      };
 
-      expect(result.isError).toBe(true);
-      expect(result.content[0]).toHaveProperty('text', expect.stringMatching(/statement/i));
-    });
+      mockGetFirefox.mockResolvedValue(mockFirefox);
 
-    it('should reject var statements', async () => {
-      const result = await handleEvaluatePrivilegedScript({ expression: 'var z = 3' });
+      const result = await handleEvaluatePrivilegedScript({
+        function: '() => Services.prefs.getBoolPref("foo")',
+      });
 
       expect(result.isError).toBe(true);
-      expect(result.content[0]).toHaveProperty('text', expect.stringMatching(/statement/i));
+      expect(result.content[0].text).toContain('Services is not defined');
     });
   });
 });
diff --git a/tests/utils/js-validation.test.ts b/tests/utils/js-validation.test.ts
index 6e021d5..d93eabf 100644
--- a/tests/utils/js-validation.test.ts
+++ b/tests/utils/js-validation.test.ts
@@ -1,5 +1,39 @@
 import { describe, it, expect } from 'vitest';
-import { isLikelyFunctionBody, isLikelyStatement } from '../../src/utils/js-validation.js';
+import {
+  isLikelyFunctionBody,
+  isLikelyStatement,
+  validateFunction,
+} from '../../src/utils/js-validation.js';
+
+describe('validateFunction', () => {
+  it('accepts arrow function', () => {
+    expect(() => validateFunction('() => document.title')).not.toThrow();
+  });
+
+  it('accepts async arrow function', () => {
+    expect(() => validateFunction('async () => { return await fetch("/api") }')).not.toThrow();
+  });
+
+  it('accepts function declaration', () => {
+    expect(() => validateFunction('function() { return 1; }')).not.toThrow();
+  });
+
+  it('accepts async function declaration', () => {
+    expect(() => validateFunction('async function() { return 1; }')).not.toThrow();
+  });
+
+  it('rejects plain expression', () => {
+    expect(() => validateFunction('document.title')).toThrow('Invalid function format');
+  });
+
+  it('rejects empty string', () => {
+    expect(() => validateFunction('')).toThrow('function parameter is required');
+  });
+
+  it('rejects oversized function', () => {
+    expect(() => validateFunction('() => ' + 'x'.repeat(16 * 1024))).toThrow('Function too large');
+  });
+});
 
 describe('isLikelyStatement', () => {
   it('should detect const declarations', () => {

From 5e8d8f5113e3dcc92fd679d4e12691c164ea5ac4 Mon Sep 17 00:00:00 2001
From: Julian Descottes <jdescottes@mozilla.com>
Date: Fri, 12 Jun 2026 11:55:16 +0200
Subject: [PATCH 4/4] Bug 2018268 - refactor: remove now unused
 isLikelyFunctionBody and isLikelyStatement helpers

---
 src/utils/js-validation.ts        | 19 ---------
 tests/utils/js-validation.test.ts | 66 +------------------------------
 2 files changed, 1 insertion(+), 84 deletions(-)

diff --git a/src/utils/js-validation.ts b/src/utils/js-validation.ts
index 0f89c4e..1efcf78 100644
--- a/src/utils/js-validation.ts
+++ b/src/utils/js-validation.ts
@@ -35,22 +35,3 @@ export function validateFunction(fnString: string): void {
     );
   }
 }
-
-/**
- * Detects if the input looks like a JavaScript statement rather than an expression.
- * Statements like const/let/var declarations cannot be used in expression contexts.
- */
-export function isLikelyStatement(input: string): boolean {
-  const trimmed = input.trim();
-  return /^(const|let|var)\s/.test(trimmed);
-}
-
-/**
- * Detects if the input looks like a function body (starts with "return").
- * Function bodies are valid in WebDriver Classic executeScript but not in BiDi script.evaluate,
- * which expects an expression.
- */
-export function isLikelyFunctionBody(input: string): boolean {
-  const trimmed = input.trim();
-  return /^return(\s|;|$)/.test(trimmed);
-}
diff --git a/tests/utils/js-validation.test.ts b/tests/utils/js-validation.test.ts
index d93eabf..ec8efe1 100644
--- a/tests/utils/js-validation.test.ts
+++ b/tests/utils/js-validation.test.ts
@@ -1,9 +1,5 @@
 import { describe, it, expect } from 'vitest';
-import {
-  isLikelyFunctionBody,
-  isLikelyStatement,
-  validateFunction,
-} from '../../src/utils/js-validation.js';
+import { validateFunction } from '../../src/utils/js-validation.js';
 
 describe('validateFunction', () => {
   it('accepts arrow function', () => {
@@ -34,63 +30,3 @@ describe('validateFunction', () => {
     expect(() => validateFunction('() => ' + 'x'.repeat(16 * 1024))).toThrow('Function too large');
   });
 });
-
-describe('isLikelyStatement', () => {
-  it('should detect const declarations', () => {
-    expect(isLikelyStatement('const x = 1')).toBe(true);
-  });
-
-  it('should detect let declarations', () => {
-    expect(isLikelyStatement('let x = 1')).toBe(true);
-  });
-
-  it('should detect var declarations', () => {
-    expect(isLikelyStatement('var x = 1')).toBe(true);
-  });
-
-  it('should allow function calls', () => {
-    expect(isLikelyStatement('Services.prefs.getBoolPref("foo")')).toBe(false);
-  });
-
-  it('should allow simple expressions', () => {
-    expect(isLikelyStatement('1 + 2')).toBe(false);
-  });
-
-  it('should allow property access', () => {
-    expect(isLikelyStatement('document.title')).toBe(false);
-  });
-
-  it('should handle leading whitespace', () => {
-    expect(isLikelyStatement('  const x = 1')).toBe(true);
-  });
-});
-
-describe('isLikelyFunctionBody', () => {
-  it('should detect return with a value', () => {
-    expect(isLikelyFunctionBody('return document.title')).toBe(true);
-  });
-
-  it('should detect return followed by semicolon', () => {
-    expect(isLikelyFunctionBody('return;')).toBe(true);
-  });
-
-  it('should detect bare return', () => {
-    expect(isLikelyFunctionBody('return')).toBe(true);
-  });
-
-  it('should handle leading whitespace', () => {
-    expect(isLikelyFunctionBody('  return x')).toBe(true);
-  });
-
-  it('should allow expressions', () => {
-    expect(isLikelyFunctionBody('document.title')).toBe(false);
-  });
-
-  it('should allow IIFEs', () => {
-    expect(isLikelyFunctionBody('(function() { return 1; })()')).toBe(false);
-  });
-
-  it('should not match identifiers starting with return', () => {
-    expect(isLikelyFunctionBody('returnValue')).toBe(false);
-  });
-});