refactor: move apiClient into state

fmenezes · fmenezes · commit 616a32eca81a · 2025-04-11T13:27:23.000+01:00
diff --git a/src/common/atlas/apiClient.ts b/src/common/atlas/apiClient.ts
@@ -1,9 +1,11 @@
 import config from "../../config.js";
-import createClient, { FetchOptions, Middleware } from "openapi-fetch";
+import createClient, { Client, FetchOptions, Middleware } from "openapi-fetch";
 import { AccessToken, ClientCredentials } from "simple-oauth2";
 
 import { paths, operations } from "./openapi.js";
 
+const ATLAS_API_VERSION = "2025-03-12";
+
 export class ApiClientError extends Error {
     response?: Response;
 
@@ -25,38 +27,32 @@ export class ApiClientError extends Error {
 }
 
 export interface ApiClientOptions {
-    credentials: {
+    credentials?: {
         clientId: string;
         clientSecret: string;
     };
     baseUrl?: string;
+    userAgent?: string;
 }
 
 export class ApiClient {
-    private client = createClient<paths>({
-        baseUrl: config.apiBaseUrl,
-        headers: {
-            "User-Agent": config.userAgent,
-            Accept: `application/vnd.atlas.${config.atlasApiVersion}+json`,
-        },
-    });
-    private oauth2Client = new ClientCredentials({
-        client: {
-            id: this.options.credentials.clientId,
-            secret: this.options.credentials.clientSecret,
-        },
-        auth: {
-            tokenHost: this.options.baseUrl || config.apiBaseUrl,
-            tokenPath: "/api/oauth/token",
-        },
-    });
+    private options: {
+        baseUrl: string;
+        userAgent: string;
+        credentials?: {
+            clientId: string;
+            clientSecret: string;
+        };
+    };
+    private client: Client<paths>;
+    private oauth2Client?: ClientCredentials;
     private accessToken?: AccessToken;
 
     private getAccessToken = async () => {
-        if (!this.accessToken || this.accessToken.expired()) {
+        if (this.oauth2Client && (!this.accessToken || this.accessToken.expired())) {
             this.accessToken = await this.oauth2Client.getToken({});
         }
-        return this.accessToken.token.access_token;
+        return this.accessToken?.token.access_token as string | undefined;
     };
 
     private authMiddleware = (apiClient: ApiClient): Middleware => ({
@@ -82,22 +78,51 @@ export class ApiClient {
         },
     });
 
-    constructor(private options: ApiClientOptions) {
-        this.client.use(this.authMiddleware(this));
+    constructor(options?: ApiClientOptions) {
+        const defaultOptions = {
+            baseUrl: "https://cloud.mongodb.com/",
+            userAgent: `AtlasMCP/${config.version} (${process.platform}; ${process.arch}; ${process.env.HOSTNAME || "unknown"})`,
+        };
+
+        this.options = {
+            ...defaultOptions,
+            ...options,
+        };
+
+        this.client = createClient<paths>({
+            baseUrl: this.options.baseUrl,
+            headers: {
+                "User-Agent": this.options.userAgent,
+                Accept: `application/vnd.atlas.${ATLAS_API_VERSION}+json`,
+            },
+        });
+        if (this.options.credentials?.clientId && this.options.credentials?.clientSecret) {
+            this.oauth2Client = new ClientCredentials({
+                client: {
+                    id: this.options.credentials.clientId,
+                    secret: this.options.credentials.clientSecret,
+                },
+                auth: {
+                    tokenHost: this.options.baseUrl,
+                    tokenPath: "/api/oauth/token",
+                },
+            });
+            this.client.use(this.authMiddleware(this));
+        }
         this.client.use(this.errorMiddleware());
     }
 
     async getIpInfo() {
         const accessToken = await this.getAccessToken();
 
         const endpoint = "api/private/ipinfo";
-        const url = new URL(endpoint, config.apiBaseUrl);
+        const url = new URL(endpoint, this.options.baseUrl || config.apiBaseUrl);
         const response = await fetch(url, {
             method: "GET",
             headers: {
                 Accept: "application/json",
                 Authorization: `Bearer ${accessToken}`,
-                "User-Agent": config.userAgent,
+                "User-Agent": this.options.userAgent,
             },
         });
 
diff --git a/src/config.ts b/src/config.ts
@@ -3,13 +3,12 @@ import os from "os";
 import argv from "yargs-parser";
 
 import packageJson from "../package.json" with { type: "json" };
-import fs from "fs";
 import { ReadConcernLevel, ReadPreferenceMode, W } from "mongodb";
 
 // If we decide to support non-string config options, we'll need to extend the mechanism for parsing
 // env variables.
 interface UserConfig {
-    apiBaseUrl: string;
+    apiBaseUrl?: string;
     apiClientId?: string;
     apiClientSecret?: string;
     logPath: string;
@@ -23,7 +22,6 @@ interface UserConfig {
 }
 
 const defaults: UserConfig = {
-    apiBaseUrl: "https://cloud.mongodb.com/",
     logPath: getLogPath(),
     connectOptions: {
         readConcern: "local",
@@ -41,17 +39,16 @@ const mergedUserConfig = {
 
 const config = {
     ...mergedUserConfig,
-    atlasApiVersion: `2025-03-12`,
     version: packageJson.version,
-    userAgent: `AtlasMCP/${packageJson.version} (${process.platform}; ${process.arch}; ${process.env.HOSTNAME || "unknown"})`,
 };
 
 export default config;
 
 function getLogPath(): string {
-    const localDataPath = (process.platform === "win32") ?
-        path.join(process.env.LOCALAPPDATA || process.env.APPDATA || os.homedir(), "mongodb")
-        : path.join(os.homedir(), ".mongodb");
+    const localDataPath =
+        process.platform === "win32"
+            ? path.join(process.env.LOCALAPPDATA || process.env.APPDATA || os.homedir(), "mongodb")
+            : path.join(os.homedir(), ".mongodb");
 
     const logPath = path.join(localDataPath, "mongodb-mcp", ".app-logs");
 
diff --git a/src/server.ts b/src/server.ts
@@ -1,5 +1,4 @@
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
-import { ApiClient } from "./common/atlas/apiClient.js";
 import defaultState, { State } from "./state.js";
 import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
 import { registerAtlasTools } from "./tools/atlas/tools.js";
@@ -10,7 +9,6 @@ import { mongoLogId } from "mongodb-log-writer";
 
 export class Server {
     state: State = defaultState;
-    apiClient?: ApiClient;
     private server?: McpServer;
 
     async connect(transport: Transport) {
@@ -21,12 +19,12 @@ export class Server {
 
         this.server.server.registerCapabilities({ logging: {} });
 
-        registerAtlasTools(this.server, this.state, this.apiClient);
+        registerAtlasTools(this.server, this.state);
         registerMongoDBTools(this.server, this.state);
 
         await this.server.connect(transport);
         await initializeLogger(this.server);
-        
+
         logger.info(mongoLogId(1_000_004), "server", `Server started with transport ${transport.constructor.name}`);
     }
 
diff --git a/src/state.ts b/src/state.ts
@@ -1,7 +1,28 @@
 import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
+import { ApiClient } from "./common/atlas/apiClient.js";
+import config from "./config.js";
 
 export class State {
     serviceProvider?: NodeDriverServiceProvider;
+    apiClient?: ApiClient;
+
+    ensureApiClient(): asserts this is { apiClient: ApiClient } {
+        if (!this.apiClient) {
+            if (!config.apiClientId || !config.apiClientSecret) {
+                throw new Error(
+                    "Not authenticated make sure to configure MCP server with MDB_MCP_API_CLIENT_ID and MDB_MCP_API_CLIENT_SECRET environment variables."
+                );
+            }
+
+            this.apiClient = new ApiClient({
+                baseUrl: config.apiBaseUrl,
+                credentials: {
+                    clientId: config.apiClientId,
+                    clientSecret: config.apiClientSecret,
+                },
+            });
+        }
+    }
 }
 
 const defaultState = new State();
diff --git a/src/tools/atlas/atlasTool.ts b/src/tools/atlas/atlasTool.ts
@@ -1,20 +1,8 @@
 import { ToolBase } from "../tool.js";
-import { ApiClient } from "../../common/atlas/apiClient.js";
 import { State } from "../../state.js";
 
 export abstract class AtlasToolBase extends ToolBase {
-    constructor(
-        state: State,
-        protected apiClient?: ApiClient
-    ) {
+    constructor(state: State) {
         super(state);
     }
-
-    protected ensureAuthenticated(): asserts this is { apiClient: ApiClient } {
-        if (!this.apiClient) {
-            throw new Error(
-                "Not authenticated make sure to configure MCP server with MDB_MCP_API_CLIENT_ID and MDB_MCP_API_CLIENT_SECRET environment variables."
-            );
-        }
-    }
 }
diff --git a/src/tools/atlas/createAccessList.ts b/src/tools/atlas/createAccessList.ts
@@ -26,7 +26,7 @@ export class CreateAccessListTool extends AtlasToolBase {
         comment,
         currentIpAddress,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        this.ensureAuthenticated();
+        this.state.ensureApiClient();
 
         if (!ipAddresses?.length && !cidrBlocks?.length && !currentIpAddress) {
             throw new Error("One of  ipAddresses, cidrBlocks, currentIpAddress must be provided.");
@@ -39,7 +39,7 @@ export class CreateAccessListTool extends AtlasToolBase {
         }));
 
         if (currentIpAddress) {
-            const currentIp = await this.apiClient.getIpInfo();
+            const currentIp = await this.state.apiClient.getIpInfo();
             const input = {
                 groupId: projectId,
                 ipAddress: currentIp.currentIpv4Address,
@@ -56,7 +56,7 @@ export class CreateAccessListTool extends AtlasToolBase {
 
         const inputs = [...ipInputs, ...cidrInputs];
 
-        await this.apiClient.createProjectIpAccessList({
+        await this.state.apiClient.createProjectIpAccessList({
             params: {
                 path: {
                     groupId: projectId,
diff --git a/src/tools/atlas/createDBUser.ts b/src/tools/atlas/createDBUser.ts
@@ -33,7 +33,7 @@ export class CreateDBUserTool extends AtlasToolBase {
         roles,
         clusters,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        this.ensureAuthenticated();
+        this.state.ensureApiClient();
 
         const input = {
             groupId: projectId,
@@ -53,7 +53,7 @@ export class CreateDBUserTool extends AtlasToolBase {
                 : undefined,
         } as CloudDatabaseUser;
 
-        await this.apiClient.createDatabaseUser({
+        await this.state.apiClient.createDatabaseUser({
             params: {
                 path: {
                     groupId: projectId,
diff --git a/src/tools/atlas/createFreeCluster.ts b/src/tools/atlas/createFreeCluster.ts
@@ -14,7 +14,7 @@ export class CreateFreeClusterTool extends AtlasToolBase {
     };
 
     protected async execute({ projectId, name, region }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        this.ensureAuthenticated();
+        this.state.ensureApiClient();
 
         const input = {
             groupId: projectId,
@@ -38,7 +38,7 @@ export class CreateFreeClusterTool extends AtlasToolBase {
             terminationProtectionEnabled: false,
         } as unknown as ClusterDescription20240805;
 
-        await this.apiClient.createCluster({
+        await this.state.apiClient.createCluster({
             params: {
                 path: {
                     groupId: projectId,
diff --git a/src/tools/atlas/inspectAccessList.ts b/src/tools/atlas/inspectAccessList.ts
@@ -11,9 +11,9 @@ export class InspectAccessListTool extends AtlasToolBase {
     };
 
     protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        this.ensureAuthenticated();
+        this.state.ensureApiClient();
 
-        const accessList = await this.apiClient.listProjectIpAccessLists({
+        const accessList = await this.state.apiClient.listProjectIpAccessLists({
             params: {
                 path: {
                     groupId: projectId,
diff --git a/src/tools/atlas/inspectCluster.ts b/src/tools/atlas/inspectCluster.ts
@@ -13,9 +13,9 @@ export class InspectClusterTool extends AtlasToolBase {
     };
 
     protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        this.ensureAuthenticated();
+        this.state.ensureApiClient();
 
-        const cluster = await this.apiClient.getCluster({
+        const cluster = await this.state.apiClient.getCluster({
             params: {
                 path: {
                     groupId: projectId,
diff --git a/src/tools/atlas/listClusters.ts b/src/tools/atlas/listClusters.ts
@@ -12,14 +12,14 @@ export class ListClustersTool extends AtlasToolBase {
     };
 
     protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        this.ensureAuthenticated();
+        this.state.ensureApiClient();
 
         if (!projectId) {
-            const data = await this.apiClient.listClustersForAllProjects();
+            const data = await this.state.apiClient.listClustersForAllProjects();
 
             return this.formatAllClustersTable(data);
         } else {
-            const project = await this.apiClient.getProject({
+            const project = await this.state.apiClient.getProject({
                 params: {
                     path: {
                         groupId: projectId,
@@ -31,7 +31,7 @@ export class ListClustersTool extends AtlasToolBase {
                 throw new Error(`Project with ID "${projectId}" not found.`);
             }
 
-            const data = await this.apiClient.listClusters({
+            const data = await this.state.apiClient.listClusters({
                 params: {
                     path: {
                         groupId: project.id || "",
diff --git a/src/tools/atlas/listDBUsers.ts b/src/tools/atlas/listDBUsers.ts
@@ -12,9 +12,9 @@ export class ListDBUsersTool extends AtlasToolBase {
     };
 
     protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        this.ensureAuthenticated();
+        this.state.ensureApiClient();
 
-        const data = await this.apiClient.listDatabaseUsers({
+        const data = await this.state.apiClient.listDatabaseUsers({
             params: {
                 path: {
                     groupId: projectId,
diff --git a/src/tools/atlas/listProjects.ts b/src/tools/atlas/listProjects.ts
@@ -7,9 +7,9 @@ export class ListProjectsTool extends AtlasToolBase {
     protected argsShape = {};
 
     protected async execute(): Promise<CallToolResult> {
-        this.ensureAuthenticated();
+        this.state.ensureApiClient();
 
-        const data = await this.apiClient.listProjects();
+        const data = await this.state.apiClient.listProjects();
 
         if (!data?.results?.length) {
             throw new Error("No projects found in your MongoDB Atlas account.");
diff --git a/src/tools/atlas/tools.ts b/src/tools/atlas/tools.ts
