chore: steer towards mcp types a bit (#880)

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 './tools/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[] {
+  async listTools(): Promise<mcpServer.Tool[]> {
-    return this._tools.map(tool => toToolDefinition(tool.schema));
+    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/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 '../tools/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[] {
+  async listTools(): Promise<mcpServer.Tool[]> {
-    return this._tools.map(tool => toToolDefinition(tool.schema));
+    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

@@ -22,7 +22,7 @@ import type { ToolSchema } from '../tools/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 MCPProvider = {
-
-export type MCPFactory = {
   name: string;
   description: string;
-  create(): Promise<Transport>;
+  connect(): 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 _contextSwitchTool: Tool;
-  private _tools: ToolDefinition[] = [];
+  private _roots: Root[] = [];
-  private _server: Server | undefined;
 
-  constructor(clientFactories: MCPFactoryList) {
+  constructor(mcpProviders: MCPProvider[]) {
-    this._mcpFactories = clientFactories;
+    this._mcpProviders = mcpProviders;
     this._contextSwitchTool = this._defineContextSwitchTool();
   }
 
   async initialize(server: Server): Promise<void> {
-    this._server = server;
+    const version = server.getClientVersion();
-    await this._setCurrentClient(this._mcpFactories[0]);
+    const capabilities = server.getClientCapabilities();
+    if (capabilities?.roots && version && clientsWithRoots.includes(version.name)) {
+      const { roots } = await server.listRoots();
+      this._roots = roots;
+    }
+
+    await this._setCurrentClient(this._mcpProviders[0]);
   }
 
-  tools(): ToolDefinition[] {
+  async listTools(): Promise<Tool[]> {
-    if (this._mcpFactories.length === 1)
+    const response = await this._currentClient!.listTools();
-      return this._tools;
+    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);
+      return this._callContextSwitchTool(args);
-    const result = await this._currentClient!.callTool({
+    return await this._currentClient!.callTool({
       name,
-      arguments: rawArguments,
+      arguments: args,
-    });
+    }) as CallToolResult;
-    return result as unknown as ToolResponse;
   }
 
   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,16 +100,16 @@ 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}`),
       ].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'),
-      }), { strictUnions: true }) as ToolDefinition['inputSchema'],
+      }), { strictUnions: true }) as Tool['inputSchema'],
       annotations: {
         title: 'Connect to a browser context',
         readOnlyHint: true,
@@ -116,7 +118,7 @@ export class ProxyBackend implements ServerBackend {
     };
   }
 
-  private async _setCurrentClient(factory: MCPFactory) {
+  private async _setCurrentClient(factory: MCPProvider) {
     await this._currentClient?.close();
     this._currentClient = undefined;
 
@@ -126,23 +128,13 @@ export class ProxyBackend implements ServerBackend {
         listRoots: true,
       },
     });
-    client.setRequestHandler(ListRootsRequestSchema, async () => {
+    client.setRequestHandler(ListRootsRequestSchema, () => ({ roots: this._roots }));
-      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(PingRequestSchema, () => ({}));
 
-    const transport = await factory.create();
+    const transport = await factory.connect();
     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[];
+  listTools(): Promise<Tool[]>;
-  callTool(name: string, rawArguments: any): Promise<ToolResponse>;
+  callTool(name: string, args: CallToolRequest['params']['arguments']): Promise<CallToolResult>;
   serverClosed?(): void;
 }
 
@@ -66,7 +54,7 @@ export function createServer(backend: ServerBackend, runHeartbeat: boolean): Ser
 
   server.setRequestHandler(ListToolsRequestSchema, async () => {
     serverDebug('listTools');
-    const tools = backend.tools();
+    const tools = await backend.listTools();
     return { tools };
   });
 
@@ -80,19 +68,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/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,12 +23,13 @@ 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 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)
@@ -78,18 +79,17 @@ program
         await mcpTransport.start(serverBackendFactory, config.server);
         return;
       }
+
       if (options.loopTools) {
         await runLoopTools(config);
         return;
       }
 
       const browserContextFactory = contextFactory(config);
-      const factories: MCPFactoryList = [
+      const providers: MCPProvider[] = [mcpProviderForBrowserContextFactory(config, browserContextFactory)];
-        new InProcessMCPFactory(browserContextFactory, config),
-      ];
       if (options.connectTool)
-        factories.push(new InProcessMCPFactory(createExtensionContextFactory(config), config));
+        providers.push(mcpProviderForBrowserContextFactory(config, createExtensionContextFactory(config)));
-      await mcpTransport.start(() => new ProxyBackend(factories), config.server);
+      await mcpTransport.start(() => new ProxyBackend(providers), config.server);
     });
 
 function setupExitWatchdog() {
@@ -112,4 +112,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

@@ -22,7 +22,7 @@ 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 * as mcpServer from '../mcp/server.js';
 
 export type FileUploadModalState = {
   type: 'fileChooser';
@@ -46,11 +46,11 @@ export type ToolSchema<Input extends z.Schema> = {
   type: 'readOnly' | 'destructive';
 };
 
-export function toToolDefinition(tool: ToolSchema<any>): ToolDefinition {
+export function toMcpTool(tool: ToolSchema<any>): mcpServer.Tool {
   return {
     name: tool.name,
     description: tool.description,
-    inputSchema: zodToJsonSchema(tool.inputSchema, { strictUnions: true }) as ToolDefinition['inputSchema'],
+    inputSchema: zodToJsonSchema(tool.inputSchema, { strictUnions: true }) as mcpServer.Tool['inputSchema'],
     annotations: {
       title: tool.title,
       readOnlyHint: tool.type === 'readOnly',

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