From e153ac3b7c0f9053b8f7a547f2a772df3b7c0332 Mon Sep 17 00:00:00 2001 From: Yury Semikhatsky Date: Thu, 24 Jul 2025 16:02:02 -0700 Subject: [PATCH] chore(extension): exit gracefully when waiting for extension connection (#754) --- src/context.ts | 3 ++- src/extension/cdpRelay.ts | 21 ++++++++++----------- src/extension/main.ts | 4 ++-- src/program.ts | 9 +++++++-- 4 files changed, 21 insertions(+), 16 deletions(-) diff --git a/src/context.ts b/src/context.ts index c8377e2..5165c8c 100644 --- a/src/context.ts +++ b/src/context.ts @@ -17,6 +17,7 @@ import debug from 'debug'; import * as playwright from 'playwright'; +import { logUnhandledError } from './log.js'; import { Tab } from './tab.js'; import type { Tool } from './tools/tool.js'; @@ -140,7 +141,7 @@ export class Context { async closeBrowserContext() { if (!this._closeBrowserContextPromise) - this._closeBrowserContextPromise = this._closeBrowserContextImpl(); + this._closeBrowserContextPromise = this._closeBrowserContextImpl().catch(logUnhandledError); await this._closeBrowserContextPromise; this._closeBrowserContextPromise = undefined; } diff --git a/src/extension/cdpRelay.ts b/src/extension/cdpRelay.ts index 022cba1..7b28b7c 100644 --- a/src/extension/cdpRelay.ts +++ b/src/extension/cdpRelay.ts @@ -30,6 +30,8 @@ import * as playwright from 'playwright'; // @ts-ignore const { registry } = await import('playwright-core/lib/server/registry/index'); import { httpAddressToString, startHttpServer } from '../httpServer.js'; +import { logUnhandledError } from '../log.js'; +import { ManualPromise } from '../manualPromise.js'; import type { BrowserContextFactory } from '../browserContextFactory.js'; import type websocket from 'ws'; @@ -65,8 +67,7 @@ export class CDPRelayServer { sessionId: string; } | undefined; private _nextSessionId: number = 1; - private _extensionConnectionPromise: Promise; - private _extensionConnectionResolve: (() => void) | null = null; + private _extensionConnectionPromise!: ManualPromise; constructor(server: http.Server, browserChannel: string) { this._wsHost = httpAddressToString(server.address()).replace(/^http/, 'ws'); @@ -76,9 +77,7 @@ export class CDPRelayServer { this._cdpPath = `/cdp/${uuid}`; this._extensionPath = `/extension/${uuid}`; - this._extensionConnectionPromise = new Promise(resolve => { - this._extensionConnectionResolve = resolve; - }); + this._resetExtensionConnection(); this._wss = new WebSocketServer({ server }); this._wss.on('connection', this._onConnection.bind(this)); } @@ -166,15 +165,15 @@ export class CDPRelayServer { private _closeExtensionConnection(reason: string) { this._extensionConnection?.close(reason); + this._extensionConnectionPromise.reject(new Error(reason)); this._resetExtensionConnection(); } private _resetExtensionConnection() { this._connectedTabInfo = undefined; this._extensionConnection = null; - this._extensionConnectionPromise = new Promise(resolve => { - this._extensionConnectionResolve = resolve; - }); + this._extensionConnectionPromise = new ManualPromise(); + void this._extensionConnectionPromise.catch(logUnhandledError); } private _closePlaywrightConnection(reason: string) { @@ -197,7 +196,7 @@ export class CDPRelayServer { this._closePlaywrightConnection(`Extension disconnected: ${reason}`); }; this._extensionConnection.onmessage = this._handleExtensionMessage.bind(this); - this._extensionConnectionResolve?.(); + this._extensionConnectionPromise.resolve(); } private _handleExtensionMessage(method: string, params: any) { @@ -323,10 +322,10 @@ class ExtensionContextFactory implements BrowserContextFactory { } } -export async function startCDPRelayServer(browserChannel: string) { +export async function startCDPRelayServer(browserChannel: string, abortController: AbortController) { const httpServer = await startHttpServer({}); const cdpRelayServer = new CDPRelayServer(httpServer, browserChannel); - process.on('exit', () => cdpRelayServer.stop()); + abortController.signal.addEventListener('abort', () => cdpRelayServer.stop()); debugLogger(`CDP relay server started, extension endpoint: ${cdpRelayServer.extensionEndpoint()}.`); return new ExtensionContextFactory(cdpRelayServer); } diff --git a/src/extension/main.ts b/src/extension/main.ts index cfabbad..760aef6 100644 --- a/src/extension/main.ts +++ b/src/extension/main.ts @@ -20,8 +20,8 @@ import * as mcpTransport from '../mcp/transport.js'; import type { FullConfig } from '../config.js'; -export async function runWithExtension(config: FullConfig) { - const contextFactory = await startCDPRelayServer(config.browser.launchOptions.channel || 'chrome'); +export async function runWithExtension(config: FullConfig, abortController: AbortController) { + const contextFactory = await startCDPRelayServer(config.browser.launchOptions.channel || 'chrome', abortController); const serverBackendFactory = () => new BrowserServerBackend(config, contextFactory); await mcpTransport.start(serverBackendFactory, config.server); } diff --git a/src/program.ts b/src/program.ts index a34205c..e9f4bc8 100644 --- a/src/program.ts +++ b/src/program.ts @@ -57,7 +57,7 @@ program .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('--vision', 'Legacy option, use --caps=vision instead').hideHelp()) .action(async options => { - setupExitWatchdog(); + const abortController = setupExitWatchdog(); if (options.vision) { // eslint-disable-next-line no-console @@ -67,7 +67,7 @@ program const config = await resolveCLIConfig(options); if (options.extension) { - await runWithExtension(config); + await runWithExtension(config, abortController); return; } @@ -85,12 +85,15 @@ program }); function setupExitWatchdog() { + const abortController = new AbortController(); + let isExiting = false; const handleExit = async () => { if (isExiting) return; isExiting = true; setTimeout(() => process.exit(0), 15000); + abortController.abort('Process exiting'); await Context.disposeAll(); process.exit(0); }; @@ -98,6 +101,8 @@ function setupExitWatchdog() { process.stdin.on('close', handleExit); process.on('SIGINT', handleExit); process.on('SIGTERM', handleExit); + + return abortController; } void program.parseAsync(process.argv);