refactor: adapt to the latest structure

gagik · gagik · commit d43beaa8502b · 2025-04-11T15:52:51.000+02:00
diff --git a/src/index.ts b/src/index.ts
@@ -3,27 +3,26 @@
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.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 { Session } from "./session.js";
 import { Server } from "./server.js";
 
 try {
-    const state = new Session();
+    const session = new Session();
     const mcpServer = new McpServer({
         name: "MongoDB Atlas",
         version: config.version,
     });
-    const transport = new StdioServerTransport();
 
-    const server = new Seriver({
+    const server = new Server({
         mcpServer,
-        state,
-        transport,
+        session,
     });
 
-    await server.registerAndConnect();
+    const transport = new StdioServerTransport();
+
+    await server.connect(transport);
 } catch (error) {
     logger.emergency(mongoLogId(1_000_004), "server", `Fatal error running server: ${error}`);
 
diff --git a/src/server.ts b/src/server.ts
@@ -9,28 +9,22 @@ import { mongoLogId } from "mongodb-log-writer";
 export class Server {
     public readonly session: Session;
     private readonly mcpServer: McpServer;
-    private readonly transport: Transport;
 
-    constructor({ mcpServer, transport, session }: { mcpServer: McpServer; session: Session; transport: Transport }) {
+    constructor({ mcpServer, session }: { mcpServer: McpServer; session: Session }) {
         this.mcpServer = mcpServer;
-        this.transport = transport;
         this.session = session;
     }
 
-    async connect() {
+    async connect(transport: Transport) {
         this.mcpServer.server.registerCapabilities({ logging: {} });
 
         this.registerTools();
 
         await initializeLogger(this.mcpServer);
 
-        await this.mcpServer.connect(this.transport);
+        await this.mcpServer.connect(transport);
 
-        logger.info(
-            mongoLogId(1_000_004),
-            "server",
-            `Server started with transport ${this.transport.constructor.name}`
-        );
+        logger.info(mongoLogId(1_000_004), "server", `Server started with transport ${transport.constructor.name}`);
     }
 
     async close(): Promise<void> {
diff --git a/src/session.ts b/src/session.ts
@@ -6,11 +6,11 @@ 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(
-                    "Not authenticated make sure to configure MCP server with MDB_MCP_API_CLIENT_ID and MDB_MCP_API_CLIENT_SECRET environment variables."
+                    "Authenticated failed. Please configure the MCP server with MDB_MCP_API_CLIENT_ID and MDB_MCP_API_CLIENT_SECRET environment variables."
                 );
             }
 
diff --git a/src/tools/atlas/atlasTool.ts b/src/tools/atlas/atlasTool.ts
@@ -5,25 +5,7 @@ import { Session } from "../../session.js";
 export abstract class AtlasToolBase extends ToolBase {
     private 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,
-                },
-            });
-        }
-    }
-
     constructor(protected readonly session: Session) {
-        super(state);
+        super(session);
     }
 }
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.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,
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.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,
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.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,
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.state.ensureApiClient();
+        this.session.ensureAuthenticated();
 
-        const accessList = await this.state.apiClient.listProjectIpAccessLists({
+        const accessList = await this.session.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.state.ensureApiClient();
+        this.session.ensureAuthenticated();
 
-        const cluster = await this.state.apiClient.getCluster({
+        const cluster = await this.session.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.state.ensureApiClient();
+        this.session.ensureAuthenticated();
 
         if (!projectId) {
-            const data = await this.state.apiClient.listClustersForAllProjects();
+            const data = await this.session.apiClient.listClustersForAllProjects();
 
             return this.formatAllClustersTable(data);
         } else {
-            const project = await this.state.apiClient.getProject({
+            const project = await this.session.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.state.apiClient.listClusters({
+            const data = await this.session.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.state.ensureApiClient();
+        this.session.ensureAuthenticated();
 
-        const data = await this.state.apiClient.listDatabaseUsers({
+        const data = await this.session.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.state.ensureApiClient();
+        this.session.ensureAuthenticated();
 
-        const data = await this.state.apiClient.listProjects();
+        const data = await this.session.apiClient.listProjects();
 
         if (!data?.results?.length) {
             throw new Error("No projects found in your MongoDB Atlas account.");
diff --git a/src/tools/mongodb/index.ts b/src/tools/mongodb/index.ts
@@ -1,5 +1,3 @@
-import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
-import { State } from "../../session.js";
 import { ConnectTool } from "./connect.js";
 import { ListCollectionsTool } from "./metadata/listCollections.js";
 import { CollectionIndexesTool } from "./collectionIndexes.js";
diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts
@@ -1,6 +1,6 @@
 import { z } from "zod";
 import { ToolBase } from "../tool.js";
-import { State } from "../../state.js";
+import { Session } from "../../session.js";
 import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { ErrorCodes, MongoDBError } from "../../errors.js";
@@ -14,16 +14,16 @@ export const DbOperationArgs = {
 export type DbOperationType = "metadata" | "read" | "create" | "update" | "delete";
 
 export abstract class MongoDBToolBase extends ToolBase {
-    constructor(state: State) {
-        super(state);
+    constructor(session: Session) {
+        super(session);
     }
 
     protected abstract operationType: DbOperationType;
 
     protected async ensureConnected(): Promise<NodeDriverServiceProvider> {
-        const provider = this.state.serviceProvider;
+        const provider = this.session.serviceProvider;
         if (!provider && config.connectionString) {
-            await this.connectToMongoDB(config.connectionString, this.state);
+            await this.connectToMongoDB(config.connectionString);
         }
 
         if (!provider) {
@@ -53,7 +53,7 @@ export abstract class MongoDBToolBase extends ToolBase {
         return super.handleError(error);
     }
 
-    protected async connectToMongoDB(connectionString: string, state: State): Promise<void> {
+    protected async connectToMongoDB(connectionString: string): Promise<void> {
         const provider = await NodeDriverServiceProvider.connect(connectionString, {
             productDocsLink: "https://docs.mongodb.com/todo-mcp",
             productName: "MongoDB MCP",
@@ -67,6 +67,6 @@ export abstract class MongoDBToolBase extends ToolBase {
             timeoutMS: config.connectOptions.timeoutMS,
         });
 
-        state.serviceProvider = provider;
+        this.session.serviceProvider = provider;
     }
 }
diff --git a/src/tools/tool.ts b/src/tools/tool.ts
@@ -1,7 +1,7 @@
 import { McpServer, ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { z, ZodNever, ZodRawShape } from "zod";
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
-import { State } from "../session.js";
+import { Session } from "../session.js";
 import logger from "../logger.js";
 import { mongoLogId } from "mongodb-log-writer";
 
@@ -16,7 +16,7 @@ export abstract class ToolBase {
 
     protected abstract execute(...args: Parameters<ToolCallback<typeof this.argsShape>>): Promise<CallToolResult>;
 
-    protected constructor(protected state: State) {}
+    protected constructor(protected session: Session) {}
 
     public register(server: McpServer): void {
         const callback: ToolCallback<typeof this.argsShape> = async (...args) => {
diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts
@@ -4,6 +4,8 @@ import { Server } from "../../src/server.js";
 import runner, { MongoCluster } from "mongodb-runner";
 import path from "path";
 import fs from "fs/promises";
+import { Session } from "../../src/session.js";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 
 export async function setupIntegrationTest(): Promise<{
     client: Client;
@@ -29,7 +31,13 @@ export async function setupIntegrationTest(): Promise<{
         }
     );
 
-    const server = new Server();
+    const server = new Server({
+        mcpServer: new McpServer({
+            name: "test-server",
+            version: "1.2.3",
+        }),
+        session: new Session(),
+    });
     await server.connect(serverTransport);
     await client.connect(clientTransport);
 
diff --git a/tests/integration/server.test.ts b/tests/integration/server.test.ts
@@ -3,7 +3,7 @@ import { setupIntegrationTest } from "./helpers.js";
 
 describe("Server integration test", () => {
     let client: Client;
-    let teardown: () => Promise<void>;
+    let teardown: () => Promise<void> = () => Promise.resolve();
 
     beforeEach(async () => {
         ({ client, teardown } = await setupIntegrationTest());
diff --git a/tests/unit/index.test.ts b/tests/unit/index.test.ts
