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),
+    atlasConnectSuccessed: mongoLogId(1_001_006),
 
     telemetryDisabled: mongoLogId(1_002_001),
     telemetryEmitFailure: mongoLogId(1_002_002),

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,32 @@ 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"> {
+        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) {
@@ -83,9 +107,20 @@ export class ConnectClusterTool extends AtlasToolBase {
         cn.searchParams.set("authSource", "admin");
         const connectionString = cn.toString();
 
+        return connectionString;
+    }
+
+    private async connectToCluster(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}`
+        );
+
+        for (let i = 0; i < 600; i++) {
+            // try for 5 minutes
             try {
                 await this.session.connectToMongoDB(connectionString, this.config.connectOptions);
                 lastError = undefined;
@@ -106,14 +141,93 @@ export class ConnectClusterTool extends AtlasToolBase {
         }
 
         if (lastError) {
+            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.atlasConnectSuccessed,
+            "atlas-connect-cluster",
+            `connected to cluster: ${this.session.connectedAtlasCluster?.clusterName}`
+        );
+    }
+
+    protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
+        try {
+            const state = await this.queryConnection(projectId, clusterName);
+            switch (state) {
+                case "connected":
+                    return {
+                        content: [
+                            {
+                                type: "text",
+                                text: "Cluster is already connected.",
+                            },
+                        ],
+                    };
+                case "connecting":
+                    return {
+                        content: [
+                            {
+                                type: "text",
+                                text: "Cluster is connecting...",
+                            },
+                        ],
+                    };
+                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);
+        process.nextTick(async () => {
+            try {
+                await this.connectToCluster(connectionString);
+            } 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}`
+                );
+            }
+        });
+
         return {
             content: [
                 {
                     type: "text",
-                    text: `Connected to cluster "${clusterName}"`,
+                    text: `Attempting to connect to cluster "${clusterName}"...`,
                 },
             ],
         };

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

@@ -189,7 +189,23 @@ describeWithAtlas("clusters", (integration) => {
                 })) as CallToolResult;
                 expect(response.content).toBeArray();
                 expect(response.content).toHaveLength(1);
-                expect(response.content[0]?.text).toContain(`Connected to cluster "${clusterName}"`);
+                expect(response.content[0]?.type).toEqual("text");
+                expect(response.content[0]?.text).toContain(`Attempting to connect 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).toHaveLength(1);
+                    expect(response.content[0]?.type).toEqual("text");
+                    const c = response.content[0] as { text: string };
+                    if (c.text.includes("Cluster is already connected.")) {
+                        break; // success
+                    }
+                    await sleep(500);
+                }
             });
         });
     });