chore: slice profile dirs by root in vscode (#814)

2025-10-12 00:25:14 +03:00 · 2025-08-01 16:59:59 -07:00
parent ffe0117456
commit a60d7b8cd1
16 changed files with 220 additions and 108 deletions
--- a/src/browserContextFactory.ts
+++ b/src/browserContextFactory.ts
@@ -14,41 +14,44 @@
 * limitations under the License.
 */

-import fs from 'node:fs';
-import net from 'node:net';
-import path from 'node:path';
-import os from 'node:os';
+import fs from 'fs';
+import net from 'net';
+import path from 'path';

 import * as playwright from 'playwright';
-
+// @ts-ignore
+import { registryDirectory } from 'playwright-core/lib/server/registry/index';
 import { logUnhandledError, testDebug } from './log.js';
+import { createHash } from './utils.js';
+import { outputFile  } from './config.js';

 import type { FullConfig } from './config.js';

-export function contextFactory(browserConfig: FullConfig['browser']): BrowserContextFactory {
-  if (browserConfig.remoteEndpoint)
-    return new RemoteContextFactory(browserConfig);
-  if (browserConfig.cdpEndpoint)
-    return new CdpContextFactory(browserConfig);
-  if (browserConfig.isolated)
-    return new IsolatedContextFactory(browserConfig);
-  return new PersistentContextFactory(browserConfig);
+export function contextFactory(config: FullConfig): BrowserContextFactory {
+  if (config.browser.remoteEndpoint)
+    return new RemoteContextFactory(config);
+  if (config.browser.cdpEndpoint)
+    return new CdpContextFactory(config);
+  if (config.browser.isolated)
+    return new IsolatedContextFactory(config);
+  return new PersistentContextFactory(config);
 }

-export type ClientInfo = { name: string, version: string };
+export type ClientInfo = { name?: string, version?: string, rootPath?: string };

 export interface BrowserContextFactory {
  createContext(clientInfo: ClientInfo, abortSignal: AbortSignal): Promise<{ browserContext: playwright.BrowserContext, close: () => Promise<void> }>;
 }

 class BaseContextFactory implements BrowserContextFactory {
-  readonly browserConfig: FullConfig['browser'];
+  readonly config: FullConfig;
  protected _browserPromise: Promise<playwright.Browser> | undefined;
+  protected _tracesDir: string | undefined;
  readonly name: string;

-  constructor(name: string, browserConfig: FullConfig['browser']) {
+  constructor(name: string, config: FullConfig) {
    this.name = name;
-    this.browserConfig = browserConfig;
+    this.config = config;
  }

  protected async _obtainBrowser(): Promise<playwright.Browser> {
@@ -70,7 +73,10 @@ class BaseContextFactory implements BrowserContextFactory {
    throw new Error('Not implemented');
  }

-  async createContext(): Promise<{ browserContext: playwright.BrowserContext, close: () => Promise<void> }> {
+  async createContext(clientInfo: ClientInfo): Promise<{ browserContext: playwright.BrowserContext, close: () => Promise<void> }> {
+    if (this.config.saveTrace)
+      this._tracesDir = await outputFile(this.config, clientInfo.rootPath, `traces-${Date.now()}`);
+
    testDebug(`create browser context (${this.name})`);
    const browser = await this._obtainBrowser();
    const browserContext = await this._doCreateContext(browser);
@@ -94,15 +100,16 @@ class BaseContextFactory implements BrowserContextFactory {
 }

 class IsolatedContextFactory extends BaseContextFactory {
-  constructor(browserConfig: FullConfig['browser']) {
-    super('isolated', browserConfig);
+  constructor(config: FullConfig) {
+    super('isolated', config);
  }

  protected override async _doObtainBrowser(): Promise<playwright.Browser> {
-    await injectCdpPort(this.browserConfig);
-    const browserType = playwright[this.browserConfig.browserName];
+    await injectCdpPort(this.config.browser);
+    const browserType = playwright[this.config.browser.browserName];
    return browserType.launch({
-      ...this.browserConfig.launchOptions,
+      tracesDir: this._tracesDir,
+      ...this.config.browser.launchOptions,
      handleSIGINT: false,
      handleSIGTERM: false,
    }).catch(error => {
@@ -113,35 +120,35 @@ class IsolatedContextFactory extends BaseContextFactory {
  }

  protected override async _doCreateContext(browser: playwright.Browser): Promise<playwright.BrowserContext> {
-    return browser.newContext(this.browserConfig.contextOptions);
+    return browser.newContext(this.config.browser.contextOptions);
  }
 }

 class CdpContextFactory extends BaseContextFactory {
-  constructor(browserConfig: FullConfig['browser']) {
-    super('cdp', browserConfig);
+  constructor(config: FullConfig) {
+    super('cdp', config);
  }

  protected override async _doObtainBrowser(): Promise<playwright.Browser> {
-    return playwright.chromium.connectOverCDP(this.browserConfig.cdpEndpoint!);
+    return playwright.chromium.connectOverCDP(this.config.browser.cdpEndpoint!);
  }

  protected override async _doCreateContext(browser: playwright.Browser): Promise<playwright.BrowserContext> {
-    return this.browserConfig.isolated ? await browser.newContext() : browser.contexts()[0];
+    return this.config.browser.isolated ? await browser.newContext() : browser.contexts()[0];
  }
 }

 class RemoteContextFactory extends BaseContextFactory {
-  constructor(browserConfig: FullConfig['browser']) {
-    super('remote', browserConfig);
+  constructor(config: FullConfig) {
+    super('remote', config);
  }

  protected override async _doObtainBrowser(): Promise<playwright.Browser> {
-    const url = new URL(this.browserConfig.remoteEndpoint!);
-    url.searchParams.set('browser', this.browserConfig.browserName);
-    if (this.browserConfig.launchOptions)
-      url.searchParams.set('launch-options', JSON.stringify(this.browserConfig.launchOptions));
-    return playwright[this.browserConfig.browserName].connect(String(url));
+    const url = new URL(this.config.browser.remoteEndpoint!);
+    url.searchParams.set('browser', this.config.browser.browserName);
+    if (this.config.browser.launchOptions)
+      url.searchParams.set('launch-options', JSON.stringify(this.config.browser.launchOptions));
+    return playwright[this.config.browser.browserName].connect(String(url));
  }

  protected override async _doCreateContext(browser: playwright.Browser): Promise<playwright.BrowserContext> {
@@ -150,27 +157,31 @@ class RemoteContextFactory extends BaseContextFactory {
 }

 class PersistentContextFactory implements BrowserContextFactory {
-  readonly browserConfig: FullConfig['browser'];
+  readonly config: FullConfig;
  private _userDataDirs = new Set<string>();

-  constructor(browserConfig: FullConfig['browser']) {
-    this.browserConfig = browserConfig;
+  constructor(config: FullConfig) {
+    this.config = config;
  }

-  async createContext(): Promise<{ browserContext: playwright.BrowserContext, close: () => Promise<void> }> {
-    await injectCdpPort(this.browserConfig);
+  async createContext(clientInfo: ClientInfo): Promise<{ browserContext: playwright.BrowserContext, close: () => Promise<void> }> {
+    await injectCdpPort(this.config.browser);
    testDebug('create browser context (persistent)');
-    const userDataDir = this.browserConfig.userDataDir ?? await this._createUserDataDir();
+    const userDataDir = this.config.browser.userDataDir ?? await this._createUserDataDir(clientInfo.rootPath);
+    let tracesDir: string | undefined;
+    if (this.config.saveTrace)
+      tracesDir = await outputFile(this.config, clientInfo.rootPath, `traces-${Date.now()}`);

    this._userDataDirs.add(userDataDir);
    testDebug('lock user data dir', userDataDir);

-    const browserType = playwright[this.browserConfig.browserName];
+    const browserType = playwright[this.config.browser.browserName];
    for (let i = 0; i < 5; i++) {
      try {
        const browserContext = await browserType.launchPersistentContext(userDataDir, {
-          ...this.browserConfig.launchOptions,
-          ...this.browserConfig.contextOptions,
+          tracesDir,
+          ...this.config.browser.launchOptions,
+          ...this.config.browser.contextOptions,
          handleSIGINT: false,
          handleSIGTERM: false,
        });
@@ -198,17 +209,12 @@ class PersistentContextFactory implements BrowserContextFactory {
    testDebug('close browser context complete (persistent)');
  }

-  private async _createUserDataDir() {
-    let cacheDirectory: string;
-    if (process.platform === 'linux')
-      cacheDirectory = process.env.XDG_CACHE_HOME || path.join(os.homedir(), '.cache');
-    else if (process.platform === 'darwin')
-      cacheDirectory = path.join(os.homedir(), 'Library', 'Caches');
-    else if (process.platform === 'win32')
-      cacheDirectory = process.env.LOCALAPPDATA || path.join(os.homedir(), 'AppData', 'Local');
-    else
-      throw new Error('Unsupported platform: ' + process.platform);
-    const result = path.join(cacheDirectory, 'ms-playwright', `mcp-${this.browserConfig.launchOptions?.channel ?? this.browserConfig?.browserName}-profile`);
+  private async _createUserDataDir(rootPath: string | undefined) {
+    const dir = process.env.PWMCP_PROFILES_DIR_FOR_TEST ?? registryDirectory;
+    const browserToken = this.config.browser.launchOptions?.channel ?? this.config.browser?.browserName;
+    // Hesitant putting hundreds of files into the user's workspace, so using it for hashing instead.
+    const rootPathToken = rootPath ? `-${createHash(rootPath)}` : '';
+    const result = path.join(dir, `mcp-${browserToken}${rootPathToken}`);
    await fs.promises.mkdir(result, { recursive: true });
    return result;
  }
--- a/src/browserServerBackend.ts
+++ b/src/browserServerBackend.ts
@@ -14,6 +14,7 @@
 * limitations under the License.
 */

+import { fileURLToPath } from 'url';
 import { FullConfig } from './config.js';
 import { Context } from './context.js';
 import { logUnhandledError } from './log.js';
@@ -44,14 +45,21 @@ export class BrowserServerBackend implements ServerBackend {
  }

  async initialize(server: mcpServer.Server): Promise<void> {
-    this._sessionLog = this._config.saveSession ? await SessionLog.create(this._config) : undefined;
+    const capabilities = server.getClientCapabilities() as mcpServer.ClientCapabilities;
+    let rootPath: string | undefined;
+    if (capabilities.roots) {
+      const { roots } = await server.listRoots();
+      const firstRootUri = roots[0]?.uri;
+      const url = firstRootUri ? new URL(firstRootUri) : undefined;
+      rootPath = url ? fileURLToPath(url) : undefined;
+    }
+    this._sessionLog = this._config.saveSession ? await SessionLog.create(this._config, rootPath) : undefined;
    this._context = new Context({
      tools: this._tools,
      config: this._config,
      browserContextFactory: this._browserContextFactory,
      sessionLog: this._sessionLog,
-      clientVersion: server.getClientVersion(),
-      capabilities: server.getClientCapabilities() as mcpServer.ClientCapabilities,
+      clientInfo: { ...server.getClientVersion(), rootPath },
    });
  }

--- a/src/config.ts
+++ b/src/config.ts
@@ -18,7 +18,8 @@ import fs from 'fs';
 import os from 'os';
 import path from 'path';
 import { devices } from 'playwright';
-import { sanitizeForFilePath } from './tools/utils.js';
+import { sanitizeForFilePath } from './utils.js';
+
 import type { Config, ToolCapability } from '../config.js';
 import type { BrowserContextOptions, LaunchOptions } from 'playwright';

@@ -67,7 +68,7 @@ const defaultConfig: FullConfig = {
    blockedOrigins: undefined,
  },
  server: {},
-  outputDir: path.join(os.tmpdir(), 'playwright-mcp-output', sanitizeForFilePath(new Date().toISOString())),
+  saveTrace: false,
 };

 type BrowserUserConfig = NonNullable<Config['browser']>;
@@ -79,7 +80,7 @@ export type FullConfig = Config & {
    contextOptions: NonNullable<BrowserUserConfig['contextOptions']>;
  },
  network: NonNullable<Config['network']>,
-  outputDir: string;
+  saveTrace: boolean;
  server: NonNullable<Config['server']>,
 };

@@ -95,9 +96,6 @@ export async function resolveCLIConfig(cliOptions: CLIOptions): Promise<FullConf
  result = mergeConfig(result, configInFile);
  result = mergeConfig(result, envOverrides);
  result = mergeConfig(result, cliOverrides);
-  // Derive artifact output directory from config.outputDir
-  if (result.saveTrace)
-    result.browser.launchOptions.tracesDir = path.join(result.outputDir, 'traces');
  return result;
 }

@@ -240,10 +238,14 @@ async function loadConfig(configFile: string | undefined): Promise<Config> {
  }
 }

-export async function outputFile(config: FullConfig, name: string): Promise<string> {
-  await fs.promises.mkdir(config.outputDir, { recursive: true });
+export async function outputFile(config: FullConfig, rootPath: string | undefined, name: string): Promise<string> {
+  const outputDir = config.outputDir
+    ?? (rootPath ? path.join(rootPath, '.playwright-mcp') : undefined)
+    ?? path.join(os.tmpdir(), 'playwright-mcp-output', sanitizeForFilePath(new Date().toISOString()));
+
+  await fs.promises.mkdir(outputDir, { recursive: true });
  const fileName = sanitizeForFilePath(name);
-  return path.join(config.outputDir, fileName);
+  return path.join(outputDir, fileName);
 }

 function pickDefined<T extends object>(obj: T | undefined): Partial<T> {
--- a/src/context.ts
+++ b/src/context.ts
@@ -19,11 +19,11 @@ import * as playwright from 'playwright';

 import { logUnhandledError } from './log.js';
 import { Tab } from './tab.js';
+import { outputFile  } from './config.js';

-import type * as mcpServer from './mcp/server.js';
-import type { Tool } from './tools/tool.js';
 import type { FullConfig } from './config.js';
-import type { BrowserContextFactory } from './browserContextFactory.js';
+import type { Tool } from './tools/tool.js';
+import type { BrowserContextFactory, ClientInfo } from './browserContextFactory.js';
 import type * as actions from './actions.js';
 import type { SessionLog } from './sessionLog.js';

@@ -34,9 +34,9 @@ type ContextOptions = {
  config: FullConfig;
  browserContextFactory: BrowserContextFactory;
  sessionLog: SessionLog | undefined;
-  clientVersion: { name: string; version: string; } | undefined;
-  capabilities: mcpServer.ClientCapabilities | undefined;
+  clientInfo: ClientInfo;
 };
+
 export class Context {
  readonly tools: Tool[];
  readonly config: FullConfig;
@@ -45,8 +45,7 @@ export class Context {
  private _browserContextFactory: BrowserContextFactory;
  private _tabs: Tab[] = [];
  private _currentTab: Tab | undefined;
-  private _clientVersion: { name: string; version: string; } | undefined;
-  private _clientCapabilities: mcpServer.ClientCapabilities;
+  private _clientInfo: ClientInfo;

  private static _allContexts: Set<Context> = new Set();
  private _closeBrowserContextPromise: Promise<void> | undefined;
@@ -58,8 +57,7 @@ export class Context {
    this.config = options.config;
    this.sessionLog = options.sessionLog;
    this._browserContextFactory = options.browserContextFactory;
-    this._clientVersion = options.clientVersion;
-    this._clientCapabilities = options.capabilities || {};
+    this._clientInfo = options.clientInfo;
    testDebug('create context');
    Context._allContexts.add(this);
  }
@@ -105,7 +103,6 @@ export class Context {
    return this._currentTab!;
  }

-
  async closeTab(index: number | undefined): Promise<string> {
    const tab = index === undefined ? this._currentTab : this._tabs[index];
    if (!tab)
@@ -115,6 +112,10 @@ export class Context {
    return url;
  }

+  async outputFile(name: string): Promise<string> {
+    return outputFile(this.config, this._clientInfo.rootPath, name);
+  }
+
  private _onPageCreated(page: playwright.Page) {
    const tab = new Tab(this, page, tab => this._onPageClosed(tab));
    this._tabs.push(tab);
@@ -199,7 +200,7 @@ export class Context {
    if (this._closeBrowserContextPromise)
      throw new Error('Another browser context is being closed.');
    // TODO: move to the browser context factory to make it based on isolation mode.
-    const result = await this._browserContextFactory.createContext(this._clientVersion!, this._abortController.signal);
+    const result = await this._browserContextFactory.createContext(this._clientInfo, this._abortController.signal);
    const { browserContext } = result;
    await this._setupRequestInterception(browserContext);
    if (this.sessionLog)
--- a/src/index.ts
+++ b/src/index.ts
@@ -26,7 +26,7 @@ import type { Server } from '@modelcontextprotocol/sdk/server/index.js';

 export async function createConnection(userConfig: Config = {}, contextGetter?: () => Promise<BrowserContext>): Promise<Server> {
  const config = await resolveConfig(userConfig);
-  const factory = contextGetter ? new SimpleBrowserContextFactory(contextGetter) : contextFactory(config.browser);
+  const factory = contextGetter ? new SimpleBrowserContextFactory(contextGetter) : contextFactory(config);
  return mcpServer.createServer(new BrowserServerBackend(config, factory), false);
 }

--- a/src/loopTools/context.ts
+++ b/src/loopTools/context.ts
@@ -45,7 +45,7 @@ export class Context {

  static async create(config: FullConfig) {
    const client = new Client({ name: 'Playwright Proxy', version: '1.0.0' });
-    const browserContextFactory = contextFactory(config.browser);
+    const browserContextFactory = contextFactory(config);
    const server = mcpServer.createServer(new BrowserServerBackend(config, browserContextFactory), false);
    await client.connect(new InProcessTransport(server));
    await client.ping();
--- a/src/program.ts
+++ b/src/program.ts
@@ -77,7 +77,7 @@ program
        return;
      }

-      const browserContextFactory = contextFactory(config.browser);
+      const browserContextFactory = contextFactory(config);
      const serverBackendFactory = () => new BrowserServerBackend(config, browserContextFactory);
      await mcpTransport.start(serverBackendFactory, config.server);

--- a/src/sessionLog.ts
+++ b/src/sessionLog.ts
@@ -17,10 +17,10 @@
 import fs from 'fs';
 import path from 'path';

-import { outputFile  } from './config.js';
 import { Response } from './response.js';
-
 import { logUnhandledError } from './log.js';
+import { outputFile  } from './config.js';
+
 import type { FullConfig } from './config.js';
 import type * as actions from './actions.js';
 import type { Tab, TabSnapshot } from './tab.js';
@@ -51,8 +51,8 @@ export class SessionLog {
    this._file = path.join(this._folder, 'session.md');
  }

-  static async create(config: FullConfig): Promise<SessionLog> {
-    const sessionFolder = await outputFile(config, `session-${Date.now()}`);
+  static async create(config: FullConfig, rootPath: string | undefined): Promise<SessionLog> {
+    const sessionFolder = await outputFile(config, rootPath, `session-${Date.now()}`);
    await fs.promises.mkdir(sessionFolder, { recursive: true });
    // eslint-disable-next-line no-console
    console.error(`Session: ${sessionFolder}`);
--- a/src/tab.ts
+++ b/src/tab.ts
@@ -20,7 +20,6 @@ import { callOnPageNoTrace, waitForCompletion } from './tools/utils.js';
 import { logUnhandledError } from './log.js';
 import { ManualPromise } from './manualPromise.js';
 import { ModalState } from './tools/tool.js';
-import { outputFile } from './config.js';

 import type { Context } from './context.js';

@@ -115,7 +114,7 @@ export class Tab extends EventEmitter<TabEventsInterface> {
    const entry = {
      download,
      finished: false,
-      outputFile: await outputFile(this.context.config, download.suggestedFilename())
+      outputFile: await this.context.outputFile(download.suggestedFilename())
    };
    this._downloads.push(entry);
    await download.saveAs(entry.outputFile);
--- a/src/tools/pdf.ts
+++ b/src/tools/pdf.ts
@@ -18,7 +18,6 @@ import { z } from 'zod';
 import { defineTabTool } from './tool.js';

 import * as javascript from '../javascript.js';
-import { outputFile } from '../config.js';

 const pdfSchema = z.object({
  filename: z.string().optional().describe('File name to save the pdf to. Defaults to `page-{timestamp}.pdf` if not specified.'),
@@ -36,7 +35,7 @@ const pdf = defineTabTool({
  },

  handle: async (tab, params, response) => {
-    const fileName = await outputFile(tab.context.config, params.filename ?? `page-${new Date().toISOString()}.pdf`);
+    const fileName = await tab.context.outputFile(params.filename ?? `page-${new Date().toISOString()}.pdf`);
    response.addCode(`await page.pdf(${javascript.formatObject({ path: fileName })});`);
    response.addResult(`Saved page as ${fileName}`);
    await tab.page.pdf({ path: fileName });
--- a/src/tools/screenshot.ts
+++ b/src/tools/screenshot.ts
@@ -18,7 +18,6 @@ import { z } from 'zod';

 import { defineTabTool } from './tool.js';
 import * as javascript from '../javascript.js';
-import { outputFile } from '../config.js';
 import { generateLocator } from './utils.js';

 import type * as playwright from 'playwright';
@@ -53,7 +52,7 @@ const screenshot = defineTabTool({

  handle: async (tab, params, response) => {
    const fileType = params.type || 'png';
-    const fileName = await outputFile(tab.context.config, params.filename ?? `page-${new Date().toISOString()}.${fileType}`);
+    const fileName = await tab.context.outputFile(params.filename ?? `page-${new Date().toISOString()}.${fileType}`);
    const options: playwright.PageScreenshotOptions = {
      type: fileType,
      quality: fileType === 'png' ? undefined : 90,
--- a/src/tools/utils.ts
+++ b/src/tools/utils.ts
@@ -71,14 +71,6 @@ export async function waitForCompletion<R>(tab: Tab, callback: () => Promise<R>)
  }
 }

-export function sanitizeForFilePath(s: string) {
-  const sanitize = (s: string) => s.replace(/[\x00-\x2C\x2E-\x2F\x3A-\x40\x5B-\x60\x7B-\x7F]+/g, '-');
-  const separator = s.lastIndexOf('.');
-  if (separator === -1)
-    return sanitize(s);
-  return sanitize(s.substring(0, separator)) + '.' + sanitize(s.substring(separator + 1));
-}
-
 export async function generateLocator(locator: playwright.Locator): Promise<string> {
  try {
    const { resolvedSelector } = await (locator as any)._resolveSelector();
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -0,0 +1,29 @@
+/**
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import crypto from 'crypto';
+
+export function createHash(data: string): string {
+  return crypto.createHash('sha256').update(data).digest('hex').slice(0, 7);
+}
+
+export function sanitizeForFilePath(s: string) {
+  const sanitize = (s: string) => s.replace(/[\x00-\x2C\x2E-\x2F\x3A-\x40\x5B-\x60\x7B-\x7F]+/g, '-');
+  const separator = s.lastIndexOf('.');
+  if (separator === -1)
+    return sanitize(s);
+  return sanitize(s.substring(0, separator)) + '.' + sanitize(s.substring(separator + 1));
+}
--- a/tests/fixtures.ts
+++ b/tests/fixtures.ts
@@ -22,6 +22,7 @@ import { chromium } from 'playwright';
 import { test as baseTest, expect as baseExpect } from '@playwright/test';
 import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
+import { ListRootsRequestSchema } from '@modelcontextprotocol/sdk/types.js';
 import { TestServer } from './testserver/index.ts';

 import type { Config } from '../config';
@@ -41,7 +42,12 @@ type CDPServer = {

 type TestFixtures = {
  client: Client;
-  startClient: (options?: { clientName?: string, args?: string[], config?: Config }) => Promise<{ client: Client, stderr: () => string }>;
+  startClient: (options?: {
+    clientName?: string,
+    args?: string[],
+    config?: Config,
+    roots?: { name: string, uri: string }[],
+  }) => Promise<{ client: Client, stderr: () => string }>;
  wsEndpoint: string;
  cdpServer: CDPServer;
  server: TestServer;
@@ -61,14 +67,11 @@ export const test = baseTest.extend<TestFixtures & TestOptions, WorkerFixtures>(
  },

  startClient: async ({ mcpHeadless, mcpBrowser, mcpMode }, use, testInfo) => {
-    const userDataDir = mcpMode !== 'docker' ? testInfo.outputPath('user-data-dir') : undefined;
    const configDir = path.dirname(test.info().config.configFile!);
    let client: Client | undefined;

    await use(async options => {
      const args: string[] = [];
-      if (userDataDir)
-        args.push('--user-data-dir', userDataDir);
      if (process.env.CI && process.platform === 'linux')
        args.push('--no-sandbox');
      if (mcpHeadless)
@@ -83,8 +86,15 @@ export const test = baseTest.extend<TestFixtures & TestOptions, WorkerFixtures>(
        args.push(`--config=${path.relative(configDir, configFile)}`);
      }

-      client = new Client({ name: options?.clientName ?? 'test', version: '1.0.0' });
-      const { transport, stderr } = await createTransport(args, mcpMode);
+      client = new Client({ name: options?.clientName ?? 'test', version: '1.0.0' }, options?.roots ? { capabilities: { roots: {} } } : undefined);
+      if (options?.roots) {
+        client.setRequestHandler(ListRootsRequestSchema, async request => {
+          return {
+            roots: options.roots,
+          };
+        });
+      }
+      const { transport, stderr } = await createTransport(args, mcpMode, testInfo.outputPath('ms-playwright'));
      let stderrBuffer = '';
      stderr?.on('data', data => {
        if (process.env.PWMCP_DEBUG)
@@ -160,7 +170,7 @@ export const test = baseTest.extend<TestFixtures & TestOptions, WorkerFixtures>(
  },
 });

-async function createTransport(args: string[], mcpMode: TestOptions['mcpMode']): Promise<{
+async function createTransport(args: string[], mcpMode: TestOptions['mcpMode'], profilesDir: string): Promise<{
  transport: Transport,
  stderr: Stream | null,
 }> {
@@ -188,6 +198,7 @@ async function createTransport(args: string[], mcpMode: TestOptions['mcpMode']):
      DEBUG: 'pw:mcp:test',
      DEBUG_COLORS: '0',
      DEBUG_HIDE_DATE: '1',
+      PWMCP_PROFILES_DIR_FOR_TEST: profilesDir,
    },
  });
  return {
--- a/tests/roots.spec.ts
+++ b/tests/roots.spec.ts
@@ -0,0 +1,66 @@
+/**
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import fs from 'fs';
+import path from 'path';
+import { pathToFileURL } from 'url';
+
+import { test, expect } from './fixtures.js';
+import { createHash } from '../src/utils.js';
+
+test('should use separate user data by root path', async ({ startClient, server }, testInfo) => {
+  const { client } = await startClient({
+    roots: [
+      {
+        name: 'test',
+        uri: 'file:///non/existent/folder',
+      },
+    ],
+  });
+
+  await client.callTool({
+    name: 'browser_navigate',
+    arguments: { url: server.HELLO_WORLD },
+  });
+
+  const hash = createHash('/non/existent/folder');
+  const [file] = await fs.promises.readdir(testInfo.outputPath('ms-playwright'));
+  expect(file).toContain(hash);
+});
+
+
+test('check that trace is saved in workspace', async ({ startClient, server, mcpMode }, testInfo) => {
+  const rootPath = testInfo.outputPath('workspace');
+  const { client } = await startClient({
+    args: ['--save-trace'],
+    roots: [
+      {
+        name: 'workspace',
+        uri: pathToFileURL(rootPath).toString(),
+      },
+    ],
+  });
+
+  expect(await client.callTool({
+    name: 'browser_navigate',
+    arguments: { url: server.HELLO_WORLD },
+  })).toHaveResponse({
+    code: expect.stringContaining(`page.goto('http://localhost`),
+  });
+
+  const [file] = await fs.promises.readdir(path.join(rootPath, '.playwright-mcp'));
+  expect(file).toContain('traces');
+});
--- a/tests/trace.spec.ts
+++ b/tests/trace.spec.ts
@@ -15,7 +15,6 @@
 */

 import fs from 'fs';
-import path from 'path';

 import { test, expect } from './fixtures.js';

@@ -33,5 +32,6 @@ test('check that trace is saved', async ({ startClient, server, mcpMode }, testI
    code: expect.stringContaining(`page.goto('http://localhost`),
  });

-  expect(fs.existsSync(path.join(outputDir, 'traces', 'trace.trace'))).toBeTruthy();
+  const [file] = await fs.promises.readdir(outputDir);
+  expect(file).toContain('traces');
 });