chore: extract abstract ConnectionManager

This commit extracts an abstract class ConnectionManager out of
MCPConnectionManager and modifies the TransportRunner interface to have
the ConnectionManager implementation injected through a factory
function.

Contains a small drive by fix for making the ConnectionManager event
emitting internal to the class itself.

Author: himanshusinghs

package-lock.json

@@ -0,0 +1,92 @@
+import { EventEmitter } from "events";
+import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
+
+export interface AtlasClusterConnectionInfo {
+    username: string;
+    projectId: string;
+    clusterName: string;
+    expiryDate: Date;
+}
+
+type ConnectionTag = "connected" | "connecting" | "disconnected" | "errored";
+export 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 interface MCPConnectParams {
+    connectionString: string;
+    atlas?: AtlasClusterConnectionInfo;
+}
+
+export abstract class ConnectionManager<ConnectParams extends MCPConnectParams = MCPConnectParams> {
+    protected clientName: string = "unknown";
+
+    protected readonly _events = new EventEmitter<ConnectionManagerEvents>();
+    readonly events: Pick<EventEmitter<ConnectionManagerEvents>, "on" | "off" | "once"> = this._events;
+
+    protected state: AnyConnectionState = { 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._events.emit(event, ...([newState] as any));
+        return newState;
+    }
+
+    setClientName(clientName: string): void {
+        this.clientName = clientName;
+    }
+
+    abstract connect(connectParams: ConnectParams): Promise<AnyConnectionState>;
+
+    abstract disconnect(): Promise<ConnectionStateDisconnected | ConnectionStateErrored>;
+}

src/common/connectionManager.ts

@@ -12,69 +12,18 @@ import type { CompositeLogger } from "./logger.js";
 import { LogId } from "./logger.js";
 import type { ConnectionInfo } from "@mongosh/arg-parser";
 import { generateConnectionInfoFromCliArgs } from "@mongosh/arg-parser";
-
-export interface AtlasClusterConnectionInfo {
-    username: string;
-    projectId: string;
-    clusterName: string;
-    expiryDate: Date;
-}
-
-export interface ConnectionSettings {
-    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 MCPConnectionManagerEvents {
-    "connection-requested": [AnyConnectionState];
-    "connection-succeeded": [ConnectionStateConnected];
-    "connection-timed-out": [ConnectionStateErrored];
-    "connection-closed": [ConnectionStateDisconnected];
-    "connection-errored": [ConnectionStateErrored];
-}
-
-export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvents> {
-    private state: AnyConnectionState;
+import {
+    ConnectionManager,
+    type AnyConnectionState,
+    type ConnectionStringAuthType,
+    type OIDCConnectionAuthType,
+    type ConnectionStateDisconnected,
+    type ConnectionStateErrored,
+    type MCPConnectParams,
+} from "./connectionManager.js";
+
+export class MCPConnectionManager extends ConnectionManager<MCPConnectParams> {
     private deviceId: DeviceId;
-    private clientName: string;
     private bus: EventEmitter;
 
     constructor(
@@ -85,23 +34,15 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
         bus?: EventEmitter
     ) {
         super();
-
         this.bus = bus ?? new EventEmitter();
-        this.state = { tag: "disconnected" };
-
         this.bus.on("mongodb-oidc-plugin:auth-failed", this.onOidcAuthFailed.bind(this));
         this.bus.on("mongodb-oidc-plugin:auth-succeeded", this.onOidcAuthSucceeded.bind(this));
-
         this.deviceId = deviceId;
         this.clientName = "unknown";
     }
 
-    setClientName(clientName: string): void {
-        this.clientName = clientName;
-    }
-
-    async connect(settings: ConnectionSettings): Promise<AnyConnectionState> {
-        this.emit("connection-requested", this.state);
+    async connect(connectParams: MCPConnectParams): Promise<AnyConnectionState> {
+        this._events.emit("connection-requested", this.state);
 
         if (this.state.tag === "connected" || this.state.tag === "connecting") {
             await this.disconnect();
@@ -111,22 +52,22 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
         let connectionInfo: ConnectionInfo;
 
         try {
-            settings = { ...settings };
+            connectParams = { ...connectParams };
             const appNameComponents: AppNameComponents = {
                 appName: `${packageInfo.mcpServerName} ${packageInfo.version}`,
                 deviceId: this.deviceId.get(),
                 clientName: this.clientName,
             };
 
-            settings.connectionString = await setAppNameParamIfMissing({
-                connectionString: settings.connectionString,
+            connectParams.connectionString = await setAppNameParamIfMissing({
+                connectionString: connectParams.connectionString,
                 components: appNameComponents,
             });
 
             connectionInfo = generateConnectionInfoFromCliArgs({
                 ...this.userConfig,
                 ...this.driverOptions,
-                connectionSpecifier: settings.connectionString,
+                connectionSpecifier: connectParams.connectionString,
             });
 
             if (connectionInfo.driverOptions.oidc) {
@@ -152,7 +93,7 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
             this.changeState("connection-errored", {
                 tag: "errored",
                 errorReason,
-                connectedAtlasCluster: settings.atlas,
+                connectedAtlasCluster: connectParams.atlas,
             });
             throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, errorReason);
         }
@@ -167,7 +108,7 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
 
                 return this.changeState("connection-requested", {
                     tag: "connecting",
-                    connectedAtlasCluster: settings.atlas,
+                    connectedAtlasCluster: connectParams.atlas,
                     serviceProvider,
                     connectionStringAuthType: connectionType,
                     oidcConnectionType: connectionType as OIDCConnectionAuthType,
@@ -178,7 +119,7 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
 
             return this.changeState("connection-succeeded", {
                 tag: "connected",
-                connectedAtlasCluster: settings.atlas,
+                connectedAtlasCluster: connectParams.atlas,
                 serviceProvider,
                 connectionStringAuthType: connectionType,
             });
@@ -187,7 +128,7 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
             this.changeState("connection-errored", {
                 tag: "errored",
                 errorReason,
-                connectedAtlasCluster: settings.atlas,
+                connectedAtlasCluster: connectParams.atlas,
             });
             throw new MongoDBError(ErrorCodes.NotConnectedToMongoDB, errorReason);
         }
@@ -211,21 +152,6 @@ export class MCPConnectionManager extends EventEmitter<MCPConnectionManagerEvent
         return { tag: "disconnected" };
     }
 
-    get currentConnectionState(): AnyConnectionState {
-        return this.state;
-    }
-
-    changeState<Event extends keyof MCPConnectionManagerEvents, State extends MCPConnectionManagerEvents[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;
-    }
-
     private onOidcAuthFailed(error: unknown): void {
         if (this.state.tag === "connecting" && this.state.connectionStringAuthType?.startsWith("oidc")) {
             void this.disconnectOnOidcError(error);

src/common/mcpConnectionManager.ts

@@ -7,10 +7,10 @@ import { LogId } from "./logger.js";
 import EventEmitter from "events";
 import type {
     AtlasClusterConnectionInfo,
-    MCPConnectionManager,
-    ConnectionSettings,
+    ConnectionManager,
     ConnectionStateConnected,
-} from "./mcpConnectionManager.js";
+    MCPConnectParams,
+} from "./connectionManager.js";
 import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { ErrorCodes, MongoDBError } from "./errors.js";
 import type { ExportsManager } from "./exportsManager.js";
@@ -21,7 +21,7 @@ export interface SessionOptions {
     apiClientSecret?: string;
     logger: CompositeLogger;
     exportsManager: ExportsManager;
-    connectionManager: MCPConnectionManager;
+    connectionManager: ConnectionManager;
 }
 
 export type SessionEvents = {
@@ -34,7 +34,7 @@ export type SessionEvents = {
 export class Session extends EventEmitter<SessionEvents> {
     readonly sessionId: string = new ObjectId().toString();
     readonly exportsManager: ExportsManager;
-    readonly connectionManager: MCPConnectionManager;
+    readonly connectionManager: ConnectionManager;
     readonly apiClient: ApiClient;
     mcpClient?: {
         name?: string;
@@ -66,10 +66,14 @@ export class Session extends EventEmitter<SessionEvents> {
         this.apiClient = new ApiClient({ baseUrl: apiBaseUrl, credentials }, logger);
         this.exportsManager = exportsManager;
         this.connectionManager = 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));
+        this.connectionManager.events.on("connection-succeeded", () => this.emit("connect"));
+        this.connectionManager.events.on("connection-timed-out", (error) =>
+            this.emit("connection-error", error.errorReason)
+        );
+        this.connectionManager.events.on("connection-closed", () => this.emit("disconnect"));
+        this.connectionManager.events.on("connection-errored", (error) =>
+            this.emit("connection-error", error.errorReason)
+        );
     }
 
     setMcpClient(mcpClient: Implementation | undefined): void {
@@ -135,9 +139,9 @@ export class Session extends EventEmitter<SessionEvents> {
         this.emit("close");
     }
 
-    async connectToMongoDB(settings: ConnectionSettings): Promise<void> {
+    async connectToMongoDB(connectParams: MCPConnectParams): Promise<void> {
         try {
-            await this.connectionManager.connect({ ...settings });
+            await this.connectionManager.connect({ ...connectParams });
         } catch (error: unknown) {
             const message = error instanceof Error ? error.message : (error as string);
             this.emit("connection-error", message);

src/common/session.ts

@@ -42,17 +42,23 @@ import { packageInfo } from "./common/packageInfo.js";
 import { StdioRunner } from "./transports/stdio.js";
 import { StreamableHttpRunner } from "./transports/streamableHttp.js";
 import { systemCA } from "@mongodb-js/devtools-proxy-support";
+import type { MCPConnectParams } from "./lib.js";
+import type { CreateConnectionManagerFn } from "./transports/base.js";
+import { MCPConnectionManager } from "./common/mcpConnectionManager.js";
 
 async function main(): Promise<void> {
     systemCA().catch(() => undefined); // load system CA asynchronously as in mongosh
 
     assertHelpMode();
     assertVersionMode();
 
+    const createConnectionManager: CreateConnectionManagerFn<MCPConnectParams> = ({ logger, deviceId }) =>
+        new MCPConnectionManager(config, driverOptions, logger, deviceId);
+
     const transportRunner =
         config.transport === "stdio"
-            ? new StdioRunner(config, driverOptions)
-            : new StreamableHttpRunner(config, driverOptions);
+            ? new StdioRunner<MCPConnectParams>(config, createConnectionManager)
+            : new StreamableHttpRunner<MCPConnectParams>(config, createConnectionManager);
     const shutdown = (): void => {
         transportRunner.logger.info({
             id: LogId.serverCloseRequested,

src/index.ts

@@ -3,5 +3,5 @@ export { Telemetry } from "./telemetry/telemetry.js";
 export { Session, type SessionOptions } from "./common/session.js";
 export { type UserConfig, defaultUserConfig } from "./common/config.js";
 export { StreamableHttpRunner } from "./transports/streamableHttp.js";
-export { LoggerBase } from "./common/logger.js";
-export type { LogPayload, LoggerType, LogLevel } from "./common/logger.js";
+export { LoggerBase, CompositeLogger, type LogPayload, type LoggerType, type LogLevel } from "./common/logger.js";
+export * from "./common/connectionManager.js";

src/lib.ts

@@ -6,7 +6,7 @@ import { generateSecurePassword } from "../../../helpers/generatePassword.js";
 import { LogId } from "../../../common/logger.js";
 import { inspectCluster } from "../../../common/atlas/cluster.js";
 import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUtils.js";
-import type { AtlasClusterConnectionInfo } from "../../../common/mcpConnectionManager.js";
+import type { AtlasClusterConnectionInfo } from "../../../common/connectionManager.js";
 import { getDefaultRoleFromConfig } from "../../../common/atlas/roles.js";
 
 const EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours