diff --git a/package.json b/package.json index 17a361a14..e70e43aa0 100644 --- a/package.json +++ b/package.json @@ -1331,12 +1331,12 @@ "mdb.mcp.server": { "type": "string", "enum": [ + "ask", "enabled", - "disabled", - "ask" + "disabled" ], "default": "ask", - "description": "Controls whether MongoDB MCP server starts automatically with the extension or if a permission is requested when connecting to a new connection. If automatic startup is disabled, the server can still be started using the 'MongoDB: Start MCP Server' command." + "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": { "type": "string", diff --git a/src/mcp/mcpController.ts b/src/mcp/mcpController.ts index 0ca90a856..733992bba 100644 --- a/src/mcp/mcpController.ts +++ b/src/mcp/mcpController.ts @@ -17,7 +17,7 @@ import type { MCPConnectParams } from './mcpConnectionManager'; import { MCPConnectionManager } from './mcpConnectionManager'; import { createMCPConnectionErrorHandler } from './mcpConnectionErrorHandler'; -type mcpServerStartupConfig = 'ask' | 'enabled' | 'disabled'; +export type McpServerStartupConfig = 'enabled' | 'disabled'; class VSCodeMCPLogger extends LoggerBase { private readonly _logger = createLogger('mcp-server'); @@ -32,10 +32,13 @@ class VSCodeMCPLogger extends LoggerBase { } } +const logger = createLogger('mcp-controller'); + export type MCPServerInfo = { runner: StreamableHttpRunner; headers: Record; }; + export class MCPController { private didChangeEmitter = new vscode.EventEmitter(); private server?: MCPServerInfo; @@ -58,7 +61,7 @@ export class MCPController { ); } - public activate(): Promise { + public async activate(): Promise { this.connectionController.addEventListener( 'ACTIVE_CONNECTION_CHANGED', () => { @@ -66,52 +69,119 @@ export class MCPController { }, ); - return Promise.resolve(); + if (this.shouldStartMCPServer()) { + await this.startServer(); + void this.notifyOnFirstStart(); + } } public async startServer(): Promise { - const headers: Record = { - authorization: `Bearer ${crypto.randomUUID()}`, - }; - - const mcpConfig: UserConfig = { - ...defaultUserConfig, - httpPort: 0, - httpHeaders: headers, - disabledTools: ['connect'], - loggers: ['mcp'], - }; - - const createConnectionManager: ConnectionManagerFactoryFn = async ({ - logger, - }) => { - const connectionManager = (this.mcpConnectionManager = - new MCPConnectionManager(logger)); - await this.switchConnectionManagerToCurrentConnection(); - return connectionManager; - }; - - const runner = new StreamableHttpRunner({ - userConfig: mcpConfig, - createConnectionManager, - connectionErrorHandler: createMCPConnectionErrorHandler( - this.connectionController, - ), - additionalLoggers: [new VSCodeMCPLogger()], - }); - await runner.start(); - - this.server = { - runner, - headers, - }; - this.didChangeEmitter.fire(); + try { + const headers: Record = { + authorization: `Bearer ${crypto.randomUUID()}`, + }; + + const mcpConfig: UserConfig = { + ...defaultUserConfig, + httpPort: 0, + httpHeaders: headers, + disabledTools: ['connect'], + loggers: ['mcp'], + }; + + const createConnectionManager: ConnectionManagerFactoryFn = async ({ + logger, + }) => { + const connectionManager = (this.mcpConnectionManager = + new MCPConnectionManager(logger)); + await this.switchConnectionManagerToCurrentConnection(); + return connectionManager; + }; + + const runner = new StreamableHttpRunner({ + userConfig: mcpConfig, + createConnectionManager, + connectionErrorHandler: createMCPConnectionErrorHandler( + this.connectionController, + ), + additionalLoggers: [new VSCodeMCPLogger()], + }); + await runner.start(); + + this.server = { + runner, + headers, + }; + this.didChangeEmitter.fire(); + } catch (error) { + // In case of errors we don't want VSCode extension process to crash so we + // silence MCP start errors and instead log them for debugging. + logger.error('Error when attempting to start MCP server', error); + } } public async stopServer(): Promise { - await this.server?.runner.close(); - this.server = undefined; - this.didChangeEmitter.fire(); + 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 async notifyOnFirstStart(): Promise { + try { + if (!this.server) { + // Server was never started so no need to notify + 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); + } catch (error) { + logger.error( + 'Error while attempting to emit MCP server started notification', + error, + ); + } } public async openServerConfig(): Promise { @@ -185,58 +255,10 @@ ${jsonConfig}`, } private async onActiveConnectionChanged(): Promise { - if (this.server) { - await this.switchConnectionManagerToCurrentConnection(); - return; - } - - if (!this.connectionController.getActiveConnectionId()) { - // No connection, don't prompt the user + if (!this.server) { return; } - - const serverStartConfig = vscode.workspace - .getConfiguration('mdb') - .get('mcp.server'); - - let shouldStartServer = false; - switch (serverStartConfig) { - case 'enabled': - shouldStartServer = true; - break; - case 'disabled': - break; - default: - const prompt = await vscode.window.showInformationMessage( - 'Do you want to start the MongoDB MCP server automatically when connected to MongoDB?', - 'Yes', - 'No', - ); - - switch (prompt) { - case 'Yes': - shouldStartServer = true; - break; - case 'No': - shouldStartServer = false; - break; - default: - // User canceled/dismissed the notification - don't do anything. - return; - } - - await vscode.workspace - .getConfiguration('mdb') - .update( - 'mcp.server', - shouldStartServer ? 'enabled' : 'disabled', - true, - ); - } - - if (shouldStartServer) { - await this.startServer(); - } + await this.switchConnectionManagerToCurrentConnection(); } private async switchConnectionManagerToCurrentConnection(): Promise { @@ -254,4 +276,22 @@ ${jsonConfig}`, : undefined; await this.mcpConnectionManager?.updateConnection(connectParams); } + + private shouldStartMCPServer(): boolean { + return this.getMCPAutoStartConfig() !== 'disabled'; + } + + private getMCPAutoStartConfig(): McpServerStartupConfig | undefined { + return vscode.workspace + .getConfiguration('mdb') + .get('mcp.server'); + } + + private async setMCPAutoStartConfig( + config: McpServerStartupConfig, + ): Promise { + await vscode.workspace + .getConfiguration('mdb') + .update('mcp.server', config, true); + } } diff --git a/src/test/suite/mcp/mcpController.test.ts b/src/test/suite/mcp/mcpController.test.ts index 02e3c20b4..c91439173 100644 --- a/src/test/suite/mcp/mcpController.test.ts +++ b/src/test/suite/mcp/mcpController.test.ts @@ -76,8 +76,8 @@ suite('MCPController test suite', function () { }); }); - suite('when mcp server start is enabled from config', function () { - test('it should start mcp server without any confirmation', async function () { + 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'); @@ -87,20 +87,15 @@ suite('MCPController test suite', function () { 'showInformationMessage', ); const startServerSpy = sandbox.spy(mcpController, 'startServer'); - // listen to connection events await mcpController.activate(); - // add a new connection to trigger connection change - await connectionController.addNewConnectionStringAndConnect({ - connectionString: TEST_DATABASE_URI, - }); expect(showInformationSpy).to.not.be.called; expect(startServerSpy).to.be.calledOnce; }); }); - suite('when mcp server start is disabled from config', function () { - test('it should not start mcp server and ask for no confirmation', async function () { + 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'); @@ -110,20 +105,31 @@ suite('MCPController test suite', function () { 'showInformationMessage', ); const startServerSpy = sandbox.spy(mcpController, 'startServer'); - // listen to connection events await mcpController.activate(); - // add a new connection to trigger connection change - await connectionController.addNewConnectionStringAndConnect({ - connectionString: TEST_DATABASE_URI, - }); expect(showInformationSpy).to.not.be.called; expect(startServerSpy).to.not.be.called; }); }); - suite('when mcp server start is not configured', function () { - test('it should ask before starting the mcp server, and update the configuration with the chosen value', async function () { + 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; + }); + 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, @@ -135,59 +141,128 @@ suite('MCPController test suite', function () { fakeGetConfiguration, ); - const showInformationStub: SinonStub = sandbox.stub( - vscode.window, - 'showInformationMessage', - ); - showInformationStub.resolves('Yes'); + // Equivalent to dismissing the popup + informationStubResolvedValue = undefined; + const startServerSpy = sandbox.spy(mcpController, 'startServer'); - // listen to connection events + const stopServerSpy = sandbox.spy(mcpController, 'stopServer'); await mcpController.activate(); - // add a new connection to trigger connection change - await connectionController.addNewConnectionStringAndConnect({ - connectionString: TEST_DATABASE_URI, - }); + + await informationStubCalledNotification; expect(showInformationStub).to.be.calledOnce; expect(updateStub).to.be.calledWith('mcp.server', 'enabled', true); expect(startServerSpy).to.be.called; + expect(stopServerSpy).to.not.be.called; }); - test('it should ask before starting the mcp server, and when denied, should not start the server', async function () { - const updateStub = sandbox.stub(); - const fakeGetConfiguration = sandbox.fake.returns({ - get: () => null, - update: updateStub, - }); - sandbox.replace( - vscode.workspace, - 'getConfiguration', - fakeGetConfiguration, - ); + suite( + 'on the notification popup, if user selects to keep auto starting', + 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, + }); + sandbox.replace( + vscode.workspace, + 'getConfiguration', + fakeGetConfiguration, + ); - const showInformationStub: SinonStub = sandbox.stub( - vscode.window, - 'showInformationMessage', - ); - showInformationStub.resolves('No'); - const startServerSpy = sandbox.spy(mcpController, 'startServer'); - // listen to connection events + informationStubResolvedValue = 'Keep'; + const startServerSpy = sandbox.spy(mcpController, 'startServer'); + const stopServerSpy = sandbox.spy(mcpController, 'stopServer'); + await mcpController.activate(); + + await informationStubCalledNotification; + expect(showInformationStub).to.be.calledOnce; + expect(updateStub).to.be.calledWith('mcp.server', 'enabled', true); + expect(startServerSpy).to.be.called; + expect(stopServerSpy).to.not.be.called; + }); + }, + ); + + suite( + 'on the notification popup, if user selects to disable auto starting', + 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; + }); + + // 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, + ); + + informationStubResolvedValue = 'Disable'; + const startServerSpy = sandbox.spy(mcpController, 'startServer'); + const stopServerSpy = sandbox.spy(mcpController, 'stopServer'); + await mcpController.activate(); + + await informationStubCalledNotification; + expect(showInformationStub).to.be.calledOnce; + + await updateCalledNotification; + expect(updateStub.lastCall).to.be.calledWith( + 'mcp.server', + 'disabled', + true, + ); + 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'); + + // Start the controller and list to events await mcpController.activate(); - // add a new connection to trigger connection change + + // Add a connection await connectionController.addNewConnectionStringAndConnect({ connectionString: TEST_DATABASE_URI, }); - expect(showInformationStub).to.be.calledOnce; - expect(updateStub).to.be.calledWith('mcp.server', 'disabled', true); - expect(startServerSpy).to.not.be.called; + + const switchConnectionManagerSpy = sandbox.spy( + mcpController as any, + 'switchConnectionManagerToCurrentConnection', + ); + + await connectionController.disconnect(); + expect(switchConnectionManagerSpy).to.be.calledOnce; }); }); - 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 + 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', 'enabled'); + .update('mcp.server', 'disabled'); // Start the controller and list to events await mcpController.activate(); @@ -197,13 +272,13 @@ suite('MCPController test suite', function () { connectionString: TEST_DATABASE_URI, }); - const connectionChangedSpy = sandbox.spy( + const switchConnectionManagerSpy = sandbox.spy( mcpController as any, - 'onActiveConnectionChanged', + 'switchConnectionManagerToCurrentConnection', ); await connectionController.disconnect(); - expect(connectionChangedSpy).to.be.calledOnce; + expect(switchConnectionManagerSpy).not.to.be.called; }); });