feat: support client roots feature

Author: OrKoN

src/McpContext.ts

@@ -18,6 +18,11 @@ import {
   type ListenerMap,
   type UncaughtError,
 } from './PageCollector.js';
+import {Locator} from './third_party/index.js';
+import {
+  PredefinedNetworkConditions,
+  fileURLToPath,
+} from './third_party/index.js';
 import type {
   Browser,
   BrowserContext,
@@ -31,8 +36,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 +94,7 @@ export class McpContext implements Context {
   #locatorClass: typeof Locator;
   #options: McpContextOptions;
   #heapSnapshotManager = new HeapSnapshotManager();
+  #roots: Root[] = [];
 
   private constructor(
     browser: Browser,
@@ -154,6 +159,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.`,
+    );
+  }
+
   resolveCdpRequestId(page: McpPage, cdpRequestId: string): number | undefined {
     if (!cdpRequestId) {
       this.logger('no network request');
@@ -643,13 +676,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 +762,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 +793,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);
   }
 }

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;
   }

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';
@@ -71,4 +75,5 @@ export const generateReport = generateReportImpl as (
   format: string,
 ) => string;
 
+export {fileURLToPath} from 'node:url';
 export * as DevTools from '../../node_modules/chrome-devtools-frontend/mcp/mcp.js';

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;

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>;

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,

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,

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,

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;
     }

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.',