fix(privileged-context): Reject statements with helpful error message

Models sometimes pass JavaScript statements (const/let/var declarations)
to evaluate_chrome_script, which wraps input as return(expression),
which causes tool aborts, resulting in slowdown and wasted tokens.

Adds isLikelyStatement() validation that detects statement keywords
and returns a clear error message suggesting the IIFE workaround.
Also updates tool description to warn against statement usage.

Author: dmose

src/tools/privileged-context.ts

@@ -35,7 +35,7 @@ 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.',
+    '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; })()',
   inputSchema: {
     type: 'object',
     properties: {
@@ -48,6 +48,15 @@ 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';
@@ -131,6 +140,16 @@ export async function handleEvaluatePrivilegedScript(args: unknown): Promise<Mcp
       throw new Error('expression parameter is required and must be a 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; })()'
+        )
+      );
+    }
+
     const { getFirefox } = await import('../index.js');
     const firefox = await getFirefox();
 

tests/tools/privileged-context-state.test.ts

@@ -0,0 +1,89 @@
+/**
+ * Tests for privileged context state consistency
+ */
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+
+describe('Privileged context state consistency', () => {
+  const mockSetContext = vi.fn();
+  const mockSwitchToWindow = vi.fn();
+  const mockExecuteScript = vi.fn();
+  const mockExecuteAsyncScript = vi.fn();
+
+  // Track currentContextId state as the real code would
+  let mockCurrentContextId: string | null;
+  const mockSetCurrentContextId = vi.fn((id: string) => {
+    mockCurrentContextId = id;
+  });
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.resetModules();
+    // Start with a content context (user has a normal tab open)
+    mockCurrentContextId = 'original-content-context';
+
+    vi.doMock('../../src/index.js', () => ({
+      getFirefox: vi.fn().mockResolvedValue({
+        getDriver: () => ({
+          switchTo: () => ({
+            window: mockSwitchToWindow,
+          }),
+          setContext: mockSetContext,
+          executeScript: mockExecuteScript,
+          executeAsyncScript: mockExecuteAsyncScript,
+        }),
+        getCurrentContextId: () => mockCurrentContextId,
+        setCurrentContextId: mockSetCurrentContextId,
+        sendBiDiCommand: vi.fn().mockResolvedValue({
+          contexts: [{ context: 'chrome-context-id', url: 'chrome://browser/content/browser.xhtml' }],
+        }),
+      }),
+    }));
+  });
+
+  it('select_privileged_context should update currentContextId (BUG: it does not)', async () => {
+    const { handleSelectPrivilegedContext } = await import(
+      '../../src/tools/privileged-context.js'
+    );
+
+    await handleSelectPrivilegedContext({ contextId: 'chrome-context-id' });
+
+    expect(mockSwitchToWindow).toHaveBeenCalledWith('chrome-context-id');
+    expect(mockSetContext).toHaveBeenCalledWith('chrome');
+
+    // BUG: select_privileged_context does NOT call setCurrentContextId
+    // so currentContextId stays as 'original-content-context'
+    expect(mockSetCurrentContextId).toHaveBeenCalledWith('chrome-context-id');
+  });
+
+  it('set_firefox_prefs after select_privileged_context should not revert to old context (BUG: it does)', async () => {
+    const { handleSelectPrivilegedContext } = await import(
+      '../../src/tools/privileged-context.js'
+    );
+    const { handleSetFirefoxPrefs } = await import(
+      '../../src/tools/firefox-prefs.js'
+    );
+
+    // User selects privileged context
+    await handleSelectPrivilegedContext({ contextId: 'chrome-context-id' });
+
+    // BUG: currentContextId is still 'original-content-context' because
+    // select_privileged_context never called setCurrentContextId.
+    // So when set_firefox_prefs saves originalContextId, it gets the wrong value.
+    mockExecuteScript.mockResolvedValue(undefined);
+    mockSwitchToWindow.mockClear();
+    mockSetContext.mockClear();
+
+    await handleSetFirefoxPrefs({ prefs: { 'browser.ml.enable': true } });
+
+    // The finally block in set_firefox_prefs restores to originalContextId.
+    // Because currentContextId was never updated, it restores to
+    // 'original-content-context' — silently undoing the privileged context selection.
+    const setContextCalls = mockSetContext.mock.calls;
+    const lastSetContext = setContextCalls[setContextCalls.length - 1];
+
+    // BUG: last setContext call is 'content' — it reverted the privileged context
+    // After fix: the tool should detect we're already in a privileged context
+    // and restore back to it (or at minimum, currentContextId should be correct)
+    expect(lastSetContext[0]).not.toBe('content');
+  });
+});