diff --git a/src/program.ts b/src/program.ts index 35f8c08..221672a 100644 --- a/src/program.ts +++ b/src/program.ts @@ -25,8 +25,7 @@ import { ProxyBackend } from './mcp/proxyBackend.js'; import { BrowserServerBackend } from './browserServerBackend.js'; import { ExtensionContextFactory } from './extension/extensionContextFactory.js'; -import { VSCodeMCPFactory } from './vscode/host.js'; -import { VSCodeProxyBackend } from './vscode/proxyBackend.js'; +import { VSCodeProxyBackend } from './vscode/host.js'; import type { MCPProvider } from './mcp/proxyBackend.js'; program @@ -86,6 +85,17 @@ program return; } + if (options.vscode) { + const serverBackendFactory: mcpServer.ServerBackendFactory = { + name: 'Playwright w/ vscode', + nameInConfig: 'playwright-vscode', + version: packageJSON.version, + create: () => new VSCodeProxyBackend(config, () => mcpServer.wrapInProcess(new BrowserServerBackend(config, browserContextFactory))) + }; + await mcpServer.start(serverBackendFactory, config.server); + return; + } + if (options.loopTools) { await runLoopTools(config); return; diff --git a/src/vscode/host.ts b/src/vscode/host.ts index 1774cb4..e22d6c7 100644 --- a/src/vscode/host.ts +++ b/src/vscode/host.ts @@ -13,27 +13,120 @@ * See the License for the specific language governing permissions and * limitations under the License. */ + +import { z } from 'zod'; +import { zodToJsonSchema } from 'zod-to-json-schema'; + +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { ListRootsRequestSchema, PingRequestSchema } from '@modelcontextprotocol/sdk/types.js'; import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'; +import { logUnhandledError } from '../utils/log.js'; +import { packageJSON } from '../utils/package.js'; + +import { FullConfig } from '../config.js'; import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'; -import type { FullConfig } from '../config.js'; -import type { MCPProvider } from './proxyBackend.js'; +import type { ClientVersion, ServerBackend } from '../mcp/server.js'; +import type { Root, Tool, CallToolResult, CallToolRequest } from '@modelcontextprotocol/sdk/types.js'; -export class VSCodeMCPFactory implements MCPProvider { - name = 'vscode'; - description = 'Connect to a browser running in the Playwright VS Code extension'; +const contextSwitchOptions = z.object({ + connectionString: z.string().optional().describe('The connection string to use to connect to the browser'), + lib: z.string().optional().describe('The library to use for the connection'), +}); - constructor(private readonly _config: FullConfig, private readonly _connectionString: string, private readonly _lib: string) {} +export class VSCodeProxyBackend implements ServerBackend { + name = 'Playwright MCP Client Switcher'; + version = packageJSON.version; - async connect(): Promise { - return new StdioClientTransport({ - command: process.execPath, - cwd: process.cwd(), - args: [ - new URL('./main.js', import.meta.url).pathname, - JSON.stringify(this._config), - this._connectionString, - this._lib, - ], + private _currentClient: Client | undefined; + private _contextSwitchTool: Tool; + private _roots: Root[] = []; + private _clientVersion?: ClientVersion; + + constructor(private readonly _config: FullConfig, private readonly _defaultTransportFactory: () => Promise) { + this._contextSwitchTool = this._defineContextSwitchTool(); + } + + async initialize(clientVersion: ClientVersion, roots: Root[]): Promise { + this._clientVersion = clientVersion; + this._roots = roots; + const transport = await this._defaultTransportFactory(); + await this._setCurrentClient(transport); + } + + async listTools(): Promise { + const response = await this._currentClient!.listTools(); + return [ + ...response.tools, + this._contextSwitchTool, + ]; + } + + async callTool(name: string, args: CallToolRequest['params']['arguments']): Promise { + if (name === this._contextSwitchTool.name) + return this._callContextSwitchTool(args as any); + return await this._currentClient!.callTool({ + name, + arguments: args, + }) as CallToolResult; + } + + serverClosed?(): void { + void this._currentClient?.close().catch(logUnhandledError); + } + + private async _callContextSwitchTool(params: z.infer): Promise { + if (!params.connectionString || !params.lib) { + const transport = await this._defaultTransportFactory(); + await this._setCurrentClient(transport); + return { + content: [{ type: 'text', text: '### Result\nSuccessfully disconnected.\n' }], + }; + } + + await this._setCurrentClient( + new StdioClientTransport({ + command: process.execPath, + cwd: process.cwd(), + args: [ + new URL('./main.js', import.meta.url).pathname, + JSON.stringify(this._config), + params.connectionString, + params.lib, + ], + }) + ); + return { + content: [{ type: 'text', text: '### Result\nSuccessfully connected.\n' }], + }; + } + + private _defineContextSwitchTool(): Tool { + return { + name: 'browser_connect', + description: 'Connect to a browser running in VS Code.', + inputSchema: zodToJsonSchema(contextSwitchOptions, { strictUnions: true }) as Tool['inputSchema'], + annotations: { + title: 'Connect to a browser running in VS Code.', + readOnlyHint: true, + openWorldHint: false, + }, + }; + } + + private async _setCurrentClient(transport: Transport) { + await this._currentClient?.close(); + this._currentClient = undefined; + + const client = new Client(this._clientVersion!); + client.registerCapabilities({ + roots: { + listRoots: true, + }, }); + client.setRequestHandler(ListRootsRequestSchema, () => ({ roots: this._roots })); + client.setRequestHandler(PingRequestSchema, () => ({})); + + await client.connect(transport); + this._currentClient = client; } } diff --git a/src/vscode/proxyBackend.ts b/src/vscode/proxyBackend.ts deleted file mode 100644 index fe63136..0000000 --- a/src/vscode/proxyBackend.ts +++ /dev/null @@ -1,139 +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 { z } from 'zod'; -import { zodToJsonSchema } from 'zod-to-json-schema'; - -import { Client } from '@modelcontextprotocol/sdk/client/index.js'; -import { ListRootsRequestSchema, PingRequestSchema } from '@modelcontextprotocol/sdk/types.js'; -import { logUnhandledError } from '../utils/log.js'; -import { packageJSON } from '../utils/package.js'; - - -import { VSCodeMCPFactory } from './host.js'; -import { FullConfig } from '../config.js'; -import type { Server } from '@modelcontextprotocol/sdk/server/index.js'; -import type { ServerBackend } from '../mcp/server.js'; -import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'; -import type { Root, Tool, CallToolResult, CallToolRequest } from '@modelcontextprotocol/sdk/types.js'; - -export type MCPProvider = { - name: string; - description: string; - connect(): Promise; -}; - -const contextSwitchOptions = z.object({ - connectionString: z.string().optional().describe('The connection string to use to connect to the browser'), - lib: z.string().optional().describe('The library to use for the connection'), -}); - -export class VSCodeProxyBackend implements ServerBackend { - name = 'Playwright MCP Client Switcher'; - version = packageJSON.version; - - private _currentClient: Client | undefined; - private _contextSwitchTool: Tool; - private _roots: Root[] = []; - - constructor(private readonly _config: FullConfig, private readonly _defaultProvider: MCPProvider) { - this._contextSwitchTool = this._defineContextSwitchTool(); - } - - async initialize(server: Server): Promise { - const version = server.getClientVersion(); - const capabilities = server.getClientCapabilities(); - if (capabilities?.roots && version && clientsWithRoots.includes(version.name)) { - const { roots } = await server.listRoots(); - this._roots = roots; - } - - await this._setCurrentClient(this._defaultProvider); - } - - async listTools(): Promise { - const response = await this._currentClient!.listTools(); - return [ - ...response.tools, - this._contextSwitchTool, - ]; - } - - async callTool(name: string, args: CallToolRequest['params']['arguments']): Promise { - if (name === this._contextSwitchTool.name) - return this._callContextSwitchTool(args as any); - return await this._currentClient!.callTool({ - name, - arguments: args, - }) as CallToolResult; - } - - serverClosed?(): void { - void this._currentClient?.close().catch(logUnhandledError); - } - - private async _callContextSwitchTool(params: z.infer): Promise { - if (!params.connectionString || !params.lib) { - await this._setCurrentClient(this._defaultProvider); - return { - content: [{ type: 'text', text: '### Result\nSuccessfully disconnected.\n' }], - }; - } - - await this._setCurrentClient(new VSCodeMCPFactory(this._config, params.connectionString, params.lib)); - return { - content: [{ type: 'text', text: '### Result\nSuccessfully connected.\n' }], - }; - } - - private _defineContextSwitchTool(): Tool { - return { - name: 'browser_connect', - description: 'Connect to a browser running in VS Code.', - inputSchema: zodToJsonSchema(contextSwitchOptions, { strictUnions: true }) as Tool['inputSchema'], - annotations: { - title: 'Connect to a browser running in VS Code.', - readOnlyHint: true, - openWorldHint: false, - }, - }; - } - - private async _setCurrentClient(factory: MCPProvider) { - await this._currentClient?.close(); - this._currentClient = undefined; - - const client = new Client({ name: 'Playwright MCP Proxy', version: packageJSON.version }); - client.registerCapabilities({ - roots: { - listRoots: true, - }, - }); - client.setRequestHandler(ListRootsRequestSchema, () => ({ roots: this._roots })); - client.setRequestHandler(PingRequestSchema, () => ({})); - - const transport = await factory.connect(); - await client.connect(transport); - this._currentClient = client; - } -} - -const clientsWithRoots = ['Visual Studio Code', 'Visual Studio Code - Insiders']; - -export async function runLoopTools(config: FullConfig) { - const serverBackendFactory = () => new LoopToolsServerBackend(config); - await mcpTransport.start(serverBackendFactory, config.server); -}