refactor: remove type assertions and simplify state setup

Author: gagik

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;
@@ -85,6 +86,18 @@ export class ApiClient {
         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> {
         this.token = token;
 

src/index.ts

@@ -1,17 +1,37 @@
 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) => {
-    logger.emergency(mongoLogId(1_000_004), "server", `Fatal error running server: ${error}`);
+        const transport = new StdioServerTransport();
+        await mcp.server.connect(transport);
+
+        registerAtlasTools(mcp, state, apiClient);
+        registerMongoDBTools(mcp, state);
+    } catch (error) {
+        logger.emergency(mongoLogId(1_000_004), "server", `Fatal error running server: ${error}`);
+
+        process.exit(1);
+    }
+}
 
-    process.exit(1);
-});
+runServer();

src/server.ts

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

src/state.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/createDBUser.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/listClusters.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/listDBUsers.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}`;
             })

src/tools/atlas/listProjects.ts

@@ -1,6 +1,5 @@
 import { describe, it } from "@jest/globals";
-import { runServer } from "../../src/index";
-import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio";
+import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
 
 // mock the StdioServerTransport
 jest.mock("@modelcontextprotocol/sdk/server/stdio");
@@ -21,7 +20,6 @@ jest.mock("../../src/server.ts", () => {
 
 describe("Server initialization", () => {
     it("should create a server instance", async () => {
-        await runServer();
-        expect(StdioServerTransport).toHaveBeenCalled();
+        await expect(StdioServerTransport).toHaveBeenCalled();
     });
 });