refactor: rename state to session, combine tool registration, and clearer dependencies (#55)

Co-authored-by: Filipe Constantinov Menezes <[email protected]>

Author: gagik

eslint.config.js

@@ -22,6 +22,7 @@ export default defineConfig([
         files,
         rules: {
             "@typescript-eslint/switch-exhaustiveness-check": "error",
+            "@typescript-eslint/no-non-null-assertion": "error",
         },
     },
     // Ignore features specific to TypeScript resolved rules

src/common/atlas/apiClient.ts

@@ -55,30 +55,31 @@ export class ApiClient {
         return this.accessToken?.token.access_token as string | undefined;
     };
 
-    private authMiddleware = (apiClient: ApiClient): Middleware => ({
-        async onRequest({ request, schemaPath }) {
+    private authMiddleware: Middleware = {
+        onRequest: async ({ request, schemaPath }) => {
             if (schemaPath.startsWith("/api/private/unauth") || schemaPath.startsWith("/api/oauth")) {
                 return undefined;
             }
 
             try {
-                const accessToken = await apiClient.getAccessToken();
+                const accessToken = await this.getAccessToken();
                 request.headers.set("Authorization", `Bearer ${accessToken}`);
                 return request;
             } catch {
                 // ignore not availble tokens, API will return 401
             }
         },
-    });
-    private errorMiddleware = (): Middleware => ({
+    };
+
+    private readonly errorMiddleware: Middleware = {
         async onResponse({ response }) {
             if (!response.ok) {
                 throw await ApiClientError.fromResponse(response);
             }
         },
-    });
+    };
 
-    constructor(options?: ApiClientOptions) {
+    constructor(options: ApiClientOptions) {
         const defaultOptions = {
             baseUrl: "https://cloud.mongodb.com/",
             userAgent: `AtlasMCP/${config.version} (${process.platform}; ${process.arch}; ${process.env.HOSTNAME || "unknown"})`,
@@ -107,9 +108,9 @@ export class ApiClient {
                     tokenPath: "/api/oauth/token",
                 },
             });
-            this.client.use(this.authMiddleware(this));
+            this.client.use(this.authMiddleware);
         }
-        this.client.use(this.errorMiddleware());
+        this.client.use(this.errorMiddleware);
     }
 
     public async getIpInfo(): Promise<{

src/config.ts

@@ -60,7 +60,10 @@ function getLogPath(): string {
 // to SNAKE_UPPER_CASE.
 function getEnvConfig(): Partial<UserConfig> {
     function setValue(obj: Record<string, unknown>, path: string[], value: string): void {
-        const currentField = path.shift()!;
+        const currentField = path.shift();
+        if (!currentField) {
+            return;
+        }
         if (path.length === 0) {
             const numberValue = Number(value);
             if (!isNaN(numberValue)) {

src/index.ts

@@ -1,19 +1,30 @@
 #!/usr/bin/env node
 
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
-import { Server } from "./server.js";
 import logger from "./logger.js";
 import { mongoLogId } from "mongodb-log-writer";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import config from "./config.js";
+import { Session } from "./session.js";
+import { Server } from "./server.js";
 
-export async function runServer() {
-    const server = new Server();
+try {
+    const session = new Session();
+    const mcpServer = new McpServer({
+        name: "MongoDB Atlas",
+        version: config.version,
+    });
+
+    const server = new Server({
+        mcpServer,
+        session,
+    });
 
     const transport = new StdioServerTransport();
-    await server.connect(transport);
-}
 
-runServer().catch((error) => {
-    logger.emergency(mongoLogId(1_000_004), "server", `Fatal error running server: ${error}`);
+    await server.connect(transport);
+} catch (error: unknown) {
+    logger.emergency(mongoLogId(1_000_004), "server", `Fatal error running server: ${error as string}`);
 
     process.exit(1);
-});
+}

src/server.ts

@@ -1,39 +1,40 @@
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
-import defaultState, { State } from "./state.js";
+import { Session } from "./session.js";
 import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
-import { registerAtlasTools } from "./tools/atlas/tools.js";
-import { registerMongoDBTools } from "./tools/mongodb/index.js";
-import config from "./config.js";
+import { AtlasTools } from "./tools/atlas/tools.js";
+import { MongoDbTools } from "./tools/mongodb/tools.js";
 import logger, { initializeLogger } from "./logger.js";
 import { mongoLogId } from "mongodb-log-writer";
 
 export class Server {
-    state: State = defaultState;
-    private server?: McpServer;
+    public readonly session: Session;
+    private readonly mcpServer: McpServer;
+
+    constructor({ mcpServer, session }: { mcpServer: McpServer; session: Session }) {
+        this.mcpServer = mcpServer;
+        this.session = session;
+    }
 
     async connect(transport: Transport) {
-        this.server = new McpServer({
-            name: "MongoDB Atlas",
-            version: config.version,
-        });
+        this.mcpServer.server.registerCapabilities({ logging: {} });
 
-        this.server.server.registerCapabilities({ logging: {} });
+        this.registerTools();
 
-        registerAtlasTools(this.server, this.state);
-        registerMongoDBTools(this.server, this.state);
+        await initializeLogger(this.mcpServer);
 
-        await initializeLogger(this.server);
-        await this.server.connect(transport);
+        await this.mcpServer.connect(transport);
 
         logger.info(mongoLogId(1_000_004), "server", `Server started with transport ${transport.constructor.name}`);
     }
 
     async close(): Promise<void> {
-        try {
-            await this.state.serviceProvider?.close(true);
-        } catch {
-            // Ignore errors during service provider close
+        await this.session.close();
+        await this.mcpServer.close();
+    }
+
+    private registerTools() {
+        for (const tool of [...AtlasTools, ...MongoDbTools]) {
+            new tool(this.session).register(this.mcpServer);
         }
-        await this.server?.close();
     }
 }

src/state.ts renamed to src/session.ts

@@ -2,11 +2,11 @@ import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver
 import { ApiClient } from "./common/atlas/apiClient.js";
 import config from "./config.js";
 
-export class State {
+export class Session {
     serviceProvider?: NodeDriverServiceProvider;
     apiClient?: ApiClient;
 
-    ensureApiClient(): asserts this is { apiClient: ApiClient } {
+    ensureAuthenticated(): asserts this is { apiClient: ApiClient } {
         if (!this.apiClient) {
             if (!config.apiClientId || !config.apiClientSecret) {
                 throw new Error(
@@ -23,7 +23,15 @@ export class State {
             });
         }
     }
-}
 
-const defaultState = new State();
-export default defaultState;
+    async close(): Promise<void> {
+        if (this.serviceProvider) {
+            try {
+                await this.serviceProvider.close(true);
+            } catch (error) {
+                console.error("Error closing service provider:", error);
+            }
+            this.serviceProvider = undefined;
+        }
+    }
+}

src/tools/atlas/atlasTool.ts

@@ -1,8 +1,8 @@
 import { ToolBase } from "../tool.js";
-import { State } from "../../state.js";
+import { Session } from "../../session.js";
 
 export abstract class AtlasToolBase extends ToolBase {
-    constructor(state: State) {
-        super(state);
+    constructor(protected readonly session: Session) {
+        super(session);
     }
 }

src/tools/atlas/createAccessList.ts

@@ -26,7 +26,7 @@ export class CreateAccessListTool extends AtlasToolBase {
         comment,
         currentIpAddress,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        this.state.ensureApiClient();
+        this.session.ensureAuthenticated();
 
         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.state.apiClient.getIpInfo();
+            const currentIp = await this.session.apiClient.getIpInfo();
             const input = {
                 groupId: projectId,
                 ipAddress: currentIp.currentIpv4Address,
@@ -56,7 +56,7 @@ export class CreateAccessListTool extends AtlasToolBase {
 
         const inputs = [...ipInputs, ...cidrInputs];
 
-        await this.state.apiClient.createProjectIpAccessList({
+        await this.session.apiClient.createProjectIpAccessList({
             params: {
                 path: {
                     groupId: projectId,

src/tools/atlas/createDBUser.ts

@@ -33,7 +33,7 @@ export class CreateDBUserTool extends AtlasToolBase {
         roles,
         clusters,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
-        this.state.ensureApiClient();
+        this.session.ensureAuthenticated();
 
         const input = {
             groupId: projectId,
@@ -53,7 +53,7 @@ export class CreateDBUserTool extends AtlasToolBase {
                 : undefined,
         } as CloudDatabaseUser;
 
-        await this.state.apiClient.createDatabaseUser({
+        await this.session.apiClient.createDatabaseUser({
             params: {
                 path: {
                     groupId: projectId,

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.state.ensureApiClient();
+        this.session.ensureAuthenticated();
 
         const input = {
             groupId: projectId,
@@ -38,7 +38,7 @@ export class CreateFreeClusterTool extends AtlasToolBase {
             terminationProtectionEnabled: false,
         } as unknown as ClusterDescription20240805;
 
-        await this.state.apiClient.createCluster({
+        await this.session.apiClient.createCluster({
             params: {
                 path: {
                     groupId: projectId,