Merge branch 'main' into vscode-client-factory

src/browserServerBackend.ts

@@ -22,7 +22,7 @@ import { Response } from './response.js';
 import { SessionLog } from './sessionLog.js';
 import { filteredTools } from './tools.js';
 import { packageJSON } from './utils/package.js';
-import { toToolDefinition } from './tools/tool.js';
+import { toMcpTool } from './mcp/tool.js';
 
 import type { Tool } from './tools/tool.js';
 import type { BrowserContextFactory } from './browserContextFactory.js';
@@ -64,12 +64,14 @@ export class BrowserServerBackend implements ServerBackend {
     });
   }
 
-  tools(): mcpServer.ToolDefinition[] {
-    return this._tools.map(tool => toToolDefinition(tool.schema));
+  async listTools(): Promise<mcpServer.Tool[]> {
+    return this._tools.map(tool => toMcpTool(tool.schema));
   }
 
-  async callTool(name: string, rawArguments: any) {
+  async callTool(name: string, rawArguments: mcpServer.CallToolRequest['params']['arguments']) {
     const tool = this._tools.find(tool => tool.schema.name === name)!;
+    if (!tool)
+      throw new Error(`Tool "${name}" not found`);
     const parsedArguments = tool.schema.inputSchema.parse(rawArguments || {});
     const context = this._context!;
     const response = new Response(context, name, parsedArguments);

src/inProcessMcpFactrory.ts

@@ -1,45 +0,0 @@
-/**
- * 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 { BrowserContextFactory } from './browserContextFactory.js';
-import { BrowserServerBackend } from './browserServerBackend.js';
-import { InProcessTransport } from './mcp/inProcessTransport.js';
-import * as mcpServer from './mcp/server.js';
-
-import type { FullConfig } from './config.js';
-import type { MCPFactory } from './mcp/proxyBackend.js';
-import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
-
-export class InProcessMCPFactory implements MCPFactory {
-  name: string;
-  description: string;
-
-  private _contextFactory: BrowserContextFactory;
-  private _config: FullConfig;
-
-  constructor(contextFactory: BrowserContextFactory, config: FullConfig) {
-    this.name = contextFactory.name;
-    this.description = contextFactory.description;
-    this._contextFactory = contextFactory;
-    this._config = config;
-  }
-
-  async create(): Promise<Transport> {
-    const delegate = mcpServer.createServer(new BrowserServerBackend(this._config, this._contextFactory), false);
-    return new InProcessTransport(delegate);
-  }
-}

src/loopTools/DEPS.list

@@ -2,5 +2,4 @@
 ../
 ../loop/
 ../mcp/
-../tools/
 ../utils/

src/loopTools/context.ts

@@ -52,7 +52,7 @@ export class Context {
     return new Context(config, client);
   }
 
-  async runTask(task: string, oneShot: boolean = false): Promise<mcpServer.ToolResponse> {
+  async runTask(task: string, oneShot: boolean = false): Promise<mcpServer.CallToolResult> {
     const messages = await runTask(this._delegate, this._client!, task, oneShot);
     const lines: string[] = [];
 

src/loopTools/main.ts

@@ -22,7 +22,7 @@ import { packageJSON } from '../utils/package.js';
 import { Context } from './context.js';
 import { perform } from './perform.js';
 import { snapshot } from './snapshot.js';
-import { toToolDefinition } from '../tools/tool.js';
+import { toMcpTool } from '../mcp/tool.js';
 
 import type { FullConfig } from '../config.js';
 import type { ServerBackend } from '../mcp/server.js';
@@ -49,13 +49,13 @@ class LoopToolsServerBackend implements ServerBackend {
     this._context = await Context.create(this._config);
   }
 
-  tools(): mcpServer.ToolDefinition[] {
-    return this._tools.map(tool => toToolDefinition(tool.schema));
+  async listTools(): Promise<mcpServer.Tool[]> {
+    return this._tools.map(tool => toMcpTool(tool.schema));
   }
 
-  async callTool(name: string, rawArguments: any): Promise<mcpServer.ToolResponse> {
+  async callTool(name: string, args: mcpServer.CallToolRequest['params']['arguments']): Promise<mcpServer.CallToolResult> {
     const tool = this._tools.find(tool => tool.schema.name === name)!;
-    const parsedArguments = tool.schema.inputSchema.parse(rawArguments || {});
+    const parsedArguments = tool.schema.inputSchema.parse(args || {});
     return await tool.handle(this._context!, parsedArguments);
   }
 

src/loopTools/tool.ts

@@ -17,12 +17,12 @@
 import type { z } from 'zod';
 import type * as mcpServer from '../mcp/server.js';
 import type { Context } from './context.js';
-import type { ToolSchema } from '../tools/tool.js';
+import type { ToolSchema } from '../mcp/tool.js';
 
 
 export type Tool<Input extends z.Schema = z.Schema> = {
   schema: ToolSchema<Input>;
-  handle: (context: Context, params: z.output<Input>) => Promise<mcpServer.ToolResponse>;
+  handle: (context: Context, params: z.output<Input>) => Promise<mcpServer.CallToolResult>;
 };
 
 export function defineTool<Input extends z.Schema>(tool: Tool<Input>): Tool<Input> {

src/mcp/proxyBackend.ts

@@ -24,65 +24,67 @@ import { packageJSON } from '../utils/package.js';
 
 
 import type { Server } from '@modelcontextprotocol/sdk/server/index.js';
-import type { ToolDefinition, ServerBackend, ToolResponse } from './server.js';
+import type { ServerBackend } from './server.js';
 import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
+import type { Root, Tool, CallToolResult, CallToolRequest } from '@modelcontextprotocol/sdk/types.js';
 
-type NonEmptyArray<T> = [T, ...T[]];
-
-export type MCPFactory = {
+export type MCPProvider = {
   name: string;
   description: string;
-  create(options: any): Promise<Transport>;
+  connect(options: any): Promise<Transport>;
 };
 
-export type MCPFactoryList = NonEmptyArray<MCPFactory>;
-
 export class ProxyBackend implements ServerBackend {
   name = 'Playwright MCP Client Switcher';
   version = packageJSON.version;
 
-  private _mcpFactories: MCPFactoryList;
+  private _mcpProviders: MCPProvider[];
   private _currentClient: Client | undefined;
-  private _contextSwitchTool: ToolDefinition;
-  private _tools: ToolDefinition[] = [];
-  private _server: Server | undefined;
+  private _contextSwitchTool: Tool;
+  private _roots: Root[] = [];
 
-  constructor(clientFactories: MCPFactoryList) {
-    this._mcpFactories = clientFactories;
+  constructor(mcpProviders: MCPProvider[]) {
+    this._mcpProviders = mcpProviders;
     this._contextSwitchTool = this._defineContextSwitchTool();
   }
 
   async initialize(server: Server): Promise<void> {
-    this._server = server;
-    await this._setCurrentClient(this._mcpFactories[0], undefined);
+    const version = server.getClientVersion();
+    const capabilities = server.getClientCapabilities();
+    if (capabilities?.roots && version && clientsWithRoots.includes(version.name)) {
+      const { roots } = await server.listRoots();
+      this._roots = roots;
     }
 
-  tools(): ToolDefinition[] {
-    if (this._mcpFactories.length === 1)
-      return this._tools;
+    await this._setCurrentClient(this._mcpProviders[0], undefined);
+  }
+
+  async listTools(): Promise<Tool[]> {
+    const response = await this._currentClient!.listTools();
+    if (this._mcpProviders.length === 1)
+      return response.tools;
     return [
-      ...this._tools,
+      ...response.tools,
       this._contextSwitchTool,
     ];
   }
 
-  async callTool(name: string, rawArguments: any): Promise<ToolResponse> {
+  async callTool(name: string, args: CallToolRequest['params']['arguments']): Promise<CallToolResult> {
     if (name === this._contextSwitchTool.name)
-      return this._callContextSwitchTool(rawArguments);
-    const result = await this._currentClient!.callTool({
+      return this._callContextSwitchTool(args);
+    return await this._currentClient!.callTool({
       name,
-      arguments: rawArguments,
-    });
-    return result as unknown as ToolResponse;
+      arguments: args,
+    }) as CallToolResult;
   }
 
   serverClosed?(): void {
     void this._currentClient?.close().catch(logUnhandledError);
   }
 
-  private async _callContextSwitchTool(params: any): Promise<ToolResponse> {
+  private async _callContextSwitchTool(params: any): Promise<CallToolResult> {
     try {
-      const factory = this._mcpFactories.find(factory => factory.name === params.name);
+      const factory = this._mcpProviders.find(factory => factory.name === params.name);
       if (!factory)
         throw new Error('Unknown connection method: ' + params.name);
 
@@ -98,18 +100,18 @@ export class ProxyBackend implements ServerBackend {
     }
   }
 
-  private _defineContextSwitchTool(): ToolDefinition {
+  private _defineContextSwitchTool(): Tool {
     return {
       name: 'browser_connect',
       description: [
         'Connect to a browser using one of the available methods:',
-        ...this._mcpFactories.map(factory => `- "${factory.name}": ${factory.description}`),
+        ...this._mcpProviders.map(factory => `- "${factory.name}": ${factory.description}`),
         `By default, you're connected to the first method. Only call this tool to change it.`,
       ].join('\n'),
       inputSchema: zodToJsonSchema(z.object({
-        name: z.enum(this._mcpFactories.map(factory => factory.name) as [string, ...string[]]).default(this._mcpFactories[0].name).describe('The method to use to connect to the browser'),
+        name: z.enum(this._mcpProviders.map(factory => factory.name) as [string, ...string[]]).default(this._mcpProviders[0].name).describe('The method to use to connect to the browser'),
         options: z.any().optional().describe('Options for the connection method'),
-      }), { strictUnions: true }) as ToolDefinition['inputSchema'],
+      }), { strictUnions: true }) as Tool['inputSchema'],
       annotations: {
         title: 'Connect to a browser context',
         readOnlyHint: true,
@@ -118,7 +120,7 @@ export class ProxyBackend implements ServerBackend {
     };
   }
 
-  private async _setCurrentClient(factory: MCPFactory, options: any) {
+  private async _setCurrentClient(factory: MCPProvider, options: any) {
     await this._currentClient?.close();
     this._currentClient = undefined;
 
@@ -128,23 +130,13 @@ export class ProxyBackend implements ServerBackend {
         listRoots: true,
       },
     });
-    client.setRequestHandler(ListRootsRequestSchema, async () => {
-      const clientName = this._server!.getClientVersion()?.name;
-      if (this._server!.getClientCapabilities()?.roots && (
-        clientName === 'Visual Studio Code' ||
-        clientName === 'Visual Studio Code - Insiders')) {
-        const { roots } = await this._server!.listRoots();
-        return { roots };
-      }
-      return { roots: [] };
-    });
+    client.setRequestHandler(ListRootsRequestSchema, () => ({ roots: this._roots }));
     client.setRequestHandler(PingRequestSchema, () => ({}));
 
-    const transport = await factory.create(options);
+    const transport = await factory.connect(options);
     await client.connect(transport);
-
     this._currentClient = client;
-    const tools = await this._currentClient.listTools();
-    this._tools = tools.tools;
   }
 }
+
+const clientsWithRoots = ['Visual Studio Code', 'Visual Studio Code - Insiders'];

src/mcp/server.ts

@@ -20,31 +20,19 @@ import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprot
 import { ManualPromise } from '../utils/manualPromise.js';
 import { logUnhandledError } from '../utils/log.js';
 
-import type { ImageContent, TextContent, Tool } from '@modelcontextprotocol/sdk/types.js';
+import type { Tool, CallToolResult, CallToolRequest } 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 { Tool, CallToolResult, CallToolRequest } from '@modelcontextprotocol/sdk/types.js';
 
 const serverDebug = debug('pw:mcp:server');
 
-export type ClientCapabilities = {
-  roots?: {
-    listRoots?: boolean
-  };
-};
-
-export type ToolResponse = {
-  content: (TextContent | ImageContent)[];
-  isError?: boolean;
-};
-
-export type ToolDefinition = Tool;
-
 export interface ServerBackend {
   name: string;
   version: string;
   initialize?(server: Server): Promise<void>;
-  tools(): ToolDefinition[];
-  callTool(name: string, rawArguments: any): Promise<ToolResponse>;
+  listTools(): Promise<Tool[]>;
+  callTool(name: string, args: CallToolRequest['params']['arguments']): Promise<CallToolResult>;
   serverClosed?(): void;
 }
 
@@ -66,7 +54,8 @@ export function createServer(backend: ServerBackend, runHeartbeat: boolean): Ser
 
   server.setRequestHandler(ListToolsRequestSchema, async () => {
     serverDebug('listTools');
-    const tools = backend.tools();
+    await initializedPromise;
+    const tools = await backend.listTools();
     return { tools };
   });
 
@@ -80,19 +69,13 @@ export function createServer(backend: ServerBackend, runHeartbeat: boolean): Ser
       startHeartbeat(server);
     }
 
-    const errorResult = (...messages: string[]) => ({
-      content: [{ type: 'text', text: '### Result\n' + messages.join('\n') }],
-      isError: true,
-    });
-    const tools = backend.tools();
-    const tool = tools.find(tool => tool.name === request.params.name);
-    if (!tool)
-      return errorResult(`Error: Tool "${request.params.name}" not found`);
-
     try {
-      return await backend.callTool(tool.name, request.params.arguments || {});
+      return await backend.callTool(request.params.name, request.params.arguments || {});
     } catch (error) {
-      return errorResult(String(error));
+      return {
+        content: [{ type: 'text', text: '### Result\n' + String(error) }],
+        isError: true,
+      };
     }
   });
   addServerListener(server, 'initialized', () => {

src/mcp/tool.ts

@@ -0,0 +1,42 @@
+/**
+ * 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 { zodToJsonSchema } from 'zod-to-json-schema';
+
+import type { z } from 'zod';
+import type * as mcpServer from './server.js';
+
+export type ToolSchema<Input extends z.Schema> = {
+  name: string;
+  title: string;
+  description: string;
+  inputSchema: Input;
+  type: 'readOnly' | 'destructive';
+};
+
+export function toMcpTool(tool: ToolSchema<any>): mcpServer.Tool {
+  return {
+    name: tool.name,
+    description: tool.description,
+    inputSchema: zodToJsonSchema(tool.inputSchema, { strictUnions: true }) as mcpServer.Tool['inputSchema'],
+    annotations: {
+      title: tool.title,
+      readOnlyHint: tool.type === 'readOnly',
+      destructiveHint: tool.type === 'destructive',
+      openWorldHint: true,
+    },
+  };
+}

src/program.ts

@@ -15,7 +15,7 @@
  */
 
 import { program, Option } from 'commander';
-
+import * as mcpServer from './mcp/server.js';
 import * as mcpTransport from './mcp/transport.js';
 import { commaSeparatedList, resolveCLIConfig, semicolonSeparatedList } from './config.js';
 import { packageJSON } from './utils/package.js';
@@ -23,13 +23,14 @@ import { Context } from './context.js';
 import { contextFactory } from './browserContextFactory.js';
 import { runLoopTools } from './loopTools/main.js';
 import { ProxyBackend } from './mcp/proxyBackend.js';
-import { InProcessMCPFactory } from './inProcessMcpFactrory.js';
 import { BrowserServerBackend } from './browserServerBackend.js';
 import { ExtensionContextFactory } from './extension/extensionContextFactory.js';
+import { InProcessTransport } from './mcp/inProcessTransport.js';
 
 import { VSCodeMCPFactory } from './vscode/host.js';
-import type { MCPFactoryList } from './mcp/proxyBackend.js';
+import type { MCPProvider } from './mcp/proxyBackend.js';
 import type { FullConfig } from './config.js';
+import type { BrowserContextFactory } from './browserContextFactory.js';
 
 program
     .version('Version ' + packageJSON.version)
@@ -79,22 +80,21 @@ program
         await mcpTransport.start(serverBackendFactory, config.server);
         return;
       }
+
       if (options.loopTools) {
         await runLoopTools(config);
         return;
       }
 
       const browserContextFactory = contextFactory(config);
-      const factories: MCPFactoryList = [
-        new InProcessMCPFactory(browserContextFactory, config),
-      ];
+      const providers: MCPProvider[] = [mcpProviderForBrowserContextFactory(config, browserContextFactory)];
       if (options.connectTool) {
-        factories.push(
-            new InProcessMCPFactory(createExtensionContextFactory(config), config),
+        providers.push(
+            mcpProviderForBrowserContextFactory(config, createExtensionContextFactory(config)),
             new VSCodeMCPFactory(config), // TODO: automatically add this based on the client name
         );
       }
-      await mcpTransport.start(() => new ProxyBackend(factories), config.server);
+      await mcpTransport.start(() => new ProxyBackend(providers), config.server);
     });
 
 function setupExitWatchdog() {
@@ -117,4 +117,15 @@ function createExtensionContextFactory(config: FullConfig) {
   return new ExtensionContextFactory(config.browser.launchOptions.channel || 'chrome', config.browser.userDataDir);
 }
 
+function mcpProviderForBrowserContextFactory(config: FullConfig, browserContextFactory: BrowserContextFactory) {
+  return {
+    name: browserContextFactory.name,
+    description: browserContextFactory.description,
+    connect: async () => {
+      const server = mcpServer.createServer(new BrowserServerBackend(config, browserContextFactory), false);
+      return new InProcessTransport(server);
+    },
+  };
+}
+
 void program.parseAsync(process.argv);

src/tools/tool.ts

@@ -14,15 +14,13 @@
  * limitations under the License.
  */
 
-import { zodToJsonSchema } from 'zod-to-json-schema';
-
 import type { z } from 'zod';
 import type { Context } from '../context.js';
 import type * as playwright from 'playwright';
 import type { ToolCapability } from '../../config.js';
 import type { Tab } from '../tab.js';
 import type { Response } from '../response.js';
-import type { ToolDefinition } from '../mcp/server.js';
+import type { ToolSchema } from '../mcp/tool.js';
 
 export type FileUploadModalState = {
   type: 'fileChooser';
@@ -38,28 +36,6 @@ export type DialogModalState = {
 
 export type ModalState = FileUploadModalState | DialogModalState;
 
-export type ToolSchema<Input extends z.Schema> = {
-  name: string;
-  title: string;
-  description: string;
-  inputSchema: Input;
-  type: 'readOnly' | 'destructive';
-};
-
-export function toToolDefinition(tool: ToolSchema<any>): ToolDefinition {
-  return {
-    name: tool.name,
-    description: tool.description,
-    inputSchema: zodToJsonSchema(tool.inputSchema, { strictUnions: true }) as ToolDefinition['inputSchema'],
-    annotations: {
-      title: tool.title,
-      readOnlyHint: tool.type === 'readOnly',
-      destructiveHint: tool.type === 'destructive',
-      openWorldHint: true,
-    },
-  };
-}
-
 export type Tool<Input extends z.Schema = z.Schema> = {
   capability: ToolCapability;
   schema: ToolSchema<Input>;

src/utils/fileUtils.ts

@@ -17,8 +17,6 @@
 import os from 'node:os';
 import path from 'node:path';
 
-import type { FullConfig } from '../config.js';
-
 export function cacheDir() {
   let cacheDirectory: string;
   if (process.platform === 'linux')
@@ -32,10 +30,6 @@ export function cacheDir() {
   return path.join(cacheDirectory, 'ms-playwright');
 }
 
-export async function userDataDir(browserConfig: FullConfig['browser']) {
-  return path.join(cacheDir(), 'ms-playwright', `mcp-${browserConfig.launchOptions?.channel ?? browserConfig?.browserName}-profile`);
-}
-
 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('.');

src/utils/guid.ts

@@ -16,6 +16,10 @@
 
 import crypto from 'crypto';
 
+export function createGuid(): string {
+  return crypto.randomBytes(16).toString('hex');
+}
+
 export function createHash(data: string): string {
   return crypto.createHash('sha256').update(data).digest('hex').slice(0, 7);
 }

src/vscode/host.ts

@@ -16,15 +16,15 @@
 import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
 import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
 import type { FullConfig } from '../config.js';
-import type { MCPFactory } from '../mcp/proxyBackend.js';
+import type { MCPProvider } from '../mcp/proxyBackend.js';
 
-export class VSCodeMCPFactory implements MCPFactory {
+export class VSCodeMCPFactory implements MCPProvider {
   name = 'vscode';
   description = 'Connect to a browser running in the Playwright VS Code extension';
 
   constructor(private readonly _config: FullConfig) {}
 
-  async create(options: any): Promise<Transport> {
+  async connect(options: any): Promise<Transport> {
     if (typeof options.connectionString !== 'string')
       throw new Error('Missing options.connectionString');
     if (typeof options.lib !== 'string')

tests/fixtures.ts

@@ -47,6 +47,7 @@ type TestFixtures = {
     args?: string[],
     config?: Config,
     roots?: { name: string, uri: string }[],
+    rootsResponseDelay?: number,
   }) => Promise<{ client: Client, stderr: () => string }>;
   wsEndpoint: string;
   cdpServer: CDPServer;
@@ -89,6 +90,8 @@ export const test = baseTest.extend<TestFixtures & TestOptions, WorkerFixtures>(
       client = new Client({ name: options?.clientName ?? 'test', version: '1.0.0' }, options?.roots ? { capabilities: { roots: {} } } : undefined);
       if (options?.roots) {
         client.setRequestHandler(ListRootsRequestSchema, async request => {
+          if (options.rootsResponseDelay)
+            await new Promise(resolve => setTimeout(resolve, options.rootsResponseDelay));
           return {
             roots: options.roots,
           };

tests/roots.spec.ts

@@ -50,7 +50,7 @@ for (const mode of ['default', 'proxy']) {
     });
 
 
-    test('check that trace is saved in workspace', async ({ startClient, server, mcpMode }, testInfo) => {
+    test('check that trace is saved in workspace', async ({ startClient, server }, testInfo) => {
       const rootPath = testInfo.outputPath('workspace');
       const { client } = await startClient({
         args: ['--save-trace', ...extraArgs],
@@ -73,5 +73,15 @@ for (const mode of ['default', 'proxy']) {
       const [file] = await fs.promises.readdir(path.join(rootPath, '.playwright-mcp'));
       expect(file).toContain('traces');
     });
+
+    test('should list all tools when listRoots is slow', async ({ startClient, server }, testInfo) => {
+      const { client } = await startClient({
+        clientName: 'Visual Studio Code', // Simulate VS Code client, roots only work with it
+        roots: [],
+        rootsResponseDelay: 1000,
+      });
+      const tools = await client.listTools();
+      expect(tools.tools.length).toBeGreaterThan(20);
+    });
   });
 }