diff --git a/package-lock.json b/package-lock.json index 19c96d7b5..a60b03a40 100644 --- a/package-lock.json +++ b/package-lock.json @@ -12,6 +12,7 @@ "@babel/core": "^7.25.8", "@babel/parser": "^7.25.8", "@babel/traverse": "^7.25.7", + "@himanshusinghs/mongodb-mcp-server": "^0.3.1", "@mongodb-js/compass-components": "^1.38.1", "@mongodb-js/connection-form": "^1.52.3", "@mongodb-js/connection-info": "^0.17.1", @@ -34,7 +35,6 @@ "mongodb-connection-string-url": "^3.0.2", "mongodb-data-service": "^22.30.1", "mongodb-log-writer": "^2.4.1", - "mongodb-mcp-server": "^0.3.0", "mongodb-query-parser": "^4.4.2", "mongodb-schema": "^12.6.2", "node-machine-id": "1.1.12", @@ -61,6 +61,7 @@ "@testing-library/user-event": "^14.5.2", "@types/babel__traverse": "^7.20.6", "@types/chai": "^4.3.20", + "@types/chai-as-promised": "^8.0.2", "@types/debug": "^4.1.12", "@types/glob": "^7.2.0", "@types/lodash": "^4.17.14", @@ -5684,6 +5685,145 @@ "integrity": "sha512-MDWhGtE+eHw5JW7lq4qhc5yRLS11ERl1c7Z6Xd0a58DozHES6EnNNwUWbMiG4J9Cgj053Bhk8zvlhFYKVhULwg==", "license": "MIT" }, + "node_modules/@himanshusinghs/mongodb-mcp-server": { + "version": "0.3.1", + "resolved": "https://registry.npmjs.org/@himanshusinghs/mongodb-mcp-server/-/mongodb-mcp-server-0.3.1.tgz", + "integrity": "sha512-H/QfueuEg7/qiUH8EI8NDP//7RsBy/iP6EwSemrz5fymovQSUnxsJnwiKlPM267HGpSDwT5QrJpo6Td6jfGA7A==", + "license": "Apache-2.0", + "dependencies": { + "@modelcontextprotocol/sdk": "^1.17.4", + "@mongodb-js/device-id": "^0.3.1", + "@mongodb-js/devtools-connect": "^3.9.3", + "@mongodb-js/devtools-proxy-support": "^0.5.2", + "@mongosh/arg-parser": "^3.14.0", + "@mongosh/service-provider-node-driver": "^3.12.0", + "@vitest/eslint-plugin": "^1.3.4", + "bson": "^6.10.4", + "express": "^5.1.0", + "lru-cache": "^11.1.0", + "mongodb": "^6.19.0", + "mongodb-connection-string-url": "^3.0.2", + "mongodb-log-writer": "^2.4.1", + "mongodb-redact": "^1.1.8", + "mongodb-schema": "^12.6.2", + "node-fetch": "^3.3.2", + "node-machine-id": "1.1.12", + "oauth4webapi": "^3.8.0", + "openapi-fetch": "^0.14.0", + "yargs-parser": "^22.0.0", + "zod": "^3.25.76" + }, + "bin": { + "mongodb-mcp-server": "dist/esm/index.js" + }, + "engines": { + "node": "^20.19.0 || ^22.12.0 || >= 23.0.0" + }, + "optionalDependencies": { + "kerberos": "^2.2.2" + } + }, + "node_modules/@himanshusinghs/mongodb-mcp-server/node_modules/data-uri-to-buffer": { + "version": "4.0.1", + "resolved": "https://registry.npmjs.org/data-uri-to-buffer/-/data-uri-to-buffer-4.0.1.tgz", + "integrity": "sha512-0R9ikRb668HB7QDxT1vkpuUBtqc53YyAwMwGeUFKRojY/NWKvdZ+9UYtRfGmhqNbRkTSVpMbmyhXipFFv2cb/A==", + "license": "MIT", + "engines": { + "node": ">= 12" + } + }, + "node_modules/@himanshusinghs/mongodb-mcp-server/node_modules/kerberos": { + "version": "2.2.2", + "resolved": "https://registry.npmjs.org/kerberos/-/kerberos-2.2.2.tgz", + "integrity": "sha512-42O7+/1Zatsc3MkxaMPpXcIl/ukIrbQaGoArZEAr6GcEi2qhfprOBYOPhj+YvSMJkEkdpTjApUx+2DuWaKwRhg==", + "hasInstallScript": true, + "license": "Apache-2.0", + "optional": true, + "dependencies": { + "node-addon-api": "^6.1.0", + "prebuild-install": "^7.1.2" + }, + "engines": { + "node": ">=12.9.0" + } + }, + "node_modules/@himanshusinghs/mongodb-mcp-server/node_modules/lru-cache": { + "version": "11.1.0", + "resolved": "https://registry.npmjs.org/lru-cache/-/lru-cache-11.1.0.tgz", + "integrity": "sha512-QIXZUBJUx+2zHUdQujWejBkcD9+cs94tLn0+YL8UrCh+D5sCXZ4c7LaEH48pNwRY3MLDgqUFyhlCyjJPf1WP0A==", + "license": "ISC", + "engines": { + "node": "20 || >=22" + } + }, + "node_modules/@himanshusinghs/mongodb-mcp-server/node_modules/napi-build-utils": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/napi-build-utils/-/napi-build-utils-2.0.0.tgz", + "integrity": "sha512-GEbrYkbfF7MoNaoh2iGG84Mnf/WZfB0GdGEsM8wz7Expx/LlWf5U8t9nvJKXSp3qr5IsEbK04cBGhol/KwOsWA==", + "license": "MIT", + "optional": true + }, + "node_modules/@himanshusinghs/mongodb-mcp-server/node_modules/node-addon-api": { + "version": "6.1.0", + "resolved": "https://registry.npmjs.org/node-addon-api/-/node-addon-api-6.1.0.tgz", + "integrity": "sha512-+eawOlIgy680F0kBzPUNFhMZGtJ1YmqM6l4+Crf4IkImjYrO/mqPwRMh352g23uIaQKFItcQ64I7KMaJxHgAVA==", + "license": "MIT", + "optional": true + }, + "node_modules/@himanshusinghs/mongodb-mcp-server/node_modules/node-fetch": { + "version": "3.3.2", + "resolved": "https://registry.npmjs.org/node-fetch/-/node-fetch-3.3.2.tgz", + "integrity": "sha512-dRB78srN/l6gqWulah9SrxeYnxeddIG30+GOqK/9OlLVyLg3HPnr6SqOWTWOXKRwC2eGYCkZ59NNuSgvSrpgOA==", + "license": "MIT", + "dependencies": { + "data-uri-to-buffer": "^4.0.0", + "fetch-blob": "^3.1.4", + "formdata-polyfill": "^4.0.10" + }, + "engines": { + "node": "^12.20.0 || ^14.13.1 || >=16.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/node-fetch" + } + }, + "node_modules/@himanshusinghs/mongodb-mcp-server/node_modules/prebuild-install": { + "version": "7.1.3", + "resolved": "https://registry.npmjs.org/prebuild-install/-/prebuild-install-7.1.3.tgz", + "integrity": "sha512-8Mf2cbV7x1cXPUILADGI3wuhfqWvtiLA1iclTDbFRZkgRQS0NqsPZphna9V+HyTEadheuPmjaJMsbzKQFOzLug==", + "license": "MIT", + "optional": true, + "dependencies": { + "detect-libc": "^2.0.0", + "expand-template": "^2.0.3", + "github-from-package": "0.0.0", + "minimist": "^1.2.3", + "mkdirp-classic": "^0.5.3", + "napi-build-utils": "^2.0.0", + "node-abi": "^3.3.0", + "pump": "^3.0.0", + "rc": "^1.2.7", + "simple-get": "^4.0.0", + "tar-fs": "^2.0.0", + "tunnel-agent": "^0.6.0" + }, + "bin": { + "prebuild-install": "bin.js" + }, + "engines": { + "node": ">=10" + } + }, + "node_modules/@himanshusinghs/mongodb-mcp-server/node_modules/yargs-parser": { + "version": "22.0.0", + "resolved": "https://registry.npmjs.org/yargs-parser/-/yargs-parser-22.0.0.tgz", + "integrity": "sha512-rwu/ClNdSMpkSrUb+d6BRsSkLUq1fmfsY6TOpYzTwvwkg1/NRG85KBy3kq++A8LKQwX6lsu+aWad+2khvuXrqw==", + "license": "ISC", + "engines": { + "node": "^20.19.0 || ^22.12.0 || >=23" + } + }, "node_modules/@humanwhocodes/config-array": { "version": "0.13.0", "resolved": "https://registry.npmjs.org/@humanwhocodes/config-array/-/config-array-0.13.0.tgz", @@ -10152,6 +10292,16 @@ "integrity": "sha512-/pC9HAB5I/xMlc5FP77qjCnI16ChlJfW0tGa0IUcFn38VJrTV6DeZ60NU5KZBtaOZqjdpwTWohz5HU1RrhiYxQ==", "dev": true }, + "node_modules/@types/chai-as-promised": { + "version": "8.0.2", + "resolved": "https://registry.npmjs.org/@types/chai-as-promised/-/chai-as-promised-8.0.2.tgz", + "integrity": "sha512-meQ1wDr1K5KRCSvG2lX7n7/5wf70BeptTKst0axGvnN6zqaVpRqegoIbugiAPSqOW9K9aL8gDVrm7a2LXOtn2Q==", + "dev": true, + "license": "MIT", + "dependencies": { + "@types/chai": "*" + } + }, "node_modules/@types/debug": { "version": "4.1.12", "resolved": "https://registry.npmjs.org/@types/debug/-/debug-4.1.12.tgz", @@ -20911,9 +21061,9 @@ } }, "node_modules/mongodb": { - "version": "6.17.0", - "resolved": "https://registry.npmjs.org/mongodb/-/mongodb-6.17.0.tgz", - "integrity": "sha512-neerUzg/8U26cgruLysKEjJvoNSXhyID3RvzvdcpsIi2COYM3FS3o9nlH7fxFtefTb942dX3W9i37oPfCVj4wA==", + "version": "6.19.0", + "resolved": "https://registry.npmjs.org/mongodb/-/mongodb-6.19.0.tgz", + "integrity": "sha512-H3GtYujOJdeKIMLKBT9PwlDhGrQfplABNF1G904w6r5ZXKWyv77aB0X9B+rhmaAwjtllHzaEkvi9mkGVZxs2Bw==", "license": "Apache-2.0", "dependencies": { "@mongodb-js/saslprep": "^1.1.9", @@ -20929,7 +21079,7 @@ "gcp-metadata": "^5.2.0", "kerberos": "^2.0.1", "mongodb-client-encryption": ">=6.0.0 <7", - "snappy": "^7.2.2", + "snappy": "^7.3.2", "socks": "^2.7.1" }, "peerDependenciesMeta": { @@ -21105,145 +21255,6 @@ "bson": "6.x" } }, - "node_modules/mongodb-mcp-server": { - "version": "0.3.0", - "resolved": "https://registry.npmjs.org/mongodb-mcp-server/-/mongodb-mcp-server-0.3.0.tgz", - "integrity": "sha512-jcJsqmbBfWSgpXE8dR+T+9WCQwigBo7OcFoYJGswm7ja7CCHqBaExOumwNsD2MRqhMcO9t3igNr9Y4kgLohmQg==", - "license": "Apache-2.0", - "dependencies": { - "@modelcontextprotocol/sdk": "^1.15.0", - "@mongodb-js/device-id": "^0.3.1", - "@mongodb-js/devtools-connect": "^3.9.2", - "@mongodb-js/devtools-proxy-support": "^0.5.1", - "@mongosh/arg-parser": "^3.14.0", - "@mongosh/service-provider-node-driver": "^3.10.2", - "@vitest/eslint-plugin": "^1.3.4", - "bson": "^6.10.4", - "express": "^5.1.0", - "lru-cache": "^11.1.0", - "mongodb": "^6.17.0", - "mongodb-connection-string-url": "^3.0.2", - "mongodb-log-writer": "^2.4.1", - "mongodb-redact": "^1.1.8", - "mongodb-schema": "^12.6.2", - "node-fetch": "^3.3.2", - "node-machine-id": "1.1.12", - "oauth4webapi": "^3.6.0", - "openapi-fetch": "^0.14.0", - "yargs-parser": "^22.0.0", - "zod": "^3.25.76" - }, - "bin": { - "mongodb-mcp-server": "dist/esm/index.js" - }, - "engines": { - "node": "^20.19.0 || ^22.12.0 || >= 23.0.0" - }, - "optionalDependencies": { - "kerberos": "^2.2.2" - } - }, - "node_modules/mongodb-mcp-server/node_modules/data-uri-to-buffer": { - "version": "4.0.1", - "resolved": "https://registry.npmjs.org/data-uri-to-buffer/-/data-uri-to-buffer-4.0.1.tgz", - "integrity": "sha512-0R9ikRb668HB7QDxT1vkpuUBtqc53YyAwMwGeUFKRojY/NWKvdZ+9UYtRfGmhqNbRkTSVpMbmyhXipFFv2cb/A==", - "license": "MIT", - "engines": { - "node": ">= 12" - } - }, - "node_modules/mongodb-mcp-server/node_modules/kerberos": { - "version": "2.2.2", - "resolved": "https://registry.npmjs.org/kerberos/-/kerberos-2.2.2.tgz", - "integrity": "sha512-42O7+/1Zatsc3MkxaMPpXcIl/ukIrbQaGoArZEAr6GcEi2qhfprOBYOPhj+YvSMJkEkdpTjApUx+2DuWaKwRhg==", - "hasInstallScript": true, - "license": "Apache-2.0", - "optional": true, - "dependencies": { - "node-addon-api": "^6.1.0", - "prebuild-install": "^7.1.2" - }, - "engines": { - "node": ">=12.9.0" - } - }, - "node_modules/mongodb-mcp-server/node_modules/lru-cache": { - "version": "11.1.0", - "resolved": "https://registry.npmjs.org/lru-cache/-/lru-cache-11.1.0.tgz", - "integrity": "sha512-QIXZUBJUx+2zHUdQujWejBkcD9+cs94tLn0+YL8UrCh+D5sCXZ4c7LaEH48pNwRY3MLDgqUFyhlCyjJPf1WP0A==", - "license": "ISC", - "engines": { - "node": "20 || >=22" - } - }, - "node_modules/mongodb-mcp-server/node_modules/napi-build-utils": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/napi-build-utils/-/napi-build-utils-2.0.0.tgz", - "integrity": "sha512-GEbrYkbfF7MoNaoh2iGG84Mnf/WZfB0GdGEsM8wz7Expx/LlWf5U8t9nvJKXSp3qr5IsEbK04cBGhol/KwOsWA==", - "license": "MIT", - "optional": true - }, - "node_modules/mongodb-mcp-server/node_modules/node-addon-api": { - "version": "6.1.0", - "resolved": "https://registry.npmjs.org/node-addon-api/-/node-addon-api-6.1.0.tgz", - "integrity": "sha512-+eawOlIgy680F0kBzPUNFhMZGtJ1YmqM6l4+Crf4IkImjYrO/mqPwRMh352g23uIaQKFItcQ64I7KMaJxHgAVA==", - "license": "MIT", - "optional": true - }, - "node_modules/mongodb-mcp-server/node_modules/node-fetch": { - "version": "3.3.2", - "resolved": "https://registry.npmjs.org/node-fetch/-/node-fetch-3.3.2.tgz", - "integrity": "sha512-dRB78srN/l6gqWulah9SrxeYnxeddIG30+GOqK/9OlLVyLg3HPnr6SqOWTWOXKRwC2eGYCkZ59NNuSgvSrpgOA==", - "license": "MIT", - "dependencies": { - "data-uri-to-buffer": "^4.0.0", - "fetch-blob": "^3.1.4", - "formdata-polyfill": "^4.0.10" - }, - "engines": { - "node": "^12.20.0 || ^14.13.1 || >=16.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/node-fetch" - } - }, - "node_modules/mongodb-mcp-server/node_modules/prebuild-install": { - "version": "7.1.3", - "resolved": "https://registry.npmjs.org/prebuild-install/-/prebuild-install-7.1.3.tgz", - "integrity": "sha512-8Mf2cbV7x1cXPUILADGI3wuhfqWvtiLA1iclTDbFRZkgRQS0NqsPZphna9V+HyTEadheuPmjaJMsbzKQFOzLug==", - "license": "MIT", - "optional": true, - "dependencies": { - "detect-libc": "^2.0.0", - "expand-template": "^2.0.3", - "github-from-package": "0.0.0", - "minimist": "^1.2.3", - "mkdirp-classic": "^0.5.3", - "napi-build-utils": "^2.0.0", - "node-abi": "^3.3.0", - "pump": "^3.0.0", - "rc": "^1.2.7", - "simple-get": "^4.0.0", - "tar-fs": "^2.0.0", - "tunnel-agent": "^0.6.0" - }, - "bin": { - "prebuild-install": "bin.js" - }, - "engines": { - "node": ">=10" - } - }, - "node_modules/mongodb-mcp-server/node_modules/yargs-parser": { - "version": "22.0.0", - "resolved": "https://registry.npmjs.org/yargs-parser/-/yargs-parser-22.0.0.tgz", - "integrity": "sha512-rwu/ClNdSMpkSrUb+d6BRsSkLUq1fmfsY6TOpYzTwvwkg1/NRG85KBy3kq++A8LKQwX6lsu+aWad+2khvuXrqw==", - "license": "ISC", - "engines": { - "node": "^20.19.0 || ^22.12.0 || >=23" - } - }, "node_modules/mongodb-ns": { "version": "2.4.2", "resolved": "https://registry.npmjs.org/mongodb-ns/-/mongodb-ns-2.4.2.tgz", @@ -23006,9 +23017,9 @@ } }, "node_modules/oauth4webapi": { - "version": "3.7.0", - "resolved": "https://registry.npmjs.org/oauth4webapi/-/oauth4webapi-3.7.0.tgz", - "integrity": "sha512-Q52wTPUWPsVLVVmTViXPQFMW2h2xv2jnDGxypjpelCFKaOjLsm7AxYuOk1oQgFm95VNDbuggasu9htXrz6XwKw==", + "version": "3.8.1", + "resolved": "https://registry.npmjs.org/oauth4webapi/-/oauth4webapi-3.8.1.tgz", + "integrity": "sha512-olkZDELNycOWQf9LrsELFq8n05LwJgV8UkrS0cburk6FOwf8GvLam+YB+Uj5Qvryee+vwWOfQVeI5Vm0MVg7SA==", "license": "MIT", "funding": { "url": "https://github.com/sponsors/panva" diff --git a/package.json b/package.json index afacf16a6..d528f0dc5 100644 --- a/package.json +++ b/package.json @@ -1396,6 +1396,7 @@ "@babel/core": "^7.25.8", "@babel/parser": "^7.25.8", "@babel/traverse": "^7.25.7", + "@himanshusinghs/mongodb-mcp-server": "^0.3.1", "@mongodb-js/compass-components": "^1.38.1", "@mongodb-js/connection-form": "^1.52.3", "@mongodb-js/connection-info": "^0.17.1", @@ -1418,7 +1419,6 @@ "mongodb-connection-string-url": "^3.0.2", "mongodb-data-service": "^22.30.1", "mongodb-log-writer": "^2.4.1", - "mongodb-mcp-server": "^0.3.0", "mongodb-query-parser": "^4.4.2", "mongodb-schema": "^12.6.2", "node-machine-id": "1.1.12", @@ -1445,6 +1445,7 @@ "@testing-library/user-event": "^14.5.2", "@types/babel__traverse": "^7.20.6", "@types/chai": "^4.3.20", + "@types/chai-as-promised": "^8.0.2", "@types/debug": "^4.1.12", "@types/glob": "^7.2.0", "@types/lodash": "^4.17.14", diff --git a/src/connectionController.ts b/src/connectionController.ts index b3ac8314c..03248ea93 100644 --- a/src/connectionController.ts +++ b/src/connectionController.ts @@ -40,8 +40,8 @@ const log = createLogger('connection controller'); const MAX_CONNECTION_NAME_LENGTH = 512; interface DataServiceEventTypes { - CONNECTIONS_DID_CHANGE: any; - ACTIVE_CONNECTION_CHANGED: any; + CONNECTIONS_DID_CHANGE: []; + ACTIVE_CONNECTION_CHANGED: []; } export enum ConnectionTypes { diff --git a/src/mcp/mcpConnectionManager.ts b/src/mcp/mcpConnectionManager.ts new file mode 100644 index 000000000..4057bd2e9 --- /dev/null +++ b/src/mcp/mcpConnectionManager.ts @@ -0,0 +1,118 @@ +import { + ConnectionManager, + type AnyConnectionState, + type ConnectionStateDisconnected, + type LoggerBase, +} from '@himanshusinghs/mongodb-mcp-server'; +import { + NodeDriverServiceProvider, + type DevtoolsConnectOptions, +} from '@mongosh/service-provider-node-driver'; +import type { ServiceProvider } from '@mongosh/service-provider-core'; +import { isAtlasStream } from 'mongodb-build-info'; +import { MCPLogIds } from './mcpLogIds'; + +export interface MCPConnectParams { + connectionId: string; + connectionString: string; + connectOptions: DevtoolsConnectOptions; +} + +export class MCPConnectionManager extends ConnectionManager { + private activeConnection: { + id: string; + provider: ServiceProvider; + } | null = null; + + constructor(private readonly logger: LoggerBase) { + super(); + } + + connect(): Promise { + return Promise.reject( + new Error( + [ + 'MongoDB MCP Server in MongoDB VSCode extension makes use of the connection that the MongoDB VSCode extension is connected to.', + "To connect, choose a connection from MongoDB VSCode extensions's sidepanel - https://www.mongodb.com/docs/mongodb-vscode/connect/#connect-to-your-mongodb-deployment", + ].join(' '), + ), + ); + } + + async connectToVSCodeConnection( + connectParams: MCPConnectParams, + ): Promise { + try { + const serviceProvider = await NodeDriverServiceProvider.connect( + connectParams.connectionString, + connectParams.connectOptions, + ); + await serviceProvider.runCommand('admin', { hello: 1 }); + this.activeConnection = { + id: connectParams.connectionId, + provider: serviceProvider, + }; + return this.changeState('connection-success', { + tag: 'connected', + serviceProvider, + }); + } catch (error) { + this.logger.error({ + id: MCPLogIds.ConnectError, + context: 'vscode-mcp-connection-manager', + message: `Error connecting to VSCode connection - ${error instanceof Error ? error.message : String(error)}`, + }); + return this.changeState('connection-error', { + tag: 'errored', + errorReason: error instanceof Error ? error.message : String(error), + }); + } + } + + async disconnect(): Promise { + try { + await this.activeConnection?.provider?.close(true); + } catch (error) { + this.logger.error({ + id: MCPLogIds.DisconnectError, + context: 'vscode-mcp-connection-manager', + message: `Error disconnecting from VSCode connection - ${error instanceof Error ? error.message : String(error)}`, + }); + } + + this.activeConnection = null; + return this.changeState('connection-close', { + tag: 'disconnected', + }); + } + + async updateConnection( + connectParams: MCPConnectParams | undefined, + ): Promise { + if (connectParams?.connectionId === this.activeConnection?.id) { + return; + } + + await this.disconnect(); + + if (!connectParams) { + return; + } + + if (isAtlasStream(connectParams.connectionString)) { + this.logger.warning({ + id: MCPLogIds.UpdateConnectionError, + context: 'vscode-mcp-connection-manager', + message: 'Attempting a connection to an AtlasStreams.', + }); + this.changeState('connection-error', { + tag: 'errored', + errorReason: + 'MongoDB MCP server does not support connecting to Atlas Streams', + }); + return; + } + + await this.connectToVSCodeConnection(connectParams); + } +} diff --git a/src/mcp/mcpController.ts b/src/mcp/mcpController.ts index 558c41c8c..c75994f72 100644 --- a/src/mcp/mcpController.ts +++ b/src/mcp/mcpController.ts @@ -1,12 +1,20 @@ import * as vscode from 'vscode'; -import type { LoggerType, LogLevel, LogPayload } from 'mongodb-mcp-server'; +import type { + LoggerType, + LogLevel, + LogPayload, + UserConfig, + ConnectionManagerFactoryFn, +} from '@himanshusinghs/mongodb-mcp-server'; import { defaultUserConfig, LoggerBase, StreamableHttpRunner, -} from 'mongodb-mcp-server'; +} from '@himanshusinghs/mongodb-mcp-server'; import type ConnectionController from '../connectionController'; import { createLogger } from '../logging'; +import type { MCPConnectParams } from './mcpConnectionManager'; +import { MCPConnectionManager } from './mcpConnectionManager'; type mcpServerStartupConfig = 'ask' | 'enabled' | 'disabled'; @@ -23,12 +31,14 @@ class VSCodeMCPLogger extends LoggerBase { } } +export type MCPServerInfo = { + runner: StreamableHttpRunner; + headers: Record; +}; export class MCPController { private didChangeEmitter = new vscode.EventEmitter(); - private server?: { - runner: StreamableHttpRunner; - headers: Record; - }; + private server?: MCPServerInfo; + private mcpConnectionManager?: MCPConnectionManager; constructor( private readonly context: vscode.ExtensionContext, @@ -62,18 +72,29 @@ export class MCPController { 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( - { - ...defaultUserConfig, - httpPort: 0, - httpHeaders: headers, - disabledTools: ['connect'], - loggers: ['mcp'], - }, - {}, + mcpConfig, + createConnectionManager, [new VSCodeMCPLogger()], ); - await runner.start(); this.server = { @@ -161,8 +182,7 @@ ${jsonConfig}`, private async onActiveConnectionChanged(): Promise { if (this.server) { - // Server is created - update the connection information - // this.server.runner.updateConnection(); + await this.switchConnectionManagerToCurrentConnection(); return; } @@ -214,4 +234,20 @@ ${jsonConfig}`, await this.startServer(); } } + + 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); + } } diff --git a/src/mcp/mcpLogIds.ts b/src/mcp/mcpLogIds.ts new file mode 100644 index 000000000..93357c0e4 --- /dev/null +++ b/src/mcp/mcpLogIds.ts @@ -0,0 +1,7 @@ +import { mongoLogId } from 'mongodb-log-writer'; + +export const MCPLogIds = { + ConnectError: mongoLogId(2_000_001), + DisconnectError: mongoLogId(2_000_002), + UpdateConnectionError: mongoLogId(2_000_003), +} as const; diff --git a/src/test/suite/mcp/mcpConnectionManager.test.ts b/src/test/suite/mcp/mcpConnectionManager.test.ts new file mode 100644 index 000000000..b977a67c6 --- /dev/null +++ b/src/test/suite/mcp/mcpConnectionManager.test.ts @@ -0,0 +1,290 @@ +import sinon from 'sinon'; +import { afterEach, beforeEach } from 'mocha'; +import chai, { expect } from 'chai'; +import chaiAsPromised from 'chai-as-promised'; +import type { LoggerBase } from '@himanshusinghs/mongodb-mcp-server'; +import type { ConnectionStateErrored } from '@himanshusinghs/mongodb-mcp-server'; +import type { DevtoolsConnectOptions } from '@mongosh/service-provider-node-driver'; +import { NodeDriverServiceProvider } from '@mongosh/service-provider-node-driver'; +import { MCPConnectionManager } from '../../../mcp/mcpConnectionManager'; + +chai.use(chaiAsPromised); + +const sandbox = sinon.createSandbox(); +suite('MCPConnectionManager Test Suite', function () { + let mcpConnectionManager: MCPConnectionManager; + let fakeServiceProvider: NodeDriverServiceProvider; + + beforeEach(() => { + mcpConnectionManager = new MCPConnectionManager({ + error: () => {}, + warning: () => {}, + } as unknown as LoggerBase); + fakeServiceProvider = { + runCommand: (() => + Promise.resolve({})) as NodeDriverServiceProvider['runCommand'], + close: (() => Promise.resolve()) as NodeDriverServiceProvider['close'], + } as NodeDriverServiceProvider; + sandbox + .stub(NodeDriverServiceProvider, 'connect') + .resolves(fakeServiceProvider); + }); + + afterEach(() => { + sandbox.restore(); + sandbox.reset(); + }); + + suite('#connect', function () { + test('should throw an error', async function () { + await expect(mcpConnectionManager.connect()).to.be.rejected; + }); + }); + + suite('#connectToVSCodeConnection', function () { + test('should connect successfully and update the state', async function () { + const newState = await mcpConnectionManager.connectToVSCodeConnection({ + connectionId: '1', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + }); + + expect(newState.tag).to.equal('connected'); + expect(mcpConnectionManager.currentConnectionState.tag).to.equal( + 'connected', + ); + }); + + test('should update the state when there is an error', async function () { + fakeServiceProvider.runCommand = (() => + Promise.reject( + new Error('Bad error'), + )) as NodeDriverServiceProvider['runCommand']; + const newState = (await mcpConnectionManager.connectToVSCodeConnection({ + connectionId: '1', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + })) as ConnectionStateErrored; + + expect(newState.tag).to.equal('errored'); + expect(newState.errorReason).to.equal('Bad error'); + expect(mcpConnectionManager.currentConnectionState.tag).to.equal( + 'errored', + ); + expect( + (mcpConnectionManager.currentConnectionState as ConnectionStateErrored) + .errorReason, + ).to.equal('Bad error'); + }); + }); + + suite('#disconnect', function () { + test('should disconnect successfully and update the state', async function () { + const newState = await mcpConnectionManager.connectToVSCodeConnection({ + connectionId: '1', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + }); + + expect(newState.tag).to.equal('connected'); + + const nextState = await mcpConnectionManager.disconnect(); + expect(nextState.tag).to.equal('disconnected'); + expect(mcpConnectionManager.currentConnectionState.tag).to.equal( + 'disconnected', + ); + expect((mcpConnectionManager as any).activeConnection).to.be.null; + }); + + test('should attempt to disconnect and on failure clear out the state', async function () { + fakeServiceProvider.close = (() => + Promise.reject( + new Error('Bad close error'), + )) as NodeDriverServiceProvider['close']; + const newState = await mcpConnectionManager.connectToVSCodeConnection({ + connectionId: '1', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + }); + + expect(newState.tag).to.equal('connected'); + + const nextState = await mcpConnectionManager.disconnect(); + expect(nextState.tag).to.equal('disconnected'); + expect(mcpConnectionManager.currentConnectionState.tag).to.equal( + 'disconnected', + ); + expect((mcpConnectionManager as any).activeConnection).to.be.null; + }); + }); + + suite('#updateConnection', function () { + suite('when not connected to any connection', function () { + test('should do nothing when invoked for a disconnected connection', async function () { + const connectSpy = sandbox.spy( + mcpConnectionManager, + 'connectToVSCodeConnection', + ); + const disconnectSpy = sandbox.spy(mcpConnectionManager, 'disconnect'); + await mcpConnectionManager.updateConnection(undefined); + + expect(connectSpy).to.not.be.called; + expect(disconnectSpy).to.not.be.called; + }); + + test('should switch to error state when invoked for an Atlas streams connection', async function () { + const connectSpy = sandbox.spy( + mcpConnectionManager, + 'connectToVSCodeConnection', + ); + const disconnectSpy = sandbox.spy(mcpConnectionManager, 'disconnect'); + await mcpConnectionManager.updateConnection({ + connectionId: '1', + connectionString: + 'mongodb://admin:catscatscats@atlas-stream-64ba1372b2a9f1545031f34d-gkumd.virginia-usa.a.query.mongodb.net/', + connectOptions: {} as DevtoolsConnectOptions, + }); + + expect(connectSpy).to.not.be.called; + expect(disconnectSpy).to.be.called; + expect(mcpConnectionManager.currentConnectionState.tag).to.equal( + 'errored', + ); + expect( + ( + mcpConnectionManager.currentConnectionState as ConnectionStateErrored + ).errorReason, + ).to.equal( + 'MongoDB MCP server does not support connecting to Atlas Streams', + ); + }); + + test('should connect to the connection when invoked for a newly connected connection', async function () { + const connectSpy = sandbox.spy( + mcpConnectionManager, + 'connectToVSCodeConnection', + ); + const disconnectSpy = sandbox.spy(mcpConnectionManager, 'disconnect'); + await mcpConnectionManager.updateConnection({ + connectionId: '1', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + }); + + expect(disconnectSpy).to.be.called; + expect(connectSpy).to.be.calledWithExactly({ + connectionId: '1', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + }); + }); + }); + + suite('when already connected to a connection', function () { + test('should do nothing when invoked for the already connected connection', async function () { + await mcpConnectionManager.connectToVSCodeConnection({ + connectionId: '1', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + }); + + // now we setup spies + const connectSpy = sandbox.spy( + mcpConnectionManager, + 'connectToVSCodeConnection', + ); + const disconnectSpy = sandbox.spy(mcpConnectionManager, 'disconnect'); + await mcpConnectionManager.updateConnection({ + connectionId: '1', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + }); + + expect(connectSpy).to.not.be.called; + expect(disconnectSpy).to.not.be.called; + }); + + test('should disconnect and return early when invoked for a disconnected connection', async function () { + await mcpConnectionManager.connectToVSCodeConnection({ + connectionId: '1', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + }); + + // now we setup spies + const connectSpy = sandbox.spy( + mcpConnectionManager, + 'connectToVSCodeConnection', + ); + const disconnectSpy = sandbox.spy(mcpConnectionManager, 'disconnect'); + await mcpConnectionManager.updateConnection(undefined); + + expect(connectSpy).to.not.be.called; + expect(disconnectSpy).to.be.called; + }); + + test('should switch to error state when invoked for an Atlas streams connection', async function () { + await mcpConnectionManager.connectToVSCodeConnection({ + connectionId: '1', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + }); + + // now we setup spies + const connectSpy = sandbox.spy( + mcpConnectionManager, + 'connectToVSCodeConnection', + ); + const disconnectSpy = sandbox.spy(mcpConnectionManager, 'disconnect'); + + // update connection + await mcpConnectionManager.updateConnection({ + connectionId: '2', + connectionString: + 'mongodb://admin:catscatscats@atlas-stream-64ba1372b2a9f1545031f34d-gkumd.virginia-usa.a.query.mongodb.net/', + connectOptions: {} as DevtoolsConnectOptions, + }); + + expect(disconnectSpy).to.be.called; + expect(connectSpy).to.not.be.called; + expect(mcpConnectionManager.currentConnectionState.tag).to.equal( + 'errored', + ); + expect( + ( + mcpConnectionManager.currentConnectionState as ConnectionStateErrored + ).errorReason, + ).to.equal( + 'MongoDB MCP server does not support connecting to Atlas Streams', + ); + }); + + test('should disconnect and attempt to connect to the new connection when invoked for a different connection', async function () { + await mcpConnectionManager.connectToVSCodeConnection({ + connectionId: '1', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + }); + + // now we setup spies + const connectSpy = sandbox.spy( + mcpConnectionManager, + 'connectToVSCodeConnection', + ); + const disconnectSpy = sandbox.spy(mcpConnectionManager, 'disconnect'); + await mcpConnectionManager.updateConnection({ + connectionId: '2', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + }); + + expect(disconnectSpy).to.be.called; + expect(connectSpy).to.be.calledWithExactly({ + connectionId: '2', + connectionString: 'mongodb://localhost:27017', + connectOptions: {} as unknown as DevtoolsConnectOptions, + }); + }); + }); + }); +}); diff --git a/src/test/suite/mcp/mcpController.test.ts b/src/test/suite/mcp/mcpController.test.ts new file mode 100644 index 000000000..02e3c20b4 --- /dev/null +++ b/src/test/suite/mcp/mcpController.test.ts @@ -0,0 +1,261 @@ +import type { SinonStub } from 'sinon'; +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 '@himanshusinghs/mongodb-mcp-server'; +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'; + +const sandbox = sinon.createSandbox(); +suite('MCPController test suite', function () { + let extensionContext: ExtensionContext; + let connectionController: ConnectionController; + let mcpController: MCPController; + + beforeEach(() => { + extensionContext = new ExtensionContextStub(); + const testStorageController = new StorageController(extensionContext); + const testTelemetryService = new TelemetryService( + testStorageController, + extensionContext, + ); + connectionController = new ConnectionController({ + statusView: new StatusView(extensionContext), + storageController: testStorageController, + telemetryService: testTelemetryService, + }); + + mcpController = new MCPController(extensionContext, connectionController); + }); + + afterEach(async () => { + 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('#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', + ); + }); + }); + + 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('when mcp server start is enabled from config', function () { + test('it should start mcp server without any confirmation', async function () { + await vscode.workspace + .getConfiguration('mdb') + .update('mcp.server', 'enabled'); + + const showInformationSpy = sandbox.spy( + vscode.window, + '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 () { + await vscode.workspace + .getConfiguration('mdb') + .update('mcp.server', 'disabled'); + + const showInformationSpy = sandbox.spy( + vscode.window, + '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 () { + 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('Yes'); + 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(showInformationStub).to.be.calledOnce; + expect(updateStub).to.be.calledWith('mcp.server', 'enabled', true); + expect(startServerSpy).to.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, + ); + + const showInformationStub: SinonStub = sandbox.stub( + vscode.window, + 'showInformationMessage', + ); + showInformationStub.resolves('No'); + 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(showInformationStub).to.be.calledOnce; + expect(updateStub).to.be.calledWith('mcp.server', 'disabled', true); + expect(startServerSpy).to.not.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 connection + await connectionController.addNewConnectionStringAndConnect({ + connectionString: TEST_DATABASE_URI, + }); + + const connectionChangedSpy = sandbox.spy( + mcpController as any, + 'onActiveConnectionChanged', + ); + + await connectionController.disconnect(); + expect(connectionChangedSpy).to.be.calledOnce; + }); + }); + + suite('#openServerConfig', function () { + suite('when the server is not running', function () { + test('should notify that server is not running', async function () { + const showErrorMessageSpy = sandbox.spy( + vscode.window, + 'showErrorMessage', + ); + expect(await mcpController.openServerConfig()).to.equal(false); + expect(showErrorMessageSpy).to.be.calledWith( + 'MongoDB MCP Server is not running. Start the server by running "MDB: Start MCP Server" in the command palette.', + ); + }); + }); + + 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(); + }); + + 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 list to events + await mcpController.activate(); + + // Add a connection + await connectionController.addNewConnectionStringAndConnect({ + connectionString: TEST_DATABASE_URI, + }); + await startServerCalled; + expect(await mcpController.openServerConfig()).to.equal(true); + expect(showTextDocumentStub).to.be.called; + }); + }); + }); +});