refactor: adapt to latest structure

gagik · gagik · commit 71a3885e7563 · 2025-04-11T15:41:49.000+02:00
diff --git a/src/index.ts b/src/index.ts
@@ -6,11 +6,11 @@ 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 { Session } from "./session.js";
 import { Server } from "./server.js";
 
 try {
-    const state = new State();
+    const state = new Session();
     const mcpServer = new McpServer({
         name: "MongoDB Atlas",
         version: config.version,
diff --git a/src/server.ts b/src/server.ts
@@ -1,44 +1,50 @@
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
-import { 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 { AtlasTools } from "./tools/atlas/tools.js";
+import { MongoDbTools } from "./tools/mongodb/index.js";
 import logger, { initializeLogger } from "./logger.js";
 import { mongoLogId } from "mongodb-log-writer";
-import { ApiClient } from "./common/atlas/apiClient.js";
 
 export class Server {
-    public readonly state: State;
+    public readonly session: Session;
     private readonly mcpServer: McpServer;
+    private readonly transport: Transport;
 
-    constructor({ mcpServer, state, transport }: { mcpServer: McpServer; state: State; transport: Transport }) {
-        this.mcpServer = server;
-        this.state = new State();
+    constructor({ mcpServer, transport, session }: { mcpServer: McpServer; session: Session; transport: Transport }) {
+        this.mcpServer = mcpServer;
+        this.transport = transport;
+        this.session = session;
     }
 
-    async connect(transport: Transport) {
-        this.server = new McpServer({
-            name: "MongoDB Atlas",
-            version: config.version,
-        });
+    async connect() {
+        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(this.transport);
 
-        logger.info(mongoLogId(1_000_004), "server", `Server started with transport ${transport.constructor.name}`);
+        logger.info(
+            mongoLogId(1_000_004),
+            "server",
+            `Server started with transport ${this.transport.constructor.name}`
+        );
     }
 
     async close(): Promise<void> {
         try {
-            await this.state.serviceProvider?.close(true);
+            await this.session.serviceProvider?.close(true);
         } catch {
             // Ignore errors during service provider close
         }
-        await this.server?.close();
+        await this.mcpServer?.close();
+    }
+
+    private registerTools() {
+        for (const tool of [...AtlasTools, ...MongoDbTools]) {
+            new tool(this.session).register(this.mcpServer);
+        }
     }
 }
diff --git a/src/session.ts b/src/session.ts
@@ -2,7 +2,8 @@ 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 } {
diff --git a/src/tools/atlas/atlasTool.ts b/src/tools/atlas/atlasTool.ts
@@ -1,8 +1,29 @@
 import { ToolBase } from "../tool.js";
-import { State } from "../../state.js";
+import { ApiClient } from "../../common/atlas/apiClient.js";
+import { Session } from "../../session.js";
 
 export abstract class AtlasToolBase extends ToolBase {
-    constructor(state: State) {
+    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);
     }
 }
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.state.ensureApiClient();
+        this.ensureApiClient();
 
         const input = {
             groupId: projectId,
diff --git a/src/tools/atlas/tools.ts b/src/tools/atlas/tools.ts
@@ -1,6 +1,3 @@
-import { ToolBase } from "../tool.js";
-import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
-import { State } from "../../state.js";
 import { ListClustersTool } from "./listClusters.js";
 import { ListProjectsTool } from "./listProjects.js";
 import { InspectClusterTool } from "./inspectCluster.js";
@@ -10,19 +7,13 @@ import { InspectAccessListTool } from "./inspectAccessList.js";
 import { ListDBUsersTool } from "./listDBUsers.js";
 import { CreateDBUserTool } from "./createDBUser.js";
 
-export function registerAtlasTools(server: McpServer, state: State) {
-    const tools: ToolBase[] = [
-        new ListClustersTool(state),
-        new ListProjectsTool(state),
-        new InspectClusterTool(state),
-        new CreateFreeClusterTool(state),
-        new CreateAccessListTool(state),
-        new InspectAccessListTool(state),
-        new ListDBUsersTool(state),
-        new CreateDBUserTool(state),
-    ];
-
-    for (const tool of tools) {
-        tool.register(server);
-    }
-}
+export const AtlasTools = [
+    ListClustersTool,
+    ListProjectsTool,
+    InspectClusterTool,
+    CreateFreeClusterTool,
+    CreateAccessListTool,
+    InspectAccessListTool,
+    ListDBUsersTool,
+    CreateDBUserTool,
+];
diff --git a/src/tools/mongodb/index.ts b/src/tools/mongodb/index.ts
@@ -1,5 +1,5 @@
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
-import { State } from "../../state.js";
+import { State } from "../../session.js";
 import { ConnectTool } from "./connect.js";
 import { ListCollectionsTool } from "./metadata/listCollections.js";
 import { CollectionIndexesTool } from "./collectionIndexes.js";
@@ -21,32 +21,25 @@ import { RenameCollectionTool } from "./update/renameCollection.js";
 import { DropDatabaseTool } from "./delete/dropDatabase.js";
 import { DropCollectionTool } from "./delete/dropCollection.js";
 
-export function registerMongoDBTools(server: McpServer, state: State) {
-    const tools = [
-        ConnectTool,
-        ListCollectionsTool,
-        ListDatabasesTool,
-        CollectionIndexesTool,
-        CreateIndexTool,
-        CollectionSchemaTool,
-        InsertOneTool,
-        FindTool,
-        InsertManyTool,
-        DeleteManyTool,
-        DeleteOneTool,
-        CollectionStorageSizeTool,
-        CountTool,
-        DbStatsTool,
-        AggregateTool,
-        UpdateOneTool,
-        UpdateManyTool,
-        RenameCollectionTool,
-        DropDatabaseTool,
-        DropCollectionTool,
-    ];
-
-    for (const tool of tools) {
-        const instance = new tool(state);
-        instance.register(server);
-    }
-}
+export const MongoDbTools = [
+    ConnectTool,
+    ListCollectionsTool,
+    ListDatabasesTool,
+    CollectionIndexesTool,
+    CreateIndexTool,
+    CollectionSchemaTool,
+    InsertOneTool,
+    FindTool,
+    InsertManyTool,
+    DeleteManyTool,
+    DeleteOneTool,
+    CollectionStorageSizeTool,
+    CountTool,
+    DbStatsTool,
+    AggregateTool,
+    UpdateOneTool,
+    UpdateManyTool,
+    RenameCollectionTool,
+    DropDatabaseTool,
+    DropCollectionTool,
+];
diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts
@@ -1,9 +1,9 @@
 import { z } from "zod";
 import { ToolBase } from "../tool.js";
+import { State } from "../../state.js";
 import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { ErrorCodes, MongoDBError } from "../../errors.js";
-import config from "../../config.js";
 
 export const DbOperationArgs = {
     database: z.string().describe("Database name"),
@@ -13,25 +13,22 @@ export const DbOperationArgs = {
 export type DbOperationType = "metadata" | "read" | "create" | "update" | "delete";
 
 export abstract class MongoDBToolBase extends ToolBase {
-    constructor(private serviceProvider: NodeDriverServiceProvider | undefined) {
-        super();
+    constructor(state: State) {
+        super(state);
     }
 
     protected abstract operationType: DbOperationType;
 
-    protected async ensureConnected(): Promise<NodeDriverServiceProvider> {
-        if (!this.serviceProvider && config.connectionString) {
-            this.serviceProvider = await this.connectToMongoDB(config.connectionString);
-        }
-
-        if (!this.serviceProvider) {
+    protected ensureConnected(): NodeDriverServiceProvider {
+        const provider = this.state.serviceProvider;
+        if (!provider) {
             throw new MongoDBError(ErrorCodes.NotConnectedToMongoDB, "Not connected to MongoDB");
         }
 
-        return this.serviceProvider;
+        return provider;
     }
 
-    protected handleError(error: unknown): Promise<CallToolResult> | CallToolResult {
+    protected handleError(error: unknown): CallToolResult | undefined {
         if (error instanceof MongoDBError && error.code === ErrorCodes.NotConnectedToMongoDB) {
             return {
                 content: [
@@ -44,25 +41,9 @@ export abstract class MongoDBToolBase extends ToolBase {
                         text: "Please use the 'connect' tool to connect to a MongoDB instance.",
                     },
                 ],
-                isError: true,
             };
         }
 
-        return super.handleError(error);
-    }
-
-    protected async connectToMongoDB(connectionString: string): Promise<NodeDriverServiceProvider> {
-        return NodeDriverServiceProvider.connect(connectionString, {
-            productDocsLink: "https://docs.mongodb.com/todo-mcp",
-            productName: "MongoDB MCP",
-            readConcern: {
-                level: config.connectOptions.readConcern,
-            },
-            readPreference: config.connectOptions.readPreference,
-            writeConcern: {
-                w: config.connectOptions.writeConcern,
-            },
-            timeoutMS: config.connectOptions.timeoutMS,
-        });
+        return undefined;
     }
 }
diff --git a/src/tools/tool.ts b/src/tools/tool.ts
@@ -1,6 +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 logger from "../logger.js";
 import { mongoLogId } from "mongodb-log-writer";
 
@@ -15,7 +16,7 @@ export abstract class ToolBase {
 
     protected abstract execute(...args: Parameters<ToolCallback<typeof this.argsShape>>): Promise<CallToolResult>;
 
-    protected constructor() {}
+    protected constructor(protected state: State) {}
 
     public register(server: McpServer): void {
         const callback: ToolCallback<typeof this.argsShape> = async (...args) => {
diff --git a/tests/unit/index.test.ts b/tests/unit/index.test.ts
@@ -1,5 +1,5 @@
 import { describe, it } from "@jest/globals";
-import { State } from "../../src/state";
+import { Session } from "../../src/session";
 
 // mock the StdioServerTransport
 jest.mock("@modelcontextprotocol/sdk/server/stdio");
@@ -20,7 +20,7 @@ jest.mock("../../src/server.ts", () => {
 
 describe("Server initialization", () => {
     it("should define a default state", async () => {
-        const state = new State();
+        const state = new Session();
 
         expect(state.credentials).toEqual({
             auth: {
