chore(extension): exit gracefully when waiting for extension connection (#754)

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

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<void>;
-  private _extensionConnectionResolve: (() => void) | null = null;
+  private _extensionConnectionPromise!: ManualPromise<void>;
 
   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);
 }

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

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