move atlas local client initialization to 'TransportRunnerBase'

Author: jeroenvervaeke

src/common/atlasLocal.ts

@@ -0,0 +1,36 @@
+import type { Client } from "@mongodb-js-preview/atlas-local";
+
+export type AtlasLocalClientFactoryFn = () => Promise<Client | undefined>;
+
+export const defaultCreateAtlasLocalClient: AtlasLocalClientFactoryFn = async () => {
+    console.log("defaultCreateAtlasLocalClient");
+    try {
+        // Import Atlas Local client asyncronously
+        // This will fail on unsupported platforms
+        // also measure the time it takes to import the client
+        const { Client: AtlasLocalClient } = await import("@mongodb-js-preview/atlas-local");
+
+        try {
+            // Connect to Atlas Local client
+            // This will fail if docker is not running
+            const client = AtlasLocalClient.connect();
+
+            // Set Atlas Local client
+            return client;
+        } catch (dockerError) {
+            console.warn(
+                "Failed to connect to Atlas Local client (Docker not available or not running), atlas-local tools will be disabled (error: ",
+                dockerError,
+                ")"
+            );
+        }
+    } catch (importError) {
+        console.warn(
+            "Failed to import Atlas Local client (platform not supported), atlas-local tools will be disabled (error: ",
+            importError,
+            ")"
+        );
+    }
+
+    return undefined;
+};

src/common/session.ts

@@ -26,6 +26,7 @@ export interface SessionOptions {
     exportsManager: ExportsManager;
     connectionManager: ConnectionManager;
     keychain: Keychain;
+    atlasLocalClient?: Client;
 }
 
 export type SessionEvents = {
@@ -40,14 +41,14 @@ export class Session extends EventEmitter<SessionEvents> {
     readonly exportsManager: ExportsManager;
     readonly connectionManager: ConnectionManager;
     readonly apiClient: ApiClient;
+    readonly atlasLocalClient?: Client;
     readonly keychain: Keychain;
 
     mcpClient?: {
         name?: string;
         version?: string;
         title?: string;
     };
-    atlasLocalClient?: Client;
 
     public logger: CompositeLogger;
 
@@ -59,6 +60,7 @@ export class Session extends EventEmitter<SessionEvents> {
         connectionManager,
         exportsManager,
         keychain,
+        atlasLocalClient,
     }: SessionOptions) {
         super();
 
@@ -73,6 +75,7 @@ export class Session extends EventEmitter<SessionEvents> {
                 : undefined;
 
         this.apiClient = new ApiClient({ baseUrl: apiBaseUrl, credentials }, logger);
+        this.atlasLocalClient = atlasLocalClient;
         this.exportsManager = exportsManager;
         this.connectionManager = connectionManager;
         this.connectionManager.events.on("connection-success", () => this.emit("connect"));

src/server.ts

@@ -87,9 +87,6 @@ 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.
@@ -220,55 +217,8 @@ export class Server {
         this.telemetry.emitEvents([event]);
     }
 
-    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");
-
-            try {
-                // Connect to Atlas Local client
-                // This will fail if docker is not running
-                const client = AtlasLocalClient.connect();
-
-                // Set Atlas Local client
-                this.session.atlasLocalClient = client;
-
-                // Register Atlas Local tools
-                for (const toolConstructor of AtlasLocalTools) {
-                    const tool = new toolConstructor({
-                        session: this.session,
-                        config: this.userConfig,
-                        telemetry: this.telemetry,
-                        elicitation: this.elicitation,
-                    });
-                    if (tool.register(this)) {
-                        this.tools.push(tool);
-                    }
-                }
-            } catch (dockerError) {
-                console.warn(
-                    "Failed to connect to Atlas Local client (Docker not available or not running), atlas-local tools will be disabled (error: ",
-                    dockerError,
-                    ")"
-                );
-            }
-        } catch (importError) {
-            console.warn(
-                "Failed to import Atlas Local client (platform not supported), atlas-local tools will be disabled (error: ",
-                importError,
-                ")"
-            );
-        }
-    }
-
     private registerTools(): void {
-        for (const toolConstructor of this.toolConstructors) {
+        for (const toolConstructor of [...this.toolConstructors, ...AtlasLocalTools]) {
             const tool = new toolConstructor({
                 session: this.session,
                 config: this.userConfig,

src/transports/base.ts

@@ -16,11 +16,15 @@ import {
 } from "../common/connectionErrorHandler.js";
 import type { CommonProperties } from "../telemetry/types.js";
 import { Elicitation } from "../elicitation.js";
+import type { AtlasLocalClientFactoryFn } from "../common/atlasLocal.js";
+import { defaultCreateAtlasLocalClient } from "../common/atlasLocal.js";
+import type { Client } from "@mongodb-js-preview/atlas-local";
 
 export type TransportRunnerConfig = {
     userConfig: UserConfig;
     createConnectionManager?: ConnectionManagerFactoryFn;
     connectionErrorHandler?: ConnectionErrorHandler;
+    createAtlasLocalClient?: AtlasLocalClientFactoryFn;
     additionalLoggers?: LoggerBase[];
     telemetryProperties?: Partial<CommonProperties>;
 };
@@ -31,18 +35,21 @@ export abstract class TransportRunnerBase {
     protected readonly userConfig: UserConfig;
     private readonly createConnectionManager: ConnectionManagerFactoryFn;
     private readonly connectionErrorHandler: ConnectionErrorHandler;
+    private readonly atlasLocalClient: Promise<Client | undefined>;
     private readonly telemetryProperties: Partial<CommonProperties>;
 
     protected constructor({
         userConfig,
         createConnectionManager = createMCPConnectionManager,
         connectionErrorHandler = defaultConnectionErrorHandler,
+        createAtlasLocalClient = defaultCreateAtlasLocalClient,
         additionalLoggers = [],
         telemetryProperties = {},
     }: TransportRunnerConfig) {
         this.userConfig = userConfig;
         this.createConnectionManager = createConnectionManager;
         this.connectionErrorHandler = connectionErrorHandler;
+        this.atlasLocalClient = createAtlasLocalClient();
         this.telemetryProperties = telemetryProperties;
         const loggers: LoggerBase[] = [...additionalLoggers];
         if (this.userConfig.loggers.includes("stderr")) {
@@ -85,6 +92,7 @@ export abstract class TransportRunnerBase {
             apiBaseUrl: this.userConfig.apiBaseUrl,
             apiClientId: this.userConfig.apiClientId,
             apiClientSecret: this.userConfig.apiClientSecret,
+            atlasLocalClient: await this.atlasLocalClient,
             logger,
             exportsManager,
             connectionManager,