diff --git a/.gitignore b/.gitignore index e5785ab14..181ee74cf 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,5 @@ constants.json .eslintcache .sbom src/test/ai-accuracy-tests/test-results.html +.yalc +yalc.lock diff --git a/package-lock.json b/package-lock.json index 8c2e8abf2..fc9f5dcba 100644 --- a/package-lock.json +++ b/package-lock.json @@ -34,7 +34,7 @@ "mongodb-connection-string-url": "^3.0.2", "mongodb-data-service": "^22.30.1", "mongodb-log-writer": "^2.4.1", - "mongodb-mcp-server": "^1.0.2-prerelease.1", + "mongodb-mcp-server": "^1.0.2-prerelease.2", "mongodb-query-parser": "^4.4.2", "mongodb-schema": "^12.6.2", "node-machine-id": "1.1.12", @@ -51,6 +51,7 @@ "devDependencies": { "@babel/preset-typescript": "^7.25.7", "@babel/types": "^7.25.8", + "@modelcontextprotocol/sdk": "^1.18.2", "@mongodb-js/oidc-mock-provider": "^0.11.4", "@mongodb-js/oidc-plugin": "^2.0.3", "@mongodb-js/prettier-config-devtools": "^1.0.2", @@ -7648,9 +7649,9 @@ "license": "BSD-2-Clause" }, "node_modules/@modelcontextprotocol/sdk": { - "version": "1.17.4", - "resolved": "https://registry.npmjs.org/@modelcontextprotocol/sdk/-/sdk-1.17.4.tgz", - "integrity": "sha512-zq24hfuAmmlNZvik0FLI58uE5sriN0WWsQzIlYnzSuKDAHFqJtBFrl/LfB1NLgJT5Y7dEBzaX4yAKqOPrcetaw==", + "version": "1.18.2", + "resolved": "https://registry.npmjs.org/@modelcontextprotocol/sdk/-/sdk-1.18.2.tgz", + "integrity": "sha512-beedclIvFcCnPrYgHsylqiYJVJ/CI47Vyc4tY8no1/Li/O8U4BTlJfy6ZwxkYwx+Mx10nrgwSVrA7VBbhh4slg==", "license": "MIT", "dependencies": { "ajv": "^6.12.6", @@ -21007,9 +21008,9 @@ } }, "node_modules/mongodb-mcp-server": { - "version": "1.0.2-prerelease.1", - "resolved": "https://registry.npmjs.org/mongodb-mcp-server/-/mongodb-mcp-server-1.0.2-prerelease.1.tgz", - "integrity": "sha512-BMe85cMXG6bKtzk7awqeOiSRO89UQSzcXL2B6ku30CXMSydeCmOqJtJw/m7MpJfgrlGaMGeg5epar9FmQWrtmg==", + "version": "1.0.2-prerelease.2", + "resolved": "https://registry.npmjs.org/mongodb-mcp-server/-/mongodb-mcp-server-1.0.2-prerelease.2.tgz", + "integrity": "sha512-6577bW2tSsrZYXzj8U/w4X01Kp1WwRjmSrBeIGof7IIalohSE9A/Tc5rxaRfLtHiQT3mDyVLndSP4VjWiVFVQw==", "license": "Apache-2.0", "dependencies": { "@modelcontextprotocol/sdk": "^1.17.4", diff --git a/package.json b/package.json index 51056aaa5..53368791c 100644 --- a/package.json +++ b/package.json @@ -1320,7 +1320,7 @@ }, "mdb.mcp.readOnly": { "type": "boolean", - "default": false, + "default": true, "description": "When set to true, MongoDB MCP server will only allow read, connect, and metadata operation types, disabling create/update/delete operations." }, "mdb.mcp.indexCheck": { @@ -1331,11 +1331,16 @@ "mdb.mcp.server": { "type": "string", "enum": [ - "ask", - "enabled", - "disabled" + "prompt", + "autoStartEnabled", + "autoStartDisabled" + ], + "enumItemLabels": [ + "Ask", + "Auto Start Enabled", + "Auto Start Disabled" ], - "default": "ask", + "default": "prompt", "description": "Controls whether MongoDB MCP server starts automatically with the extension and connects to the active connection. If automatic startup is disabled, the server can still be started using the 'MongoDB: Start MCP Server' command." }, "mdb.mcp.exportsPath": { @@ -1418,7 +1423,7 @@ "mongodb-connection-string-url": "^3.0.2", "mongodb-data-service": "^22.30.1", "mongodb-log-writer": "^2.4.1", - "mongodb-mcp-server": "^1.0.2-prerelease.1", + "mongodb-mcp-server": "^1.0.2-prerelease.2", "mongodb-query-parser": "^4.4.2", "mongodb-schema": "^12.6.2", "node-machine-id": "1.1.12", @@ -1435,6 +1440,7 @@ "devDependencies": { "@babel/preset-typescript": "^7.25.7", "@babel/types": "^7.25.8", + "@modelcontextprotocol/sdk": "^1.18.2", "@mongodb-js/oidc-mock-provider": "^0.11.4", "@mongodb-js/oidc-plugin": "^2.0.3", "@mongodb-js/prettier-config-devtools": "^1.0.2", diff --git a/src/connectionController.ts b/src/connectionController.ts index e0f7a58ba..b68a5debf 100644 --- a/src/connectionController.ts +++ b/src/connectionController.ts @@ -732,7 +732,9 @@ export default class ConnectionController { return false; } - const originalDisconnect = this._activeDataService.disconnect.bind(this); + const originalDisconnect = this._activeDataService.disconnect.bind( + this._activeDataService, + ); this._activeDataService = null; try { diff --git a/src/mcp/mcpConnectionManager.ts b/src/mcp/mcpConnectionManager.ts index 20e762006..508f96d52 100644 --- a/src/mcp/mcpConnectionManager.ts +++ b/src/mcp/mcpConnectionManager.ts @@ -41,7 +41,7 @@ export class MCPConnectionManager extends ConnectionManager { this.getTelemetryAnonymousId = getTelemetryAnonymousId; } - connect(): Promise { + override connect(): Promise { return Promise.reject( new Error( // eslint-disable-next-line no-multi-str @@ -83,7 +83,7 @@ To connect, choose a connection from MongoDB VSCode extensions's sidepanel - htt } } - async disconnect(): Promise { + override async disconnect(): Promise { try { await this.activeConnection?.provider?.close(); } catch (error) { @@ -100,6 +100,11 @@ To connect, choose a connection from MongoDB VSCode extensions's sidepanel - htt }); } + override async close(): Promise { + await this.disconnect(); + this._events.emit('close', this.currentConnectionState); + } + async updateConnection( connectParams: MCPConnectParams | undefined, ): Promise { diff --git a/src/mcp/mcpController.ts b/src/mcp/mcpController.ts index 06ca2f165..5be87b362 100644 --- a/src/mcp/mcpController.ts +++ b/src/mcp/mcpController.ts @@ -5,6 +5,7 @@ import type { LogPayload, UserConfig, ConnectionManagerFactoryFn, + ConnectionManager, } from 'mongodb-mcp-server'; import { defaultUserConfig, @@ -19,8 +20,12 @@ import type { MCPConnectParams } from './mcpConnectionManager'; import { MCPConnectionManager } from './mcpConnectionManager'; import { createMCPConnectionErrorHandler } from './mcpConnectionErrorHandler'; import { getMCPConfigFromVSCodeSettings } from './mcpConfig'; +import { DEFAULT_TELEMETRY_APP_NAME } from '../connectionController'; -export type McpServerStartupConfig = 'enabled' | 'disabled'; +export type MCPServerStartupConfig = + | 'prompt' + | 'autoStartEnabled' + | 'autoStartDisabled'; class VSCodeMCPLogger extends LoggerBase { private readonly _logger = createLogger('mcp-server'); @@ -52,10 +57,10 @@ export class MCPController { private context: vscode.ExtensionContext; private connectionController: ConnectionController; private getTelemetryAnonymousId: () => string; + private mcpConnectionManagers: MCPConnectionManager[] = []; private didChangeEmitter = new vscode.EventEmitter(); private server?: MCPServerInfo; - private mcpConnectionManager?: MCPConnectionManager; constructor({ context, @@ -68,6 +73,13 @@ export class MCPController { } public async activate(): Promise { + await this.migrateOldConfigToNewConfig( + // At this point we don't know for certain if the "mdb.mcp.server" holds + // one of the old values or something totally unknown so to keep cases + // covered we consider the retrieved value unknown. + this.getMCPAutoStartConfig(), + ); + this.context.subscriptions.push( vscode.lm.registerMcpServerDefinitionProvider('mongodb', { onDidChangeMcpServerDefinitions: this.didChangeEmitter.event, @@ -87,16 +99,104 @@ export class MCPController { }, ); - if (this.shouldStartMCPServer()) { + if (this.getMCPAutoStartConfig() === 'autoStartEnabled') { await this.startServer(); - void this.notifyOnFirstStart(); + } + } + + private async migrateOldConfigToNewConfig(oldConfig: unknown): Promise { + try { + switch (oldConfig) { + // The previous logic would set the mdb.mcp.server to 'enabled' on + // extension activate (with a notification) so we're assuming that this + // value is not the result of explicit user action and hence mapping it + // to 'prompt'. + case 'ask': + case 'enabled': { + await this.setMCPAutoStartConfig('prompt'); + break; + } + + // In the previous logic only 'disabled' value would've represented an + // explicit user action which is why we preserve that and map it to new + // disabled value. + case 'disabled': { + await this.setMCPAutoStartConfig('autoStartDisabled'); + break; + } + + // Any other value is possible only if: + // 1. user explicitly did the modification or, + // 2. the old values were already migrated to the new values. + // So we don't migrate in this case. + default: { + break; + } + } + } catch (error) { + logger.error('Error when migrating old config to the new config', error); + } + } + + private async promptForMCPAutoStart(): Promise { + try { + const autoStartConfig = this.getMCPAutoStartConfig(); + const shouldPrompt = autoStartConfig === 'prompt'; + + logger.debug('Prompt to configure MCP auto start requested.', { + autoStartConfig, + shouldPrompt, + serverRunning: !!this.server, + }); + + if (!shouldPrompt) { + return; + } + + // 'Start Once' action might confuse users if the server is already + // running so we skip exposing this action in this particular case. + const notificationActions = this.server + ? (['Auto-Start', 'Never'] as const) + : (['Auto-Start', 'Start Once', 'Never'] as const); + + const promptResponse = await vscode.window.showInformationMessage( + 'Would you like to automatically start the MongoDB MCP server for a streamlined experience? When started, the server will automatically connect to your active MongoDB instance.', + ...notificationActions, + ); + + switch (promptResponse) { + case 'Auto-Start': { + await this.setMCPAutoStartConfig('autoStartEnabled'); + await this.startServer(); + break; + } + + case 'Start Once': { + await this.startServer(); + break; + } + + case 'Never': { + await this.setMCPAutoStartConfig('autoStartDisabled'); + break; + } + + default: + break; + } + } catch (error) { + logger.error('Error when prompting for MCP auto start', error); } } public async startServer(): Promise { try { - // Stop an already running server if any - await this.stopServer(); + if (this.server) { + logger.info( + 'MCP server start requested. An MCP server is already running, will not start a new server.', + ); + return; + } const token = crypto.randomUUID(); const headers: Record = { @@ -104,24 +204,7 @@ export class MCPController { }; registerGlobalSecretToRedact(token, 'password'); - const vscodeConfiguredMCPConfig = getMCPConfigFromVSCodeSettings(); - - const mcpConfig: UserConfig = { - ...defaultUserConfig, - ...vscodeConfiguredMCPConfig, - transport: 'http', - httpPort: 0, - httpHeaders: headers, - disabledTools: Array.from( - new Set([ - 'connect', - ...(vscodeConfiguredMCPConfig.disabledTools ?? []), - ]), - ), - loggers: Array.from( - new Set(['mcp', ...(vscodeConfiguredMCPConfig.loggers ?? [])]), - ), - }; + const mcpConfig = this.getMCPServerConfig(headers); logger.info('Starting MCP server with config', { ...mcpConfig, @@ -130,27 +213,16 @@ export class MCPController { apiClientSecret: '', }); - const createConnectionManager: ConnectionManagerFactoryFn = async ({ - logger, - }) => { - const connectionManager = (this.mcpConnectionManager = - new MCPConnectionManager({ - logger, - getTelemetryAnonymousId: this.getTelemetryAnonymousId, - })); - await this.switchConnectionManagerToCurrentConnection(); - return connectionManager; - }; - const runner = new StreamableHttpRunner({ userConfig: mcpConfig, - createConnectionManager, + createConnectionManager: (...params) => + MCPController.createConnectionManager(this, ...params), connectionErrorHandler: createMCPConnectionErrorHandler( this.connectionController, ), additionalLoggers: [new VSCodeMCPLogger(Keychain.root)], telemetryProperties: { - hosting_mode: 'vscode-extension', + hosting_mode: DEFAULT_TELEMETRY_APP_NAME, }, }); await runner.start(); @@ -167,67 +239,74 @@ export class MCPController { } } - public async stopServer(): Promise { - try { - await this.server?.runner.close(); - this.server = undefined; - this.didChangeEmitter.fire(); - } catch (error) { - logger.error('Error when attempting to close the MCP server', error); - } + private getMCPServerConfig(headers: Record): UserConfig { + const vscodeConfiguredMCPConfig = getMCPConfigFromVSCodeSettings(); + + return { + ...defaultUserConfig, + ...vscodeConfiguredMCPConfig, + transport: 'http', + httpPort: 0, + httpHeaders: headers, + disabledTools: Array.from( + new Set([ + 'connect', + ...(vscodeConfiguredMCPConfig.disabledTools ?? []), + ]), + ), + loggers: Array.from( + new Set(['mcp', ...(vscodeConfiguredMCPConfig.loggers ?? [])]), + ), + }; } - private async notifyOnFirstStart(): Promise { + private static async createConnectionManager( + mcpController: MCPController, + ...params: Parameters + ): Promise { + const [{ logger: mcpLogger }] = params; + const connectionManager = new MCPConnectionManager({ + logger: mcpLogger, + getTelemetryAnonymousId: mcpController.getTelemetryAnonymousId, + }); + + // Track this ConnectionManager instance for future connection updates + mcpController.mcpConnectionManagers.push(connectionManager); + + // Also set up listener on close event to perform a cleanup when the Client + // closes connection to MCP server and eventually ConnectionManager shuts + // down. + connectionManager.events.on('close', (): void => { + logger.debug('MCPConnectionManager closed. Performing cleanup', { + connectionManagerClientName: connectionManager.clientName, + }); + mcpController.mcpConnectionManagers = + mcpController.mcpConnectionManagers.filter( + (manager) => manager !== connectionManager, + ); + }); + + // The newly created ConnectionManager need to be brought up to date with + // the current connection state. + await mcpController.switchConnectionManagerToCurrentConnection( + connectionManager, + ); + return connectionManager; + } + + public async stopServer(): Promise { try { if (!this.server) { - // Server was never started so no need to notify + logger.info( + 'MCP server stop requested. No MCP server running, nothing to stop.', + ); return; } - - const serverStartConfig = this.getMCPAutoStartConfig(); - - // If the config value is one of the following values means they are - // intentional (either set by user or by this function itself) and we - // should not notify in that case. - const shouldNotNotify = - serverStartConfig === 'enabled' || serverStartConfig === 'disabled'; - - if (shouldNotNotify) { - return; - } - - // We set the auto start already to enabled to not prompt user again for - // this on the next boot. We do it this way because chances are that the - // user might not act on the notification in which case the final update - // will never happen. - await this.setMCPAutoStartConfig('enabled'); - let selectedServerStartConfig: McpServerStartupConfig = 'enabled'; - - const prompt = await vscode.window.showInformationMessage( - 'MongoDB MCP server started automatically and will connect to your active connection. Would you like to keep or disable automatic startup?', - 'Keep', - 'Disable', - ); - - switch (prompt) { - case 'Keep': - default: - // The default happens only when users explicity dismiss the - // notification. - selectedServerStartConfig = 'enabled'; - break; - case 'Disable': { - selectedServerStartConfig = 'disabled'; - await this.stopServer(); - } - } - - await this.setMCPAutoStartConfig(selectedServerStartConfig); + await this.server.runner.close(); + this.server = undefined; + this.didChangeEmitter.fire(); } catch (error) { - logger.error( - 'Error while attempting to emit MCP server started notification', - error, - ); + logger.error('Error when attempting to close the MCP server', error); } } @@ -302,43 +381,63 @@ ${jsonConfig}`, } private async onActiveConnectionChanged(): Promise { - if (!this.server) { - return; + logger.debug( + 'Active connection changed, will switch connection manager to new connection', + { + connectionId: this.connectionController.getActiveConnectionId(), + serverStarted: !!this.server, + }, + ); + + if (this.connectionController.getActiveConnectionId()) { + void this.promptForMCPAutoStart(); } - await this.switchConnectionManagerToCurrentConnection(); - } - private async switchConnectionManagerToCurrentConnection(): Promise { - const connectionId = this.connectionController.getActiveConnectionId(); - const mongoClientOptions = - this.connectionController.getMongoClientConnectionOptions(); - - const connectParams: MCPConnectParams | undefined = - connectionId && mongoClientOptions - ? { - connectionId: connectionId, - connectionString: mongoClientOptions.url, - connectOptions: mongoClientOptions.options, - } - : undefined; - await this.mcpConnectionManager?.updateConnection(connectParams); + await Promise.all( + this.mcpConnectionManagers.map((manager) => + this.switchConnectionManagerToCurrentConnection(manager), + ), + ); } - private shouldStartMCPServer(): boolean { - return this.getMCPAutoStartConfig() !== 'disabled'; + private async switchConnectionManagerToCurrentConnection( + connectionManager: MCPConnectionManager, + ): Promise { + try { + const connectionId = this.connectionController.getActiveConnectionId(); + const mongoClientOptions = + this.connectionController.getMongoClientConnectionOptions(); + + const connectParams: MCPConnectParams | undefined = + connectionId && mongoClientOptions + ? { + connectionId: connectionId, + connectionString: mongoClientOptions.url, + connectOptions: mongoClientOptions.options, + } + : undefined; + await connectionManager.updateConnection(connectParams); + } catch (error) { + logger.error( + 'Error when attempting to switch connection for connection manager', + error, + ); + } } - private getMCPAutoStartConfig(): McpServerStartupConfig | undefined { + private getMCPAutoStartConfig(): + | ConfigValue + | undefined { return vscode.workspace - .getConfiguration('mdb') - .get('mcp.server'); + .getConfiguration() + .get('mdb.mcp.server'); } private async setMCPAutoStartConfig( - config: McpServerStartupConfig, + config: MCPServerStartupConfig, ): Promise { await vscode.workspace - .getConfiguration('mdb') - .update('mcp.server', config, true); + .getConfiguration() + .update('mdb.mcp.server', config, true); } } diff --git a/src/test/suite/connectionController.test.ts b/src/test/suite/connectionController.test.ts index 8a443fef2..01eb20102 100644 --- a/src/test/suite/connectionController.test.ts +++ b/src/test/suite/connectionController.test.ts @@ -518,8 +518,10 @@ suite('Connection Controller Test Suite', function () { // The number of times we expect to re-render connections on the sidebar: // - connection attempt started // - connection attempt finished - // - disconnect - const expectedTimesToFire = 3; + // - disconnect from our call in the tests + // - disconnect from on('close') listener on DataService that gets called as + // a result of our disconnect call + const expectedTimesToFire = 4; let connectionsDidChangeEventFiredCount = 0; testConnectionController.addEventListener('CONNECTIONS_DID_CHANGE', () => { diff --git a/src/test/suite/mcp/mcpConnectionManager.test.ts b/src/test/suite/mcp/mcpConnectionManager.test.ts index e531eb2df..e3bee4578 100644 --- a/src/test/suite/mcp/mcpConnectionManager.test.ts +++ b/src/test/suite/mcp/mcpConnectionManager.test.ts @@ -124,6 +124,20 @@ suite('MCPConnectionManager Test Suite', function () { }); }); + suite('#close', function () { + test('should call disconnect and emit close event', async function () { + const disconnectSpy = sandbox.spy(mcpConnectionManager, 'disconnect'); + let resolveWhenEmitted: (() => void) | undefined; + const closeEventEmitted = new Promise((resolve) => { + resolveWhenEmitted = resolve; + }); + mcpConnectionManager.events.on('close', () => resolveWhenEmitted?.()); + await mcpConnectionManager.close(); + expect(disconnectSpy).to.have.been.calledOnce; + await closeEventEmitted; + }); + }); + suite('#updateConnection', function () { suite('when not connected to any connection', function () { test('should do nothing when invoked for a disconnected connection', async function () { diff --git a/src/test/suite/mcp/mcpController.test.ts b/src/test/suite/mcp/mcpController.test.ts index 98ca74a12..3e53ad780 100644 --- a/src/test/suite/mcp/mcpController.test.ts +++ b/src/test/suite/mcp/mcpController.test.ts @@ -3,25 +3,73 @@ import sinon from 'sinon'; import { expect } from 'chai'; import { afterEach, beforeEach } from 'mocha'; import * as vscode from 'vscode'; -import type { ExtensionContext } from 'vscode'; -import * as MCPServer from 'mongodb-mcp-server'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; import { ExtensionContextStub } from '../stubs'; -import type { MCPServerInfo } from '../../../mcp/mcpController'; import { MCPController } from '../../../mcp/mcpController'; import ConnectionController from '../../../connectionController'; import { StatusView } from '../../../views'; import { StorageController } from '../../../storage'; import { TelemetryService } from '../../../telemetry'; import { TEST_DATABASE_URI } from '../dbTestHelper'; +import { waitFor } from '../waitFor'; + +function timeout(ms: number): Promise { + return new Promise((resolve) => { + setTimeout(resolve, ms); + }); +} + +async function createConnectedMCPClient( + clientName: string, + mcpController: MCPController, +): Promise<{ + client: Client; + transport: StreamableHTTPClientTransport; + closeClient: () => Promise; +}> { + const httpServerDefinition: vscode.McpHttpServerDefinition = ( + mcpController as any + ).getServerConfig(); + expect(httpServerDefinition).to.not.be.undefined; + const { uri, headers } = httpServerDefinition; + const transport = new StreamableHTTPClientTransport(new URL(uri.toString()), { + requestInit: { + headers: headers, + }, + }); + const client = new Client({ name: clientName, version: '1.0.0' }); + await client.connect(transport); + return { + client, + transport, + async closeClient(): Promise { + await transport.terminateSession(); + await client.close(); + }, + }; +} const sandbox = sinon.createSandbox(); suite('MCPController test suite', function () { - let extensionContext: ExtensionContext; let connectionController: ConnectionController; let mcpController: MCPController; + let mcpAutoStartValue: string | null | undefined; + let getConfigurationStub: SinonStub; + let updateConfigurationStub: SinonStub; + + let showInformationSelection: string | undefined; + let showInformationMessageStub: SinonStub; + let showInformationCalledNotification: Promise; + + let startServerStub: SinonStub; + let startServerCalledNotification: Promise; + + let stopServerStub: SinonStub; + beforeEach(() => { - extensionContext = new ExtensionContextStub(); + const extensionContext = new ExtensionContextStub(); const testStorageController = new StorageController(extensionContext); const testTelemetryService = new TelemetryService( testStorageController, @@ -38,247 +86,755 @@ suite('MCPController test suite', function () { connectionController: connectionController, getTelemetryAnonymousId: (): string => '1FOO', }); + + // GetConfiguration Stubs + mcpAutoStartValue = undefined; + getConfigurationStub = sandbox.stub().callsFake((key) => { + if (key === 'mdb.mcp.server') { + return mcpAutoStartValue; + } + }); + updateConfigurationStub = sandbox.stub().callsFake((key, value) => { + mcpAutoStartValue = value; + }); + const fakeGetConfiguration = sandbox.fake.returns({ + get: getConfigurationStub, + update: updateConfigurationStub, + }); + sandbox.replace(vscode.workspace, 'getConfiguration', fakeGetConfiguration); + + // Show information message stubs + showInformationSelection = undefined; + let notifyInformationMessageVisible: (() => void) | undefined; + showInformationCalledNotification = new Promise((resolve) => { + notifyInformationMessageVisible = resolve; + }); + showInformationMessageStub = sandbox + .stub(vscode.window, 'showInformationMessage') + .callsFake(((): Promise => { + notifyInformationMessageVisible?.(); + return Promise.resolve(showInformationSelection); + }) as unknown as any); + + // Other spies + let notifyStartServerCalled: (() => void) | undefined; + startServerCalledNotification = new Promise((resolve) => { + notifyStartServerCalled = resolve; + }); + const originalStartServer = mcpController.startServer.bind(mcpController); + startServerStub = sandbox + .stub(mcpController, 'startServer') + .callsFake((...args) => { + notifyStartServerCalled?.(); + return originalStartServer(...args); + }); + + const originalStopServer = mcpController.stopServer.bind(mcpController); + stopServerStub = sandbox + .stub(mcpController, 'stopServer') + .callsFake((...args) => { + return originalStopServer(...args); + }); }); - afterEach(async () => { + afterEach(() => { sandbox.restore(); sandbox.reset(); connectionController.clearAllConnections(); - await vscode.workspace.getConfiguration('mdb').update('mcp.server', null); }); - test('should register mcp server definition provider', function () { - // At-least one from our mcp controller - expect(extensionContext.subscriptions.length).to.be.greaterThanOrEqual(1); - }); + suite('on extension activate', function () { + for (const { storedValue, migratedValue, expectMigration } of [ + { storedValue: 'ask', migratedValue: 'prompt', expectMigration: true }, + { + storedValue: 'enabled', + migratedValue: 'prompt', + expectMigration: true, + }, + { + storedValue: 'prompt', + migratedValue: 'prompt', + expectMigration: false, + }, + ]) { + // eslint-disable-next-line no-loop-func + suite(`if stored config "${storedValue}"`, function () { + const testName = expectMigration + ? `should migrate the stored value to "${migratedValue}", not show any auto start config popups and not start the server` + : 'should keep the stored value as it is, not show any auto-start config popups and not start the server'; + + test(testName, async function () { + mcpAutoStartValue = storedValue; + await mcpController.activate(); - suite('#activate', function () { - test('should subscribe to ACTIVE_CONNECTION_CHANGED event', async function () { - const addEventListenerSpy = sandbox.spy( - connectionController, - 'addEventListener', - ); - await mcpController.activate(); - expect(addEventListenerSpy).to.be.called; - expect(addEventListenerSpy.args[0]).to.contain( - 'ACTIVE_CONNECTION_CHANGED', - ); - }); - }); + // Update configuration is called anyways + if (expectMigration) { + expect(updateConfigurationStub).to.be.calledWithExactly( + 'mdb.mcp.server', + migratedValue, + true, + ); + expect(mcpAutoStartValue).to.equal(migratedValue); + } else { + expect(updateConfigurationStub).to.not.be.called; + expect(mcpAutoStartValue).to.equal(storedValue); + } + + expect(showInformationMessageStub).to.not.be.called; + + // A small timeout to ensure the background tasks did happen + await timeout(10); + expect(startServerStub).to.not.be.called; + // Open server config will return false when server is not already running + expect(await mcpController.openServerConfig()).to.be.false; + }); + }); + } - suite('#startServer', function () { - test('should initialize HTTP transport and start it', async function () { - await mcpController.startServer(); - const serverInfo = (mcpController as any).server as - | MCPServerInfo - | undefined; - expect(serverInfo).to.not.be.undefined; - expect(serverInfo?.runner).to.be.instanceOf( - MCPServer.StreamableHttpRunner, - ); - expect(serverInfo?.headers?.authorization).to.not.be.undefined; + suite('if stored config "autoStartEnabled"', function () { + test('should keep the stored value as it is, not show any auto-start config popups and start the server', async function () { + mcpAutoStartValue = 'autoStartEnabled'; + await mcpController.activate(); + + // A small timeout to ensure the background tasks did happen + await timeout(10); + expect(startServerStub).to.be.called; + // Open server config will return true when server is running + expect(await mcpController.openServerConfig()).to.be.true; + + // no popup shown + expect(showInformationMessageStub).to.not.be.called; + }); }); + + for (const storedValue of ['anything-else', null]) { + // eslint-disable-next-line no-loop-func + suite(`if stored config "${storedValue}"`, function () { + test('should keep the stored value as it is, not show any auto-start config popups and not start the server', async function () { + mcpAutoStartValue = storedValue; + await mcpController.activate(); + + expect(updateConfigurationStub).to.not.be.called; + + // A small timeout to ensure the background tasks did happen + await timeout(10); + expect(startServerStub).to.not.be.called; + // Open server config will return true when server is running + expect(await mcpController.openServerConfig()).to.be.false; + + // no popup shown + expect(showInformationMessageStub).to.not.be.called; + }); + }); + } + + for (const { storedValue, migratedValue, expectMigration } of [ + { + storedValue: 'disabled', + migratedValue: 'autoStartDisabled', + expectMigration: true, + }, + { + storedValue: 'autoStartDisabled', + migratedValue: 'autoStartDisabled', + expectMigration: false, + }, + ]) { + // eslint-disable-next-line no-loop-func + suite(`if stored config "${storedValue}"`, function () { + const testName = expectMigration + ? `should migrate the stored value to "${migratedValue}", not show any auto start config popups and not start the server` + : 'should keep the stored value as it is, not show any auto-start config popups and not start the server'; + + test(testName, async function () { + mcpAutoStartValue = storedValue; + await mcpController.activate(); + + if (expectMigration) { + expect(updateConfigurationStub).to.be.calledWithExactly( + 'mdb.mcp.server', + migratedValue, + true, + ); + expect(mcpAutoStartValue).to.equal(migratedValue); + } + + // A small timeout to ensure the background tasks did happen + await timeout(10); + expect(startServerStub).to.not.be.called; + // Open server config will return true when server is running + expect(await mcpController.openServerConfig()).to.be.false; + + // no popup shown + expect(showInformationMessageStub).to.not.be.called; + }); + }); + } }); - suite('when mcp server auto start is enabled in the config', function () { - test('it should start mcp server without any notification', async function () { - await vscode.workspace - .getConfiguration('mdb') - .update('mcp.server', 'enabled'); + suite('autostart config popup', function () { + suite('popup visibility', function () { + suite('if server is running and a connection is connected', function () { + for (const storedValue of ['ask', 'enabled', 'prompt']) { + // eslint-disable-next-line no-loop-func + suite(`if stored config is "${storedValue}"`, function () { + test('should show the auto start config prompt with two action buttons', async function () { + mcpAutoStartValue = storedValue; + await mcpController.activate(); + await timeout(10); + + await mcpController.startServer(); + + await connectionController.addNewConnectionStringAndConnect({ + connectionString: TEST_DATABASE_URI, + }); + await timeout(10); + + await showInformationCalledNotification; + expect(showInformationMessageStub).to.be.calledWithExactly( + 'Would you like to automatically start the MongoDB MCP server for a streamlined experience? When started, the server will automatically connect to your active MongoDB instance.', + 'Auto-Start', + 'Never', + ); + }); + }); + } + }); - const showInformationSpy = sandbox.spy( - vscode.window, - 'showInformationMessage', + suite( + 'if server is not running and a connection is connected', + function () { + for (const storedValue of ['ask', 'enabled', 'prompt']) { + // eslint-disable-next-line no-loop-func + suite(`if stored config is "${storedValue}"`, function () { + test('should show the auto start config prompt with three action buttons', async function () { + mcpAutoStartValue = storedValue; + await mcpController.activate(); + await timeout(10); + + await connectionController.addNewConnectionStringAndConnect({ + connectionString: TEST_DATABASE_URI, + }); + await timeout(10); + + await showInformationCalledNotification; + expect(showInformationMessageStub).to.be.calledWithExactly( + 'Would you like to automatically start the MongoDB MCP server for a streamlined experience? When started, the server will automatically connect to your active MongoDB instance.', + 'Auto-Start', + 'Start Once', + 'Never', + ); + }); + }); + } + }, ); - const startServerSpy = sandbox.spy(mcpController, 'startServer'); - await mcpController.activate(); - expect(showInformationSpy).to.not.be.called; - expect(startServerSpy).to.be.calledOnce; + suite( + 'regardless of server state, if a connection is disconnected', + function () { + for (const storedValue of ['ask', 'enabled', 'prompt']) { + // eslint-disable-next-line no-loop-func + suite(`if stored config is "${storedValue}"`, function () { + test('should not show the auto start config prompt', async function () { + mcpAutoStartValue = storedValue; + await mcpController.activate(); + await timeout(10); + + await mcpController.startServer(); + + await connectionController.addNewConnectionStringAndConnect({ + connectionString: TEST_DATABASE_URI, + }); + await timeout(10); + + expect(showInformationMessageStub).to.be.calledWithExactly( + 'Would you like to automatically start the MongoDB MCP server for a streamlined experience? When started, the server will automatically connect to your active MongoDB instance.', + 'Auto-Start', + 'Never', + ); + + await connectionController.disconnect(); + await timeout(10); + // Only once and that too from before. + expect(showInformationMessageStub).to.be.calledWithExactly( + 'Would you like to automatically start the MongoDB MCP server for a streamlined experience? When started, the server will automatically connect to your active MongoDB instance.', + 'Auto-Start', + 'Never', + ); + }); + }); + } + }, + ); }); - }); - suite('when mcp server auto start is disabled from config', function () { - test('it should not start mcp server and show no notification', async function () { - await vscode.workspace - .getConfiguration('mdb') - .update('mcp.server', 'disabled'); + suite('popup actions', function () { + suite('when user clicks on "Auto-Start"', function () { + test('should start the server and configure MCP server to auto-start', async function () { + mcpAutoStartValue = 'prompt'; + showInformationSelection = 'Auto-Start'; + await mcpController.activate(); - const showInformationSpy = sandbox.spy( - vscode.window, - 'showInformationMessage', - ); - const startServerSpy = sandbox.spy(mcpController, 'startServer'); - await mcpController.activate(); + await connectionController.addNewConnectionStringAndConnect({ + connectionString: TEST_DATABASE_URI, + }); + + await showInformationCalledNotification; + expect(showInformationMessageStub).to.be.calledWithExactly( + 'Would you like to automatically start the MongoDB MCP server for a streamlined experience? When started, the server will automatically connect to your active MongoDB instance.', + 'Auto-Start', + 'Start Once', + 'Never', + ); + + await timeout(10); + // Server should've been started by now after popup selection + expect(startServerStub).to.have.been.called; + expect(await mcpController.openServerConfig()).to.be.true; - expect(showInformationSpy).to.not.be.called; - expect(startServerSpy).to.not.be.called; + expect(updateConfigurationStub).to.have.been.calledWithExactly( + 'mdb.mcp.server', + 'autoStartEnabled', + true, + ); + expect(mcpAutoStartValue).to.equal('autoStartEnabled'); + }); + }); + + suite('when user clicks on "Start Once"', function () { + test('should start the server once and not configure MCP server auto-start', async function () { + mcpAutoStartValue = 'prompt'; + showInformationSelection = 'Start Once'; + await mcpController.activate(); + + await connectionController.addNewConnectionStringAndConnect({ + connectionString: TEST_DATABASE_URI, + }); + + await showInformationCalledNotification; + expect(showInformationMessageStub).to.be.calledWithExactly( + 'Would you like to automatically start the MongoDB MCP server for a streamlined experience? When started, the server will automatically connect to your active MongoDB instance.', + 'Auto-Start', + 'Start Once', + 'Never', + ); + + await timeout(10); + // Server should've been started by now after popup selection + expect(startServerStub).to.have.been.called; + expect(await mcpController.openServerConfig()).to.be.true; + + expect(updateConfigurationStub).to.not.have.been.called; + expect(mcpAutoStartValue).to.equal('prompt'); + }); + }); + + suite('when user clicks on "Never"', function () { + test('should not start the server and configure MCP server to never auto-start', async function () { + mcpAutoStartValue = 'prompt'; + showInformationSelection = 'Never'; + await mcpController.activate(); + + await connectionController.addNewConnectionStringAndConnect({ + connectionString: TEST_DATABASE_URI, + }); + + await showInformationCalledNotification; + expect(showInformationMessageStub).to.be.calledWithExactly( + 'Would you like to automatically start the MongoDB MCP server for a streamlined experience? When started, the server will automatically connect to your active MongoDB instance.', + 'Auto-Start', + 'Start Once', + 'Never', + ); + + await timeout(10); + // Server should've been started by now + expect(startServerStub).to.not.have.been.called; + expect(await mcpController.openServerConfig()).to.be.false; + + expect(updateConfigurationStub).to.have.been.calledWithExactly( + 'mdb.mcp.server', + 'autoStartDisabled', + true, + ); + expect(mcpAutoStartValue).to.equal('autoStartDisabled'); + }); + }); }); }); - suite('when mcp server auto start is not configured', function () { - let showInformationStub: SinonStub; - let informationStubCalledNotification: Promise; - let informationStubResolvedValue: any; - beforeEach(() => { - informationStubResolvedValue = undefined; - let notifyInformationStubCalled: () => void; - informationStubCalledNotification = new Promise((resolve) => { - notifyInformationStubCalled = resolve; + suite('MCP commands handlers', function () { + suite('when MCP server is not running', function () { + test('"startServer" should start the server', async function () { + mcpAutoStartValue = 'autoStartDisabled'; + await mcpController.activate(); + // Server is not running + expect(await mcpController.openServerConfig()).to.be.false; + + await mcpController.startServer(); + // Assert no side effects or attempt to close previous runs + expect(stopServerStub).to.not.have.been.called; + expect(await mcpController.openServerConfig()).to.be.true; + }); + + test('"stopServer" should do nothing', async function () { + mcpAutoStartValue = 'autoStartDisabled'; + await mcpController.activate(); + // Server not running + expect(await mcpController.openServerConfig()).to.be.false; + + await mcpController.stopServer(); + // Server not running + expect(await mcpController.openServerConfig()).to.be.false; }); - showInformationStub = sandbox - .stub(vscode.window, 'showInformationMessage') - .callsFake(() => { - notifyInformationStubCalled(); - return Promise.resolve(informationStubResolvedValue); - }); }); - test('it start the mcp server, set auto start to enabled and, notify the user with an information message', async function () { - const updateStub = sandbox.stub(); - const fakeGetConfiguration = sandbox.fake.returns({ - get: () => null, - update: updateStub, + + suite('when MCP server is running', function () { + test('"startServer" should do nothing', async function () { + mcpAutoStartValue = 'autoStartEnabled'; + await mcpController.activate(); + // Assert no side effects or attempt to close previous runs + expect(stopServerStub).to.not.have.been.called; + expect(await mcpController.openServerConfig()).to.be.true; + + await mcpController.startServer(); + // Assert no side effects or attempt to close previous runs + expect(stopServerStub).to.not.have.been.called; + expect(await mcpController.openServerConfig()).to.be.true; }); - sandbox.replace( - vscode.workspace, - 'getConfiguration', - fakeGetConfiguration, - ); - // Equivalent to dismissing the popup - informationStubResolvedValue = undefined; + test('"stopServer" should stop the server', async function () { + mcpAutoStartValue = 'autoStartEnabled'; + await mcpController.activate(); + // Server is running + expect(await mcpController.openServerConfig()).to.be.true; - const startServerSpy = sandbox.spy(mcpController, 'startServer'); - await mcpController.activate(); + await mcpController.stopServer(); + // Not anymore + expect(await mcpController.openServerConfig()).to.be.false; + }); + }); + }); - await informationStubCalledNotification; - expect(showInformationStub).to.be.calledOnce; - expect(updateStub).to.be.calledWith('mcp.server', 'enabled', true); - expect(startServerSpy).to.be.called; + suite('clients connection handling', function () { + suite('when there is an active connection in VSCode', function () { + test('connected clients should be able to query mongodb', async function () { + mcpAutoStartValue = 'autoStartEnabled'; + // Connect already + await connectionController.addNewConnectionStringAndConnect({ + connectionString: TEST_DATABASE_URI, + }); + await mcpController.activate(); + + const { client: firstClient } = await createConnectedMCPClient( + 'firstClient', + mcpController, + ); + const { client: secondClient } = await createConnectedMCPClient( + 'secondClient', + mcpController, + ); + + const [firstResponse, secondResponse] = await Promise.all([ + firstClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + secondClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + ]); + + expect(JSON.stringify(firstResponse.content)).to.contain( + 'Found 3 databases', + ); + expect(JSON.stringify(secondResponse.content)).to.contain( + 'Found 3 databases', + ); + }); }); suite( - 'on the notification popup, if user selects to keep auto starting', + 'when connection state changes from connected to disconnected', function () { - test('it should keep the config set to auto start and continue running the MCP server', async function () { - const updateStub = sandbox.stub(); - const fakeGetConfiguration = sandbox.fake.returns({ - get: () => null, - update: updateStub, + test('connected clients should get responded with no connection response', async function () { + mcpAutoStartValue = 'autoStartEnabled'; + // Connect already + await connectionController.addNewConnectionStringAndConnect({ + connectionString: TEST_DATABASE_URI, }); - sandbox.replace( - vscode.workspace, - 'getConfiguration', - fakeGetConfiguration, + await mcpController.activate(); + + const { client: firstClient } = await createConnectedMCPClient( + 'firstClient', + mcpController, + ); + const { client: secondClient } = await createConnectedMCPClient( + 'secondClient', + mcpController, ); - informationStubResolvedValue = 'Keep'; - const startServerSpy = sandbox.spy(mcpController, 'startServer'); - await mcpController.activate(); + let [firstResponse, secondResponse] = await Promise.all([ + firstClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + secondClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + ]); + + expect(JSON.stringify(firstResponse.content)).to.contain( + 'Found 3 databases', + ); + expect(JSON.stringify(secondResponse.content)).to.contain( + 'Found 3 databases', + ); - await informationStubCalledNotification; - expect(showInformationStub).to.be.calledOnce; - expect(updateStub).to.be.calledWith('mcp.server', 'enabled', true); - expect(startServerSpy).to.be.called; + // Now disconnect + await connectionController.disconnect(); + + // Next call should respond back with disconnected content + [firstResponse, secondResponse] = await Promise.all([ + firstClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + secondClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + ]); + + expect(JSON.stringify(firstResponse.content)).to.contain( + 'You need to connect to a MongoDB instance before you can access its data.', + ); + expect(JSON.stringify(secondResponse.content)).to.contain( + 'You need to connect to a MongoDB instance before you can access its data.', + ); }); }, ); suite( - 'on the notification popup, if user selects to disable auto starting', + 'when connection state changes from disconnected to connected', function () { - test('it should set the config to disable auto start and stop the MCP server', async function () { - let notifyUpdateCalled: () => void; - const updateCalledNotification = new Promise((resolve) => { - notifyUpdateCalled = resolve; - }); + test('connected clients should be able to query mongodb', async function () { + mcpAutoStartValue = 'autoStartEnabled'; - // There will be two calls to update, one which we do by default and - // second to update the config to disabled. - let callCount = 0; - const updateStub = sandbox.stub().callsFake(() => { - if (++callCount === 2) { - notifyUpdateCalled(); - } - }); - const fakeGetConfiguration = sandbox.fake.returns({ - get: () => null, - update: updateStub, - }); - sandbox.replace( - vscode.workspace, - 'getConfiguration', - fakeGetConfiguration, + await mcpController.activate(); + + const { client: firstClient } = await createConnectedMCPClient( + 'firstClient', + mcpController, + ); + const { client: secondClient } = await createConnectedMCPClient( + 'secondClient', + mcpController, ); - informationStubResolvedValue = 'Disable'; - const startServerSpy = sandbox.spy(mcpController, 'startServer'); - const stopServerSpy = sandbox.spy(mcpController, 'stopServer'); - await mcpController.activate(); + let [firstResponse, secondResponse] = await Promise.all([ + firstClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + secondClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + ]); + + expect(JSON.stringify(firstResponse.content)).to.contain( + 'You need to connect to a MongoDB instance before you can access its data.', + ); + expect(JSON.stringify(secondResponse.content)).to.contain( + 'You need to connect to a MongoDB instance before you can access its data.', + ); - await informationStubCalledNotification; - expect(showInformationStub).to.be.calledOnce; + // Now connect + await connectionController.addNewConnectionStringAndConnect({ + connectionString: TEST_DATABASE_URI, + }); - await updateCalledNotification; - expect(updateStub.lastCall).to.be.calledWith( - 'mcp.server', - 'disabled', - true, + // A little timeout + await timeout(100); + + // Next call should respond back with disconnected content + [firstResponse, secondResponse] = await Promise.all([ + firstClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + secondClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + ]); + + expect(JSON.stringify(firstResponse.content)).to.contain( + 'Found 3 databases', + ); + expect(JSON.stringify(secondResponse.content)).to.contain( + 'Found 3 databases', ); - expect(startServerSpy).to.be.called; - expect(stopServerSpy).to.be.called; }); }, ); - }); - suite('when an MCP server is already running', function () { - test('it should notify the connection manager of the connection changed event', async function () { - // We want to connect as soon as connection changes - await vscode.workspace - .getConfiguration('mdb') - .update('mcp.server', 'enabled'); + suite('when MCP server shuts down', function () { + test('should terminate individual client connections and clear up the internal connection manager state', async function () { + mcpAutoStartValue = 'autoStartEnabled'; + // Connect already in VSCode + await connectionController.addNewConnectionStringAndConnect({ + connectionString: TEST_DATABASE_URI, + }); + await mcpController.activate(); + await timeout(100); - // Start the controller and list to events - await mcpController.activate(); + // Construct our MCP clients connected to the VSCode MCP + const { client: firstClient } = await createConnectedMCPClient( + 'firstClient', + mcpController, + ); + const { client: secondClient } = await createConnectedMCPClient( + 'secondClient', + mcpController, + ); - // Add a connection + // Both clients are connected so both should be able to query MCP server + let [firstClientResponse, secondClientResponse] = await Promise.all([ + firstClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + secondClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + ]); + + expect(JSON.stringify(firstClientResponse.content)).to.contain( + 'Found 3 databases', + ); + expect(JSON.stringify(secondClientResponse.content)).to.contain( + 'Found 3 databases', + ); + // There should be 2 connection managers for the two clients we created + await waitFor( + () => (mcpController as any).mcpConnectionManagers.length === 2, + ); + + // MCP server shuts down + await mcpController.stopServer(); + + [firstClientResponse, secondClientResponse] = await Promise.all([ + firstClient + .callTool({ name: 'list-databases', arguments: {} }) + .catch((error) => error.message), + secondClient + .callTool({ name: 'list-databases', arguments: {} }) + .catch((error) => error.message), + ]); + // fetch would fail because server is not running + expect(firstClientResponse).to.contain('fetch failed'); + expect(secondClientResponse).to.contain('fetch failed'); + + // Cleanup that we are expecting + await waitFor( + () => (mcpController as any).mcpConnectionManagers.length === 0, + ); + }); + }); + + test('different clients should have their own connection state and not overstep each other', async function () { + mcpAutoStartValue = 'autoStartEnabled'; + // Connect already in VSCode await connectionController.addNewConnectionStringAndConnect({ connectionString: TEST_DATABASE_URI, }); + await mcpController.activate(); - const switchConnectionManagerSpy = sandbox.spy( - mcpController as any, - 'switchConnectionManagerToCurrentConnection', + // Construct our MCP clients connected to the VSCode MCP + const { client: firstClient } = await createConnectedMCPClient( + 'firstClient', + mcpController, + ); + const { client: secondClient, closeClient: closeSecondClient } = + await createConnectedMCPClient('secondClient', mcpController); + + // Both clients are connected so both should be able to query MCP server + let [firstResponse, secondResponse] = await Promise.all([ + firstClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + secondClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + ]); + + expect(JSON.stringify(firstResponse.content)).to.contain( + 'Found 3 databases', + ); + expect(JSON.stringify(secondResponse.content)).to.contain( + 'Found 3 databases', ); - await connectionController.disconnect(); - expect(switchConnectionManagerSpy).to.be.calledOnce; - }); - }); - - suite('when an MCP server is not running', function () { - test('it should not notify the connection manager of the connection changed event', async function () { - // Disable connecting - await vscode.workspace - .getConfiguration('mdb') - .update('mcp.server', 'disabled'); - - // Start the controller and list to events - await mcpController.activate(); - - // Add a connection - await connectionController.addNewConnectionStringAndConnect({ - connectionString: TEST_DATABASE_URI, + // Closing second client to test that it clears up and not affect the + // first client + await closeSecondClient(); + await waitFor(() => { + return (mcpController as any).mcpConnectionManagers.length === 1; }); - const switchConnectionManagerSpy = sandbox.spy( - mcpController as any, - 'switchConnectionManagerToCurrentConnection', + // Second client is closed but the first should still get a response + [firstResponse, secondResponse] = await Promise.all([ + firstClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + secondClient + .callTool({ + name: 'list-databases', + arguments: {}, + }) + .catch((error) => error.message), + ]); + // Only first client responds with actual tool response + expect(JSON.stringify(firstResponse.content)).to.contain( + 'Found 3 databases', ); + expect(secondResponse).to.contain('Not connected'); + // Another state change from VSCode await connectionController.disconnect(); - expect(switchConnectionManagerSpy).not.to.be.called; + + // A small timeout + await timeout(10); + + // Second client is closed so that should respond with the error message + // but the first client should get the disconnected response. + [firstResponse, secondResponse] = await Promise.all([ + firstClient.callTool({ + name: 'list-databases', + arguments: {}, + }), + secondClient + .callTool({ + name: 'list-databases', + arguments: {}, + }) + .catch((error) => error.message), + ]); + // Only first client responds with actual tool response + expect(JSON.stringify(firstResponse.content)).to.contain( + 'You need to connect to a MongoDB instance before you can access its data.', + ); + expect(secondResponse).to.contain('Not connected'); }); }); @@ -298,28 +854,13 @@ suite('MCPController test suite', function () { suite('when the server is running', function () { test('should open the document with the server config', async function () { - const startServer = mcpController.startServer.bind(mcpController); - let notifyStartServerCalled: () => void = () => {}; - const startServerCalled: Promise = new Promise( - (resolve) => { - notifyStartServerCalled = resolve; - }, - ); - sandbox.replace(mcpController, 'startServer', async () => { - await startServer(); - notifyStartServerCalled(); - }); + mcpAutoStartValue = 'autoStartEnabled'; const showTextDocumentStub = sandbox.spy( vscode.window, 'showTextDocument', ); - // We want to connect as soon as connection changes - await vscode.workspace - .getConfiguration('mdb') - .update('mcp.server', 'enabled'); - // Start the controller and listen to events await mcpController.activate(); @@ -327,7 +868,7 @@ suite('MCPController test suite', function () { await connectionController.addNewConnectionStringAndConnect({ connectionString: TEST_DATABASE_URI, }); - await startServerCalled; + await startServerCalledNotification; expect(await mcpController.openServerConfig()).to.equal(true); expect(showTextDocumentStub).to.be.called; });