chore: allow providing Voyage AI API key from config and tests to ensure unusable tools are not registered MCP-233 MCP-237 (#622)

Author: himanshusinghs

README.md

@@ -360,6 +360,7 @@ The MongoDB MCP Server can be configured using multiple methods, with the follow
 | `exportTimeoutMs`                      | `MDB_MCP_EXPORT_TIMEOUT_MS`                         | 300000                                                                      | Time in milliseconds after which an export is considered expired and eligible for cleanup.                                                                                                              |
 | `exportCleanupIntervalMs`              | `MDB_MCP_EXPORT_CLEANUP_INTERVAL_MS`                | 120000                                                                      | Time in milliseconds between export cleanup cycles that remove expired export files.                                                                                                                    |
 | `atlasTemporaryDatabaseUserLifetimeMs` | `MDB_MCP_ATLAS_TEMPORARY_DATABASE_USER_LIFETIME_MS` | 14400000                                                                    | Time in milliseconds that temporary database users created when connecting to MongoDB Atlas clusters will remain active before being automatically deleted.                                             |
+| `voyageApiKey`                         | `MDB_VOYAGE_API_KEY`                                | <not set>                                                                   | API key for communicating with Voyage AI. Used for generating embeddings for Vector search.                                                                                                             |
 
 #### Logger Options
 

src/common/config.ts

@@ -53,6 +53,7 @@ const OPTIONS = {
         "exportsPath",
         "exportTimeoutMs",
         "exportCleanupIntervalMs",
+        "voyageApiKey",
     ],
     boolean: [
         "apiDeprecationErrors",
@@ -181,6 +182,7 @@ export interface UserConfig extends CliOptions {
     maxDocumentsPerQuery: number;
     maxBytesPerQuery: number;
     atlasTemporaryDatabaseUserLifetimeMs: number;
+    voyageApiKey: string;
 }
 
 export const defaultUserConfig: UserConfig = {
@@ -210,6 +212,7 @@ export const defaultUserConfig: UserConfig = {
     maxDocumentsPerQuery: 100, // By default, we only fetch a maximum 100 documents per query / aggregation
     maxBytesPerQuery: 16 * 1024 * 1024, // By default, we only return ~16 mb of data per query / aggregation
     atlasTemporaryDatabaseUserLifetimeMs: 4 * 60 * 60 * 1000, // 4 hours
+    voyageApiKey: "",
 };
 
 export const config = setupUserConfig({

src/server.ts

@@ -18,7 +18,7 @@ import {
     UnsubscribeRequestSchema,
 } from "@modelcontextprotocol/sdk/types.js";
 import assert from "assert";
-import type { ToolBase } from "./tools/tool.js";
+import type { ToolBase, ToolConstructorParams } from "./tools/tool.js";
 import { validateConnectionString } from "./helpers/connectionOptions.js";
 import { packageInfo } from "./common/packageInfo.js";
 import { type ConnectionErrorHandler } from "./common/connectionErrorHandler.js";
@@ -31,6 +31,7 @@ export interface ServerOptions {
     telemetry: Telemetry;
     elicitation: Elicitation;
     connectionErrorHandler: ConnectionErrorHandler;
+    toolConstructors?: (new (params: ToolConstructorParams) => ToolBase)[];
 }
 
 export class Server {
@@ -39,6 +40,7 @@ export class Server {
     private readonly telemetry: Telemetry;
     public readonly userConfig: UserConfig;
     public readonly elicitation: Elicitation;
+    private readonly toolConstructors: (new (params: ToolConstructorParams) => ToolBase)[];
     public readonly tools: ToolBase[] = [];
     public readonly connectionErrorHandler: ConnectionErrorHandler;
 
@@ -51,14 +53,23 @@ export class Server {
     private readonly startTime: number;
     private readonly subscriptions = new Set<string>();
 
-    constructor({ session, mcpServer, userConfig, telemetry, connectionErrorHandler, elicitation }: ServerOptions) {
+    constructor({
+        session,
+        mcpServer,
+        userConfig,
+        telemetry,
+        connectionErrorHandler,
+        elicitation,
+        toolConstructors,
+    }: ServerOptions) {
         this.startTime = Date.now();
         this.session = session;
         this.telemetry = telemetry;
         this.mcpServer = mcpServer;
         this.userConfig = userConfig;
         this.elicitation = elicitation;
         this.connectionErrorHandler = connectionErrorHandler;
+        this.toolConstructors = toolConstructors ?? [...AtlasTools, ...MongoDbTools];
     }
 
     async connect(transport: Transport): Promise<void> {
@@ -206,7 +217,7 @@ export class Server {
     }
 
     private registerTools(): void {
-        for (const toolConstructor of [...AtlasTools, ...MongoDbTools]) {
+        for (const toolConstructor of this.toolConstructors) {
             const tool = new toolConstructor({
                 session: this.session,
                 config: this.userConfig,

tests/integration/tools/mongodb/mongodbTool.test.ts

@@ -1,9 +1,9 @@
-import { vi, it, describe, beforeEach, afterEach, type MockedFunction, afterAll, expect } from "vitest";
+import { vi, it, describe, beforeEach, afterEach, afterAll, expect } from "vitest";
 import { type CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { MongoDBToolBase } from "../../../../src/tools/mongodb/mongodbTool.js";
-import { type OperationType } from "../../../../src/tools/tool.js";
+import { type ToolBase, type ToolConstructorParams, type OperationType } from "../../../../src/tools/tool.js";
 import { defaultDriverOptions, type UserConfig } from "../../../../src/common/config.js";
 import { MCPConnectionManager } from "../../../../src/common/connectionManager.js";
 import { Session } from "../../../../src/common/session.js";
@@ -19,6 +19,7 @@ import { setupMongoDBIntegrationTest } from "./mongodbHelpers.js";
 import { ErrorCodes } from "../../../../src/common/errors.js";
 import { Keychain } from "../../../../src/common/keychain.js";
 import { Elicitation } from "../../../../src/elicitation.js";
+import { MongoDbTools } from "../../../../src/tools/mongodb/tools.js";
 
 const injectedErrorHandler: ConnectionErrorHandler = (error) => {
     switch (error.code) {
@@ -51,29 +52,45 @@ const injectedErrorHandler: ConnectionErrorHandler = (error) => {
     }
 };
 
-describe("MongoDBTool implementations", () => {
-    const mdbIntegration = setupMongoDBIntegrationTest({ enterprise: false }, []);
-    const executeStub: MockedFunction<() => Promise<CallToolResult>> = vi
-        .fn()
-        .mockResolvedValue({ content: [{ type: "text", text: "Something" }] });
-    class RandomTool extends MongoDBToolBase {
-        name = "Random";
-        operationType: OperationType = "read";
-        protected description = "This is a tool.";
-        protected argsShape = {};
-        public async execute(): Promise<CallToolResult> {
-            await this.ensureConnected();
-            return executeStub();
+class RandomTool extends MongoDBToolBase {
+    name = "Random";
+    operationType: OperationType = "read";
+    protected description = "This is a tool.";
+    protected argsShape = {};
+    public async execute(): Promise<CallToolResult> {
+        await this.ensureConnected();
+        return { content: [{ type: "text", text: "Something" }] };
+    }
+}
+
+class UnusableVoyageTool extends MongoDBToolBase {
+    name = "UnusableVoyageTool";
+    operationType: OperationType = "read";
+    protected description = "This is a Voyage tool.";
+    protected argsShape = {};
+
+    override verifyAllowed(): boolean {
+        if (this.config.voyageApiKey.trim()) {
+            return super.verifyAllowed();
         }
+        return false;
+    }
+    public async execute(): Promise<CallToolResult> {
+        await this.ensureConnected();
+        return { content: [{ type: "text", text: "Something" }] };
     }
+}
+
+describe("MongoDBTool implementations", () => {
+    const mdbIntegration = setupMongoDBIntegrationTest({ enterprise: false }, []);
 
-    let tool: RandomTool | undefined;
     let mcpClient: Client | undefined;
     let mcpServer: Server | undefined;
     let deviceId: DeviceId | undefined;
 
     async function cleanupAndStartServer(
         config: Partial<UserConfig> | undefined = {},
+        toolConstructors: (new (params: ToolConstructorParams) => ToolBase)[] = [...MongoDbTools, RandomTool],
         errorHandler: ConnectionErrorHandler | undefined = connectionErrorHandler
     ): Promise<void> {
         await cleanup();
@@ -126,16 +143,9 @@ describe("MongoDBTool implementations", () => {
             mcpServer: internalMcpServer,
             connectionErrorHandler: errorHandler,
             elicitation,
+            toolConstructors,
         });
 
-        tool = new RandomTool({
-            session,
-            config: userConfig,
-            telemetry,
-            elicitation,
-        });
-        tool.register(mcpServer);
-
         await mcpServer.connect(serverTransport);
         await mcpClient.connect(clientTransport);
     }
@@ -150,8 +160,6 @@ describe("MongoDBTool implementations", () => {
 
         deviceId?.close();
         deviceId = undefined;
-
-        tool = undefined;
     }
 
     beforeEach(async () => {
@@ -232,7 +240,7 @@ describe("MongoDBTool implementations", () => {
 
     describe("when MCP is using injected connection error handler", () => {
         beforeEach(async () => {
-            await cleanupAndStartServer(defaultTestConfig, injectedErrorHandler);
+            await cleanupAndStartServer(defaultTestConfig, [...MongoDbTools, RandomTool], injectedErrorHandler);
         });
 
         describe("and comes across a MongoDB Error - NotConnectedToMongoDB", () => {
@@ -256,7 +264,11 @@ describe("MongoDBTool implementations", () => {
         describe("and comes across a MongoDB Error - MisconfiguredConnectionString", () => {
             it("should handle the error", async () => {
                 // This is a misconfigured connection string
-                await cleanupAndStartServer({ connectionString: "mongodb://localhost:1234" }, injectedErrorHandler);
+                await cleanupAndStartServer(
+                    { connectionString: "mongodb://localhost:1234" },
+                    [...MongoDbTools, RandomTool],
+                    injectedErrorHandler
+                );
                 const toolResponse = await mcpClient?.callTool({
                     name: "Random",
                     arguments: {},
@@ -278,6 +290,7 @@ describe("MongoDBTool implementations", () => {
                 // This is a misconfigured connection string
                 await cleanupAndStartServer(
                     { connectionString: mdbIntegration.connectionString(), indexCheck: true },
+                    [...MongoDbTools, RandomTool],
                     injectedErrorHandler
                 );
                 const toolResponse = await mcpClient?.callTool({
@@ -299,4 +312,17 @@ describe("MongoDBTool implementations", () => {
             });
         });
     });
+
+    describe("when a tool is not usable", () => {
+        it("should not even be registered", async () => {
+            await cleanupAndStartServer(
+                { connectionString: mdbIntegration.connectionString(), indexCheck: true },
+                [RandomTool, UnusableVoyageTool],
+                injectedErrorHandler
+            );
+            const tools = await mcpClient?.listTools({});
+            expect(tools?.tools).toHaveLength(1);
+            expect(tools?.tools.find((tool) => tool.name === "UnusableVoyageTool")).toBeUndefined();
+        });
+    });
 });