fix: turn atlas-connect-cluster async

src/logger.ts

@@ -17,6 +17,8 @@ export const LogId = {
     atlasDeleteDatabaseUserFailure: mongoLogId(1_001_002),
     atlasConnectFailure: mongoLogId(1_001_003),
     atlasInspectFailure: mongoLogId(1_001_004),
+    atlasConnectAttempt: mongoLogId(1_001_005),
+    atlasConnectSucceeded: mongoLogId(1_001_006),
 
     telemetryDisabled: mongoLogId(1_002_001),
     telemetryEmitFailure: mongoLogId(1_002_002),

src/session.ts

@@ -67,30 +67,27 @@ export class Session extends EventEmitter<{
             }
             this.serviceProvider = undefined;
         }
-        if (!this.connectedAtlasCluster) {
-            this.emit("disconnect");
-            return;
-        }
-        void this.apiClient
-            .deleteDatabaseUser({
-                params: {
-                    path: {
-                        groupId: this.connectedAtlasCluster.projectId,
-                        username: this.connectedAtlasCluster.username,
-                        databaseName: "admin",
+        if (this.connectedAtlasCluster?.username && this.connectedAtlasCluster?.projectId) {
+            void this.apiClient
+                .deleteDatabaseUser({
+                    params: {
+                        path: {
+                            groupId: this.connectedAtlasCluster.projectId,
+                            username: this.connectedAtlasCluster.username,
+                            databaseName: "admin",
+                        },
                     },
-                },
-            })
-            .catch((err: unknown) => {
-                const error = err instanceof Error ? err : new Error(String(err));
-                logger.error(
-                    LogId.atlasDeleteDatabaseUserFailure,
-                    "atlas-connect-cluster",
-                    `Error deleting previous database user: ${error.message}`
-                );
-            });
-        this.connectedAtlasCluster = undefined;
-
+                })
+                .catch((err: unknown) => {
+                    const error = err instanceof Error ? err : new Error(String(err));
+                    logger.error(
+                        LogId.atlasDeleteDatabaseUserFailure,
+                        "atlas-connect-cluster",
+                        `Error deleting previous database user: ${error.message}`
+                    );
+                });
+            this.connectedAtlasCluster = undefined;
+        }
         this.emit("disconnect");
     }
 

src/tools/atlas/metadata/connectCluster.ts

@@ -11,6 +11,7 @@ const EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours
 function sleep(ms: number): Promise<void> {
     return new Promise((resolve) => setTimeout(resolve, ms));
 }
+
 export class ConnectClusterTool extends AtlasToolBase {
     protected name = "atlas-connect-cluster";
     protected description = "Connect to MongoDB Atlas cluster";
@@ -20,9 +21,46 @@ export class ConnectClusterTool extends AtlasToolBase {
         clusterName: z.string().describe("Atlas cluster name"),
     };
 
-    protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        await this.session.disconnect();
+    private async queryConnection(
+        projectId: string,
+        clusterName: string
+    ): Promise<"connected" | "disconnected" | "connecting" | "connected-to-other-cluster" | "unknown"> {
+        if (!this.session.connectedAtlasCluster) {
+            if (this.session.serviceProvider) {
+                return "connected-to-other-cluster";
+            }
+            return "disconnected";
+        }
+
+        if (
+            this.session.connectedAtlasCluster.projectId !== projectId ||
+            this.session.connectedAtlasCluster.clusterName !== clusterName
+        ) {
+            return "connected-to-other-cluster";
+        }
+
+        if (!this.session.serviceProvider) {
+            return "connecting";
+        }
 
+        try {
+            await this.session.serviceProvider.runCommand("admin", {
+                ping: 1,
+            });
+
+            return "connected";
+        } catch (err: unknown) {
+            const error = err instanceof Error ? err : new Error(String(err));
+            logger.debug(
+                LogId.atlasConnectFailure,
+                "atlas-connect-cluster",
+                `error querying cluster: ${error.message}`
+            );
+            return "unknown";
+        }
+    }
+
+    private async prepareClusterConnection(projectId: string, clusterName: string): Promise<string> {
         const cluster = await inspectCluster(this.session.apiClient, projectId, clusterName);
 
         if (!cluster.connectionString) {
@@ -81,14 +119,32 @@ export class ConnectClusterTool extends AtlasToolBase {
         cn.username = username;
         cn.password = password;
         cn.searchParams.set("authSource", "admin");
-        const connectionString = cn.toString();
+        return cn.toString();
+    }
 
+    private async connectToCluster(projectId: string, clusterName: string, connectionString: string): Promise<void> {
         let lastError: Error | undefined = undefined;
 
-        for (let i = 0; i < 20; i++) {
+        logger.debug(
+            LogId.atlasConnectAttempt,
+            "atlas-connect-cluster",
+            `attempting to connect to cluster: ${this.session.connectedAtlasCluster?.clusterName}`
+        );
+
+        // try to connect for about 5 minutes
+        for (let i = 0; i < 600; i++) {
+            if (
+                !this.session.connectedAtlasCluster ||
+                this.session.connectedAtlasCluster.projectId != projectId ||
+                this.session.connectedAtlasCluster.clusterName != clusterName
+            ) {
+                throw new Error("Cluster connection aborted");
+            }
+
             try {
-                await this.session.connectToMongoDB(connectionString, this.config.connectOptions);
                 lastError = undefined;
+
+                await this.session.connectToMongoDB(connectionString, this.config.connectOptions);
                 break;
             } catch (err: unknown) {
                 const error = err instanceof Error ? err : new Error(String(err));
@@ -106,14 +162,94 @@ export class ConnectClusterTool extends AtlasToolBase {
         }
 
         if (lastError) {
+            if (
+                this.session.connectedAtlasCluster?.projectId == projectId &&
+                this.session.connectedAtlasCluster?.clusterName == clusterName &&
+                this.session.connectedAtlasCluster?.username
+            ) {
+                void this.session.apiClient
+                    .deleteDatabaseUser({
+                        params: {
+                            path: {
+                                groupId: this.session.connectedAtlasCluster.projectId,
+                                username: this.session.connectedAtlasCluster.username,
+                                databaseName: "admin",
+                            },
+                        },
+                    })
+                    .catch((err: unknown) => {
+                        const error = err instanceof Error ? err : new Error(String(err));
+                        logger.debug(
+                            LogId.atlasConnectFailure,
+                            "atlas-connect-cluster",
+                            `error deleting database user: ${error.message}`
+                        );
+                    });
+            }
+            this.session.connectedAtlasCluster = undefined;
             throw lastError;
         }
 
+        logger.debug(
+            LogId.atlasConnectSucceeded,
+            "atlas-connect-cluster",
+            `connected to cluster: ${this.session.connectedAtlasCluster?.clusterName}`
+        );
+    }
+
+    protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
+        for (let i = 0; i < 60; i++) {
+            const state = await this.queryConnection(projectId, clusterName);
+            switch (state) {
+                case "connected": {
+                    return {
+                        content: [
+                            {
+                                type: "text",
+                                text: `Connected to cluster "${clusterName}".`,
+                            },
+                        ],
+                    };
+                }
+                case "connecting": {
+                    break;
+                }
+                case "connected-to-other-cluster":
+                case "disconnected":
+                case "unknown":
+                default: {
+                    await this.session.disconnect();
+                    const connectionString = await this.prepareClusterConnection(projectId, clusterName);
+
+                    // try to connect for about 5 minutes asynchronously
+                    void this.connectToCluster(projectId, clusterName, connectionString).catch((err: unknown) => {
+                        const error = err instanceof Error ? err : new Error(String(err));
+                        logger.error(
+                            LogId.atlasConnectFailure,
+                            "atlas-connect-cluster",
+                            `error connecting to cluster: ${error.message}`
+                        );
+                    });
+                    break;
+                }
+            }
+
+            await sleep(500);
+        }
+
         return {
             content: [
                 {
-                    type: "text",
-                    text: `Connected to cluster "${clusterName}"`,
+                    type: "text" as const,
+                    text: `Attempting to connect to cluster "${clusterName}"...`,
+                },
+                {
+                    type: "text" as const,
+                    text: `Warning: Provisioning a user and connecting to the cluster may take more time, please check again in a few seconds.`,
+                },
+                {
+                    type: "text" as const,
+                    text: `Warning: Make sure your IP address was enabled in the allow list setting of the Atlas cluster.`,
                 },
             ],
         };

src/tools/mongodb/mongodbTool.ts

@@ -14,16 +14,25 @@ export abstract class MongoDBToolBase extends ToolBase {
     protected category: ToolCategory = "mongodb";
 
     protected async ensureConnected(): Promise<NodeDriverServiceProvider> {
-        if (!this.session.serviceProvider && this.config.connectionString) {
-            try {
-                await this.connectToMongoDB(this.config.connectionString);
-            } catch (error) {
-                logger.error(
-                    LogId.mongodbConnectFailure,
-                    "mongodbTool",
-                    `Failed to connect to MongoDB instance using the connection string from the config: ${error as string}`
+        if (!this.session.serviceProvider) {
+            if (this.session.connectedAtlasCluster) {
+                throw new MongoDBError(
+                    ErrorCodes.NotConnectedToMongoDB,
+                    `Attempting to connect to Atlas cluster "${this.session.connectedAtlasCluster.clusterName}", try again in a few seconds.`
                 );
-                throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, "Not connected to MongoDB.");
+            }
+
+            if (this.config.connectionString) {
+                try {
+                    await this.connectToMongoDB(this.config.connectionString);
+                } catch (error) {
+                    logger.error(
+                        LogId.mongodbConnectFailure,
+                        "mongodbTool",
+                        `Failed to connect to MongoDB instance using the connection string from the config: ${error as string}`
+                    );
+                    throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, "Not connected to MongoDB.");
+                }
             }
         }
 

tests/integration/tools/atlas/clusters.test.ts

@@ -183,13 +183,27 @@ describeWithAtlas("clusters", (integration) => {
             it("connects to cluster", async () => {
                 const projectId = getProjectId();
 
-                const response = (await integration.mcpClient().callTool({
-                    name: "atlas-connect-cluster",
-                    arguments: { projectId, clusterName },
-                })) as CallToolResult;
-                expect(response.content).toBeArray();
-                expect(response.content).toHaveLength(1);
-                expect(response.content[0]?.text).toContain(`Connected to cluster "${clusterName}"`);
+                for (let i = 0; i < 10; i++) {
+                    const response = (await integration.mcpClient().callTool({
+                        name: "atlas-connect-cluster",
+                        arguments: { projectId, clusterName },
+                    })) as CallToolResult;
+                    expect(response.content).toBeArray();
+                    expect(response.content.length).toBeGreaterThanOrEqual(1);
+                    expect(response.content[0]?.type).toEqual("text");
+                    const c = response.content[0] as { text: string };
+                    if (
+                        c.text.includes("Cluster is already connected.") ||
+                        c.text.includes(`Connected to cluster "${clusterName}"`)
+                    ) {
+                        break; // success
+                    } else {
+                        expect(response.content[0]?.text).toContain(
+                            `Attempting to connect to cluster "${clusterName}"...`
+                        );
+                    }
+                    await sleep(500);
+                }
             });
         });
     });