chore: refactor connections to use the new ConnectionManager to isolate long running processes like OIDC connections MCP-81 (#423)

Author: kmruiz

src/common/connectionManager.ts

@@ -0,0 +1,192 @@
+import { ConnectOptions } from "./config.js";
+import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
+import EventEmitter from "events";
+import { setAppNameParamIfMissing } from "../helpers/connectionOptions.js";
+import { packageInfo } from "./packageInfo.js";
+import ConnectionString from "mongodb-connection-string-url";
+import { MongoClientOptions } from "mongodb";
+import { ErrorCodes, MongoDBError } from "./errors.js";
+
+export interface AtlasClusterConnectionInfo {
+    username: string;
+    projectId: string;
+    clusterName: string;
+    expiryDate: Date;
+}
+
+export interface ConnectionSettings extends ConnectOptions {
+    connectionString: string;
+    atlas?: AtlasClusterConnectionInfo;
+}
+
+type ConnectionTag = "connected" | "connecting" | "disconnected" | "errored";
+type OIDCConnectionAuthType = "oidc-auth-flow" | "oidc-device-flow";
+export type ConnectionStringAuthType = "scram" | "ldap" | "kerberos" | OIDCConnectionAuthType | "x.509";
+
+export interface ConnectionState {
+    tag: ConnectionTag;
+    connectionStringAuthType?: ConnectionStringAuthType;
+    connectedAtlasCluster?: AtlasClusterConnectionInfo;
+}
+
+export interface ConnectionStateConnected extends ConnectionState {
+    tag: "connected";
+    serviceProvider: NodeDriverServiceProvider;
+}
+
+export interface ConnectionStateConnecting extends ConnectionState {
+    tag: "connecting";
+    serviceProvider: NodeDriverServiceProvider;
+    oidcConnectionType: OIDCConnectionAuthType;
+    oidcLoginUrl?: string;
+    oidcUserCode?: string;
+}
+
+export interface ConnectionStateDisconnected extends ConnectionState {
+    tag: "disconnected";
+}
+
+export interface ConnectionStateErrored extends ConnectionState {
+    tag: "errored";
+    errorReason: string;
+}
+
+export type AnyConnectionState =
+    | ConnectionStateConnected
+    | ConnectionStateConnecting
+    | ConnectionStateDisconnected
+    | ConnectionStateErrored;
+
+export interface ConnectionManagerEvents {
+    "connection-requested": [AnyConnectionState];
+    "connection-succeeded": [ConnectionStateConnected];
+    "connection-timed-out": [ConnectionStateErrored];
+    "connection-closed": [ConnectionStateDisconnected];
+    "connection-errored": [ConnectionStateErrored];
+}
+
+export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
+    private state: AnyConnectionState;
+
+    constructor() {
+        super();
+        this.state = { tag: "disconnected" };
+    }
+
+    async connect(settings: ConnectionSettings): Promise<AnyConnectionState> {
+        this.emit("connection-requested", this.state);
+
+        if (this.state.tag === "connected" || this.state.tag === "connecting") {
+            await this.disconnect();
+        }
+
+        let serviceProvider: NodeDriverServiceProvider;
+        try {
+            settings = { ...settings };
+            settings.connectionString = setAppNameParamIfMissing({
+                connectionString: settings.connectionString,
+                defaultAppName: `${packageInfo.mcpServerName} ${packageInfo.version}`,
+            });
+
+            serviceProvider = await NodeDriverServiceProvider.connect(settings.connectionString, {
+                productDocsLink: "https://github.com/mongodb-js/mongodb-mcp-server/",
+                productName: "MongoDB MCP",
+                readConcern: {
+                    level: settings.readConcern,
+                },
+                readPreference: settings.readPreference,
+                writeConcern: {
+                    w: settings.writeConcern,
+                },
+                timeoutMS: settings.timeoutMS,
+                proxy: { useEnvironmentVariableProxies: true },
+                applyProxyToOIDC: true,
+            });
+        } catch (error: unknown) {
+            const errorReason = error instanceof Error ? error.message : `${error as string}`;
+            this.changeState("connection-errored", {
+                tag: "errored",
+                errorReason,
+                connectedAtlasCluster: settings.atlas,
+            });
+            throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, errorReason);
+        }
+
+        try {
+            await serviceProvider?.runCommand?.("admin", { hello: 1 });
+
+            return this.changeState("connection-succeeded", {
+                tag: "connected",
+                connectedAtlasCluster: settings.atlas,
+                serviceProvider,
+                connectionStringAuthType: ConnectionManager.inferConnectionTypeFromSettings(settings),
+            });
+        } catch (error: unknown) {
+            const errorReason = error instanceof Error ? error.message : `${error as string}`;
+            this.changeState("connection-errored", {
+                tag: "errored",
+                errorReason,
+                connectedAtlasCluster: settings.atlas,
+            });
+            throw new MongoDBError(ErrorCodes.NotConnectedToMongoDB, errorReason);
+        }
+    }
+
+    async disconnect(): Promise<ConnectionStateDisconnected | ConnectionStateErrored> {
+        if (this.state.tag === "disconnected" || this.state.tag === "errored") {
+            return this.state;
+        }
+
+        if (this.state.tag === "connected" || this.state.tag === "connecting") {
+            try {
+                await this.state.serviceProvider?.close(true);
+            } finally {
+                this.changeState("connection-closed", {
+                    tag: "disconnected",
+                });
+            }
+        }
+
+        return { tag: "disconnected" };
+    }
+
+    get currentConnectionState(): AnyConnectionState {
+        return this.state;
+    }
+
+    changeState<Event extends keyof ConnectionManagerEvents, State extends ConnectionManagerEvents[Event][0]>(
+        event: Event,
+        newState: State
+    ): State {
+        this.state = newState;
+        // TypeScript doesn't seem to be happy with the spread operator and generics
+        // eslint-disable-next-line
+        this.emit(event, ...([newState] as any));
+        return newState;
+    }
+
+    static inferConnectionTypeFromSettings(settings: ConnectionSettings): ConnectionStringAuthType {
+        const connString = new ConnectionString(settings.connectionString);
+        const searchParams = connString.typedSearchParams<MongoClientOptions>();
+
+        switch (searchParams.get("authMechanism")) {
+            case "MONGODB-OIDC": {
+                return "oidc-auth-flow"; // TODO: depending on if we don't have a --browser later it can be oidc-device-flow
+            }
+            case "MONGODB-X509":
+                return "x.509";
+            case "GSSAPI":
+                return "kerberos";
+            case "PLAIN":
+                if (searchParams.get("authSource") === "$external") {
+                    return "ldap";
+                }
+                return "scram";
+            // default should catch also null, but eslint complains
+            // about it.
+            case null:
+            default:
+                return "scram";
+        }
+    }
+}

src/common/session.ts

@@ -1,16 +1,21 @@
-import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { ApiClient, ApiClientCredentials } from "./atlas/apiClient.js";
 import { Implementation } from "@modelcontextprotocol/sdk/types.js";
 import logger, { LogId } from "./logger.js";
 import EventEmitter from "events";
-import { ConnectOptions } from "./config.js";
-import { setAppNameParamIfMissing } from "../helpers/connectionOptions.js";
-import { packageInfo } from "./packageInfo.js";
+import {
+    AtlasClusterConnectionInfo,
+    ConnectionManager,
+    ConnectionSettings,
+    ConnectionStateConnected,
+} from "./connectionManager.js";
+import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
+import { ErrorCodes, MongoDBError } from "./errors.js";
 
 export interface SessionOptions {
     apiBaseUrl: string;
     apiClientId?: string;
     apiClientSecret?: string;
+    connectionManager?: ConnectionManager;
 }
 
 export type SessionEvents = {
@@ -22,20 +27,14 @@ export type SessionEvents = {
 
 export class Session extends EventEmitter<SessionEvents> {
     sessionId?: string;
-    serviceProvider?: NodeDriverServiceProvider;
+    connectionManager: ConnectionManager;
     apiClient: ApiClient;
     agentRunner?: {
         name: string;
         version: string;
     };
-    connectedAtlasCluster?: {
-        username: string;
-        projectId: string;
-        clusterName: string;
-        expiryDate: Date;
-    };
 
-    constructor({ apiBaseUrl, apiClientId, apiClientSecret }: SessionOptions) {
+    constructor({ apiBaseUrl, apiClientId, apiClientSecret, connectionManager }: SessionOptions) {
         super();
 
         const credentials: ApiClientCredentials | undefined =
@@ -46,10 +45,13 @@ export class Session extends EventEmitter<SessionEvents> {
                   }
                 : undefined;
 
-        this.apiClient = new ApiClient({
-            baseUrl: apiBaseUrl,
-            credentials,
-        });
+        this.apiClient = new ApiClient({ baseUrl: apiBaseUrl, credentials });
+
+        this.connectionManager = connectionManager ?? new ConnectionManager();
+        this.connectionManager.on("connection-succeeded", () => this.emit("connect"));
+        this.connectionManager.on("connection-timed-out", (error) => this.emit("connection-error", error.errorReason));
+        this.connectionManager.on("connection-closed", () => this.emit("disconnect"));
+        this.connectionManager.on("connection-errored", (error) => this.emit("connection-error", error.errorReason));
     }
 
     setAgentRunner(agentRunner: Implementation | undefined) {
@@ -62,22 +64,22 @@ export class Session extends EventEmitter<SessionEvents> {
     }
 
     async disconnect(): Promise<void> {
-        if (this.serviceProvider) {
-            try {
-                await this.serviceProvider.close(true);
-            } catch (err: unknown) {
-                const error = err instanceof Error ? err : new Error(String(err));
-                logger.error(LogId.mongodbDisconnectFailure, "Error closing service provider:", error.message);
-            }
-            this.serviceProvider = undefined;
+        const atlasCluster = this.connectedAtlasCluster;
+
+        try {
+            await this.connectionManager.disconnect();
+        } catch (err: unknown) {
+            const error = err instanceof Error ? err : new Error(String(err));
+            logger.error(LogId.mongodbDisconnectFailure, "Error closing service provider:", error.message);
         }
-        if (this.connectedAtlasCluster?.username && this.connectedAtlasCluster?.projectId) {
+
+        if (atlasCluster?.username && atlasCluster?.projectId) {
             void this.apiClient
                 .deleteDatabaseUser({
                     params: {
                         path: {
-                            groupId: this.connectedAtlasCluster.projectId,
-                            username: this.connectedAtlasCluster.username,
+                            groupId: atlasCluster.projectId,
+                            username: atlasCluster.username,
                             databaseName: "admin",
                         },
                     },
@@ -90,9 +92,7 @@ export class Session extends EventEmitter<SessionEvents> {
                         `Error deleting previous database user: ${error.message}`
                     );
                 });
-            this.connectedAtlasCluster = undefined;
         }
-        this.emit("disconnect");
     }
 
     async close(): Promise<void> {
@@ -101,35 +101,30 @@ export class Session extends EventEmitter<SessionEvents> {
         this.emit("close");
     }
 
-    async connectToMongoDB(connectionString: string, connectOptions: ConnectOptions): Promise<void> {
-        connectionString = setAppNameParamIfMissing({
-            connectionString,
-            defaultAppName: `${packageInfo.mcpServerName} ${packageInfo.version}`,
-        });
-
+    async connectToMongoDB(settings: ConnectionSettings): Promise<void> {
         try {
-            this.serviceProvider = await NodeDriverServiceProvider.connect(connectionString, {
-                productDocsLink: "https://github.com/mongodb-js/mongodb-mcp-server/",
-                productName: "MongoDB MCP",
-                readConcern: {
-                    level: connectOptions.readConcern,
-                },
-                readPreference: connectOptions.readPreference,
-                writeConcern: {
-                    w: connectOptions.writeConcern,
-                },
-                timeoutMS: connectOptions.timeoutMS,
-                proxy: { useEnvironmentVariableProxies: true },
-                applyProxyToOIDC: true,
-            });
-
-            await this.serviceProvider?.runCommand?.("admin", { hello: 1 });
+            await this.connectionManager.connect({ ...settings });
         } catch (error: unknown) {
-            const message = error instanceof Error ? error.message : `${error as string}`;
+            const message = error instanceof Error ? error.message : (error as string);
             this.emit("connection-error", message);
             throw error;
         }
+    }
+
+    get isConnectedToMongoDB(): boolean {
+        return this.connectionManager.currentConnectionState.tag === "connected";
+    }
+
+    get serviceProvider(): NodeDriverServiceProvider {
+        if (this.isConnectedToMongoDB) {
+            const state = this.connectionManager.currentConnectionState as ConnectionStateConnected;
+            return state.serviceProvider;
+        }
+
+        throw new MongoDBError(ErrorCodes.NotConnectedToMongoDB, "Not connected to MongoDB");
+    }
 
-        this.emit("connect");
+    get connectedAtlasCluster(): AtlasClusterConnectionInfo | undefined {
+        return this.connectionManager.currentConnectionState.connectedAtlasCluster;
     }
 }

src/resources/common/debug.ts

@@ -10,10 +10,11 @@ type ConnectionStateDebuggingInformation = {
 
 export class DebugResource extends ReactiveResource(
     {
-        name: "debug-mongodb-connectivity",
-        uri: "debug://mongodb-connectivity",
+        name: "debug-mongodb",
+        uri: "debug://mongodb",
         config: {
-            description: "Debugging information for connectivity issues.",
+            description:
+                "Debugging information for MongoDB connectivity issues. Tracks the last connectivity error and attempt information.",
         },
     },
     {