fix: wait for connection string when creating a deployment (#1291)

Author: nirinchev

scripts/cleanupAtlasTestLeftovers.test.ts

@@ -3,10 +3,7 @@ import { ApiClient } from "../src/common/atlas/apiClient.js";
 import { ConsoleLogger } from "../src/common/logging/index.js";
 import { Keychain } from "../src/lib.js";
 import { describe, it } from "vitest";
-
-function sleep(ms: number): Promise<void> {
-    return new Promise((resolve) => setTimeout(resolve, ms));
-}
+import { sleep } from "../src/common/managedTimeout.js";
 
 function isOlderThanTwoHours(date: string): boolean {
     const twoHoursInMs = 2 * 60 * 60 * 1000;

src/common/atlasLocal/connectionString.ts

@@ -0,0 +1,47 @@
+import type { Client } from "@mongodb-js/atlas-local";
+import { sleep } from "../managedTimeout.js";
+
+/** Keep well under the MCP client's default ~60s request timeout. */
+export const DEFAULT_MAX_ATTEMPTS = 60;
+export const DEFAULT_INTERVAL_MS = 500;
+
+export class AtlasLocalDeploymentNotReadyError extends Error {
+    constructor(public readonly deploymentName: string) {
+        super(`Atlas Local deployment "${deploymentName}" is still starting up`);
+        this.name = "AtlasLocalDeploymentNotReadyError";
+    }
+}
+
+function isMissingPortBindingError(error: unknown): boolean {
+    const message = error instanceof Error ? error.message : String(error);
+    return message.includes("Missing port binding information");
+}
+
+/**
+ * atlas-local-create-deployment can return before Docker publishes port bindings.
+ * Retry briefly so connect usually works without exceeding MCP request timeouts.
+ */
+export async function waitForConnectionString(
+    client: Client,
+    deploymentName: string,
+    {
+        maxAttempts = DEFAULT_MAX_ATTEMPTS,
+        intervalMs = DEFAULT_INTERVAL_MS,
+    }: { maxAttempts?: number; intervalMs?: number } = {}
+): Promise<string> {
+    for (let attempt = 0; attempt < maxAttempts; attempt++) {
+        try {
+            return await client.getConnectionString(deploymentName);
+        } catch (error: unknown) {
+            if (!isMissingPortBindingError(error)) {
+                throw error;
+            }
+
+            if (attempt < maxAttempts - 1) {
+                await sleep(intervalMs);
+            }
+        }
+    }
+
+    throw new AtlasLocalDeploymentNotReadyError(deploymentName);
+}

src/common/managedTimeout.ts

@@ -25,3 +25,7 @@ export function setManagedTimeout(callback: () => Promise<void> | void, timeoutM
         restart,
     };
 }
+
+export function sleep(ms: number): Promise<void> {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+}

src/tools/atlas/connect/connectCluster.ts

@@ -12,17 +12,14 @@ import { getDefaultRoleFromConfig } from "../../../common/atlas/roles.js";
 import { AtlasArgs } from "../../args.js";
 import { SHARED_TIER_METRIC_NAMES } from "../../../telemetry/types.js";
 import type { ConnectionMetadata, SharedTierTier, SharedTierMetricName } from "../../../telemetry/types.js";
+import { sleep } from "../../../common/managedTimeout.js";
 
 const addedIpAccessListMessage =
     "Note: Your current IP address has been added to the Atlas project's IP access list to enable secure connection.";
 
 const createdUserMessage =
     "Note: A temporary user has been created to enable secure connection to the cluster. For more information, see https://dochub.mongodb.org/core/mongodb-mcp-server-tools-considerations\n\nNote to LLM Agent: it is important to include the following link in your response to the user in case they want to get more information about the temporary user created: https://dochub.mongodb.org/core/mongodb-mcp-server-tools-considerations";
 
-function sleep(ms: number): Promise<void> {
-    return new Promise((resolve) => setTimeout(resolve, ms));
-}
-
 export const ConnectClusterArgs = {
     projectId: AtlasArgs.projectId().describe("Atlas project ID"),
     clusterName: AtlasArgs.clusterName().describe("Atlas cluster name"),

src/tools/atlasLocal/connect/connectDeployment.ts

@@ -2,6 +2,10 @@ import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { AtlasLocalToolBase } from "../atlasLocalTool.js";
 import type { OperationType, ToolArgs } from "../../tool.js";
 import type { Client } from "@mongodb-js/atlas-local";
+import {
+    AtlasLocalDeploymentNotReadyError,
+    waitForConnectionString,
+} from "../../../common/atlasLocal/connectionString.js";
 import { CommonArgs } from "../../args.js";
 import type { ConnectionMetadata } from "../../../telemetry/types.js";
 
@@ -17,10 +21,24 @@ export class ConnectDeploymentTool extends AtlasLocalToolBase {
         { deploymentName }: ToolArgs<typeof this.argsShape>,
         { client }: { client: Client }
     ): Promise<CallToolResult> {
-        // Get the connection string for the deployment
-        const connectionString = await client.getConnectionString(deploymentName);
+        let connectionString: string;
+        try {
+            connectionString = await waitForConnectionString(client, deploymentName);
+        } catch (error: unknown) {
+            if (error instanceof AtlasLocalDeploymentNotReadyError) {
+                return {
+                    content: [
+                        {
+                            type: "text",
+                            text: `Atlas Local deployment "${deploymentName}" is still starting up. Wait a few seconds and call atlas-local-connect-deployment again with the same deployment name.`,
+                        },
+                    ],
+                    isError: true,
+                };
+            }
+            throw error;
+        }
 
-        // Connect to the deployment
         await this.session.connectToMongoDB({ connectionString });
 
         return {

src/tools/atlasLocal/create/createDeployment.ts

@@ -2,6 +2,10 @@ import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { AtlasLocalToolBase } from "../atlasLocalTool.js";
 import type { OperationType, ToolArgs } from "../../tool.js";
 import type { Client, CreateDeploymentOptions } from "@mongodb-js/atlas-local";
+import {
+    AtlasLocalDeploymentNotReadyError,
+    waitForConnectionString,
+} from "../../../common/atlasLocal/connectionString.js";
 import { CommonArgs } from "../../args.js";
 import z from "zod";
 
@@ -35,14 +39,33 @@ export class CreateDeploymentTool extends AtlasLocalToolBase {
             ...(this.config.voyageApiKey ? { voyageApiKey: this.config.voyageApiKey } : {}),
             doNotTrack: !this.telemetry.isTelemetryEnabled(),
         };
-        // Create the deployment
         const deployment = await client.createDeployment(deploymentOptions);
 
+        // createDeployment returns once the container is healthy, but Docker may
+        // not have published port bindings yet. Block until connect can succeed.
+        const resolvedDeploymentName = deployment.name ?? deploymentName;
+        let stillStarting = false;
+        if (resolvedDeploymentName) {
+            try {
+                await waitForConnectionString(client, resolvedDeploymentName);
+            } catch (error: unknown) {
+                if (error instanceof AtlasLocalDeploymentNotReadyError) {
+                    stillStarting = true;
+                } else {
+                    throw error;
+                }
+            }
+        }
+
+        const startingNote = stillStarting
+            ? " The deployment is still initializing; if atlas-local-connect-deployment fails, wait a few seconds and try connecting again."
+            : "";
+
         return {
             content: [
                 {
                     type: "text",
-                    text: `Deployment with container ID "${deployment.containerId}" and name "${deployment.name}" created (imageTag: ${imageTag}).`,
+                    text: `Deployment with container ID "${deployment.containerId}" and name "${deployment.name}" created (imageTag: ${imageTag}).${startingNote}`,
                 },
             ],
             _meta: {

tests/eval/lib/seeding.ts

@@ -1,6 +1,7 @@
 import type { MongoClient } from "mongodb";
 import type { DbSeedEntry, SeedClassicIndex } from "./datasetTypes.js";
 import { getSeedDocuments, parseSeedEntry } from "./datasetHelpers.js";
+import { sleep } from "../../../src/common/managedTimeout.js";
 
 // ╭──────────────────────────────────────────────╮
 // │   ↘️ Seeding Constants                       │
@@ -47,8 +48,11 @@ async function waitForIndexesQueryable(
             }
         }
 
-        if (pending.size === 0) return;
-        await new Promise((resolve) => setTimeout(resolve, intervalMs));
+        if (pending.size === 0) {
+            return;
+        }
+
+        await sleep(intervalMs);
     }
 
     throw new Error(

tests/integration/common/connectionManager.oidc.test.ts

@@ -4,13 +4,14 @@ import semver from "semver";
 import process from "process";
 import type { MongoDBIntegrationTestCase } from "../tools/mongodb/mongodbHelpers.js";
 import { describeWithMongoDB, isCommunityServer, getServerVersion } from "../tools/mongodb/mongodbHelpers.js";
-import { connect, defaultTestConfig, responseAsText, timeout, waitUntil } from "../helpers.js";
+import { connect, defaultTestConfig, responseAsText, waitUntil } from "../helpers.js";
 import type { ConnectionStateConnected, ConnectionStateConnecting } from "../../../src/common/connectionManager.js";
 import type { UserConfig } from "../../../src/common/config/userConfig.js";
 import path from "path";
 import type { OIDCMockProviderConfig } from "@mongodb-js/oidc-mock-provider";
 import { OIDCMockProvider } from "@mongodb-js/oidc-mock-provider";
 import { type TestConnectionManager } from "../../utils/index.js";
+import { sleep } from "../../../src/common/managedTimeout.js";
 
 const DEFAULT_TIMEOUT = 60_000;
 const DEFAULT_RETRIES = 5;
@@ -242,7 +243,7 @@ describe.skipIf(process.platform !== "linux")("ConnectionManager OIDC Tests", as
                     signal
                 );
 
-                await timeout(2000);
+                await sleep(2000);
                 await state.serviceProvider.listDatabases("admin");
                 expect(tokenFetches).toBeGreaterThan(1);
             });

tests/integration/helpers.ts

@@ -392,10 +392,6 @@ function validateToolAnnotations(tool: ToolInfo, name: string, operationType: Op
     }
 }
 
-export function timeout(ms: number): Promise<void> {
-    return new Promise((resolve) => setTimeout(resolve, ms));
-}
-
 /**
  * Subscribes to the resources changed notification for the provided URI
  */
@@ -458,10 +454,6 @@ export function getDataFromUntrustedContent(content: string): string {
     return match.groups.data.trim();
 }
 
-export function sleep(ms: number): Promise<void> {
-    return new Promise((resolve) => setTimeout(resolve, ms));
-}
-
 export class InMemoryLogger extends LoggerBase {
     protected type?: LoggerType = "console";
     public messages: { level: LogLevel; payload: LogPayload }[] = [];

tests/integration/resources/exportedData.test.ts

@@ -3,10 +3,11 @@ import fs from "fs/promises";
 import { EJSON, Long, ObjectId } from "bson";
 import { describe, expect, it, beforeEach, afterAll } from "vitest";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
-import { defaultTestConfig, getDataFromUntrustedContent, resourceChangedNotification, timeout } from "../helpers.js";
+import { defaultTestConfig, getDataFromUntrustedContent, resourceChangedNotification } from "../helpers.js";
 import { describeWithMongoDB } from "../tools/mongodb/mongodbHelpers.js";
 import { contentWithResourceURILink } from "../tools/mongodb/read/export.test.js";
 import type { UserConfig } from "../../../src/lib.js";
+import { sleep } from "../../../src/common/managedTimeout.js";
 
 const userConfig: UserConfig = {
     ...defaultTestConfig,
@@ -83,7 +84,7 @@ describeWithMongoDB(
 
                 // wait for export expired
                 for (let tries = 0; tries < 10; tries++) {
-                    await timeout(300);
+                    await sleep(300);
                     const response = await integration.mcpClient().readResource({
                         uri: exportedResourceURI as string,
                     });