Add initial MCP server (#261898)

test/mcp/src/inMemoryEventStore.ts

@@ -0,0 +1,81 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+import { EventStore } from '@modelcontextprotocol/sdk/server/streamableHttp';
+import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types';
+
+/**
+ * Simple in-memory implementation of the EventStore interface for resumability
+ * This is primarily intended for examples and testing, not for production use
+ * where a persistent storage solution would be more appropriate.
+ */
+export class InMemoryEventStore implements EventStore {
+	private events: Map<string, { streamId: string; message: JSONRPCMessage }> = new Map();
+
+	/**
+	 * Generates a unique event ID for a given stream ID
+	 */
+	private generateEventId(streamId: string): string {
+		return `${streamId}_${Date.now()}_${Math.random().toString(36).substring(2, 10)}`;
+	}
+
+	/**
+	 * Extracts the stream ID from an event ID
+	 */
+	private getStreamIdFromEventId(eventId: string): string {
+		const parts = eventId.split('_');
+		return parts.length > 0 ? parts[0] : '';
+	}
+
+	/**
+	 * Stores an event with a generated event ID
+	 * Implements EventStore.storeEvent
+	 */
+	async storeEvent(streamId: string, message: JSONRPCMessage): Promise<string> {
+		const eventId = this.generateEventId(streamId);
+		this.events.set(eventId, { streamId, message });
+		return eventId;
+	}
+
+	/**
+	 * Replays events that occurred after a specific event ID
+	 * Implements EventStore.replayEventsAfter
+	 */
+	async replayEventsAfter(lastEventId: string,
+		{ send }: { send: (eventId: string, message: JSONRPCMessage) => Promise<void> }
+	): Promise<string> {
+		if (!lastEventId || !this.events.has(lastEventId)) {
+			return '';
+		}
+
+		// Extract the stream ID from the event ID
+		const streamId = this.getStreamIdFromEventId(lastEventId);
+		if (!streamId) {
+			return '';
+		}
+
+		let foundLastEvent = false;
+
+		// Sort events by eventId for chronological ordering
+		const sortedEvents = [...this.events.entries()].sort((a, b) => a[0].localeCompare(b[0]));
+
+		for (const [eventId, { streamId: eventStreamId, message }] of sortedEvents) {
+			// Only include events from the same stream
+			if (eventStreamId !== streamId) {
+				continue;
+			}
+
+			// Start sending events after we find the lastEventId
+			if (eventId === lastEventId) {
+				foundLastEvent = true;
+				continue;
+			}
+
+			if (foundLastEvent) {
+				await send(eventId, message);
+			}
+		}
+		return streamId;
+	}
+}

test/mcp/src/main.ts

@@ -0,0 +1,186 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+import * as express from 'express';
+import { Request, Response } from 'express';
+import * as cors from 'cors';
+import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
+import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js';
+import { randomUUID } from 'node:crypto';
+import { InMemoryEventStore } from './inMemoryEventStore';
+import { getServer } from './playwright';
+
+const MCP_PORT = 33418;
+
+const app = express();
+app.use(express.json());
+
+// Allow CORS all domains, expose the Mcp-Session-Id header
+app.use(cors({
+	origin: '*', // Allow all origins
+	exposedHeaders: ['Mcp-Session-Id']
+}));
+
+// Map to store transports by session ID
+const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {};
+
+// MCP POST endpoint with optional auth
+const mcpPostHandler = async (req: Request, res: Response) => {
+	const sessionId = req.headers['mcp-session-id'] as string | undefined;
+	if (sessionId) {
+		console.log(`Received MCP request for session: ${sessionId}`);
+	} else {
+		console.log('Request body:', req.body);
+	}
+
+	try {
+		let transport: StreamableHTTPServerTransport;
+		if (sessionId && transports[sessionId]) {
+			// Reuse existing transport
+			transport = transports[sessionId];
+		} else if (!sessionId && isInitializeRequest(req.body)) {
+			// New initialization request
+			const eventStore = new InMemoryEventStore();
+			transport = new StreamableHTTPServerTransport({
+				sessionIdGenerator: () => randomUUID(),
+				eventStore, // Enable resumability
+				onsessioninitialized: (sessionId: any) => {
+					// Store the transport by session ID when session is initialized
+					// This avoids race conditions where requests might come in before the session is stored
+					console.log(`Session initialized with ID: ${sessionId}`);
+					transports[sessionId] = transport;
+				}
+			});
+
+			// Set up onclose handler to clean up transport when closed
+			transport.onclose = () => {
+				const sid = transport.sessionId;
+				if (sid && transports[sid]) {
+					console.log(`Transport closed for session ${sid}, removing from transports map`);
+					delete transports[sid];
+				}
+			};
+
+			// Connect the transport to the MCP server BEFORE handling the request
+			// so responses can flow back through the same transport
+			// NOTE: the shape of this is very similar to a Server but not quite
+			// I think it's because it's electron while playwright expects browser.
+			// TODO: Fix that.
+			const server = await getServer() as any;
+			server.onclose = () => {
+				const sid = transport.sessionId;
+				if (sid && transports[sid]) {
+					console.log(`Transport closed for session ${sid}, removing from transports map`);
+					delete transports[sid];
+				}
+			};
+			await server.connect(transport);
+
+			await transport.handleRequest(req, res, req.body);
+			return; // Already handled
+		} else {
+			// Invalid request - no session ID or not initialization request
+			res.status(400).json({
+				jsonrpc: '2.0',
+				error: {
+					code: -32000,
+					message: 'Bad Request: No valid session ID provided',
+				},
+				id: null,
+			});
+			return;
+		}
+
+		// Handle the request with existing transport - no need to reconnect
+		// The existing transport is already connected to the server
+		await transport.handleRequest(req, res, req.body);
+	} catch (error) {
+		console.error('Error handling MCP request:', error);
+		if (!res.headersSent) {
+			res.status(500).json({
+				jsonrpc: '2.0',
+				error: {
+					code: -32603,
+					message: 'Internal server error',
+				},
+				id: null,
+			});
+		}
+	}
+};
+
+// Set up routes with conditional auth middleware
+app.post('/mcp', mcpPostHandler);
+
+// Handle GET requests for SSE streams (using built-in support from StreamableHTTP)
+const mcpGetHandler = async (req: Request, res: Response) => {
+	const sessionId = req.headers['mcp-session-id'] as string | undefined;
+	if (!sessionId || !transports[sessionId]) {
+		res.status(400).send('Invalid or missing session ID');
+		return;
+	}
+
+	// Check for Last-Event-ID header for resumability
+	const lastEventId = req.headers['last-event-id'] as string | undefined;
+	if (lastEventId) {
+		console.log(`Client reconnecting with Last-Event-ID: ${lastEventId}`);
+	} else {
+		console.log(`Establishing new SSE stream for session ${sessionId}`);
+	}
+
+	const transport = transports[sessionId];
+	await transport.handleRequest(req, res);
+};
+
+// Set up GET route with conditional auth middleware
+app.get('/mcp', mcpGetHandler);
+
+// Handle DELETE requests for session termination (according to MCP spec)
+const mcpDeleteHandler = async (req: Request, res: Response) => {
+	const sessionId = req.headers['mcp-session-id'] as string | undefined;
+	if (!sessionId || !transports[sessionId]) {
+		res.status(400).send('Invalid or missing session ID');
+		return;
+	}
+
+	console.log(`Received session termination request for session ${sessionId}`);
+
+	try {
+		const transport = transports[sessionId];
+		await transport.handleRequest(req, res);
+	} catch (error) {
+		console.error('Error handling session termination:', error);
+		if (!res.headersSent) {
+			res.status(500).send('Error processing session termination');
+		}
+	}
+};
+
+app.delete('/mcp', mcpDeleteHandler);
+
+app.listen(MCP_PORT, (error: any) => {
+	if (error) {
+		console.error('Failed to start server:', error);
+		process.exit(1);
+	}
+	console.log(`MCP available at http://localhost:${MCP_PORT}/mcp`);
+});
+
+// Handle server shutdown
+process.on('SIGINT', async () => {
+	console.log('Shutting down server...');
+
+	// Close all active transports to properly clean up resources
+	for (const sessionId in transports) {
+		try {
+			console.log(`Closing transport for session ${sessionId}`);
+			await transports[sessionId].close();
+			delete transports[sessionId];
+		} catch (error) {
+			console.error(`Error closing transport for session ${sessionId}:`, error);
+		}
+	}
+	console.log('Server shutdown complete');
+	process.exit(0);
+});

test/mcp/src/playwright.ts

@@ -0,0 +1,169 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { createConnection } from '@playwright/mcp';
+import { getDevElectronPath, Quality, ConsoleLogger, FileLogger, Logger, MultiLogger } from '../../automation';
+import * as path from 'path';
+import * as fs from 'fs';
+import * as cp from 'child_process';
+import * as os from 'os';
+import { createApp, parseVersion } from './utils';
+import * as minimist from 'minimist';
+
+const rootPath = path.join(__dirname, '..', '..', '..');
+
+const [, , ...args] = process.argv;
+const opts = minimist(args, {
+	string: [
+		'browser',
+		'build',
+		'stable-build',
+		'wait-time',
+		'test-repo',
+		'electronArgs'
+	],
+	boolean: [
+		'verbose',
+		'remote',
+		'web',
+		'headless',
+		'tracing'
+	],
+	default: {
+		verbose: false
+	}
+}) as {
+	verbose?: boolean;
+	remote?: boolean;
+	headless?: boolean;
+	web?: boolean;
+	tracing?: boolean;
+	build?: string;
+	'stable-build'?: string;
+	browser?: 'chromium' | 'webkit' | 'firefox' | 'chromium-msedge' | 'chromium-chrome' | undefined;
+	electronArgs?: string;
+};
+
+const testDataPath = path.join(os.tmpdir(), 'vscsmoke');
+if (fs.existsSync(testDataPath)) {
+	fs.rmSync(testDataPath, { recursive: true, force: true, maxRetries: 10 });
+}
+fs.mkdirSync(testDataPath, { recursive: true });
+process.once('exit', () => {
+	try {
+		fs.rmSync(testDataPath, { recursive: true, force: true, maxRetries: 10 });
+	} catch {
+		// noop
+	}
+});
+
+const testRepoUrl = 'https://github.com/microsoft/vscode-smoketest-express';
+const workspacePath = path.join(testDataPath, 'vscode-smoketest-express');
+const extensionsPath = path.join(testDataPath, 'extensions-dir');
+fs.mkdirSync(extensionsPath, { recursive: true });
+
+const logsRootPath = (() => {
+	const logsParentPath = path.join(rootPath, '.build', 'logs');
+
+	let logsName: string;
+	if (opts.web) {
+		logsName = 'mcp-browser';
+	} else if (opts.remote) {
+		logsName = 'mcp-remote';
+	} else {
+		logsName = 'mcp-electron';
+	}
+
+	return path.join(logsParentPath, logsName);
+})();
+
+const crashesRootPath = (() => {
+	const crashesParentPath = path.join(rootPath, '.build', 'crashes');
+
+	let crashesName: string;
+	if (opts.web) {
+		crashesName = 'mcp-browser';
+	} else if (opts.remote) {
+		crashesName = 'mcp-remote';
+	} else {
+		crashesName = 'mcp-electron';
+	}
+
+	return path.join(crashesParentPath, crashesName);
+})();
+
+const logger = createLogger();
+
+function createLogger(): Logger {
+	const loggers: Logger[] = [];
+
+	// Log to console if verbose
+	if (opts.verbose) {
+		loggers.push(new ConsoleLogger());
+	}
+
+	// Prepare logs rot path
+	fs.rmSync(logsRootPath, { recursive: true, force: true, maxRetries: 3 });
+	fs.mkdirSync(logsRootPath, { recursive: true });
+
+	// Always log to log file
+	loggers.push(new FileLogger(path.join(logsRootPath, 'smoke-test-runner.log')));
+
+	return new MultiLogger(loggers);
+}
+
+async function setupRepository(): Promise<void> {
+	if (!fs.existsSync(workspacePath)) {
+		logger.log('Cloning test project repository...');
+		const res = cp.spawnSync('git', ['clone', testRepoUrl, workspacePath], { stdio: 'inherit' });
+		if (!fs.existsSync(workspacePath)) {
+			throw new Error(`Clone operation failed: ${res.stderr.toString()}`);
+		}
+	} else {
+		logger.log('Cleaning test project repository...');
+		cp.spawnSync('git', ['fetch'], { cwd: workspacePath, stdio: 'inherit' });
+		cp.spawnSync('git', ['reset', '--hard', 'FETCH_HEAD'], { cwd: workspacePath, stdio: 'inherit' });
+		cp.spawnSync('git', ['clean', '-xdf'], { cwd: workspacePath, stdio: 'inherit' });
+	}
+}
+
+export async function getServer() {
+	const testCodePath = getDevElectronPath();
+	const electronPath = testCodePath;
+	if (!fs.existsSync(electronPath || '')) {
+		throw new Error(`Cannot find VSCode at ${electronPath}. Please run VSCode once first (scripts/code.sh, scripts\\code.bat) and try again.`);
+	}
+	process.env.VSCODE_REPOSITORY = rootPath;
+	process.env.VSCODE_DEV = '1';
+	process.env.VSCODE_CLI = '1';
+	const quality = Quality.Dev;
+	const userDataDir = path.join(testDataPath, 'd');
+
+	await setupRepository();
+	const application = createApp({
+		quality,
+		version: parseVersion('0.0.0'),
+		codePath: opts.build,
+		workspacePath,
+		userDataDir,
+		extensionsPath,
+		logger,
+		logsPath: path.join(logsRootPath, 'suite_unknown'),
+		crashesPath: path.join(crashesRootPath, 'suite_unknown'),
+		verbose: opts.verbose,
+		remote: opts.remote,
+		web: opts.web,
+		tracing: opts.tracing,
+		headless: opts.headless,
+		browser: opts.browser,
+		extraArgs: (opts.electronArgs || '').split(' ').map(arg => arg.trim()).filter(arg => !!arg),
+	}, opts => ({ ...opts, userDataDir: path.join(opts.userDataDir, 'ø') }));
+	await application.start();
+	const connection = await createConnection(undefined, () => Promise.resolve(application.code.driver.browserContext as any));
+	application.code.driver.browserContext.on('close', () => {
+		connection.close();
+	});
+	return connection;
+}

test/mcp/src/utils.ts

@@ -0,0 +1,46 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { dirname, join } from 'path';
+import { Application, ApplicationOptions } from '../../automation';
+
+let logsCounter = 1;
+let crashCounter = 1;
+
+export function parseVersion(version: string): { major: number; minor: number; patch: number } {
+	const [, major, minor, patch] = /^(\d+)\.(\d+)\.(\d+)/.exec(version)!;
+	return { major: parseInt(major), minor: parseInt(minor), patch: parseInt(patch) };
+}
+
+export function suiteLogsPath(options: ApplicationOptions, suiteName: string): string {
+	return join(dirname(options.logsPath), `${logsCounter++}_suite_${suiteName.replace(/[^a-z0-9\-]/ig, '_')}`);
+}
+
+export function suiteCrashPath(options: ApplicationOptions, suiteName: string): string {
+	return join(dirname(options.crashesPath), `${crashCounter++}_suite_${suiteName.replace(/[^a-z0-9\-]/ig, '_')}`);
+}
+
+export function getRandomUserDataDir(options: ApplicationOptions): string {
+
+	// Pick a random user data dir suffix that is not
+	// too long to not run into max path length issues
+	// https://github.com/microsoft/vscode/issues/34988
+	const userDataPathSuffix = [...Array(8)].map(() => Math.random().toString(36)[3]).join('');
+
+	return options.userDataDir.concat(`-${userDataPathSuffix}`);
+}
+
+export function createApp(options: ApplicationOptions, optionsTransform?: (opts: ApplicationOptions) => ApplicationOptions): Application {
+	if (optionsTransform) {
+		options = optionsTransform({ ...options });
+	}
+
+	const app = new Application({
+		...options,
+		userDataDir: getRandomUserDataDir(options)
+	});
+
+	return app;
+}