From 5a0cfb9e65a377ea8ad72e4db3682128d653d447 Mon Sep 17 00:00:00 2001
From: Simon Knott <info@simonknott.de>
Date: Wed, 13 Aug 2025 11:26:01 +0200
Subject: [PATCH] address

---
 src/vscode/host.ts |  2 +-
 src/vscode/main.ts | 46 +++++++++++++++++++++++++++++-----------------
 2 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/src/vscode/host.ts b/src/vscode/host.ts
index 1ceb815..76a80bc 100644
--- a/src/vscode/host.ts
+++ b/src/vscode/host.ts
@@ -28,7 +28,7 @@ export class VSCodeMCPFactory implements MCPFactory {
     if (typeof options.connectionString !== 'string')
       throw new Error('Missing options.connectionString');
     if (typeof options.lib !== 'string')
-      throw new Error('Missing options.library');
+      throw new Error('Missing options.lib');
 
     return new StdioClientTransport({
       command: process.execPath,
diff --git a/src/vscode/main.ts b/src/vscode/main.ts
index 4e44a6e..65966f3 100644
--- a/src/vscode/main.ts
+++ b/src/vscode/main.ts
@@ -21,26 +21,28 @@ import * as mcpServer from '../mcp/server.js';
 import { BrowserServerBackend } from '../browserServerBackend.js';
 import { BrowserContextFactory, ClientInfo } from '../browserContextFactory.js';
 
-const config: FullConfig = JSON.parse(process.argv[2]);
-const connectionString = new URL(process.argv[3]);
-const lib = process.argv[4];
-
-const playwright = await import(lib).then(mod => mod.default ?? mod) as typeof import('playwright');
-
 class VSCodeBrowserContextFactory implements BrowserContextFactory {
   name = 'vscode';
   description = 'Connect to a browser running in the Playwright VS Code extension';
 
+  constructor(private _config: FullConfig, private _playwright: typeof import('playwright'), private _connectionString: string) {}
+
   async createContext(clientInfo: ClientInfo, abortSignal: AbortSignal): Promise<{ browserContext: BrowserContext; close: () => Promise<void>; }> {
-    connectionString.searchParams.set('launch-options', JSON.stringify({
-      ...config.browser.launchOptions,
-      ...config.browser.contextOptions,
-      userDataDir: config.browser.userDataDir,
-    }));
+    let launchOptions: any = this._config.browser.launchOptions;
+    if (this._config.browser.userDataDir) {
+      launchOptions = {
+        ...launchOptions,
+        ...this._config.browser.contextOptions,
+        userDataDir: this._config.browser.userDataDir,
+      };
+    }
+    const connectionString = new URL(this._connectionString);
+    connectionString.searchParams.set('launch-options', JSON.stringify(launchOptions));
 
-    const browser = await playwright.chromium.connect(connectionString.toString());
+    const browserType = this._playwright.chromium; // it could also be firefox or webkit, we just need some browser type to call `connect` on
+    const browser = await browserType.connect(connectionString.toString());
 
-    const context = browser.contexts()[0] ?? await browser.newContext(config.browser.contextOptions);
+    const context = browser.contexts()[0] ?? await browser.newContext(this._config.browser.contextOptions);
 
     return {
       browserContext: context,
@@ -51,8 +53,18 @@ class VSCodeBrowserContextFactory implements BrowserContextFactory {
   }
 }
 
-await mcpServer.connect(
-    () => new BrowserServerBackend(config, new VSCodeBrowserContextFactory()),
-    new StdioServerTransport(),
-    false
+async function main(config: FullConfig, connectionString: string, lib: string) {
+  const playwright = await import(lib).then(mod => mod.default ?? mod);
+  const factory = new VSCodeBrowserContextFactory(config, playwright, connectionString);
+  await mcpServer.connect(
+      () => new BrowserServerBackend(config, factory),
+      new StdioServerTransport(),
+      false
+  );
+}
+
+await main(
+    JSON.parse(process.argv[2]),
+    process.argv[3],
+    process.argv[4]
 );