Skip to content

Commit 3977aa7

Browse files
kmruizkmruiz
committed
chore: refactor to include the new ConnectionManager
Now the connection itself and the user flow is managed by the new connection manager. There are tests missing for this class, that will be added in a following commit.
1 parent 0083493 commit 3977aa7

File tree

11 files changed

+269
-63
lines changed

11 files changed

+269
-63
lines changed

‎src/common/connectionManager.ts

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
import { ConnectOptions } from "./config.js";
2+
import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
3+
import EventEmitter from "events";
4+
import { setAppNameParamIfMissing } from "../helpers/connectionOptions.js";
5+
import { packageInfo } from "./packageInfo.js";
6+
import ConnectionString from "mongodb-connection-string-url";
7+
import { MongoClientOptions } from "mongodb";
8+
import { ErrorCodes, MongoDBError } from "./errors.js";
9+
10+
export interface AtlasClusterConnectionInfo {
11+
username: string;
12+
projectId: string;
13+
clusterName: string;
14+
expiryDate: Date;
15+
}
16+
17+
export interface ConnectionSettings extends ConnectOptions {
18+
connectionString: string;
19+
atlas?: AtlasClusterConnectionInfo;
20+
}
21+
22+
type ConnectionTag = "connected" | "connecting" | "disconnected" | "errored";
23+
type OIDCConnectionAuthType = "oidc-auth-flow" | "oidc-device-flow";
24+
type ConnectionStringAuthType = "scram" | "ldap" | "kerberos" | OIDCConnectionAuthType | "x.509";
25+
26+
export interface ConnectionState {
27+
tag: ConnectionTag;
28+
connectionStringAuthType?: ConnectionStringAuthType;
29+
}
30+
31+
export interface ConnectionStateConnected extends ConnectionState {
32+
tag: "connected";
33+
serviceProvider: NodeDriverServiceProvider;
34+
connectedAtlasCluster?: AtlasClusterConnectionInfo;
35+
}
36+
37+
export interface ConnectionStateConnecting extends ConnectionState {
38+
tag: "connecting";
39+
serviceProvider: NodeDriverServiceProvider;
40+
oidcConnectionType: OIDCConnectionAuthType;
41+
oidcLoginUrl?: string;
42+
oidcUserCode?: string;
43+
}
44+
45+
export interface ConnectionStateDisconnected extends ConnectionState {
46+
tag: "disconnected";
47+
}
48+
49+
export interface ConnectionStateErrored extends ConnectionState {
50+
tag: "errored";
51+
errorReason: string;
52+
}
53+
54+
export type AnyConnectionState =
55+
| ConnectionState
56+
| ConnectionStateConnected
57+
| ConnectionStateConnecting
58+
| ConnectionStateDisconnected
59+
| ConnectionStateErrored;
60+
61+
export interface ConnectionManagerEvents {
62+
"connection-requested": [AnyConnectionState];
63+
"connection-succeeded": [ConnectionStateConnected];
64+
"connection-timed-out": [ConnectionStateErrored];
65+
"connection-closed": [ConnectionStateDisconnected];
66+
"connection-errored": [ConnectionStateErrored];
67+
}
68+
69+
export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
70+
private state: AnyConnectionState;
71+
72+
constructor() {
73+
super();
74+
this.state = { tag: "disconnected" };
75+
}
76+
77+
async connect(settings: ConnectionSettings): Promise<AnyConnectionState> {
78+
if (this.state.tag == "connected" || this.state.tag == "connecting") {
79+
await this.disconnect();
80+
}
81+
82+
let serviceProvider: NodeDriverServiceProvider;
83+
try {
84+
settings = { ...settings };
85+
settings.connectionString = setAppNameParamIfMissing({
86+
connectionString: settings.connectionString,
87+
defaultAppName: `${packageInfo.mcpServerName} ${packageInfo.version}`,
88+
});
89+
90+
serviceProvider = await NodeDriverServiceProvider.connect(settings.connectionString, {
91+
productDocsLink: "https://github.com/mongodb-js/mongodb-mcp-server/",
92+
productName: "MongoDB MCP",
93+
readConcern: {
94+
level: settings.readConcern,
95+
},
96+
readPreference: settings.readPreference,
97+
writeConcern: {
98+
w: settings.writeConcern,
99+
},
100+
timeoutMS: settings.timeoutMS,
101+
proxy: { useEnvironmentVariableProxies: true },
102+
applyProxyToOIDC: true,
103+
});
104+
} catch (error: unknown) {
105+
const errorReason = error instanceof Error ? error.message : `${error as string}`;
106+
this.changeState("connection-errored", { tag: "errored", errorReason });
107+
throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, errorReason);
108+
}
109+
110+
try {
111+
await serviceProvider?.runCommand?.("admin", { hello: 1 });
112+
113+
return this.changeState("connection-succeeded", {
114+
tag: "connected",
115+
connectedAtlasCluster: settings.atlas,
116+
serviceProvider,
117+
connectionStringAuthType: this.inferConnectionTypeFromSettings(settings),
118+
});
119+
} catch (error: unknown) {
120+
const errorReason = error instanceof Error ? error.message : `${error as string}`;
121+
this.changeState("connection-errored", { tag: "errored", errorReason });
122+
throw new MongoDBError(ErrorCodes.NotConnectedToMongoDB, errorReason);
123+
}
124+
}
125+
126+
async disconnect(): Promise<ConnectionStateDisconnected | ConnectionStateErrored> {
127+
if (this.state.tag == "disconnected") {
128+
return this.state as ConnectionStateDisconnected;
129+
}
130+
131+
if (this.state.tag == "errored") {
132+
return this.state as ConnectionStateErrored;
133+
}
134+
135+
if (this.state.tag == "connected" || this.state.tag == "connecting") {
136+
const state = this.state as ConnectionStateConnecting | ConnectionStateConnected;
137+
try {
138+
await state.serviceProvider?.close(true);
139+
} finally {
140+
this.changeState("connection-closed", { tag: "disconnected" });
141+
}
142+
}
143+
144+
return { tag: "disconnected" };
145+
}
146+
147+
get currentConnectionState(): AnyConnectionState {
148+
return this.state;
149+
}
150+
151+
changeState<State extends AnyConnectionState>(event: keyof ConnectionManagerEvents, newState: State): State {
152+
this.state = newState;
153+
this.emit(event, newState);
154+
return newState;
155+
}
156+
157+
private inferConnectionTypeFromSettings(settings: ConnectionSettings): ConnectionStringAuthType {
158+
const connString = new ConnectionString(settings.connectionString);
159+
const searchParams = connString.typedSearchParams<MongoClientOptions>();
160+
161+
switch (searchParams.get("authMechanism")) {
162+
case "MONGODB-OIDC": {
163+
return "oidc-auth-flow"; // TODO: depending on if we don't have a --browser later it can be oidc-device-flow
164+
}
165+
case "MONGODB-X509":
166+
return "x.509";
167+
case "GSSAPI":
168+
return "kerberos";
169+
case "PLAIN":
170+
if (searchParams.get("authSource") == "$external") {
171+
return "ldap";
172+
}
173+
break;
174+
default:
175+
return "scram";
176+
}
177+
return "scram";
178+
}
179+
}

‎src/common/session.ts

Lines changed: 39 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,21 @@
1-
import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
21
import { ApiClient, ApiClientCredentials } from "./atlas/apiClient.js";
32
import { Implementation } from "@modelcontextprotocol/sdk/types.js";
43
import logger, { LogId } from "./logger.js";
54
import EventEmitter from "events";
6-
import { ConnectOptions } from "./config.js";
7-
import { setAppNameParamIfMissing } from "../helpers/connectionOptions.js";
8-
import { packageInfo } from "./packageInfo.js";
5+
import {
6+
AnyConnectionState,
7+
ConnectionManager,
8+
ConnectionSettings,
9+
ConnectionStateConnected,
10+
} from "./connectionManager.js";
11+
import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
12+
import { ErrorCodes, MongoDBError } from "./errors.js";
913

1014
export interface SessionOptions {
1115
apiBaseUrl: string;
1216
apiClientId?: string;
1317
apiClientSecret?: string;
18+
connectionManager?: ConnectionManager;
1419
}
1520

1621
export type SessionEvents = {
@@ -22,7 +27,7 @@ export type SessionEvents = {
2227

2328
export class Session extends EventEmitter<SessionEvents> {
2429
sessionId?: string;
25-
serviceProvider?: NodeDriverServiceProvider;
30+
connectionManager: ConnectionManager;
2631
apiClient: ApiClient;
2732
agentRunner?: {
2833
name: string;
@@ -35,7 +40,7 @@ export class Session extends EventEmitter<SessionEvents> {
3540
expiryDate: Date;
3641
};
3742

38-
constructor({ apiBaseUrl, apiClientId, apiClientSecret }: SessionOptions) {
43+
constructor({ apiBaseUrl, apiClientId, apiClientSecret, connectionManager }: SessionOptions) {
3944
super();
4045

4146
const credentials: ApiClientCredentials | undefined =
@@ -46,10 +51,13 @@ export class Session extends EventEmitter<SessionEvents> {
4651
}
4752
: undefined;
4853

49-
this.apiClient = new ApiClient({
50-
baseUrl: apiBaseUrl,
51-
credentials,
52-
});
54+
this.apiClient = new ApiClient({ baseUrl: apiBaseUrl, credentials });
55+
56+
this.connectionManager = connectionManager ?? new ConnectionManager();
57+
this.connectionManager.on("connection-succeeded", () => this.emit("connect"));
58+
this.connectionManager.on("connection-timed-out", (error) => this.emit("connection-error", error.errorReason));
59+
this.connectionManager.on("connection-closed", () => this.emit("disconnect"));
60+
this.connectionManager.on("connection-errored", (error) => this.emit("connection-error", error.errorReason));
5361
}
5462

5563
setAgentRunner(agentRunner: Implementation | undefined) {
@@ -62,15 +70,13 @@ export class Session extends EventEmitter<SessionEvents> {
6270
}
6371

6472
async disconnect(): Promise<void> {
65-
if (this.serviceProvider) {
66-
try {
67-
await this.serviceProvider.close(true);
68-
} catch (err: unknown) {
69-
const error = err instanceof Error ? err : new Error(String(err));
70-
logger.error(LogId.mongodbDisconnectFailure, "Error closing service provider:", error.message);
71-
}
72-
this.serviceProvider = undefined;
73+
try {
74+
await this.connectionManager.disconnect();
75+
} catch (err: unknown) {
76+
const error = err instanceof Error ? err : new Error(String(err));
77+
logger.error(LogId.mongodbDisconnectFailure, "Error closing service provider:", error.message);
7378
}
79+
7480
if (this.connectedAtlasCluster?.username && this.connectedAtlasCluster?.projectId) {
7581
void this.apiClient
7682
.deleteDatabaseUser({
@@ -92,44 +98,33 @@ export class Session extends EventEmitter<SessionEvents> {
9298
});
9399
this.connectedAtlasCluster = undefined;
94100
}
95-
this.emit("disconnect");
96101
}
97102

98103
async close(): Promise<void> {
99104
await this.disconnect();
100105
await this.apiClient.close();
101-
this.emit("close");
102106
}
103107

104-
async connectToMongoDB(connectionString: string, connectOptions: ConnectOptions): Promise<void> {
105-
connectionString = setAppNameParamIfMissing({
106-
connectionString,
107-
defaultAppName: `${packageInfo.mcpServerName} ${packageInfo.version}`,
108-
});
109-
108+
async connectToMongoDB(settings: ConnectionSettings): Promise<AnyConnectionState> {
110109
try {
111-
this.serviceProvider = await NodeDriverServiceProvider.connect(connectionString, {
112-
productDocsLink: "https://github.com/mongodb-js/mongodb-mcp-server/",
113-
productName: "MongoDB MCP",
114-
readConcern: {
115-
level: connectOptions.readConcern,
116-
},
117-
readPreference: connectOptions.readPreference,
118-
writeConcern: {
119-
w: connectOptions.writeConcern,
120-
},
121-
timeoutMS: connectOptions.timeoutMS,
122-
proxy: { useEnvironmentVariableProxies: true },
123-
applyProxyToOIDC: true,
124-
});
125-
126-
await this.serviceProvider?.runCommand?.("admin", { hello: 1 });
110+
return await this.connectionManager.connect({ ...settings });
127111
} catch (error: unknown) {
128-
const message = error instanceof Error ? error.message : `${error as string}`;
112+
const message = error instanceof Error ? error.message : (error as string);
129113
this.emit("connection-error", message);
130114
throw error;
131115
}
116+
}
117+
118+
isConnectedToMongoDB(): boolean {
119+
return this.connectionManager.currentConnectionState.tag == "connected";
120+
}
121+
122+
get serviceProvider(): NodeDriverServiceProvider {
123+
if (this.isConnectedToMongoDB()) {
124+
const state = this.connectionManager.currentConnectionState as ConnectionStateConnected;
125+
return state.serviceProvider;
126+
}
132127

133-
this.emit("connect");
128+
throw new MongoDBError(ErrorCodes.NotConnectedToMongoDB, "Not connected to MongoDB");
134129
}
135130
}

‎src/resources/common/debug.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,10 @@ type ConnectionStateDebuggingInformation = {
1010

1111
export class DebugResource extends ReactiveResource(
1212
{
13-
name: "debug-mongodb-connectivity",
14-
uri: "debug://mongodb-connectivity",
13+
name: "debug-mongodb",
14+
uri: "debug://mongodb",
1515
config: {
16-
description: "Debugging information for connectivity issues.",
16+
description: "Debugging information for MongoDB connectivity issues.",
1717
},
1818
},
1919
{

‎src/server.ts

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,15 @@ export class Server {
3838
}
3939

4040
async connect(transport: Transport): Promise<void> {
41+
// Resources are now reactive, so we register them ASAP so they can listen to events like
42+
// connection events.
43+
this.registerResources();
4144
await this.validateConfig();
4245

4346
this.mcpServer.server.registerCapabilities({ logging: {} });
4447

48+
// TODO: Eventually we might want to make tools reactive too instead of relying on custom logic.
4549
this.registerTools();
46-
this.registerResources();
4750

4851
// This is a workaround for an issue we've seen with some models, where they'll see that everything in the `arguments`
4952
// object is optional, and then not pass it at all. However, the MCP server expects the `arguments` object to be if
@@ -194,7 +197,10 @@ export class Server {
194197

195198
if (this.userConfig.connectionString) {
196199
try {
197-
await this.session.connectToMongoDB(this.userConfig.connectionString, this.userConfig.connectOptions);
200+
await this.session.connectToMongoDB({
201+
connectionString: this.userConfig.connectionString,
202+
...this.userConfig.connectOptions,
203+
});
198204
} catch (error) {
199205
console.error(
200206
"Failed to connect to MongoDB instance using the connection string from the config: ",

‎src/tools/atlas/connect/connectCluster.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ export class ConnectClusterTool extends AtlasToolBase {
2727
clusterName: string
2828
): Promise<"connected" | "disconnected" | "connecting" | "connected-to-other-cluster" | "unknown"> {
2929
if (!this.session.connectedAtlasCluster) {
30-
if (this.session.serviceProvider) {
30+
if (this.session.isConnectedToMongoDB()) {
3131
return "connected-to-other-cluster";
3232
}
3333
return "disconnected";
@@ -40,7 +40,7 @@ export class ConnectClusterTool extends AtlasToolBase {
4040
return "connected-to-other-cluster";
4141
}
4242

43-
if (!this.session.serviceProvider) {
43+
if (!this.session.isConnectedToMongoDB()) {
4444
return "connecting";
4545
}
4646

@@ -145,7 +145,7 @@ export class ConnectClusterTool extends AtlasToolBase {
145145
try {
146146
lastError = undefined;
147147

148-
await this.session.connectToMongoDB(connectionString, this.config.connectOptions);
148+
await this.session.connectToMongoDB({ connectionString, ...this.config.connectOptions });
149149
break;
150150
} catch (err: unknown) {
151151
const error = err instanceof Error ? err : new Error(String(err));

0 commit comments

Comments
 (0)