refactor: rename state to session, combine tool registration, and clearer dependencies

eslint.config.js

@@ -9,5 +9,11 @@ export default defineConfig([
     { files: ["src/**/*.ts"], languageOptions: { globals: globals.node } },
     tseslint.configs.recommended,
     eslintConfigPrettier,
+    {
+        files: ["src/**/*.ts"],
+        rules: {
+            "@typescript-eslint/no-non-null-assertion": "error",
+        },
+    },
     globalIgnores(["node_modules", "dist"]),
 ]);

src/common/atlas/apiClient.ts

@@ -2,6 +2,7 @@ import config from "../../config.js";
 import createClient, { FetchOptions, Middleware } from "openapi-fetch";
 
 import { paths, operations } from "./openapi.js";
+import { State } from "../../state.js";
 
 export interface OAuthToken {
     access_token: string;
@@ -50,39 +51,52 @@ export interface ApiClientOptions {
 
 export class ApiClient {
     private token?: OAuthToken;
-    private saveToken?: saveTokenFunction;
-    private client = createClient<paths>({
+    private readonly saveToken?: saveTokenFunction;
+    private readonly client = createClient<paths>({
         baseUrl: config.apiBaseUrl,
         headers: {
             "User-Agent": config.userAgent,
             Accept: `application/vnd.atlas.${config.atlasApiVersion}+json`,
         },
     });
-    private authMiddleware = (apiClient: ApiClient): Middleware => ({
-        async onRequest({ request, schemaPath }) {
+
+    private readonly authMiddleware: Middleware = {
+        onRequest: async ({ request, schemaPath }) => {
             if (schemaPath.startsWith("/api/private/unauth") || schemaPath.startsWith("/api/oauth")) {
                 return undefined;
             }
-            if (await apiClient.validateToken()) {
-                request.headers.set("Authorization", `Bearer ${apiClient.token!.access_token}`);
+            if (this.token && (await this.validateToken())) {
+                request.headers.set("Authorization", `Bearer ${this.token.access_token}`);
                 return request;
             }
         },
-    });
-    private errorMiddleware = (): Middleware => ({
+    };
+    private readonly errorMiddleware: Middleware = {
         async onResponse({ response }) {
             if (!response.ok) {
                 throw await ApiClientError.fromResponse(response);
             }
         },
-    });
+    };
 
     constructor(options: ApiClientOptions) {
         const { token, saveToken } = options;
         this.token = token;
         this.saveToken = saveToken;
-        this.client.use(this.authMiddleware(this));
-        this.client.use(this.errorMiddleware());
+        this.client.use(this.authMiddleware);
+        this.client.use(this.errorMiddleware);
+    }
+
+    static fromState(state: State): ApiClient {
+        return new ApiClient({
+            token: state.credentials.auth.token,
+            saveToken: async (token) => {
+                state.credentials.auth.code = undefined;
+                state.credentials.auth.token = token;
+                state.credentials.auth.status = "issued";
+                await state.persistCredentials();
+            },
+        });
     }
 
     async storeToken(token: OAuthToken): Promise<OAuthToken> {
@@ -160,7 +174,7 @@ export class ApiClient {
         }
     }
 
-    async refreshToken(token?: OAuthToken): Promise<OAuthToken | null> {
+    async refreshToken(token: OAuthToken): Promise<OAuthToken> {
         const endpoint = "api/private/unauth/account/device/token";
         const url = new URL(endpoint, config.apiBaseUrl);
         const response = await fetch(url, {
@@ -171,7 +185,7 @@ export class ApiClient {
             },
             body: new URLSearchParams({
                 client_id: config.clientId,
-                refresh_token: (token || this.token)?.refresh_token || "",
+                refresh_token: token.refresh_token,
                 grant_type: "refresh_token",
                 scope: "openid profile offline_access",
             }).toString(),
@@ -194,7 +208,7 @@ export class ApiClient {
         return await this.storeToken(tokenToStore);
     }
 
-    async revokeToken(token?: OAuthToken): Promise<void> {
+    async revokeToken(token: OAuthToken): Promise<void> {
         const endpoint = "api/private/unauth/account/device/token";
         const url = new URL(endpoint, config.apiBaseUrl);
         const response = await fetch(url, {
@@ -206,7 +220,7 @@ export class ApiClient {
             },
             body: new URLSearchParams({
                 client_id: config.clientId,
-                token: (token || this.token)?.access_token || "",
+                token: token.access_token || "",
                 token_type_hint: "refresh_token",
             }).toString(),
         });
@@ -222,9 +236,8 @@ export class ApiClient {
         return;
     }
 
-    private checkTokenExpiry(token?: OAuthToken): boolean {
+    private checkTokenExpiry(token: OAuthToken): boolean {
         try {
-            token = token || this.token;
             if (!token || !token.access_token) {
                 return false;
             }
@@ -239,21 +252,25 @@ export class ApiClient {
         }
     }
 
-    async validateToken(token?: OAuthToken): Promise<boolean> {
-        if (this.checkTokenExpiry(token)) {
+    async validateToken(): Promise<boolean> {
+        if (!this.token) {
+            return false;
+        }
+
+        if (this.checkTokenExpiry(this.token)) {
             return true;
         }
 
         try {
-            await this.refreshToken(token);
+            await this.refreshToken(this.token);
             return true;
         } catch {
             return false;
         }
     }
 
     async getIpInfo() {
-        if (!(await this.validateToken())) {
+        if (!this.token || !(await this.validateToken())) {
             throw new Error("Not Authenticated");
         }
 
@@ -263,7 +280,7 @@ export class ApiClient {
             method: "GET",
             headers: {
                 Accept: "application/json",
-                Authorization: `Bearer ${this.token!.access_token}`,
+                Authorization: `Bearer ${this.token.access_token}`,
                 "User-Agent": config.userAgent,
             },
         });

src/index.ts

@@ -1,17 +1,33 @@
 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 { ApiClient } from "./common/atlas/apiClient.js";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import config from "./config.js";
+import { State } from "./state.js";
+import { registerAtlasTools } from "./tools/atlas/tools.js";
+import { registerMongoDBTools } from "./tools/mongodb/index.js";
 
-export async function runServer() {
-    const server = new Server();
+try {
+    const state = new State();
+    await state.loadCredentials();
 
-    const transport = new StdioServerTransport();
-    await server.connect(transport);
-}
+    const apiClient = ApiClient.fromState(state);
+
+    const mcp = new McpServer({
+        name: "MongoDB Atlas",
+        version: config.version,
+    });
+
+    mcp.server.registerCapabilities({ logging: {} });
 
-runServer().catch((error) => {
+    registerAtlasTools(mcp, state, apiClient);
+    registerMongoDBTools(mcp, state);
+
+    const transport = new StdioServerTransport();
+    await mcp.server.connect(transport);
+} catch (error) {
     logger.emergency(mongoLogId(1_000_004), "server", `Fatal error running server: ${error}`);
 
     process.exit(1);
-});
+}

src/state.ts

@@ -40,6 +40,3 @@ export class State {
         }
     }
 }
-
-const defaultState = new State();
-export default defaultState;

src/tools/atlas/createDBUser.ts

@@ -53,7 +53,7 @@ export class CreateDBUserTool extends AtlasToolBase {
                 : undefined,
         } as CloudDatabaseUser;
 
-        await this.apiClient!.createDatabaseUser({
+        await this.apiClient.createDatabaseUser({
             params: {
                 path: {
                     groupId: projectId,

src/tools/atlas/listClusters.ts

@@ -47,8 +47,8 @@ export class ListClustersTool extends AtlasToolBase {
         if (!clusters?.results?.length) {
             throw new Error("No clusters found.");
         }
-        const rows = clusters
-            .results!.map((result) => {
+        const rows = clusters.results
+            .map((result) => {
                 return (result.clusters || []).map((cluster) => {
                     return { ...result, ...cluster, clusters: undefined };
                 });
@@ -75,8 +75,8 @@ ${rows}`,
         if (!clusters?.results?.length) {
             throw new Error("No clusters found.");
         }
-        const rows = clusters
-            .results!.map((cluster) => {
+        const rows = clusters.results
+            .map((cluster) => {
                 const connectionString = cluster.connectionStrings?.standard || "N/A";
                 const mongoDBVersion = cluster.mongoDBVersion || "N/A";
                 return `${cluster.name} | ${cluster.stateName} | ${mongoDBVersion} | ${connectionString}`;

src/tools/atlas/listDBUsers.ts

@@ -14,7 +14,7 @@ export class ListDBUsersTool extends AtlasToolBase {
     protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         await this.ensureAuthenticated();
 
-        const data = await this.apiClient!.listDatabaseUsers({
+        const data = await this.apiClient.listDatabaseUsers({
             params: {
                 path: {
                     groupId: projectId,

src/tools/atlas/listProjects.ts

@@ -9,15 +9,15 @@ export class ListProjectsTool extends AtlasToolBase {
     protected async execute(): Promise<CallToolResult> {
         await this.ensureAuthenticated();
 
-        const data = await this.apiClient!.listProjects();
+        const data = await this.apiClient.listProjects();
 
         if (!data?.results?.length) {
             throw new Error("No projects found in your MongoDB Atlas account.");
         }
 
         // Format projects as a table
-        const rows = data!
-            .results!.map((project) => {
+        const rows = data.results
+            .map((project) => {
                 const createdAt = project.created ? new Date(project.created).toLocaleString() : "N/A";
                 return `${project.name} | ${project.id} | ${createdAt}`;
             })

tests/unit/index.test.ts

@@ -1,27 +1,14 @@
 import { describe, it } from "@jest/globals";
-import { runServer } from "../../src/index";
-import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio";
-
-// mock the StdioServerTransport
-jest.mock("@modelcontextprotocol/sdk/server/stdio");
-// mock Server class and its methods
-jest.mock("../../src/server.ts", () => {
-    return {
-        Server: jest.fn().mockImplementation(() => {
-            return {
-                connect: jest.fn().mockImplementation((transport) => {
-                    return new Promise((resolve) => {
-                        resolve(transport);
-                    });
-                }),
-            };
-        }),
-    };
-});
+import { State } from "../../src/state";
 
 describe("Server initialization", () => {
-    it("should create a server instance", async () => {
-        await runServer();
-        expect(StdioServerTransport).toHaveBeenCalled();
+    it("should define a default state", async () => {
+        const state = new State();
+
+        expect(state.credentials).toEqual({
+            auth: {
+                status: "not_auth",
+            },
+        });
     });
 });

tsconfig.json

@@ -6,6 +6,7 @@
     "rootDir": "./src",
     "outDir": "./dist",
     "strict": true,
+    "strictNullChecks": true,
     "esModuleInterop": true,
     "types": ["node"],
     "sourceMap": true,