chore: PR feedback

Author: himanshusinghs

src/mcp/mcpConnectionManager.ts

@@ -19,8 +19,10 @@ export interface MCPConnectParams {
 }
 
 export class MCPConnectionManager extends ConnectionManager {
-  private activeConnectionId: string | null = null;
-  private activeConnectionProvider: ServiceProvider | null = null;
+  private activeConnection: {
+    id: string;
+    provider: ServiceProvider;
+  } | null = null;
 
   constructor(private readonly logger: LoggerBase) {
     super();
@@ -46,17 +48,19 @@ export class MCPConnectionManager extends ConnectionManager {
         connectParams.connectOptions,
       );
       await serviceProvider.runCommand('admin', { hello: 1 });
-      this.activeConnectionId = connectParams.connectionId;
-      this.activeConnectionProvider = serviceProvider;
+      this.activeConnection = {
+        id: connectParams.connectionId,
+        provider: serviceProvider,
+      };
       return this.changeState('connection-success', {
         tag: 'connected',
         serviceProvider,
       });
     } catch (error) {
       this.logger.error({
         id: MCPLogIds.ConnectError,
-        context: 'VSCodeMCPConnectionManager.connect',
-        message: error instanceof Error ? error.message : String(error),
+        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',
@@ -67,72 +71,48 @@ export class MCPConnectionManager extends ConnectionManager {
 
   async disconnect(): Promise<ConnectionStateDisconnected> {
     try {
-      await this.activeConnectionProvider?.close(true);
+      await this.activeConnection?.provider?.close(true);
     } catch (error) {
       this.logger.error({
         id: MCPLogIds.DisconnectError,
-        context: 'VSCodeMCPConnectionManager.disconnect',
-        message: error instanceof Error ? error.message : String(error),
+        context: 'vscode-mcp-connection-manager',
+        message: `Error disconnecting from VSCode connection - ${error instanceof Error ? error.message : String(error)}`,
       });
     }
 
-    this.activeConnectionId = null;
-    this.activeConnectionProvider = null;
+    this.activeConnection = null;
     return this.changeState('connection-close', {
       tag: 'disconnected',
     });
   }
 
-  async updateConnection({
-    connectionId,
-    connectionString,
-    connectOptions,
-  }: {
-    connectionId: string | undefined;
-    connectionString: string | undefined;
-    connectOptions: DevtoolsConnectOptions | undefined;
-  }): Promise<void> {
-    try {
-      if (this.activeConnectionId && this.activeConnectionId !== connectionId) {
-        await this.disconnect();
-      }
-
-      const connectionWasDisconnected =
-        !connectionId || !connectionString || !connectOptions;
+  async updateConnection(
+    connectParams: MCPConnectParams | undefined,
+  ): Promise<void> {
+    if (connectParams?.connectionId === this.activeConnection?.id) {
+      return;
+    }
 
-      if (
-        this.activeConnectionId === connectionId ||
-        connectionWasDisconnected
-      ) {
-        return;
-      }
+    await this.disconnect();
 
-      if (isAtlasStream(connectionString)) {
-        this.logger.warning({
-          id: MCPLogIds.UpdateConnectionError,
-          context: 'VSCodeMCPConnectionManager.updateConnection',
-          message: 'updateConnection called for an AtlasStreams connection',
-        });
-        this.changeState('connection-error', {
-          tag: 'errored',
-          errorReason:
-            'MongoDB MCP server do not support connecting to Atlas Streams',
-        });
-        return;
-      }
+    if (!connectParams) {
+      return;
+    }
 
-      await this.connectToVSCodeConnection({
-        connectionString,
-        connectOptions,
-        connectionId,
-      });
-    } catch (error) {
-      this.logger.error({
+    if (isAtlasStream(connectParams.connectionString)) {
+      this.logger.warning({
         id: MCPLogIds.UpdateConnectionError,
-        context: 'VSCodeMCPConnectionManager.updateConnection',
-        message: error instanceof Error ? error.message : String(error),
+        context: 'vscode-mcp-connection-manager',
+        message: 'Attempting a connection to an AtlasStreams.',
       });
-      throw error;
+      this.changeState('connection-error', {
+        tag: 'errored',
+        errorReason:
+          'MongoDB MCP server does not support connecting to Atlas Streams',
+      });
+      return;
     }
+
+    await this.connectToVSCodeConnection(connectParams);
   }
 }

src/mcp/mcpController.ts

@@ -13,6 +13,7 @@ import {
 } 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';
@@ -238,10 +239,15 @@ ${jsonConfig}`,
     const connectionId = this.connectionController.getActiveConnectionId();
     const mongoClientOptions =
       this.connectionController.getMongoClientConnectionOptions();
-    await this.mcpConnectionManager?.updateConnection({
-      connectionId: connectionId ?? undefined,
-      connectionString: mongoClientOptions?.url,
-      connectOptions: mongoClientOptions?.options,
-    });
+
+    const connectParams: MCPConnectParams | undefined =
+      connectionId && mongoClientOptions
+        ? {
+            connectionId: connectionId,
+            connectionString: mongoClientOptions.url,
+            connectOptions: mongoClientOptions.options,
+          }
+        : undefined;
+    await this.mcpConnectionManager?.updateConnection(connectParams);
   }
 }

src/test/suite/mcp/mcpConnectionManager.test.ts

@@ -11,7 +11,7 @@ import { MCPConnectionManager } from '../../../mcp/mcpConnectionManager';
 chai.use(chaiAsPromised);
 
 const sandbox = sinon.createSandbox();
-suite('MCPConnectionManager Test Suite', function () {
+suite.only('MCPConnectionManager Test Suite', function () {
   let mcpConnectionManager: MCPConnectionManager;
   let fakeServiceProvider: NodeDriverServiceProvider;
 
@@ -93,8 +93,7 @@ suite('MCPConnectionManager Test Suite', function () {
       expect(mcpConnectionManager.currentConnectionState.tag).to.equal(
         'disconnected',
       );
-      expect((mcpConnectionManager as any).activeConnectionId).to.be.null;
-      expect((mcpConnectionManager as any).activeConnectionProvider).to.be.null;
+      expect((mcpConnectionManager as any).activeConnection).to.be.null;
     });
 
     test('should attempt to disconnect and on failure clear out the state', async function () {
@@ -115,8 +114,7 @@ suite('MCPConnectionManager Test Suite', function () {
       expect(mcpConnectionManager.currentConnectionState.tag).to.equal(
         'disconnected',
       );
-      expect((mcpConnectionManager as any).activeConnectionId).to.be.null;
-      expect((mcpConnectionManager as any).activeConnectionProvider).to.be.null;
+      expect((mcpConnectionManager as any).activeConnection).to.be.null;
     });
   });
 
@@ -128,11 +126,7 @@ suite('MCPConnectionManager Test Suite', function () {
           'connectToVSCodeConnection',
         );
         const disconnectSpy = sandbox.spy(mcpConnectionManager, 'disconnect');
-        await mcpConnectionManager.updateConnection({
-          connectionId: undefined,
-          connectionString: undefined,
-          connectOptions: undefined,
-        });
+        await mcpConnectionManager.updateConnection(undefined);
 
         expect(connectSpy).to.not.be.called;
         expect(disconnectSpy).to.not.be.called;
@@ -152,7 +146,7 @@ suite('MCPConnectionManager Test Suite', function () {
         });
 
         expect(connectSpy).to.not.be.called;
-        expect(disconnectSpy).to.not.be.called;
+        expect(disconnectSpy).to.be.called;
         expect(mcpConnectionManager.currentConnectionState.tag).to.equal(
           'errored',
         );
@@ -161,7 +155,7 @@ suite('MCPConnectionManager Test Suite', function () {
             mcpConnectionManager.currentConnectionState as ConnectionStateErrored
           ).errorReason,
         ).to.equal(
-          'MongoDB MCP server do not support connecting to Atlas Streams',
+          'MongoDB MCP server does not support connecting to Atlas Streams',
         );
       });
 
@@ -177,12 +171,12 @@ suite('MCPConnectionManager Test Suite', function () {
           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,
         });
-        expect(disconnectSpy).to.not.be.called;
       });
     });
 
@@ -223,11 +217,7 @@ suite('MCPConnectionManager Test Suite', function () {
           'connectToVSCodeConnection',
         );
         const disconnectSpy = sandbox.spy(mcpConnectionManager, 'disconnect');
-        await mcpConnectionManager.updateConnection({
-          connectionId: undefined,
-          connectionString: undefined,
-          connectOptions: undefined,
-        });
+        await mcpConnectionManager.updateConnection(undefined);
 
         expect(connectSpy).to.not.be.called;
         expect(disconnectSpy).to.be.called;
@@ -265,7 +255,7 @@ suite('MCPConnectionManager Test Suite', function () {
             mcpConnectionManager.currentConnectionState as ConnectionStateErrored
           ).errorReason,
         ).to.equal(
-          'MongoDB MCP server do not support connecting to Atlas Streams',
+          'MongoDB MCP server does not support connecting to Atlas Streams',
         );
       });
 

src/test/suite/mcp/mcpController.test.ts

@@ -15,7 +15,7 @@ import { TelemetryService } from '../../../telemetry';
 import { TEST_DATABASE_URI } from '../dbTestHelper';
 
 const sandbox = sinon.createSandbox();
-suite('MCPController test suite', function () {
+suite.only('MCPController test suite', function () {
   let extensionContext: ExtensionContext;
   let connectionController: ConnectionController;
   let mcpController: MCPController;