Merge branch 'main' into ni/connect-guidance

* main:
  chore: revoke access tokens on server shutdown [MCP-53] (#352)
  fix: turn atlas-connect-cluster async (#343)

Author: nirinchev

.vscode/launch.json

@@ -4,6 +4,15 @@
   // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
   "version": "0.2.0",
   "configurations": [
+    {
+      "type": "node",
+      "request": "launch",
+      "name": "Launch Tests",
+      "runtimeExecutable": "npm",
+      "runtimeArgs": ["test"],
+      "cwd": "${workspaceFolder}",
+      "envFile": "${workspaceFolder}/.env"
+    },
     {
       "type": "node",
       "request": "launch",

src/common/atlas/apiClient.ts

@@ -5,6 +5,7 @@ import { ApiClientError } from "./apiClientError.js";
 import { paths, operations } from "./openapi.js";
 import { CommonProperties, TelemetryEvent } from "../../telemetry/types.js";
 import { packageInfo } from "../../helpers/packageInfo.js";
+import logger, { LogId } from "../../logger.js";
 
 const ATLAS_API_VERSION = "2025-03-12";
 
@@ -34,9 +35,7 @@ export class ApiClient {
 
     private getAccessToken = async () => {
         if (this.oauth2Client && (!this.accessToken || this.accessToken.expired())) {
-            this.accessToken = await this.oauth2Client.getToken({
-                agent: this.options.userAgent,
-            });
+            this.accessToken = await this.oauth2Client.getToken({});
         }
         return this.accessToken?.token.access_token as string | undefined;
     };
@@ -49,7 +48,9 @@ export class ApiClient {
 
             try {
                 const accessToken = await this.getAccessToken();
-                request.headers.set("Authorization", `Bearer ${accessToken}`);
+                if (accessToken) {
+                    request.headers.set("Authorization", `Bearer ${accessToken}`);
+                }
                 return request;
             } catch {
                 // ignore not availble tokens, API will return 401
@@ -81,20 +82,38 @@ export class ApiClient {
                 auth: {
                     tokenHost: this.options.baseUrl,
                     tokenPath: "/api/oauth/token",
+                    revokePath: "/api/oauth/revoke",
+                },
+                http: {
+                    headers: {
+                        "User-Agent": this.options.userAgent,
+                    },
                 },
             });
             this.client.use(this.authMiddleware);
         }
     }
 
     public hasCredentials(): boolean {
-        return !!(this.oauth2Client && this.accessToken);
+        return !!this.oauth2Client;
     }
 
     public async validateAccessToken(): Promise<void> {
         await this.getAccessToken();
     }
 
+    public async close(): Promise<void> {
+        if (this.accessToken) {
+            try {
+                await this.accessToken.revoke("access_token");
+            } catch (error: unknown) {
+                const err = error instanceof Error ? error : new Error(String(error));
+                logger.error(LogId.atlasApiRevokeFailure, "apiClient", `Failed to revoke access token: ${err.message}`);
+            }
+            this.accessToken = undefined;
+        }
+    }
+
     public async getIpInfo(): Promise<{
         currentIpv4Address: string;
     }> {

src/logger.ts

@@ -17,6 +17,9 @@ 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),
+    atlasApiRevokeFailure: mongoLogId(1_001_007),
 
     telemetryDisabled: mongoLogId(1_002_001),
     telemetryEmitFailure: mongoLogId(1_002_002),

src/session.ts

@@ -67,35 +67,33 @@ 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");
     }
 
     async close(): Promise<void> {
         await this.disconnect();
+        await this.apiClient.close();
         this.emit("close");
     }
 

src/tools/atlas/connect/connectCluster.ts

@@ -21,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) {
@@ -82,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));
@@ -107,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/connect/connect.ts

@@ -46,7 +46,7 @@ export class ConnectTool extends MongoDBToolBase {
 
     constructor(session: Session, config: UserConfig, telemetry: Telemetry) {
         super(session, config, telemetry);
-        session.on("close", () => {
+        session.on("disconnect", () => {
             this.updateMetadata();
         });
     }

src/tools/mongodb/mongodbTool.ts

@@ -16,16 +16,25 @@ export abstract class MongoDBToolBase extends ToolBase {
     public 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/helpers.ts

@@ -84,8 +84,8 @@ export function setupIntegrationTest(getUserConfig: () => UserConfig): Integrati
     });
 
     afterEach(async () => {
-        if (mcpServer) {
-            await mcpServer.session.close();
+        if (mcpServer && !mcpServer.session.connectedAtlasCluster) {
+            await mcpServer.session.disconnect();
         }
     });