feat: support client roots feature

OrKoN · OrKoN · commit 51e3291dde15 · 2026-04-23T13:30:49.000+02:00
diff --git a/src/McpContext.ts b/src/McpContext.ts
@@ -6,6 +6,7 @@
 
 import fs from 'node:fs/promises';
 import path from 'node:path';
+import {fileURLToPath} from 'node:url';
 
 import type {TargetUniverse} from './DevtoolsUtils.js';
 import {UniverseManager} from './DevtoolsUtils.js';
@@ -18,6 +19,8 @@ import {
   type ListenerMap,
   type UncaughtError,
 } from './PageCollector.js';
+import {Locator} from './third_party/index.js';
+import {PredefinedNetworkConditions} from './third_party/index.js';
 import type {
   Browser,
   BrowserContext,
@@ -31,8 +34,7 @@ import type {
   Extension,
 } from './third_party/index.js';
 import type {DevTools} from './third_party/index.js';
-import {Locator} from './third_party/index.js';
-import {PredefinedNetworkConditions} from './third_party/index.js';
+import type {Root} from './third_party/index.js';
 import {listPages} from './tools/pages.js';
 import {CLOSE_PAGE_ERROR} from './tools/ToolDefinition.js';
 import type {Context, SupportedExtensions} from './tools/ToolDefinition.js';
@@ -90,6 +92,7 @@ export class McpContext implements Context {
   #locatorClass: typeof Locator;
   #options: McpContextOptions;
   #heapSnapshotManager = new HeapSnapshotManager();
+  #roots: Root[] = [];
 
   private constructor(
     browser: Browser,
@@ -154,6 +157,34 @@ export class McpContext implements Context {
     return context;
   }
 
+  roots(): Root[] {
+    return this.#roots;
+  }
+
+  setRoots(roots: Root[]): void {
+    this.#roots = roots;
+  }
+
+  validatePath(filePath: string): void {
+    const roots = this.roots();
+    if (roots.length === 0) {
+      return;
+    }
+    const absolutePath = path.resolve(filePath);
+    for (const root of roots) {
+      const rootPath = path.resolve(fileURLToPath(root.uri));
+      if (
+        absolutePath === rootPath ||
+        absolutePath.startsWith(rootPath + path.sep)
+      ) {
+        return;
+      }
+    }
+    throw new Error(
+      `Access denied: path ${filePath} is not within any of the workspace roots ${JSON.stringify(roots)}.`,
+    );
+  }
+
   resolveCdpRequestId(page: McpPage, cdpRequestId: string): number | undefined {
     if (!cdpRequestId) {
       this.logger('no network request');
@@ -643,13 +674,21 @@ export class McpContext implements Context {
     data: Uint8Array<ArrayBufferLike>,
     filename: string,
   ): Promise<{filepath: string}> {
+    if (
+      filename.includes('/') ||
+      filename.includes('\\') ||
+      filename.includes('..')
+    ) {
+      throw new Error(`Invalid filename: ${filename}`);
+    }
     return await saveTemporaryFile(data, filename);
   }
   async saveFile(
     data: Uint8Array<ArrayBufferLike>,
     clientProvidedFilePath: string,
     extension: SupportedExtensions,
   ): Promise<{filename: string}> {
+    this.validatePath(clientProvidedFilePath);
     try {
       const filePath = ensureExtension(
         path.resolve(clientProvidedFilePath),
@@ -721,6 +760,7 @@ export class McpContext implements Context {
   }
 
   async installExtension(extensionPath: string): Promise<string> {
+    this.validatePath(extensionPath);
     const id = await this.browser.installExtension(extensionPath);
     return id;
   }
@@ -751,25 +791,29 @@ export class McpContext implements Context {
   async getHeapSnapshotAggregates(
     filePath: string,
   ): Promise<Record<string, AggregatedInfoWithUid>> {
+    this.validatePath(filePath);
     return await this.#heapSnapshotManager.getAggregates(filePath);
   }
 
   async getHeapSnapshotStats(
     filePath: string,
   ): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.Statistics> {
+    this.validatePath(filePath);
     return await this.#heapSnapshotManager.getStats(filePath);
   }
 
   async getHeapSnapshotStaticData(
     filePath: string,
   ): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.StaticData | null> {
+    this.validatePath(filePath);
     return await this.#heapSnapshotManager.getStaticData(filePath);
   }
 
   async getHeapSnapshotNodesByUid(
     filePath: string,
     uid: number,
   ): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.ItemsRange> {
+    this.validatePath(filePath);
     return await this.#heapSnapshotManager.getNodesByUid(filePath, uid);
   }
 }
diff --git a/src/index.ts b/src/index.ts
@@ -21,6 +21,8 @@ import {
   McpServer,
   type CallToolResult,
   SetLevelRequestSchema,
+  ListRootsResultSchema,
+  RootsListChangedNotificationSchema,
 } from './third_party/index.js';
 import {ToolCategory} from './tools/categories.js';
 import type {DefinedPageTool, ToolDefinition} from './tools/ToolDefinition.js';
@@ -62,6 +64,24 @@ export async function createMcpServer(
     if (clientName) {
       clearcutLogger?.setClientName(clientName);
     }
+    const updateRoots = async () => {
+      try {
+        const roots = await server.server.request(
+          {method: 'roots/list'},
+          ListRootsResultSchema,
+        );
+        context.setRoots(roots.roots);
+      } catch (e) {
+        logger('Failed to list roots', e);
+      }
+    };
+    void updateRoots();
+    server.server.setNotificationHandler(
+      RootsListChangedNotificationSchema,
+      () => {
+        void updateRoots();
+      },
+    );
   };
 
   let context: McpContext;
@@ -104,11 +124,13 @@ export async function createMcpServer(
           });
 
     if (context?.browser !== browser) {
+      const roots = context?.roots() ?? [];
       context = await McpContext.from(browser, logger, {
         experimentalDevToolsDebugging: devtools,
         experimentalIncludeAllPages: serverArgs.experimentalIncludeAllPages,
         performanceCrux: serverArgs.performanceCrux,
       });
+      context.setRoots(roots);
     }
     return context;
   }
diff --git a/src/third_party/index.ts b/src/third_party/index.ts
@@ -30,6 +30,10 @@ export {
   SetLevelRequestSchema,
   type ImageContent,
   type TextContent,
+  type Root,
+  ListRootsRequestSchema,
+  RootsListChangedNotificationSchema,
+  ListRootsResultSchema,
 } from '@modelcontextprotocol/sdk/types.js';
 export {z as zod} from 'zod';
 export {default as ajv} from 'ajv';
diff --git a/src/tools/ToolDefinition.ts b/src/tools/ToolDefinition.ts
@@ -16,6 +16,7 @@ import type {
   ScreenRecorder,
   Viewport,
   DevTools,
+  Root,
 } from '../third_party/index.js';
 import type {InsightName, TraceResult} from '../trace-processing/parse.js';
 import type {
@@ -170,6 +171,9 @@ export type SupportedExtensions =
  * Only add methods required by tools/*.
  */
 export type Context = Readonly<{
+  roots(): Root[];
+  setRoots(roots: Root[]): void;
+  validatePath(filePath: string): void;
   isRunningPerformanceTrace(): boolean;
   setIsRunningPerformanceTrace(x: boolean): void;
   isCruxEnabled(): boolean;
diff --git a/src/tools/input.ts b/src/tools/input.ts
@@ -365,8 +365,9 @@ export const uploadFile = definePageTool({
     filePath: zod.string().describe('The local path of the file to upload'),
     includeSnapshot: includeSnapshotSchema,
   },
-  handler: async (request, response) => {
+  handler: async (request, response, context) => {
     const {uid, filePath} = request.params;
+    context.validatePath(filePath);
     const handle = (await request.page.getElementByUid(
       uid,
     )) as ElementHandle<HTMLInputElement>;
diff --git a/src/tools/lighthouse.ts b/src/tools/lighthouse.ts
@@ -53,6 +53,10 @@ export const lighthouseAudit = definePageTool({
       outputDirPath,
     } = request.params;
 
+    if (outputDirPath) {
+      context.validatePath(outputDirPath);
+    }
+
     const flags: Flags = {
       onlyCategories: categories,
       output: formats,
diff --git a/src/tools/memory.ts b/src/tools/memory.ts
@@ -22,8 +22,9 @@ export const takeMemorySnapshot = definePageTool({
       .string()
       .describe('A path to a .heapsnapshot file to save the heapsnapshot to.'),
   },
-  handler: async (request, response, _context) => {
+  handler: async (request, response, context) => {
     const page = request.page;
+    context.validatePath(request.params.filePath);
 
     await page.pptrPage.captureHeapSnapshot({
       path: ensureExtension(request.params.filePath, '.heapsnapshot'),
@@ -48,6 +49,7 @@ export const exploreMemorySnapshot = defineTool({
     filePath: zod.string().describe('A path to a .heapsnapshot file to read.'),
   },
   handler: async (request, response, context) => {
+    context.validatePath(request.params.filePath);
     const stats = await context.getHeapSnapshotStats(request.params.filePath);
     const staticData = await context.getHeapSnapshotStaticData(
       request.params.filePath,
@@ -78,6 +80,7 @@ export const getMemorySnapshotDetails = defineTool({
       .describe('The page size for pagination of aggregates.'),
   },
   handler: async (request, response, context) => {
+    context.validatePath(request.params.filePath);
     const aggregates = await context.getHeapSnapshotAggregates(
       request.params.filePath,
     );
@@ -109,6 +112,7 @@ export const getNodesByClass = defineTool({
     pageSize: zod.number().optional().describe('The page size for pagination.'),
   },
   handler: async (request, response, context) => {
+    context.validatePath(request.params.filePath);
     const nodes = await context.getHeapSnapshotNodesByUid(
       request.params.filePath,
       request.params.uid,
diff --git a/src/tools/network.ts b/src/tools/network.ts
@@ -114,6 +114,12 @@ export const getNetworkRequest = definePageTool({
       ),
   },
   handler: async (request, response, context) => {
+    if (request.params.requestFilePath) {
+      context.validatePath(request.params.requestFilePath);
+    }
+    if (request.params.responseFilePath) {
+      context.validatePath(request.params.responseFilePath);
+    }
     if (request.params.reqid) {
       response.attachNetworkRequest(request.params.reqid, {
         requestFilePath: request.params.requestFilePath,
diff --git a/src/tools/performance.ts b/src/tools/performance.ts
@@ -49,6 +49,9 @@ export const startTrace = definePageTool({
     filePath: filePathSchema,
   },
   handler: async (request, response, context) => {
+    if (request.params.filePath) {
+      context.validatePath(request.params.filePath);
+    }
     if (context.isRunningPerformanceTrace()) {
       response.appendResponseLine(
         'Error: a performance trace is already running. Use performance_stop_trace to stop it. Only one trace can be running at any given time.',
@@ -126,6 +129,9 @@ export const stopTrace = definePageTool({
     filePath: filePathSchema,
   },
   handler: async (request, response, context) => {
+    if (request.params.filePath) {
+      context.validatePath(request.params.filePath);
+    }
     if (!context.isRunningPerformanceTrace()) {
       return;
     }
diff --git a/src/tools/screencast.ts b/src/tools/screencast.ts
@@ -40,6 +40,9 @@ export const startScreencast = definePageTool(args => ({
       ),
   },
   handler: async (request, response, context) => {
+    if (request.params.filePath) {
+      context.validatePath(request.params.filePath);
+    }
     if (context.getScreenRecorder() !== null) {
       response.appendResponseLine(
         'Error: a screencast recording is already in progress. Use screencast_stop to stop it before starting a new one.',
diff --git a/src/tools/screenshot.ts b/src/tools/screenshot.ts
@@ -51,6 +51,9 @@ export const screenshot = definePageTool({
       ),
   },
   handler: async (request, response, context) => {
+    if (request.params.filePath) {
+      context.validatePath(request.params.filePath);
+    }
     if (request.params.uid && request.params.fullPage) {
       throw new Error('Providing both "uid" and "fullPage" is not allowed.');
     }
diff --git a/src/tools/snapshot.ts b/src/tools/snapshot.ts
@@ -33,7 +33,10 @@ in the DevTools Elements panel (if any).`,
         'The absolute path, or a path relative to the current working directory, to save the snapshot to instead of attaching it to the response.',
       ),
   },
-  handler: async (request, response) => {
+  handler: async (request, response, context) => {
+    if (request.params.filePath) {
+      context.validatePath(request.params.filePath);
+    }
     response.includeSnapshot({
       verbose: request.params.verbose ?? false,
       filePath: request.params.filePath,
diff --git a/tests/McpContext.test.ts b/tests/McpContext.test.ts
@@ -5,7 +5,9 @@
  */
 
 import assert from 'node:assert';
+import path from 'node:path';
 import {afterEach, describe, it} from 'node:test';
+import {pathToFileURL} from 'node:url';
 
 import sinon from 'sinon';
 
@@ -213,4 +215,36 @@ describe('McpContext', () => {
       fromStub.restore();
     });
   });
+
+  it('can store and retrieve roots', async () => {
+    await withMcpContext(async (_response, context) => {
+      const roots = [{uri: 'file:///test', name: 'test'}];
+      context.setRoots(roots);
+      assert.deepEqual(context.roots(), roots);
+    });
+  });
+
+  it('validatePath allows paths within roots', async () => {
+    await withMcpContext(async (_response, context) => {
+      const workspacePath = path.resolve('/tmp/workspace');
+      const roots = [
+        {uri: pathToFileURL(workspacePath).href, name: 'workspace'},
+      ];
+      context.setRoots(roots);
+      // Valid path within root
+      context.validatePath(path.join(workspacePath, 'test.txt'));
+      context.validatePath(workspacePath);
+
+      // Invalid path outside root
+      const outsidePath = path.resolve('/tmp/outside.txt');
+      assert.throws(() => context.validatePath(outsidePath), /Access denied/);
+    });
+  });
+
+  it('validatePath allows all paths if no roots are set', async () => {
+    await withMcpContext(async (_response, context) => {
+      context.setRoots([]);
+      context.validatePath(path.resolve('/tmp/anywhere.txt'));
+    });
+  });
 });
diff --git a/tests/index.test.ts b/tests/index.test.ts
