feat(atlas-local): Added Atlas Local List Deployments tool (#520)

Author: jeroenvervaeke

package-lock.json

@@ -121,6 +121,7 @@
     "node": "^20.19.0 || ^22.12.0 || >= 23.0.0"
   },
   "optionalDependencies": {
+    "@mongodb-js-preview/atlas-local": "^0.0.0-preview.1",
     "kerberos": "^2.2.2"
   }
 }

package.json

@@ -15,6 +15,7 @@ import type {
 import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { ErrorCodes, MongoDBError } from "./errors.js";
 import type { ExportsManager } from "./exportsManager.js";
+import type { Client } from "@mongodb-js-preview/atlas-local";
 import type { Keychain } from "./keychain.js";
 
 export interface SessionOptions {
@@ -46,6 +47,7 @@ export class Session extends EventEmitter<SessionEvents> {
         version?: string;
         title?: string;
     };
+    atlasLocalClient?: Client;
 
     public logger: CompositeLogger;
 
@@ -99,6 +101,10 @@ export class Session extends EventEmitter<SessionEvents> {
         this.connectionManager.setClientName(this.mcpClient.name || "unknown");
     }
 
+    setAtlasLocalClient(atlasLocalClient: Client): void {
+        this.atlasLocalClient = atlasLocalClient;
+    }
+
     async disconnect(): Promise<void> {
         const atlasCluster = this.connectedAtlasCluster;
 

src/common/session.ts

@@ -19,7 +19,7 @@ import {
     UnsubscribeRequestSchema,
 } from "@modelcontextprotocol/sdk/types.js";
 import assert from "assert";
-import type { ToolBase } from "./tools/tool.js";
+import type { ToolBase, ToolConstructor } from "./tools/tool.js";
 import { validateConnectionString } from "./helpers/connectionOptions.js";
 import { packageInfo } from "./common/packageInfo.js";
 import { type ConnectionErrorHandler } from "./common/connectionErrorHandler.js";
@@ -69,6 +69,9 @@ export class Server {
         // TODO: Eventually we might want to make tools reactive too instead of relying on custom logic.
         this.registerTools();
 
+        // Atlas Local tools are optional and require async initialization
+        void this.registerAtlasLocalTools();
+
         // This is a workaround for an issue we've seen with some models, where they'll see that everything in the `arguments`
         // object is optional, and then not pass it at all. However, the MCP server expects the `arguments` object to be if
         // the tool accepts any arguments, even if they're all optional.
@@ -197,8 +200,41 @@ export class Server {
         this.telemetry.emitEvents([event]).catch(() => {});
     }
 
+    private async registerAtlasLocalTools(): Promise<void> {
+        // If Atlas Local tools are disabled, don't attempt to connect to the client
+        if (this.userConfig.disabledTools.includes("atlas-local")) {
+            return;
+        }
+
+        try {
+            // Import Atlas Local client asyncronously
+            // This will fail on unsupported platforms
+            const { Client: AtlasLocalClient } = await import("@mongodb-js-preview/atlas-local");
+
+            // Connect to Atlas Local client
+            // This will fail if docker is not running
+            const client = AtlasLocalClient.connect();
+
+            // Set Atlas Local client
+            this.session.setAtlasLocalClient(client);
+
+            // Register Atlas Local tools
+            this.registerToolInstances(AtlasLocalTools);
+        } catch (error) {
+            console.warn(
+                "Failed to initialize Atlas Local client, atlas-local tools will be disabled (error: ",
+                error,
+                ")"
+            );
+        }
+    }
+
     private registerTools(): void {
-        for (const toolConstructor of [...AtlasTools, ...AtlasLocalTools, ...MongoDbTools]) {
+        this.registerToolInstances([...AtlasTools, ...MongoDbTools]);
+    }
+
+    private registerToolInstances(tools: Array<ToolConstructor>): void {
+        for (const toolConstructor of tools) {
             const tool = new toolConstructor(this.session, this.userConfig, this.telemetry);
             if (tool.register(this)) {
                 this.tools.push(tool);

src/server.ts

@@ -1,10 +1,45 @@
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
-import type { ToolArgs, ToolCategory } from "../tool.js";
+import type { TelemetryToolMetadata, ToolArgs, ToolCategory } from "../tool.js";
 import { ToolBase } from "../tool.js";
+import type { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
+import type { Client } from "@mongodb-js-preview/atlas-local";
 
 export abstract class AtlasLocalToolBase extends ToolBase {
     public category: ToolCategory = "atlas-local";
 
+    protected verifyAllowed(): boolean {
+        return this.session.atlasLocalClient !== undefined && super.verifyAllowed();
+    }
+
+    protected async execute(): Promise<CallToolResult> {
+        // Get the client
+        const client = this.session.atlasLocalClient;
+
+        // If the client is not found, throw an error
+        // This should never happen:
+        // - atlas-local tools are only added after the client is set
+        //   this means that if we were unable to get the client, the tool will not be registered
+        // - in case the tool was registered by accident
+        //   verifyAllowed in the base class would still return false preventing the tool from being registered,
+        //   preventing the tool from being executed
+        if (!client) {
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: `Something went wrong on our end, this tool should have been disabled but it was not.
+please log a ticket here: https://github.com/mongodb-js/mongodb-mcp-server/issues/new?template=bug_report.yml`,
+                    },
+                ],
+                isError: true,
+            };
+        }
+
+        return this.executeWithAtlasLocalClient(client);
+    }
+
+    protected abstract executeWithAtlasLocalClient(client: Client): Promise<CallToolResult>;
+
     protected handleError(
         error: unknown,
         args: ToolArgs<typeof this.argsShape>
@@ -14,4 +49,12 @@ export abstract class AtlasLocalToolBase extends ToolBase {
         // For other types of errors, use the default error handling from the base class
         return super.handleError(error, args);
     }
+
+    protected resolveTelemetryMetadata(
+        ...args: Parameters<ToolCallback<typeof this.argsShape>>
+    ): TelemetryToolMetadata {
+        // TODO: include deployment id in the metadata where possible
+        void args; // this shuts up the eslint rule until we implement the TODO above
+        return {};
+    }
 }

src/tools/atlasLocal/atlasLocalTool.ts

@@ -0,0 +1,46 @@
+import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { AtlasLocalToolBase } from "../atlasLocalTool.js";
+import type { OperationType } from "../../tool.js";
+import { formatUntrustedData } from "../../tool.js";
+import type { Deployment } from "@mongodb-js-preview/atlas-local";
+import type { Client } from "@mongodb-js-preview/atlas-local";
+
+export class ListDeploymentsTool extends AtlasLocalToolBase {
+    public name = "atlas-local-list-deployments";
+    protected description = "List MongoDB Atlas local deployments";
+    public operationType: OperationType = "read";
+    protected argsShape = {};
+
+    protected async executeWithAtlasLocalClient(client: Client): Promise<CallToolResult> {
+        // List the deployments
+        const deployments = await client.listDeployments();
+
+        // Format the deployments
+        return this.formatDeploymentsTable(deployments);
+    }
+
+    private formatDeploymentsTable(deployments: Deployment[]): CallToolResult {
+        // Check if deployments are absent
+        if (!deployments?.length) {
+            return {
+                content: [{ type: "text", text: "No deployments found." }],
+            };
+        }
+
+        // Turn the deployments into a markdown table
+        const rows = deployments
+            .map((deployment) => {
+                return `${deployment.name || "Unknown"} | ${deployment.state} | ${deployment.mongodbVersion}`;
+            })
+            .join("\n");
+
+        return {
+            content: formatUntrustedData(
+                `Found ${deployments.length} deployments:`,
+                `Deployment Name | State | MongoDB Version
+----------------|----------------|----------------
+${rows}`
+            ),
+        };
+    }
+}

src/tools/atlasLocal/read/listDeployments.ts

@@ -1 +1,3 @@
-export const AtlasLocalTools = [];
+import { ListDeploymentsTool } from "./read/listDeployments.js";
+
+export const AtlasLocalTools = [ListDeploymentsTool];

src/tools/atlasLocal/tools.ts

@@ -18,6 +18,8 @@ export type TelemetryToolMetadata = {
     orgId?: string;
 };
 
+export type ToolConstructor = new (session: Session, config: UserConfig, telemetry: Telemetry) => ToolBase;
+
 export abstract class ToolBase {
     public abstract name: string;
 

src/tools/tool.ts

@@ -15,6 +15,7 @@ import { MCPConnectionManager } from "../../src/common/connectionManager.js";
 import { DeviceId } from "../../src/helpers/deviceId.js";
 import { connectionErrorHandler } from "../../src/common/connectionErrorHandler.js";
 import { Keychain } from "../../src/common/keychain.js";
+import type { Client as AtlasLocalClient } from "@mongodb-js-preview/atlas-local";
 
 interface ParameterInfo {
     name: string;
@@ -345,6 +346,37 @@ export function waitUntil<T extends ConnectionState>(
     });
 }
 
+export function waitUntilMcpClientIsSet(
+    mcpServer: Server,
+    signal: AbortSignal,
+    timeout: number = 5000
+): Promise<AtlasLocalClient> {
+    let ts: NodeJS.Timeout | undefined;
+
+    const timeoutSignal = AbortSignal.timeout(timeout);
+    const combinedSignal = AbortSignal.any([signal, timeoutSignal]);
+
+    return new Promise<AtlasLocalClient>((resolve, reject) => {
+        ts = setInterval(() => {
+            if (combinedSignal.aborted) {
+                return reject(new Error(`Aborted: ${combinedSignal.reason}`));
+            }
+
+            // wait until session.client != undefined
+            // do not wait more than 1 second, should take a few milliseconds at most
+            // try every 50ms to see if the client is set, if it's not set after 1 second, throw an error
+            const client = mcpServer.session.atlasLocalClient;
+            if (client) {
+                return resolve(client);
+            }
+        }, 100);
+    }).finally(() => {
+        if (ts !== undefined) {
+            clearInterval(ts);
+        }
+    });
+}
+
 export function getDataFromUntrustedContent(content: string): string {
     const regex = /^[ \t]*<untrusted-user-data-[0-9a-f\\-]*>(?<data>.*)^[ \t]*<\/untrusted-user-data-[0-9a-f\\-]*>/gms;
     const match = regex.exec(content);

tests/integration/helpers.ts

@@ -11,7 +11,9 @@ describe("Server integration test", () => {
                 expectDefined(tools);
                 expect(tools.tools.length).toBeGreaterThan(0);
 
-                const atlasTools = tools.tools.filter((tool) => tool.name.startsWith("atlas-"));
+                const atlasTools = tools.tools.filter(
+                    (tool) => tool.name.startsWith("atlas-") && !tool.name.startsWith("atlas-local-")
+                );
                 expect(atlasTools.length).toBeLessThanOrEqual(0);
             });
         },