Merge remote-tracking branch 'origin/main' into copilot/fix-759

# Conflicts: # src/browserContextFactory.ts
2025-10-12 00:25:14 +03:00 · 2025-08-02 01:27:58 +00:00
parent d4667f865d 3c6eac9b21
commit 41ba5ba0eb
20 changed files with 343 additions and 151 deletions
--- a/extension/src/relayConnection.ts
+++ b/extension/src/relayConnection.ts
@@ -59,10 +59,6 @@ export class RelayConnection {
    this._ws.close(1000, message);
  }

-  private async _detachDebugger(): Promise<void> {
-    await chrome.debugger.detach(this._debuggee);
-  }
-
  private _onDebuggerEvent(source: chrome.debugger.DebuggerSession, method: string, params: any): void {
    if (source.tabId !== this._debuggee.tabId)
      return;
@@ -81,13 +77,7 @@ export class RelayConnection {
  private _onDebuggerDetach(source: chrome.debugger.Debuggee, reason: string): void {
    if (source.tabId !== this._debuggee.tabId)
      return;
-    this._sendMessage({
-      method: 'detachedFromTab',
-      params: {
-        tabId: this._debuggee.tabId,
-        reason,
-      },
-    });
+    this.close(`Debugger detached: ${reason}`);
    this._debuggee = { };
  }

@@ -131,10 +121,6 @@ export class RelayConnection {
        targetInfo: result?.targetInfo,
      };
    }
-    if (message.method === 'detachFromTab') {
-      debugLog('Detaching debugger from tab:', this._debuggee);
-      return await this._detachDebugger();
-    }
    if (message.method === 'forwardCDPCommand') {
      const { sessionId, method, params } = message.params;
      debugLog('CDP command:', method, params);
--- a/src/browserContextFactory.ts
+++ b/src/browserContextFactory.ts
@@ -14,41 +14,48 @@
 * 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 {
+  readonly name: string;
+  readonly description: string;
  createContext(clientInfo: ClientInfo, abortSignal: AbortSignal): Promise<{ browserContext: playwright.BrowserContext, close: () => Promise<void> }>;
 }

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

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

  protected async _obtainBrowser(): Promise<playwright.Browser> {
@@ -70,7 +77,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 +104,16 @@ class BaseContextFactory implements BrowserContextFactory {
 }

 class IsolatedContextFactory extends BaseContextFactory {
-  constructor(browserConfig: FullConfig['browser']) {
-    super('isolated', browserConfig);
+  constructor(config: FullConfig) {
+    super('isolated', 'Create a new isolated browser context', 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,71 +124,74 @@ 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', 'Connect to a browser over 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> {
-    if (this.browserConfig.isolated) {
-      return browser.newContext(this.browserConfig.contextOptions);
-    }
-    return browser.contexts()[0];
+    return this.config.browser.isolated ? await browser.newContext(this.config.browser.contextOptions) : browser.contexts()[0];
  }
 }

 class RemoteContextFactory extends BaseContextFactory {
-  constructor(browserConfig: FullConfig['browser']) {
-    super('remote', browserConfig);
+  constructor(config: FullConfig) {
+    super('remote', 'Connect to a browser using a remote endpoint', 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> {
-    return browser.newContext(this.browserConfig.contextOptions);
+    return browser.newContext(this.config.browser.contextOptions);
  }
 }

 class PersistentContextFactory implements BrowserContextFactory {
-  readonly browserConfig: FullConfig['browser'];
+  readonly config: FullConfig;
+  readonly name = 'persistent';
+  readonly description = 'Create a new persistent browser context';
+
  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 contextOptions = {
-          ...this.browserConfig.launchOptions,
-          ...this.browserConfig.contextOptions,
+        const browserContext = await browserType.launchPersistentContext(userDataDir, {
+          tracesDir,
+          ...this.config.browser.launchOptions,
+          ...this.config.browser.contextOptions,
          handleSIGINT: false,
          handleSIGTERM: false,
-        };
-        const browserContext = await browserType.launchPersistentContext(userDataDir, contextOptions);
+        });
        const close = () => this._closeBrowserContext(browserContext, userDataDir);
        return { browserContext, close };
      } catch (error: any) {
@@ -202,17 +216,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,8 @@
 * limitations under the License.
 */

+import { fileURLToPath } from 'url';
+import { z } from 'zod';
 import { FullConfig } from './config.js';
 import { Context } from './context.js';
 import { logUnhandledError } from './log.js';
@@ -21,11 +23,16 @@ import { Response } from './response.js';
 import { SessionLog } from './sessionLog.js';
 import { filteredTools } from './tools.js';
 import { packageJSON } from './package.js';
+import { defineTool  } from './tools/tool.js';

+import type { Tool } from './tools/tool.js';
 import type { BrowserContextFactory } from './browserContextFactory.js';
 import type * as mcpServer from './mcp/server.js';
 import type { ServerBackend } from './mcp/server.js';
-import type { Tool } from './tools/tool.js';
+
+type NonEmptyArray<T> = [T, ...T[]];
+
+export type FactoryList = NonEmptyArray<BrowserContextFactory>;

 export class BrowserServerBackend implements ServerBackend {
  name = 'Playwright';
@@ -37,15 +44,31 @@ export class BrowserServerBackend implements ServerBackend {
  private _config: FullConfig;
  private _browserContextFactory: BrowserContextFactory;

-  constructor(config: FullConfig, browserContextFactory: BrowserContextFactory) {
+  constructor(config: FullConfig, factories: FactoryList) {
    this._config = config;
-    this._browserContextFactory = browserContextFactory;
+    this._browserContextFactory = factories[0];
    this._tools = filteredTools(config);
+    if (factories.length > 1)
+      this._tools.push(this._defineContextSwitchTool(factories));
  }

-  async initialize() {
-    this._sessionLog = this._config.saveSession ? await SessionLog.create(this._config) : undefined;
-    this._context = new Context(this._tools, this._config, this._browserContextFactory, this._sessionLog);
+  async initialize(server: mcpServer.Server): Promise<void> {
+    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,
+      clientInfo: { ...server.getClientVersion(), rootPath },
+    });
  }

  tools(): mcpServer.ToolSchema<any>[] {
@@ -69,11 +92,49 @@ export class BrowserServerBackend implements ServerBackend {
    return response.serialize();
  }

-  serverInitialized(version: mcpServer.ClientVersion | undefined) {
-    this._context!.clientVersion = version;
-  }
-
  serverClosed() {
    void this._context!.dispose().catch(logUnhandledError);
  }
+
+  private _defineContextSwitchTool(factories: FactoryList): Tool<any> {
+    const self = this;
+    return defineTool({
+      capability: 'core',
+
+      schema: {
+        name: 'browser_connect',
+        title: 'Connect to a browser context',
+        description: [
+          'Connect to a browser using one of the available methods:',
+          ...factories.map(factory => `- "${factory.name}": ${factory.description}`),
+        ].join('\n'),
+        inputSchema: z.object({
+          method: z.enum(factories.map(factory => factory.name) as [string, ...string[]]).default(factories[0].name).describe('The method to use to connect to the browser'),
+        }),
+        type: 'readOnly',
+      },
+
+      async handle(context, params, response) {
+        const factory = factories.find(factory => factory.name === params.method);
+        if (!factory) {
+          response.addError('Unknown connection method: ' + params.method);
+          return;
+        }
+        await self._setContextFactory(factory);
+        response.addResult('Successfully changed connection method.');
+      }
+    });
+  }
+
+  private async _setContextFactory(newFactory: BrowserContextFactory) {
+    if (this._context) {
+      const options = {
+        ...this._context.options,
+        browserContextFactory: newFactory,
+      };
+      await this._context.dispose();
+      this._context = new Context(options);
+    }
+    this._browserContextFactory = newFactory;
+  }
 }
--- 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';

@@ -68,7 +69,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']>;
@@ -80,7 +81,7 @@ export type FullConfig = Config & {
    contextOptions: NonNullable<BrowserUserConfig['contextOptions']>;
  },
  network: NonNullable<Config['network']>,
-  outputDir: string;
+  saveTrace: boolean;
  server: NonNullable<Config['server']>,
 };

@@ -96,9 +97,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;
 }

@@ -245,10 +243,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,36 +19,47 @@ import * as playwright from 'playwright';

 import { logUnhandledError } from './log.js';
 import { Tab } from './tab.js';
+import { outputFile  } from './config.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';

 const testDebug = debug('pw:mcp:test');

+type ContextOptions = {
+  tools: Tool[];
+  config: FullConfig;
+  browserContextFactory: BrowserContextFactory;
+  sessionLog: SessionLog | undefined;
+  clientInfo: ClientInfo;
+};
+
 export class Context {
  readonly tools: Tool[];
  readonly config: FullConfig;
  readonly sessionLog: SessionLog | undefined;
+  readonly options: ContextOptions;
  private _browserContextPromise: Promise<{ browserContext: playwright.BrowserContext, close: () => Promise<void> }> | undefined;
  private _browserContextFactory: BrowserContextFactory;
  private _tabs: Tab[] = [];
  private _currentTab: Tab | undefined;
-
-  clientVersion: { name: string; version: string; } | undefined;
+  private _clientInfo: ClientInfo;

  private static _allContexts: Set<Context> = new Set();
  private _closeBrowserContextPromise: Promise<void> | undefined;
  private _isRunningTool: boolean = false;
  private _abortController = new AbortController();

-  constructor(tools: Tool[], config: FullConfig, browserContextFactory: BrowserContextFactory, sessionLog: SessionLog | undefined) {
-    this.tools = tools;
-    this.config = config;
-    this._browserContextFactory = browserContextFactory;
-    this.sessionLog = sessionLog;
+  constructor(options: ContextOptions) {
+    this.tools = options.tools;
+    this.config = options.config;
+    this.sessionLog = options.sessionLog;
+    this.options = options;
+    this._browserContextFactory = options.browserContextFactory;
+    this._clientInfo = options.clientInfo;
    testDebug('create context');
    Context._allContexts.add(this);
  }
@@ -94,7 +105,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)
@@ -104,6 +114,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);
@@ -188,7 +202,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/extension/extensionContextFactory.ts
+++ b/src/extension/extensionContextFactory.ts
@@ -24,6 +24,9 @@ import type { BrowserContextFactory, ClientInfo } from '../browserContextFactory
 const debugLogger = debug('pw:mcp:relay');

 export class ExtensionContextFactory implements BrowserContextFactory {
+  name = 'extension';
+  description = 'Connect to a browser using the Playwright MCP extension';
+
  private _browserChannel: string;
  private _relayPromise: Promise<CDPRelayServer> | undefined;
  private _browserPromise: Promise<playwright.Browser> | undefined;
--- a/src/extension/main.ts
+++ b/src/extension/main.ts
@@ -22,6 +22,10 @@ import type { FullConfig } from '../config.js';

 export async function runWithExtension(config: FullConfig) {
  const contextFactory = new ExtensionContextFactory(config.browser.launchOptions.channel || 'chrome');
-  const serverBackendFactory = () => new BrowserServerBackend(config, contextFactory);
+  const serverBackendFactory = () => new BrowserServerBackend(config, [contextFactory]);
  await mcpTransport.start(serverBackendFactory, config.server);
 }
+
+export function createExtensionContextFactory(config: FullConfig) {
+  return new ExtensionContextFactory(config.browser.launchOptions.channel || 'chrome');
+}
--- a/src/index.ts
+++ b/src/index.ts
@@ -26,11 +26,14 @@ 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);
-  return mcpServer.createServer(new BrowserServerBackend(config, factory), false);
+  const factory = contextGetter ? new SimpleBrowserContextFactory(contextGetter) : contextFactory(config);
+  return mcpServer.createServer(new BrowserServerBackend(config, [factory]), false);
 }

 class SimpleBrowserContextFactory implements BrowserContextFactory {
+  name = 'custom';
+  description = 'Connect to a browser using a custom context getter';
+
  private readonly _contextGetter: () => Promise<BrowserContext>;

  constructor(contextGetter: () => Promise<BrowserContext>) {
--- a/src/loopTools/context.ts
+++ b/src/loopTools/context.ts
@@ -45,8 +45,8 @@ 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 server = mcpServer.createServer(new BrowserServerBackend(config, browserContextFactory), false);
+    const browserContextFactory = contextFactory(config);
+    const server = mcpServer.createServer(new BrowserServerBackend(config, [browserContextFactory]), false);
    await client.connect(new InProcessTransport(server));
    await client.ping();
    return new Context(config, client);
--- a/src/mcp/server.ts
+++ b/src/mcp/server.ts
@@ -18,11 +18,18 @@ import { z } from 'zod';
 import { Server } from '@modelcontextprotocol/sdk/server/index.js';
 import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';
 import { zodToJsonSchema } from 'zod-to-json-schema';
+import { ManualPromise } from '../manualPromise.js';
+import { logUnhandledError } from '../log.js';

-import type { ImageContent, Implementation, TextContent } from '@modelcontextprotocol/sdk/types.js';
+import type { ImageContent, TextContent } from '@modelcontextprotocol/sdk/types.js';
 import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
+export type { Server } from '@modelcontextprotocol/sdk/server/index.js';

-export type ClientVersion = Implementation;
+export type ClientCapabilities = {
+  roots?: {
+    listRoots?: boolean
+  };
+};

 export type ToolResponse = {
  content: (TextContent | ImageContent)[];
@@ -42,10 +49,9 @@ export type ToolHandler = (toolName: string, params: any) => Promise<ToolRespons
 export interface ServerBackend {
  name: string;
  version: string;
-  initialize?(): Promise<void>;
+  initialize?(server: Server): Promise<void>;
  tools(): ToolSchema<any>[];
  callTool(schema: ToolSchema<any>, parsedArguments: any): Promise<ToolResponse>;
-  serverInitialized?(version: ClientVersion | undefined): void;
  serverClosed?(): void;
 }

@@ -53,12 +59,12 @@ export type ServerBackendFactory = () => ServerBackend;

 export async function connect(serverBackendFactory: ServerBackendFactory, transport: Transport, runHeartbeat: boolean) {
  const backend = serverBackendFactory();
-  await backend.initialize?.();
  const server = createServer(backend, runHeartbeat);
  await server.connect(transport);
 }

 export function createServer(backend: ServerBackend, runHeartbeat: boolean): Server {
+  const initializedPromise = new ManualPromise<void>();
  const server = new Server({ name: backend.name, version: backend.version }, {
    capabilities: {
      tools: {},
@@ -82,6 +88,8 @@ export function createServer(backend: ServerBackend, runHeartbeat: boolean): Ser

  let heartbeatRunning = false;
  server.setRequestHandler(CallToolRequestSchema, async request => {
+    await initializedPromise;
+
    if (runHeartbeat && !heartbeatRunning) {
      heartbeatRunning = true;
      startHeartbeat(server);
@@ -101,8 +109,9 @@ export function createServer(backend: ServerBackend, runHeartbeat: boolean): Ser
      return errorResult(String(error));
    }
  });
-
-  addServerListener(server, 'initialized', () => backend.serverInitialized?.(server.getClientVersion()));
+  addServerListener(server, 'initialized', () => {
+    backend.initialize?.(server).then(() => initializedPromise.resolve()).catch(logUnhandledError);
+  });
  addServerListener(server, 'close', () => backend.serverClosed?.());
  return server;
 }
--- a/src/program.ts
+++ b/src/program.ts
@@ -21,8 +21,8 @@ import { startTraceViewerServer } from 'playwright-core/lib/server';
 import * as mcpTransport from './mcp/transport.js';
 import { commaSeparatedList, resolveCLIConfig, semicolonSeparatedList } from './config.js';
 import { packageJSON } from './package.js';
-import { runWithExtension } from './extension/main.js';
-import { BrowserServerBackend } from './browserServerBackend.js';
+import { createExtensionContextFactory, runWithExtension } from './extension/main.js';
+import { BrowserServerBackend, FactoryList } from './browserServerBackend.js';
 import { Context } from './context.js';
 import { contextFactory } from './browserContextFactory.js';
 import { runLoopTools } from './loopTools/main.js';
@@ -57,6 +57,7 @@ program
    .option('--user-data-dir <path>', 'path to the user data directory. If not specified, a temporary directory will be created.')
    .option('--viewport-size <size>', 'specify browser viewport size in pixels, for example "1280, 720"')
    .addOption(new Option('--extension', 'Connect to a running browser instance (Edge/Chrome only). Requires the "Playwright MCP Bridge" browser extension to be installed.').hideHelp())
+    .addOption(new Option('--connect-tool', 'Allow to switch between different browser connection methods.').hideHelp())
    .addOption(new Option('--loop-tools', 'Run loop tools').hideHelp())
    .addOption(new Option('--vision', 'Legacy option, use --caps=vision instead').hideHelp())
    .action(async options => {
@@ -78,8 +79,11 @@ program
        return;
      }

-      const browserContextFactory = contextFactory(config.browser);
-      const serverBackendFactory = () => new BrowserServerBackend(config, browserContextFactory);
+      const browserContextFactory = contextFactory(config);
+      const factories: FactoryList = [browserContextFactory];
+      if (options.connectTool)
+        factories.push(createExtensionContextFactory(config));
+      const serverBackendFactory = () => new BrowserServerBackend(config, factories);
      await mcpTransport.start(serverBackendFactory, config.server);

      if (config.saveTrace) {
--- 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,68 @@
+/**
+ * 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';
+
+const p = process.platform === 'win32' ? 'c:\\non\\existent\\folder' : '/non/existent/folder';
+
+test('should use separate user data by root path', async ({ startClient, server }, testInfo) => {
+  const { client } = await startClient({
+    roots: [
+      {
+        name: 'test',
+        uri: 'file://' + p.replace(/\\/g, '/'),
+      }
+    ],
+  });
+
+  await client.callTool({
+    name: 'browser_navigate',
+    arguments: { url: server.HELLO_WORLD },
+  });
+
+  const hash = createHash(p);
+  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');
 });