diff --git a/packages/playwright-core/src/tools/backend/browserBackend.ts b/packages/playwright-core/src/tools/backend/browserBackend.ts index b8c8fb55d78f7..2fc6d82944af2 100644 --- a/packages/playwright-core/src/tools/backend/browserBackend.ts +++ b/packages/playwright-core/src/tools/backend/browserBackend.ts @@ -14,6 +14,8 @@ * limitations under the License. */ +import { EventEmitter } from 'events'; + import * as z from 'zod'; import debug from 'debug'; import { Context } from './context'; @@ -25,26 +27,35 @@ import type { Tool } from './tool'; import type * as mcpServer from '../utils/mcp/server'; import type { ClientInfo, ServerBackend } from '../utils/mcp/server'; -export class BrowserBackend implements ServerBackend { +export class BrowserBackend extends EventEmitter<{ disconnected: [], disposed: [] }> implements ServerBackend { private _tools: Tool[]; private _context: Context | undefined; private _sessionLog: SessionLog | undefined; private _config: ContextConfig; private _disconnected = false; - readonly browserContext: playwright.BrowserContext; + private _disposed = false; + private _browserContext: playwright.BrowserContext; + private _disposeCallback: (() => Promise) | undefined; - constructor(config: ContextConfig, browserContext: playwright.BrowserContext, tools: Tool[]) { + constructor(config: ContextConfig, browserContext: playwright.BrowserContext, tools: Tool[], disposeCallback?: () => Promise) { + super(); this._config = config; this._tools = tools; - this.browserContext = browserContext; - const markDisconnected = () => { this._disconnected = true; }; - this.browserContext.once('close', markDisconnected); - this.browserContext.browser()?.once('disconnected', markDisconnected); + this._browserContext = browserContext; + this._disposeCallback = disposeCallback; + const markDisconnected = () => { + if (this._disconnected) + return; + this._disconnected = true; + this.emit('disconnected'); + }; + this._browserContext.once('close', markDisconnected); + this._browserContext.browser()?.once('disconnected', markDisconnected); } async initialize(clientInfo: ClientInfo): Promise { this._sessionLog = this._config.saveSession ? await SessionLog.create(this._config, clientInfo.cwd) : undefined; - this._context = new Context(this.browserContext, { + this._context = new Context(this._browserContext, { config: this._config, sessionLog: this._sessionLog, cwd: clientInfo.cwd, @@ -52,10 +63,15 @@ export class BrowserBackend implements ServerBackend { } async dispose() { + if (this._disposed) + return; + this._disposed = true; await this._context?.dispose().catch(e => debug('pw:tools:error')(e)); + await this._disposeCallback?.().catch(e => debug('pw:tools:error')(e)); + this.emit('disposed'); } - async callTool(name: string, rawArguments: mcpServer.CallToolRequest['params']['arguments'] & { _meta?: Record } = {}, signal?: AbortSignal): Promise { + async callTool(name: string, rawArguments: mcpServer.CallToolRequest['params']['arguments'] & { _meta?: Record } = {}, signal?: AbortSignal): Promise { const json = !!rawArguments._meta?.json; const formatError = (message: string): mcpServer.CallToolResult => ({ content: [{ type: 'text' as const, text: json ? JSON.stringify({ isError: true, error: message }, null, 2) : `### Error\n${message}` }], @@ -90,8 +106,10 @@ export class BrowserBackend implements ServerBackend { } finally { context.setRunningTool(undefined); } - if (this._disconnected) - responseObject.isClose = true; + if (this._disconnected || responseObject.isClose) { + delete responseObject.isClose; + await this.dispose(); + } return responseObject; } } diff --git a/packages/playwright-core/src/tools/mcp/index.ts b/packages/playwright-core/src/tools/mcp/index.ts index a1a1c8c4b5a0c..d720e1c93890a 100644 --- a/packages/playwright-core/src/tools/mcp/index.ts +++ b/packages/playwright-core/src/tools/mcp/index.ts @@ -41,7 +41,6 @@ export async function createConnection(userConfig: Config = {}, contextGetter?: const context = config.browser.isolated ? await browser.newContext(config.browser.contextOptions) : browser.contexts()[0]; return new BrowserBackend(config, context, tools); }, - disposed: async () => { } }; return createServer('api', packageJSON.version, backendFactory, Promise.resolve(), false); } diff --git a/packages/playwright-core/src/tools/mcp/program.ts b/packages/playwright-core/src/tools/mcp/program.ts index 2643131dd085d..2c748848b31c7 100644 --- a/packages/playwright-core/src/tools/mcp/program.ts +++ b/packages/playwright-core/src/tools/mcp/program.ts @@ -129,25 +129,23 @@ export function decorateMCPCommand(command: Command) { await browser.bind(sessionName, { workspaceDir: clientInfo.cwd }); } const browserContext = config.browser.isolated ? await browser.newContext(config.browser.contextOptions) : browser.contexts()[0]; - return new BrowserBackend(config, browserContext, tools); - }, - disposed: async backend => { - clientCount--; - const browserContext = (backend as BrowserBackend).browserContext; + return new BrowserBackend(config, browserContext, tools, async () => { + clientCount--; - if (sharedBrowserPromise && clientCount > 0) { - if (config.browser.isolated) { - testDebug('close context'); - await browserContext.close().catch(() => { }); + if (sharedBrowserPromise && clientCount > 0) { + if (config.browser.isolated) { + testDebug('close context'); + await browserContext.close().catch(() => { }); + } + return; } - return; - } - testDebug('close browser'); - sharedBrowserPromise = undefined; - await browserContext.close().catch(() => { }); - await browserContext.browser()?.close().catch(() => { }); - } + testDebug('close browser'); + sharedBrowserPromise = undefined; + await browserContext.close().catch(() => { }); + await browserContext.browser()?.close().catch(() => { }); + }); + }, }; await mcpServer.start(factory, config.server); }); diff --git a/packages/playwright-core/src/tools/utils/mcp/server.ts b/packages/playwright-core/src/tools/utils/mcp/server.ts index 15cf6ebe1b5ed..ecdd1ae6ca11b 100644 --- a/packages/playwright-core/src/tools/utils/mcp/server.ts +++ b/packages/playwright-core/src/tools/utils/mcp/server.ts @@ -38,32 +38,11 @@ export type ClientInfo = { clientName: string; }; -class BackendManager { - private _backends = new Map(); - - async createBackend(factory: ServerBackendFactory, clientInfo: ClientInfo): Promise { - const backend = await factory.create(clientInfo); - await backend.initialize?.(clientInfo); - this._backends.set(backend, factory); - return backend; - } - - async disposeBackend(backend: ServerBackend) { - const factory = this._backends.get(backend); - if (!factory) - return; - await backend.dispose?.(); - await factory.disposed(backend).catch(serverDebug); - this._backends.delete(backend); - } -} - -const backendManager = new BackendManager(); - export interface ServerBackend { initialize?(clientInfo: ClientInfo): Promise; - callTool(name: string, args: CallToolRequest['params']['arguments'], signal: AbortSignal): Promise; + callTool(name: string, args: CallToolRequest['params']['arguments'], signal: AbortSignal): Promise; dispose?(): Promise; + once(event: 'disposed', listener: () => void): void; } export type ServerBackendFactory = { @@ -72,7 +51,6 @@ export type ServerBackendFactory = { version: string; toolSchemas: ToolSchema[]; create: (clientInfo: ClientInfo) => Promise; - disposed: (backend: ServerBackend) => Promise; }; export async function connect(factory: ServerBackendFactory, transport: Transport, transportInitialized: Promise, runHeartbeat: boolean) { @@ -94,7 +72,7 @@ export function createServer(name: string, version: string, factory: ServerBacke let backendPromise: Promise | undefined; - const onClose = () => backendPromise?.then(b => backendManager.disposeBackend(b)).catch(serverDebug); + const onClose = () => backendPromise?.then(b => b.dispose?.()).catch(serverDebug); addServerListener(server, 'close', onClose); server.setRequestHandler(CallToolRequestSchema, async (request, extra) => { @@ -102,7 +80,10 @@ export function createServer(name: string, version: string, factory: ServerBacke try { if (!backendPromise) { - backendPromise = initializeServer(server, factory, transportInitialized, runHeartbeat).catch(e => { + backendPromise = initializeServer(server, factory, transportInitialized, runHeartbeat).then(backend => { + backend.once('disposed', () => { backendPromise = undefined; }); + return backend; + }).catch(e => { backendPromise = undefined; throw e; }); @@ -110,12 +91,6 @@ export function createServer(name: string, version: string, factory: ServerBacke const backend = await backendPromise; const toolResult = await backend.callTool(request.params.name, request.params.arguments || {}, extra.signal); - if (toolResult.isClose) { - await backendManager.disposeBackend(backend).catch(serverDebug); - backendPromise = undefined; - delete toolResult.isClose; - } - const mergedResult = mergeTextParts(toolResult); serverDebugResponse('callResult', mergedResult); return mergedResult; @@ -146,7 +121,8 @@ const initializeServer = async (server: ServerType, factory: ServerBackendFactor clientName: server.getClientVersion()?.name ?? 'Playwright MCP', }; - const backend = await backendManager.createBackend(factory, clientInfo); + const backend = await factory.create(clientInfo); + await backend.initialize?.(clientInfo); if (runHeartbeat) startHeartbeat(server); return backend; diff --git a/packages/playwright/src/program.ts b/packages/playwright/src/program.ts index eea299a9db04f..2384f49a8852e 100644 --- a/packages/playwright/src/program.ts +++ b/packages/playwright/src/program.ts @@ -152,7 +152,6 @@ function addTestMCPServerCommand(program: Command) { version: packageJSON.version, toolSchemas: testServerBackendTools.map(tool => tool.schema), create: async () => new TestServerBackend(options.config, { muteConsole: options.port === undefined, headless: options.headless }), - disposed: async () => { } }; // TODO: add all options from mcp.startHttpServer. await tools.start(factory, { port: options.port === undefined ? undefined : +options.port, host: options.host });