mongodb-mcp-server 0.1.2 → 0.2.0

package/src/tools/mongodb/mongodbTool.ts

@@ -2,8 +2,9 @@ import { z } from "zod";
 import { ToolArgs, ToolBase, ToolCategory, TelemetryToolMetadata } from "../tool.js";
 import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
-import { ErrorCodes, MongoDBError } from "../../errors.js";
-import logger, { LogId } from "../../logger.js";
+import { ErrorCodes, MongoDBError } from "../../common/errors.js";
+import logger, { LogId } from "../../common/logger.js";
+import { Server } from "../../server.js";
 
 export const DbOperationArgs = {
     database: z.string().describe("Database name"),
@@ -11,19 +12,29 @@ export const DbOperationArgs = {
 };
 
 export abstract class MongoDBToolBase extends ToolBase {
-    protected category: ToolCategory = "mongodb";
+    private server?: Server;
+    public category: ToolCategory = "mongodb";
 
     protected async ensureConnected(): Promise<NodeDriverServiceProvider> {
-        if (!this.session.serviceProvider && this.config.connectionString) {
-            try {
-                await this.connectToMongoDB(this.config.connectionString);
-            } catch (error) {
-                logger.error(
-                    LogId.mongodbConnectFailure,
-                    "mongodbTool",
-                    `Failed to connect to MongoDB instance using the connection string from the config: ${error as string}`
+        if (!this.session.serviceProvider) {
+            if (this.session.connectedAtlasCluster) {
+                throw new MongoDBError(
+                    ErrorCodes.NotConnectedToMongoDB,
+                    `Attempting to connect to Atlas cluster "${this.session.connectedAtlasCluster.clusterName}", try again in a few seconds.`
                 );
-                throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, "Not connected to MongoDB.");
+            }
+
+            if (this.config.connectionString) {
+                try {
+                    await this.connectToMongoDB(this.config.connectionString);
+                } catch (error) {
+                    logger.error(
+                        LogId.mongodbConnectFailure,
+                        "mongodbTool",
+                        `Failed to connect to MongoDB instance using the connection string from the config: ${error as string}`
+                    );
+                    throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, "Not connected to MongoDB.");
+                }
             }
         }
 
@@ -34,11 +45,28 @@ export abstract class MongoDBToolBase extends ToolBase {
         return this.session.serviceProvider;
     }
 
+    public register(server: Server): boolean {
+        this.server = server;
+        return super.register(server);
+    }
+
     protected handleError(
         error: unknown,
         args: ToolArgs<typeof this.argsShape>
     ): Promise<CallToolResult> | CallToolResult {
         if (error instanceof MongoDBError) {
+            const connectTools = this.server?.tools
+                .filter((t) => t.operationType === "connect")
+                .sort((a, b) => a.category.localeCompare(b.category)); // Sort Altas tools before MongoDB tools
+
+            // Find the first Atlas connect tool if available and suggest to the LLM to use it.
+            // Note: if we ever have multiple Atlas connect tools, we may want to refine this logic to select the most appropriate one.
+            const atlasConnectTool = connectTools?.find((t) => t.category === "atlas");
+            const llmConnectHint = atlasConnectTool
+                ? `Note to LLM: prefer using the "${atlasConnectTool.name}" tool to connect to an Atlas cluster over using a connection string. Make sure to ask the user to specify a cluster name they want to connect to or ask them if they want to use the "list-clusters" tool to list all their clusters. Do not invent cluster names or connection strings unless the user has explicitly specified them. If they've previously connected to MongoDB using MCP, you can ask them if they want to reconnect using the same cluster/connection.`
+                : "Note to LLM: do not invent connection strings and explicitly ask the user to provide one. If they have previously connected to MongoDB using MCP, you can ask them if they want to reconnect using the same connection string.";
+
+            const connectToolsNames = connectTools?.map((t) => `"${t.name}"`).join(", ");
             switch (error.code) {
                 case ErrorCodes.NotConnectedToMongoDB:
                     return {
@@ -49,7 +77,9 @@ export abstract class MongoDBToolBase extends ToolBase {
                             },
                             {
                                 type: "text",
-                                text: "Please use the 'connect' or 'switch-connection' tool to connect to a MongoDB instance.",
+                                text: connectToolsNames
+                                    ? `Please use one of the following tools: ${connectToolsNames} to connect to a MongoDB instance or update the MCP server configuration to include a connection string. ${llmConnectHint}`
+                                    : "There are no tools available to connect. Please update the configuration to include a connection string and restart the server.",
                             },
                         ],
                         isError: true,
@@ -59,7 +89,23 @@ export abstract class MongoDBToolBase extends ToolBase {
                         content: [
                             {
                                 type: "text",
-                                text: "The configured connection string is not valid. Please check the connection string and confirm it points to a valid MongoDB instance. Alternatively, use the 'switch-connection' tool to connect to a different instance.",
+                                text: "The configured connection string is not valid. Please check the connection string and confirm it points to a valid MongoDB instance.",
+                            },
+                            {
+                                type: "text",
+                                text: connectTools
+                                    ? `Alternatively, you can use one of the following tools: ${connectToolsNames} to connect to a MongoDB instance. ${llmConnectHint}`
+                                    : "Please update the configuration to use a valid connection string and restart the server.",
+                            },
+                        ],
+                        isError: true,
+                    };
+                case ErrorCodes.ForbiddenCollscan:
+                    return {
+                        content: [
+                            {
+                                type: "text",
+                                text: error.message,
                             },
                         ],
                         isError: true,

package/src/tools/mongodb/read/aggregate.ts

@@ -3,19 +3,20 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs, OperationType } from "../../tool.js";
 import { EJSON } from "bson";
+import { checkIndexUsage } from "../../../helpers/indexCheck.js";
 
 export const AggregateArgs = {
-    pipeline: z.array(z.record(z.string(), z.unknown())).describe("An array of aggregation stages to execute"),
+    pipeline: z.array(z.object({}).passthrough()).describe("An array of aggregation stages to execute"),
 };
 
 export class AggregateTool extends MongoDBToolBase {
-    protected name = "aggregate";
+    public name = "aggregate";
     protected description = "Run an aggregation against a MongoDB collection";
     protected argsShape = {
         ...DbOperationArgs,
         ...AggregateArgs,
     };
-    protected operationType: OperationType = "read";
+    public operationType: OperationType = "read";
 
     protected async execute({
         database,
@@ -23,6 +24,16 @@ export class AggregateTool extends MongoDBToolBase {
         pipeline,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
+
+        // Check if aggregate operation uses an index if enabled
+        if (this.config.indexCheck) {
+            await checkIndexUsage(provider, database, collection, "aggregate", async () => {
+                return provider
+                    .aggregate(database, collection, pipeline, {}, { writeConcern: undefined })
+                    .explain("queryPlanner");
+            });
+        }
+
         const documents = await provider.aggregate(database, collection, pipeline).toArray();
 
         const content: Array<{ text: string; type: "text" }> = [

package/src/tools/mongodb/read/collectionIndexes.ts

@@ -3,10 +3,10 @@ import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs, OperationType } from "../../tool.js";
 
 export class CollectionIndexesTool extends MongoDBToolBase {
-    protected name = "collection-indexes";
+    public name = "collection-indexes";
     protected description = "Describe the indexes for a collection";
     protected argsShape = DbOperationArgs;
-    protected operationType: OperationType = "read";
+    public operationType: OperationType = "read";
 
     protected async execute({ database, collection }: ToolArgs<typeof DbOperationArgs>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();

package/src/tools/mongodb/read/count.ts

@@ -2,10 +2,12 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs, OperationType } from "../../tool.js";
 import { z } from "zod";
+import { checkIndexUsage } from "../../../helpers/indexCheck.js";
 
 export const CountArgs = {
     query: z
-        .record(z.string(), z.unknown())
+        .object({})
+        .passthrough()
         .optional()
         .describe(
             "A filter/query parameter. Allows users to filter the documents to count. Matches the syntax of the filter argument of db.collection.count()."
@@ -13,7 +15,7 @@ export const CountArgs = {
 };
 
 export class CountTool extends MongoDBToolBase {
-    protected name = "count";
+    public name = "count";
     protected description =
         "Gets the number of documents in a MongoDB collection using db.collection.count() and query as an optional filter parameter";
     protected argsShape = {
@@ -21,10 +23,24 @@ export class CountTool extends MongoDBToolBase {
         ...CountArgs,
     };
 
-    protected operationType: OperationType = "read";
+    public operationType: OperationType = "read";
 
     protected async execute({ database, collection, query }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
+
+        // Check if count operation uses an index if enabled
+        if (this.config.indexCheck) {
+            await checkIndexUsage(provider, database, collection, "count", async () => {
+                return provider.runCommandWithCheck(database, {
+                    explain: {
+                        count: collection,
+                        query,
+                    },
+                    verbosity: "queryPlanner",
+                });
+            });
+        }
+
         const count = await provider.count(database, collection, query);
 
         return {

package/src/tools/mongodb/read/find.ts

@@ -4,31 +4,37 @@ import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs, OperationType } from "../../tool.js";
 import { SortDirection } from "mongodb";
 import { EJSON } from "bson";
+import { checkIndexUsage } from "../../../helpers/indexCheck.js";
 
 export const FindArgs = {
     filter: z
-        .record(z.string(), z.unknown())
+        .object({})
+        .passthrough()
         .optional()
         .describe("The query filter, matching the syntax of the query argument of db.collection.find()"),
     projection: z
-        .record(z.string(), z.unknown())
+        .object({})
+        .passthrough()
         .optional()
         .describe("The projection, matching the syntax of the projection argument of db.collection.find()"),
     limit: z.number().optional().default(10).describe("The maximum number of documents to return"),
     sort: z
-        .record(z.string(), z.custom<SortDirection>())
+        .object({})
+        .catchall(z.custom<SortDirection>())
         .optional()
-        .describe("A document, describing the sort order, matching the syntax of the sort argument of cursor.sort()"),
+        .describe(
+            "A document, describing the sort order, matching the syntax of the sort argument of cursor.sort(). The keys of the object are the fields to sort on, while the values are the sort directions (1 for ascending, -1 for descending)."
+        ),
 };
 
 export class FindTool extends MongoDBToolBase {
-    protected name = "find";
+    public name = "find";
     protected description = "Run a find query against a MongoDB collection";
     protected argsShape = {
         ...DbOperationArgs,
         ...FindArgs,
     };
-    protected operationType: OperationType = "read";
+    public operationType: OperationType = "read";
 
     protected async execute({
         database,
@@ -39,6 +45,14 @@ export class FindTool extends MongoDBToolBase {
         sort,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
+
+        // Check if find operation uses an index if enabled
+        if (this.config.indexCheck) {
+            await checkIndexUsage(provider, database, collection, "find", async () => {
+                return provider.find(database, collection, filter, { projection, limit, sort }).explain("queryPlanner");
+            });
+        }
+
         const documents = await provider.find(database, collection, filter, { projection, limit, sort }).toArray();
 
         const content: Array<{ text: string; type: "text" }> = [

package/src/tools/mongodb/tools.ts

@@ -1,4 +1,4 @@
-import { ConnectTool } from "./metadata/connect.js";
+import { ConnectTool } from "./connect/connect.js";
 import { ListCollectionsTool } from "./metadata/listCollections.js";
 import { CollectionIndexesTool } from "./read/collectionIndexes.js";
 import { ListDatabasesTool } from "./metadata/listDatabases.js";

package/src/tools/mongodb/update/renameCollection.ts

@@ -4,14 +4,14 @@ import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs, OperationType } from "../../tool.js";
 
 export class RenameCollectionTool extends MongoDBToolBase {
-    protected name = "rename-collection";
+    public name = "rename-collection";
     protected description = "Renames a collection in a MongoDB database";
     protected argsShape = {
         ...DbOperationArgs,
         newName: z.string().describe("The new name for the collection"),
         dropTarget: z.boolean().optional().default(false).describe("If true, drops the target collection if it exists"),
     };
-    protected operationType: OperationType = "update";
+    public operationType: OperationType = "update";
 
     protected async execute({
         database,

package/src/tools/mongodb/update/updateMany.ts

@@ -2,27 +2,30 @@ import { z } from "zod";
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs, OperationType } from "../../tool.js";
+import { checkIndexUsage } from "../../../helpers/indexCheck.js";
 
 export class UpdateManyTool extends MongoDBToolBase {
-    protected name = "update-many";
+    public name = "update-many";
     protected description = "Updates all documents that match the specified filter for a collection";
     protected argsShape = {
         ...DbOperationArgs,
         filter: z
-            .record(z.string(), z.unknown())
+            .object({})
+            .passthrough()
             .optional()
             .describe(
                 "The selection criteria for the update, matching the syntax of the filter argument of db.collection.updateOne()"
             ),
         update: z
-            .record(z.string(), z.unknown())
+            .object({})
+            .passthrough()
             .describe("An update document describing the modifications to apply using update operator expressions"),
         upsert: z
             .boolean()
             .optional()
             .describe("Controls whether to insert a new document if no documents match the filter"),
     };
-    protected operationType: OperationType = "update";
+    public operationType: OperationType = "update";
 
     protected async execute({
         database,
@@ -32,6 +35,27 @@ export class UpdateManyTool extends MongoDBToolBase {
         upsert,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
+
+        // Check if update operation uses an index if enabled
+        if (this.config.indexCheck) {
+            await checkIndexUsage(provider, database, collection, "updateMany", async () => {
+                return provider.runCommandWithCheck(database, {
+                    explain: {
+                        update: collection,
+                        updates: [
+                            {
+                                q: filter || {},
+                                u: update,
+                                upsert: upsert || false,
+                                multi: true,
+                            },
+                        ],
+                    },
+                    verbosity: "queryPlanner",
+                });
+            });
+        }
+
         const result = await provider.updateMany(database, collection, filter, update, {
             upsert,
         });

package/src/tools/tool.ts

@@ -1,15 +1,16 @@
 import { z, type ZodRawShape, type ZodNever, AnyZodObject } from "zod";
-import type { McpServer, RegisteredTool, ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
+import type { RegisteredTool, ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
 import type { CallToolResult, ToolAnnotations } from "@modelcontextprotocol/sdk/types.js";
-import { Session } from "../session.js";
-import logger, { LogId } from "../logger.js";
+import { Session } from "../common/session.js";
+import logger, { LogId } from "../common/logger.js";
 import { Telemetry } from "../telemetry/telemetry.js";
 import { type ToolEvent } from "../telemetry/types.js";
-import { UserConfig } from "../config.js";
+import { UserConfig } from "../common/config.js";
+import { Server } from "../server.js";
 
 export type ToolArgs<Args extends ZodRawShape> = z.objectOutputType<Args, ZodNever>;
 
-export type OperationType = "metadata" | "read" | "create" | "delete" | "update";
+export type OperationType = "metadata" | "read" | "create" | "delete" | "update" | "connect";
 export type ToolCategory = "mongodb" | "atlas";
 export type TelemetryToolMetadata = {
     projectId?: string;
@@ -17,11 +18,11 @@ export type TelemetryToolMetadata = {
 };
 
 export abstract class ToolBase {
-    protected abstract name: string;
+    public abstract name: string;
 
-    protected abstract category: ToolCategory;
+    public abstract category: ToolCategory;
 
-    protected abstract operationType: OperationType;
+    public abstract operationType: OperationType;
 
     protected abstract description: string;
 
@@ -36,6 +37,7 @@ export abstract class ToolBase {
         switch (this.operationType) {
             case "read":
             case "metadata":
+            case "connect":
                 annotations.readOnlyHint = true;
                 annotations.destructiveHint = false;
                 break;
@@ -63,9 +65,9 @@ export abstract class ToolBase {
         protected readonly telemetry: Telemetry
     ) {}
 
-    public register(server: McpServer): void {
+    public register(server: Server): boolean {
         if (!this.verifyAllowed()) {
-            return;
+            return false;
         }
 
         const callback: ToolCallback<typeof this.argsShape> = async (...args) => {
@@ -74,24 +76,25 @@ export abstract class ToolBase {
                 logger.debug(LogId.toolExecute, "tool", `Executing tool ${this.name}`);
 
                 const result = await this.execute(...args);
-                this.emitToolEvent(startTime, result, ...args);
+                await this.emitToolEvent(startTime, result, ...args).catch(() => {});
                 return result;
             } catch (error: unknown) {
                 logger.error(LogId.toolExecuteFailure, "tool", `Error executing ${this.name}: ${error as string}`);
                 const toolResult = await this.handleError(error, args[0] as ToolArgs<typeof this.argsShape>);
-                this.emitToolEvent(startTime, toolResult, ...args);
+                await this.emitToolEvent(startTime, toolResult, ...args).catch(() => {});
                 return toolResult;
             }
         };
 
-        server.tool(this.name, this.description, this.argsShape, this.annotations, callback);
+        server.mcpServer.tool(this.name, this.description, this.argsShape, this.annotations, callback);
 
         // This is very similar to RegisteredTool.update, but without the bugs around the name.
         // In the upstream update method, the name is captured in the closure and not updated when
         // the tool name changes. This means that you only get one name update before things end up
         // in a broken state.
+        // See https://github.com/modelcontextprotocol/typescript-sdk/issues/414 for more details.
         this.update = (updates: { name?: string; description?: string; inputSchema?: AnyZodObject }) => {
-            const tools = server["_registeredTools"] as { [toolName: string]: RegisteredTool };
+            const tools = server.mcpServer["_registeredTools"] as { [toolName: string]: RegisteredTool };
             const existingTool = tools[this.name];
 
             if (!existingTool) {
@@ -118,8 +121,10 @@ export abstract class ToolBase {
                 existingTool.inputSchema = updates.inputSchema;
             }
 
-            server.sendToolListChanged();
+            server.mcpServer.sendToolListChanged();
         };
+
+        return true;
     }
 
     protected update?: (updates: { name?: string; description?: string; inputSchema?: AnyZodObject }) => void;
@@ -179,11 +184,11 @@ export abstract class ToolBase {
      * @param result - Whether the command succeeded or failed
      * @param args - The arguments passed to the tool
      */
-    private emitToolEvent(
+    private async emitToolEvent(
         startTime: number,
         result: CallToolResult,
         ...args: Parameters<ToolCallback<typeof this.argsShape>>
-    ): void {
+    ): Promise<void> {
         if (!this.telemetry.isTelemetryEnabled()) {
             return;
         }
@@ -209,6 +214,6 @@ export abstract class ToolBase {
             event.properties.project_id = metadata.projectId;
         }
 
-        this.telemetry.emitEvents([event]);
+        await this.telemetry.emitEvents([event]);
     }
 }

package/src/transports/base.ts

@@ -0,0 +1,34 @@
+import { UserConfig } from "../common/config.js";
+import { packageInfo } from "../common/packageInfo.js";
+import { Server } from "../server.js";
+import { Session } from "../common/session.js";
+import { Telemetry } from "../telemetry/telemetry.js";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+
+export abstract class TransportRunnerBase {
+    protected setupServer(userConfig: UserConfig): Server {
+        const session = new Session({
+            apiBaseUrl: userConfig.apiBaseUrl,
+            apiClientId: userConfig.apiClientId,
+            apiClientSecret: userConfig.apiClientSecret,
+        });
+
+        const telemetry = Telemetry.create(session, userConfig);
+
+        const mcpServer = new McpServer({
+            name: packageInfo.mcpServerName,
+            version: packageInfo.version,
+        });
+
+        return new Server({
+            mcpServer,
+            session,
+            telemetry,
+            userConfig,
+        });
+    }
+
+    abstract start(): Promise<void>;
+
+    abstract close(): Promise<void>;
+}

package/src/{helpers/EJsonTransport.ts → transports/stdio.ts}

@@ -1,6 +1,10 @@
+import logger, { LogId } from "../common/logger.js";
+import { Server } from "../server.js";
+import { TransportRunnerBase } from "./base.js";
 import { JSONRPCMessage, JSONRPCMessageSchema } from "@modelcontextprotocol/sdk/types.js";
 import { EJSON } from "bson";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import { UserConfig } from "../common/config.js";
 
 // This is almost a copy of ReadBuffer from @modelcontextprotocol/sdk
 // but it uses EJSON.parse instead of JSON.parse to handle BSON types
@@ -39,9 +43,34 @@ export class EJsonReadBuffer {
 //
 // This function creates a StdioServerTransport and replaces the internal readBuffer with EJsonReadBuffer
 // that uses EJson.parse instead.
-export function createEJsonTransport(): StdioServerTransport {
+export function createStdioTransport(): StdioServerTransport {
     const server = new StdioServerTransport();
     server["_readBuffer"] = new EJsonReadBuffer();
 
     return server;
 }
+
+export class StdioRunner extends TransportRunnerBase {
+    private server: Server | undefined;
+
+    constructor(private userConfig: UserConfig) {
+        super();
+    }
+
+    async start() {
+        try {
+            this.server = this.setupServer(this.userConfig);
+
+            const transport = createStdioTransport();
+
+            await this.server.connect(transport);
+        } catch (error: unknown) {
+            logger.emergency(LogId.serverStartFailure, "server", `Fatal error running server: ${error as string}`);
+            process.exit(1);
+        }
+    }
+
+    async close(): Promise<void> {
+        await this.server?.close();
+    }
+}