chore: vscode connection manager implementation

Author: himanshusinghs

src/mcp/mcpController.ts

@@ -1,17 +1,21 @@
 import * as vscode from 'vscode';
 import { defaultUserConfig, StreamableHttpRunner } from 'mongodb-mcp-server';
 import type ConnectionController from '../connectionController';
-// import * as os from 'os';
-// import * as path from 'path';
+import {
+  VSCodeMCPConnectionManager,
+  type VSCodeMCPConnectParams,
+} from './vsCodeMCPConnectionManager';
+import type { CreateConnectionManagerFn, UserConfig } from 'mongodb-mcp-server';
 
 type mcpServerStartupConfig = 'ask' | 'enabled' | 'disabled';
 
 export class MCPController {
   private didChangeEmitter = new vscode.EventEmitter<void>();
   private server?: {
-    runner: StreamableHttpRunner;
+    runner: StreamableHttpRunner<VSCodeMCPConnectParams>;
     headers: Record<string, string>;
   };
+  private mcpConnectionManager?: VSCodeMCPConnectionManager;
 
   constructor(
     private readonly context: vscode.ExtensionContext,
@@ -45,13 +49,24 @@ export class MCPController {
     const headers: Record<string, string> = {
       authorization: `Bearer ${crypto.randomUUID()}`,
     };
-    const runner = new StreamableHttpRunner({
+
+    const mcpConfig: UserConfig = {
       ...defaultUserConfig,
       httpPort: 0,
       httpHeaders: headers,
       disabledTools: ['connect'],
-    });
+    };
+
+    const createConnectionManager: CreateConnectionManagerFn<
+      VSCodeMCPConnectParams
+    > = async ({ logger }) => {
+      const connectionManager = (this.mcpConnectionManager =
+        new VSCodeMCPConnectionManager(logger));
+      await this.switchConnectionManagerToCurrentConnection();
+      return connectionManager;
+    };
 
+    const runner = new StreamableHttpRunner(mcpConfig, createConnectionManager);
     await runner.start();
 
     this.server = {
@@ -138,8 +153,7 @@ ${jsonConfig}`,
 
   private async onActiveConnectionChanged(): Promise<void> {
     if (this.server) {
-      // Server is created - update the connection information
-      // this.server.runner.updateConnection();
+      await this.switchConnectionManagerToCurrentConnection();
       return;
     }
 
@@ -191,4 +205,15 @@ ${jsonConfig}`,
       await this.startServer();
     }
   }
+
+  private async switchConnectionManagerToCurrentConnection(): Promise<void> {
+    const connectionId = this.connectionController.getActiveConnectionId();
+    const mongoClientOptions =
+      this.connectionController.getMongoClientConnectionOptions();
+    await this.mcpConnectionManager?.updateConnection({
+      connectionId,
+      connectionString: mongoClientOptions?.url,
+      connectOptions: mongoClientOptions?.options,
+    });
+  }
 }

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;

src/mcp/vsCodeMCPConnectionManager.ts

@@ -0,0 +1,123 @@
+import {
+  ConnectionManager,
+  type AnyConnectionState,
+  type CompositeLogger,
+  type MCPConnectParams,
+  type ConnectionStateDisconnected,
+} from '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 VSCodeMCPConnectParams extends MCPConnectParams {
+  connectionId: string;
+  connectOptions: DevtoolsConnectOptions;
+}
+
+export class VSCodeMCPConnectionManager extends ConnectionManager<VSCodeMCPConnectParams> {
+  private activeConnectionId: string | null = null;
+  private activeConnectionProvider: ServiceProvider | null = null;
+
+  constructor(private readonly logger: CompositeLogger) {
+    super();
+  }
+
+  async connect(
+    connectParams: VSCodeMCPConnectParams,
+  ): Promise<AnyConnectionState> {
+    try {
+      const serviceProvider = (this.activeConnectionProvider =
+        await NodeDriverServiceProvider.connect(
+          connectParams.connectionString,
+          connectParams.connectOptions,
+        ));
+      this.activeConnectionId = connectParams.connectionId;
+      return this.changeState('connection-succeeded', {
+        tag: 'connected',
+        serviceProvider,
+      });
+    } catch (error) {
+      this.logger.error({
+        id: MCPLogIds.ConnectError,
+        context: 'VSCodeMCPConnectionManager.connect',
+        message: error instanceof Error ? error.message : String(error),
+      });
+      return this.changeState('connection-errored', {
+        tag: 'errored',
+        errorReason: error instanceof Error ? error.message : String(error),
+      });
+    }
+  }
+
+  async disconnect(): Promise<ConnectionStateDisconnected> {
+    try {
+      await this.activeConnectionProvider?.close(true);
+    } catch (error) {
+      this.logger.error({
+        id: MCPLogIds.DisconnectError,
+        context: 'VSCodeMCPConnectionManager.disconnect',
+        message: error instanceof Error ? error.message : String(error),
+      });
+      throw error;
+    } finally {
+      this.activeConnectionId = null;
+      this.activeConnectionProvider = null;
+    }
+    return Promise.resolve({ tag: 'disconnected' });
+  }
+
+  async updateConnection({
+    connectionId,
+    connectionString,
+    connectOptions,
+  }: {
+    connectionId: string | null;
+    connectionString: string | undefined;
+    connectOptions: DevtoolsConnectOptions | undefined;
+  }): Promise<void> {
+    try {
+      if (this.activeConnectionId && this.activeConnectionId !== connectionId) {
+        await this.disconnect();
+      }
+
+      if (this.activeConnectionId === connectionId) {
+        return;
+      }
+
+      if (!connectionString || !connectOptions || !connectionId) {
+        this.changeState('connection-errored', {
+          tag: 'errored',
+          errorReason:
+            'MongoDB MCP server cannot establish connection without the required connection string and MongoClientOptions',
+        });
+        return;
+      }
+
+      if (isAtlasStream(connectionString)) {
+        this.changeState('connection-errored', {
+          tag: 'errored',
+          errorReason:
+            'MongoDB MCP server do not support connecting to Atlas Streams',
+        });
+        return;
+      }
+
+      await this.connect({
+        connectionString,
+        connectOptions,
+        connectionId,
+      });
+    } catch (error) {
+      this.logger.error({
+        id: MCPLogIds.UpdateConnectionError,
+        context: 'VSCodeMCPConnectionManager.updateConnection',
+        message: error instanceof Error ? error.message : String(error),
+      });
+      throw error;
+    }
+  }
+}