Merge branch 'main' of github.com:mongodb-js/mongodb-mcp-server into gagik/accuracy-embeddings

gagik · gagik · commit 2bb10e879f93 · 2025-10-29T15:42:13.000+01:00
diff --git a/README.md b/README.md
diff --git a/src/common/config.ts b/src/common/config.ts
@@ -10,6 +10,9 @@ import type { Similarity } from "./search/vectorSearchEmbeddingsManager.js";
 import { z } from "zod";
 const levenshtein = levenshteinModule.default;
 
+const previewFeatures = z.enum(["vectorSearch"]);
+export type PreviewFeature = z.infer<typeof previewFeatures>;
+
 // From: https://github.com/mongodb-js/mongosh/blob/main/packages/cli-repl/src/arg-parser.ts
 export const OPTIONS = {
     number: ["maxDocumentsPerQuery", "maxBytesPerQuery"],
@@ -81,7 +84,7 @@ export const OPTIONS = {
         "tlsFIPSMode",
         "version",
     ],
-    array: ["disabledTools", "loggers", "confirmationRequiredTools"],
+    array: ["disabledTools", "loggers", "confirmationRequiredTools", "previewFeatures"],
     alias: {
         h: "help",
         p: "password",
@@ -119,7 +122,7 @@ export const ALL_CONFIG_KEYS = new Set(
         .concat(Object.keys(OPTIONS.alias))
 );
 
-export function validateConfigKey(key: string): { valid: boolean; suggestion?: string } {
+function validateConfigKey(key: string): { valid: boolean; suggestion?: string } {
     if (ALL_CONFIG_KEYS.has(key)) {
         return { valid: true };
     }
@@ -282,6 +285,7 @@ export const UserConfigSchema = z.object({
         .optional()
         .default("euclidean")
         .describe("Default similarity function for vector search: 'euclidean', 'cosine', or 'dotProduct'."),
+    previewFeatures: z.array(previewFeatures).default([]).describe("An array of preview features that are enabled."),
 });
 
 export type UserConfig = z.infer<typeof UserConfigSchema> & CliOptions;
@@ -318,6 +322,7 @@ export const defaultUserConfig: UserConfig = {
     disableEmbeddingsValidation: false,
     vectorSearchDimensions: 1024,
     vectorSearchSimilarityFunction: "euclidean",
+    previewFeatures: [],
 };
 
 export const config = setupUserConfig({
@@ -554,13 +559,13 @@ export function setupUserConfig({
 }: {
     cli: string[];
     env: Record<string, unknown>;
-    defaults: Partial<UserConfig>;
+    defaults: UserConfig;
 }): UserConfig {
-    const userConfig: UserConfig = {
+    const userConfig = {
         ...defaults,
         ...parseEnvConfig(env),
         ...parseCliConfig(cli),
-    } as UserConfig;
+    } satisfies UserConfig;
 
     userConfig.disabledTools = commaSeparatedToArray(userConfig.disabledTools);
     userConfig.loggers = commaSeparatedToArray(userConfig.loggers);
diff --git a/src/common/search/embeddingsProvider.ts b/src/common/search/embeddingsProvider.ts
@@ -72,8 +72,8 @@ class VoyageEmbeddingsProvider implements EmbeddingsProvider<VoyageModels, Voyag
         this.voyage = createVoyage({ apiKey: voyageApiKey, fetch: customFetch });
     }
 
-    static isConfiguredIn({ voyageApiKey }: UserConfig): boolean {
-        return !!voyageApiKey;
+    static isConfiguredIn({ voyageApiKey, previewFeatures }: UserConfig): boolean {
+        return previewFeatures.includes("vectorSearch") && !!voyageApiKey;
     }
 
     async embed<Model extends VoyageModels>(
diff --git a/src/common/search/vectorSearchEmbeddingsManager.ts b/src/common/search/vectorSearchEmbeddingsManager.ts
@@ -49,6 +49,25 @@ export class VectorSearchEmbeddingsManager {
         this.embeddings.delete(embeddingDefKey);
     }
 
+    async indexExists({
+        database,
+        collection,
+        indexName,
+    }: {
+        database: string;
+        collection: string;
+        indexName: string;
+    }): Promise<boolean> {
+        const provider = await this.atlasSearchEnabledProvider();
+        if (!provider) {
+            return false;
+        }
+
+        const searchIndexesWithName = await provider.getSearchIndexes(database, collection, indexName);
+
+        return searchIndexesWithName.length >= 1;
+    }
+
     async embeddingsForNamespace({
         database,
         collection,
diff --git a/src/tools/mongodb/create/createIndex.ts b/src/tools/mongodb/create/createIndex.ts
@@ -1,7 +1,7 @@
 import { z } from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
-import { type ToolArgs, type OperationType, FeatureFlags } from "../../tool.js";
+import { type ToolArgs, type OperationType } from "../../tool.js";
 import type { IndexDirection } from "mongodb";
 import { quantizationEnum, similarityEnum } from "../../../common/search/vectorSearchEmbeddingsManager.js";
 
@@ -74,7 +74,7 @@ export class CreateIndexTool extends MongoDBToolBase {
                         type: z.literal("classic"),
                         keys: z.object({}).catchall(z.custom<IndexDirection>()).describe("The index definition"),
                     }),
-                    ...(this.isFeatureFlagEnabled(FeatureFlags.VectorSearch) ? [this.vectorSearchIndexDefinition] : []),
+                    ...(this.isFeatureEnabled("vectorSearch") ? [this.vectorSearchIndexDefinition] : []),
                 ])
             )
             .describe(
diff --git a/src/tools/mongodb/delete/dropIndex.ts b/src/tools/mongodb/delete/dropIndex.ts
@@ -2,15 +2,15 @@ import z from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
-import { type ToolArgs, type OperationType, formatUntrustedData, FeatureFlags } from "../../tool.js";
+import { type ToolArgs, type OperationType, formatUntrustedData } from "../../tool.js";
 
 export class DropIndexTool extends MongoDBToolBase {
     public name = "drop-index";
     protected description = "Drop an index for the provided database and collection.";
     protected argsShape = {
         ...DbOperationArgs,
         indexName: z.string().nonempty().describe("The name of the index to be dropped."),
-        type: this.isFeatureFlagEnabled(FeatureFlags.VectorSearch)
+        type: this.isFeatureEnabled("vectorSearch")
             ? z
                   .enum(["classic", "search"])
                   .describe(
diff --git a/src/tools/mongodb/metadata/collectionIndexes.ts b/src/tools/mongodb/metadata/collectionIndexes.ts
@@ -1,7 +1,7 @@
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import type { ToolArgs, OperationType } from "../../tool.js";
-import { FeatureFlags, formatUntrustedData } from "../../tool.js";
+import { formatUntrustedData } from "../../tool.js";
 
 type SearchIndexStatus = {
     name: string;
@@ -31,7 +31,7 @@ export class CollectionIndexesTool extends MongoDBToolBase {
         }));
 
         const searchIndexDefinitions: SearchIndexStatus[] = [];
-        if (this.isFeatureFlagEnabled(FeatureFlags.VectorSearch) && (await this.session.isSearchSupported())) {
+        if (this.isFeatureEnabled("vectorSearch") && (await this.session.isSearchSupported())) {
             const searchIndexes = await provider.getSearchIndexes(database, collection);
             searchIndexDefinitions.push(...this.extractSearchIndexDetails(searchIndexes));
         }
diff --git a/src/tools/mongodb/read/aggregate.ts b/src/tools/mongodb/read/aggregate.ts
@@ -100,11 +100,27 @@ export class AggregateTool extends MongoDBToolBase {
 
             // 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 [usesVectorSearchIndex, indexName] = await this.isVectorSearchIndexUsed({
+                    database,
+                    collection,
+                    pipeline,
                 });
+                switch (usesVectorSearchIndex) {
+                    case "not-vector-search-query":
+                        await checkIndexUsage(provider, database, collection, "aggregate", async () => {
+                            return provider
+                                .aggregate(database, collection, pipeline, {}, { writeConcern: undefined })
+                                .explain("queryPlanner");
+                        });
+                        break;
+                    case "non-existent-index":
+                        throw new MongoDBError(
+                            ErrorCodes.AtlasVectorSearchIndexNotFound,
+                            `Could not find an index with name "${indexName}" in namespace "${database}.${collection}".`
+                        );
+                    case "valid-index":
+                    // nothing to do, everything is correct so ready to run the query
+                }
             }
 
             pipeline = await this.replaceRawValuesWithEmbeddingsIfNecessary({
@@ -294,6 +310,41 @@ export class AggregateTool extends MongoDBToolBase {
         return pipeline;
     }
 
+    private async isVectorSearchIndexUsed({
+        database,
+        collection,
+        pipeline,
+    }: {
+        database: string;
+        collection: string;
+        pipeline: Document[];
+    }): Promise<["valid-index" | "non-existent-index" | "not-vector-search-query", string?]> {
+        // check if the pipeline contains a $vectorSearch stage
+        let usesVectorSearch = false;
+        let indexName: string = "default";
+
+        for (const stage of pipeline) {
+            if ("$vectorSearch" in stage) {
+                const { $vectorSearch: vectorSearchStage } = stage as z.infer<typeof VectorSearchStage>;
+                usesVectorSearch = true;
+                indexName = vectorSearchStage.index;
+                break;
+            }
+        }
+
+        if (!usesVectorSearch) {
+            return ["not-vector-search-query"];
+        }
+
+        const indexExists = await this.session.vectorSearchEmbeddingsManager.indexExists({
+            database,
+            collection,
+            indexName,
+        });
+
+        return [indexExists ? "valid-index" : "non-existent-index", indexName];
+    }
+
     private generateMessage({
         aggResultsCount,
         documents,
diff --git a/src/tools/tool.ts b/src/tools/tool.ts
@@ -6,7 +6,7 @@ import type { Session } from "../common/session.js";
 import { LogId } from "../common/logger.js";
 import type { Telemetry } from "../telemetry/telemetry.js";
 import { type ToolEvent } from "../telemetry/types.js";
-import type { UserConfig } from "../common/config.js";
+import type { PreviewFeature, UserConfig } from "../common/config.js";
 import type { Server } from "../server.js";
 import type { Elicitation } from "../elicitation.js";
 
@@ -15,10 +15,6 @@ export type ToolCallbackArgs<Args extends ZodRawShape> = Parameters<ToolCallback
 
 export type ToolExecutionContext<Args extends ZodRawShape = ZodRawShape> = Parameters<ToolCallback<Args>>[1];
 
-export const enum FeatureFlags {
-    VectorSearch = "vectorSearch",
-}
-
 /**
  * The type of operation the tool performs. This is used when evaluating if a tool is allowed to run based on
  * the config's `disabledTools` and `readOnly` settings.
@@ -325,14 +321,8 @@ export abstract class ToolBase {
         this.telemetry.emitEvents([event]);
     }
 
-    // TODO: Move this to a separate file
-    protected isFeatureFlagEnabled(flag: FeatureFlags): boolean {
-        switch (flag) {
-            case FeatureFlags.VectorSearch:
-                return this.config.voyageApiKey !== "";
-            default:
-                return false;
-        }
+    protected isFeatureEnabled(feature: PreviewFeature): boolean {
+        return this.config.previewFeatures.includes(feature);
     }
 }
 
diff --git a/tests/accuracy/createIndex.test.ts b/tests/accuracy/createIndex.test.ts
@@ -137,7 +137,7 @@ describeAccuracyTests(
         },
     ],
     {
-        userConfig: { voyageApiKey: "valid-key" },
+        userConfig: { previewFeatures: "vectorSearch" },
         clusterConfig: {
             search: true,
         },
diff --git a/tests/accuracy/createIndex.vectorSearchDisabled.test.ts b/tests/accuracy/createIndex.vectorSearchDisabled.test.ts
@@ -52,6 +52,6 @@ describeAccuracyTests(
         },
     ],
     {
-        userConfig: { voyageApiKey: "" },
+        userConfig: { previewFeatures: "" },
     }
 );
diff --git a/tests/accuracy/dropIndex.test.ts b/tests/accuracy/dropIndex.test.ts
@@ -127,7 +127,7 @@ describeAccuracyTests(
     ],
     {
         userConfig: {
-            voyageApiKey: "voyage-api-key",
+            previewFeatures: "vectorSearch",
         },
         clusterConfig: { search: true },
     }
diff --git a/tests/accuracy/dropIndex.vectorSearchDisabled.test.ts b/tests/accuracy/dropIndex.vectorSearchDisabled.test.ts
@@ -90,7 +90,7 @@ describeAccuracyTests(
     ],
     {
         userConfig: {
-            voyageApiKey: "",
+            previewFeatures: "",
         },
     }
 );
diff --git a/tests/accuracy/sdk/accuracyTestingClient.ts b/tests/accuracy/sdk/accuracyTestingClient.ts
@@ -91,11 +91,15 @@ export class AccuracyTestingClient {
             return [`--${key}`, value];
         });
 
-        const args = [MCP_SERVER_CLI_SCRIPT, "--connectionString", mdbConnectionString, ...additionalArgs];
+        const args = [MCP_SERVER_CLI_SCRIPT, mdbConnectionString, ...additionalArgs];
 
         const clientTransport = new StdioClientTransport({
             command: process.execPath,
             args,
+            env: {
+                ...process.env,
+                DO_NOT_TRACK: "1",
+            },
         });
 
         const client = await createMCPClient({
diff --git a/tests/integration/elicitation.test.ts b/tests/integration/elicitation.test.ts
@@ -9,7 +9,6 @@ import { describeWithMongoDB } from "./tools/mongodb/mongodbHelpers.js";
 function createTestConfig(config: Partial<UserConfig> = {}): UserConfig {
     return {
         ...defaultTestConfig,
-        telemetry: "disabled",
         // Add fake API credentials so Atlas tools get registered
         apiClientId: "test-client-id",
         apiClientSecret: "test-client-secret",
diff --git a/tests/integration/tools/mongodb/create/createIndex.test.ts b/tests/integration/tools/mongodb/create/createIndex.test.ts
@@ -14,7 +14,7 @@ import { afterEach, beforeEach, describe, expect, it } from "vitest";
 
 describeWithMongoDB("createIndex tool when search is not enabled", (integration) => {
     it("doesn't allow creating vector search indexes", async () => {
-        expect(integration.mcpServer().userConfig.voyageApiKey).toEqual("");
+        expect(integration.mcpServer().userConfig.previewFeatures).to.not.include("vectorSearch");
 
         const { tools } = await integration.mcpClient().listTools();
         const createIndexTool = tools.find((tool) => tool.name === "create-index");
@@ -38,7 +38,7 @@ describeWithMongoDB(
     "createIndex tool when search is enabled",
     (integration) => {
         it("allows creating vector search indexes", async () => {
-            expect(integration.mcpServer().userConfig.voyageApiKey).not.toEqual("");
+            expect(integration.mcpServer().userConfig.previewFeatures).includes("vectorSearch");
 
             const { tools } = await integration.mcpClient().listTools();
             const createIndexTool = tools.find((tool) => tool.name === "create-index");
@@ -84,7 +84,7 @@ describeWithMongoDB(
         getUserConfig: () => {
             return {
                 ...defaultTestConfig,
-                voyageApiKey: "valid_key",
+                previewFeatures: ["vectorSearch"],
             };
         },
     }
@@ -392,7 +392,7 @@ describeWithMongoDB(
         getUserConfig: () => {
             return {
                 ...defaultTestConfig,
-                voyageApiKey: "valid_key",
+                previewFeatures: ["vectorSearch"],
             };
         },
     }
@@ -613,7 +613,7 @@ describeWithMongoDB(
     {
         getUserConfig: () => ({
             ...defaultTestConfig,
-            voyageApiKey: "valid_key",
+            previewFeatures: ["vectorSearch"],
         }),
         downloadOptions: {
             search: true,
diff --git a/tests/integration/tools/mongodb/delete/dropIndex.test.ts b/tests/integration/tools/mongodb/delete/dropIndex.test.ts
diff --git a/tests/integration/tools/mongodb/metadata/collectionIndexes.test.ts b/tests/integration/tools/mongodb/metadata/collectionIndexes.test.ts
diff --git a/tests/integration/tools/mongodb/read/aggregate.test.ts b/tests/integration/tools/mongodb/read/aggregate.test.ts
diff --git a/tests/unit/toolBase.test.ts b/tests/unit/toolBase.test.ts

Original file line number	Diff line number	Diff line change
`@@ -72,8 +72,8 @@ class VoyageEmbeddingsProvider implements EmbeddingsProvider<VoyageModels, Voyag`
`72`	`72`	`this.voyage = createVoyage({ apiKey: voyageApiKey, fetch: customFetch });`
`73`	`73`	`}`
`74`	`74`
`75`		`- static isConfiguredIn({ voyageApiKey }: UserConfig): boolean {`
`76`		`- return !!voyageApiKey;`
	`75`	`+ static isConfiguredIn({ voyageApiKey, previewFeatures }: UserConfig): boolean {`
	`76`	`+ return previewFeatures.includes("vectorSearch") && !!voyageApiKey;`
`77`	`77`	`}`
`78`	`78`
`79`	`79`	`async embed<Model extends VoyageModels>(`
Original file line number	Diff line number	Diff line change
`@@ -137,7 +137,7 @@ describeAccuracyTests(`
`137`	`137`	`},`
`138`	`138`	`],`
`139`	`139`	`{`
`140`		`- userConfig: { voyageApiKey: "valid-key" },`
	`140`	`+ userConfig: { previewFeatures: "vectorSearch" },`
`141`	`141`	`clusterConfig: {`
`142`	`142`	`search: true,`
`143`	`143`	`},`