chore: NodeDriverServiceProvider.connect blocks, breaking previous contract

NodeDriverServiceProvider.connect would return a Promise that was
lazily connected. However, now, blocks until there is a ping from the
server, forcing a connection.

This would usually be fine, however, in OIDC it would mean blocking
the whole MCP flow until the OIDC flow is connected. With the previous
version, this was already handled, but now we need to split the logic
a bit more to ensure that OIDC flows stay async.

Author: kmruiz

src/common/connectionManager.ts

@@ -39,7 +39,7 @@ export interface ConnectionStateConnected extends ConnectionState {
 
 export interface ConnectionStateConnecting extends ConnectionState {
     tag: "connecting";
-    serviceProvider: NodeDriverServiceProvider;
+    serviceProvider: Promise<NodeDriverServiceProvider>;
     oidcConnectionType: OIDCConnectionAuthType;
     oidcLoginUrl?: string;
     oidcUserCode?: string;
@@ -141,9 +141,10 @@ export class MCPConnectionManager extends ConnectionManager {
             await this.disconnect();
         }
 
-        let serviceProvider: NodeDriverServiceProvider;
+        let serviceProvider: Promise<NodeDriverServiceProvider>;
         let connectionInfo: ConnectionInfo;
         let connectionStringAuthType: ConnectionStringAuthType = "scram";
+        let isOidcConnection: boolean = false;
 
         try {
             settings = { ...settings };
@@ -177,7 +178,8 @@ export class MCPConnectionManager extends ConnectionManager {
                 connectionInfo
             );
 
-            serviceProvider = await NodeDriverServiceProvider.connect(
+            isOidcConnection = connectionStringAuthType.startsWith("oidc");
+            serviceProvider = NodeDriverServiceProvider.connect(
                 connectionInfo.connectionString,
                 {
                     productDocsLink: "https://github.com/mongodb-js/mongodb-mcp-server/",
@@ -187,6 +189,16 @@ export class MCPConnectionManager extends ConnectionManager {
                 undefined,
                 this.bus
             );
+
+            if (isOidcConnection) {
+                this.changeState("connection-request", {
+                    tag: "connecting",
+                    serviceProvider,
+                    connectedAtlasCluster: settings.atlas,
+                    connectionStringAuthType,
+                    oidcConnectionType: connectionStringAuthType as OIDCConnectionAuthType,
+                });
+            }
         } catch (error: unknown) {
             const errorReason = error instanceof Error ? error.message : `${error as string}`;
             this.changeState("connection-error", {
@@ -199,26 +211,16 @@ export class MCPConnectionManager extends ConnectionManager {
         }
 
         try {
-            if (connectionStringAuthType.startsWith("oidc")) {
-                void this.pingAndForget(serviceProvider);
-
-                return this.changeState("connection-request", {
-                    tag: "connecting",
+            if (!isOidcConnection) {
+                return this.changeState("connection-success", {
+                    tag: "connected",
                     connectedAtlasCluster: settings.atlas,
-                    serviceProvider,
+                    serviceProvider: await serviceProvider,
                     connectionStringAuthType,
-                    oidcConnectionType: connectionStringAuthType as OIDCConnectionAuthType,
                 });
+            } else {
+                return this.currentConnectionState;
             }
-
-            await serviceProvider?.runCommand?.("admin", { hello: 1 });
-
-            return this.changeState("connection-success", {
-                tag: "connected",
-                connectedAtlasCluster: settings.atlas,
-                serviceProvider,
-                connectionStringAuthType,
-            });
         } catch (error: unknown) {
             const errorReason = error instanceof Error ? error.message : `${error as string}`;
             this.changeState("connection-error", {
@@ -238,7 +240,13 @@ export class MCPConnectionManager extends ConnectionManager {
 
         if (this.currentConnectionState.tag === "connected" || this.currentConnectionState.tag === "connecting") {
             try {
-                await this.currentConnectionState.serviceProvider?.close();
+                if (this.currentConnectionState.tag === "connected") {
+                    await this.currentConnectionState.serviceProvider?.close();
+                }
+                if (this.currentConnectionState.tag === "connecting") {
+                    const serviceProvider = await this.currentConnectionState.serviceProvider;
+                    await serviceProvider.close();
+                }
             } finally {
                 this.changeState("connection-close", {
                     tag: "disconnected",
@@ -258,12 +266,16 @@ export class MCPConnectionManager extends ConnectionManager {
         }
     }
 
-    private onOidcAuthSucceeded(): void {
+    private async onOidcAuthSucceeded(): Promise<void> {
         if (
             this.currentConnectionState.tag === "connecting" &&
             this.currentConnectionState.connectionStringAuthType?.startsWith("oidc")
         ) {
-            this.changeState("connection-success", { ...this.currentConnectionState, tag: "connected" });
+            this.changeState("connection-success", {
+                ...this.currentConnectionState,
+                tag: "connected",
+                serviceProvider: await this.currentConnectionState.serviceProvider,
+            });
         }
 
         this.logger.info({
@@ -330,18 +342,6 @@ export class MCPConnectionManager extends ConnectionManager {
         }
     }
 
-    private async pingAndForget(serviceProvider: NodeDriverServiceProvider): Promise<void> {
-        try {
-            await serviceProvider?.runCommand?.("admin", { hello: 1 });
-        } catch (error: unknown) {
-            this.logger.warning({
-                id: LogId.oidcFlow,
-                context: "pingAndForget",
-                message: String(error),
-            });
-        }
-    }
-
     private async disconnectOnOidcError(error: unknown): Promise<void> {
         try {
             await this.disconnect();