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,18 +11,42 @@ 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";
+    protected description = "Connect to / Inspect connection of MongoDB Atlas cluster";
     protected operationType: OperationType = "metadata";
     protected argsShape = {
         projectId: z.string().describe("Atlas project ID"),
         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"> {
+        if (!this.session.connectedAtlasCluster) {
+            return "disconnected";
+        }
+
+        if (
+            this.session.connectedAtlasCluster.projectId !== projectId ||
+            this.session.connectedAtlasCluster.clusterName !== clusterName
+        ) {
+            return "connected-to-other-cluster";
+        }
+
+        if (!this.session.serviceProvider) {
+            return "connecting";
+        }
+
+        await this.session.serviceProvider.runCommand("admin", {
+            ping: 1,
+        });
+        return "connected";
+    }
 
+    private async prepareClusterConnection(projectId: string, clusterName: string): Promise<string> {
         const cluster = await inspectCluster(this.session.apiClient, projectId, clusterName);
 
         if (!cluster.connectionString) {
@@ -81,14 +105,37 @@ 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,
+        tryCount: number
+    ): 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}`
+        );
+
+        for (let i = 0; i < tryCount; i++) {
+            if (
+                !this.session.connectedAtlasCluster ||
+                this.session.connectedAtlasCluster.projectId != projectId ||
+                this.session.connectedAtlasCluster.clusterName != clusterName
+            ) {
+                lastError = new Error("Cluster connection aborted");
+                break;
+            }
+
             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,16 +153,121 @@ export class ConnectClusterTool extends AtlasToolBase {
         }
 
         if (lastError) {
+            if (this.session.connectedAtlasCluster?.projectId && 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;
         }
 
-        return {
+        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> {
+        const connectingResult = {
             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.`,
                 },
             ],
         };
+
+        try {
+            const state = await this.queryConnection(projectId, clusterName);
+            switch (state) {
+                case "connected":
+                    return {
+                        content: [
+                            {
+                                type: "text",
+                                text: "Cluster is already connected.",
+                            },
+                        ],
+                    };
+                case "connecting":
+                    return connectingResult;
+                case "connected-to-other-cluster":
+                case "disconnected":
+                default:
+                    // fall through to create new connection
+                    break;
+            }
+        } 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}`
+            );
+            // fall through to create new connection
+        }
+
+        await this.session.disconnect();
+        const connectionString = await this.prepareClusterConnection(projectId, clusterName);
+
+        try {
+            // First, try to connect to the cluster within the current tool call.
+            // We give it 60 attempts with 500 ms delay between each, so ~30 seconds
+            await this.connectToCluster(projectId, clusterName, connectionString, 60);
+
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: `Connected to cluster "${clusterName}".`,
+                    },
+                ],
+            };
+        } catch (err: unknown) {
+            const error = err instanceof Error ? err : new Error(String(err));
+            logger.debug(
+                LogId.atlasConnectFailure,
+                "atlas-connect-cluster",
+                `error connecting to cluster: ${error.message}`
+            );
+
+            // We couldn't connect in ~30 seconds, likely because user creation is taking longer.
+            // Retry the connection with longer timeout (~5 minutes), while also returning a response
+            // to the client. Many clients will have a 1 minute timeout for tool calls, so we want to
+            // return well before that.
+            //
+            // Once we add support for streamable http, we'd want to use progress notifications here.
+            void this.connectToCluster(projectId, clusterName, connectionString, 600).catch((err) => {
+                const error = err instanceof Error ? err : new Error(String(err));
+                logger.debug(
+                    LogId.atlasConnectFailure,
+                    "atlas-connect-cluster",
+                    `error connecting to cluster: ${error.message}`
+                );
+            });
+
+            return connectingResult;
+        }
     }
 }

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 < 600; 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);
+                }
             });
         });
     });