feat: support client roots feature

Author: OrKoN

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,21 +19,22 @@ import {
   type ListenerMap,
   type UncaughtError,
 } from './PageCollector.js';
-import type {
-  Browser,
-  BrowserContext,
-  ConsoleMessage,
-  Debugger,
-  HTTPRequest,
-  Page,
-  ScreenRecorder,
-  Viewport,
-  Target,
-  Extension,
+import {
+  Locator,
+  PredefinedNetworkConditions,
+  type Browser,
+  type BrowserContext,
+  type ConsoleMessage,
+  type Debugger,
+  type HTTPRequest,
+  type Page,
+  type ScreenRecorder,
+  type Viewport,
+  type Target,
+  type Extension,
+  type Root,
+  type DevTools,
 } 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 {listPages} from './tools/pages.js';
 import {CLOSE_PAGE_ERROR} from './tools/ToolDefinition.js';
 import type {Context, SupportedExtensions} from './tools/ToolDefinition.js';
@@ -42,7 +44,7 @@ import type {
   GeolocationOptions,
   ExtensionServiceWorker,
 } from './types.js';
-import {ensureExtension, saveTemporaryFile} from './utils/files.js';
+import {ensureExtension, getTempFilePath} from './utils/files.js';
 import {getNetworkMultiplierFromString} from './WaitForHelper.js';
 
 interface McpContextOptions {
@@ -90,6 +92,7 @@ export class McpContext implements Context {
   #locatorClass: typeof Locator;
   #options: McpContextOptions;
   #heapSnapshotManager = new HeapSnapshotManager();
+  #roots: Root[] | undefined = undefined;
 
   private constructor(
     browser: Browser,
@@ -154,6 +157,34 @@ export class McpContext implements Context {
     return context;
   }
 
+  roots(): Root[] | undefined {
+    return this.#roots;
+  }
+
+  setRoots(roots: Root[] | undefined): void {
+    this.#roots = roots;
+  }
+
+  validatePath(filePath: string): void {
+    const roots = this.roots();
+    if (roots === undefined) {
+      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,18 @@ export class McpContext implements Context {
     data: Uint8Array<ArrayBufferLike>,
     filename: string,
   ): Promise<{filepath: string}> {
-    return await saveTemporaryFile(data, filename);
+    const filepath = await getTempFilePath(filename);
+    this.validatePath(filepath);
+    await fs.writeFile(filepath, data);
+    return {filepath};
   }
+
   async saveFile(
     data: Uint8Array<ArrayBufferLike>,
     clientProvidedFilePath: string,
     extension: SupportedExtensions,
   ): Promise<{filename: string}> {
+    this.validatePath(clientProvidedFilePath);
     try {
       const filePath = ensureExtension(
         path.resolve(clientProvidedFilePath),
@@ -721,6 +757,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 +788,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/daemon/client.ts

@@ -11,7 +11,7 @@ import net from 'node:net';
 import {logger} from '../logger.js';
 import type {CallToolResult} from '../third_party/index.js';
 import {PipeTransport} from '../third_party/index.js';
-import {saveTemporaryFile} from '../utils/files.js';
+import {getTempFilePath} from '../utils/files.js';
 
 import type {DaemonMessage, DaemonResponse} from './types.js';
 import {
@@ -179,7 +179,8 @@ export async function handleResponse(
       }
       const data = Buffer.from(imageData, 'base64');
       const name = crypto.randomUUID();
-      const {filepath} = await saveTemporaryFile(data, `${name}${extension}`);
+      const filepath = await getTempFilePath(`${name}${extension}`);
+      fs.writeFileSync(filepath, data);
       chunks.push(`Saved to ${filepath}.`);
     } else {
       throw new Error('Not supported response content type');

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';
@@ -57,11 +59,35 @@ export async function createMcpServer(
     return {};
   });
 
+  const updateRoots = async () => {
+    if (!server.server.getClientCapabilities()?.roots) {
+      return;
+    }
+    try {
+      const roots = await server.server.request(
+        {method: 'roots/list'},
+        ListRootsResultSchema,
+      );
+      context?.setRoots(roots.roots);
+    } catch (e) {
+      logger('Failed to list roots', e);
+    }
+  };
+
   server.server.oninitialized = () => {
     const clientName = server.server.getClientVersion()?.name;
     if (clientName) {
       clearcutLogger?.setClientName(clientName);
     }
+    if (server.server.getClientCapabilities()?.roots) {
+      void updateRoots();
+      server.server.setNotificationHandler(
+        RootsListChangedNotificationSchema,
+        () => {
+          void updateRoots();
+        },
+      );
+    }
   };
 
   let context: McpContext;
@@ -109,6 +135,7 @@ export async function createMcpServer(
         experimentalIncludeAllPages: serverArgs.experimentalIncludeAllPages,
         performanceCrux: serverArgs.performanceCrux,
       });
+      await updateRoots();
     }
     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';

src/tools/ToolDefinition.ts

@@ -167,9 +167,10 @@ export type SupportedExtensions =
   | '.json.gz';
 
 /**
- * Only add methods required by tools/*.
+ * Only add methods used by tools/*.
  */
 export type Context = Readonly<{
+  validatePath(filePath: string): void;
   isRunningPerformanceTrace(): boolean;
   setIsRunningPerformanceTrace(x: boolean): void;
   isCruxEnabled(): boolean;
@@ -244,6 +245,9 @@ export type Context = Readonly<{
   ): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.ItemsRange>;
 }>;
 
+/**
+ * Only add methods used by tools/*.
+ */
 export type ContextPage = Readonly<{
   readonly pptrPage: Page;
   getAXNodeByUid(uid: string): TextSnapshotNode | undefined;

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