chore: check that a vector search index exists with indexCheck (#693)

Co-authored-by: Copilot <[email protected]>

Author: kmruiz

src/common/search/vectorSearchEmbeddingsManager.ts

@@ -48,6 +48,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,

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({
@@ -279,6 +295,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,

tests/integration/tools/mongodb/read/aggregate.test.ts

@@ -405,6 +405,47 @@ describeWithMongoDB(
             await integration.mongoClient().db(integration.randomDbName()).collection("databases").drop();
         });
 
+        it("should throw an exception when using an index that does not exist", async () => {
+            await waitUntilSearchIsReady(integration.mongoClient());
+
+            const collection = integration.mongoClient().db(integration.randomDbName()).collection("databases");
+
+            await collection.insertOne({ name: "mongodb", description_embedding: [1, 2, 3, 4] });
+            await integration.connectMcpClient();
+            const response = await integration.mcpClient().callTool({
+                name: "aggregate",
+                arguments: {
+                    database: integration.randomDbName(),
+                    collection: "databases",
+                    pipeline: [
+                        {
+                            $vectorSearch: {
+                                index: "non_existing",
+                                path: "description_embedding",
+                                queryVector: "example",
+                                numCandidates: 10,
+                                limit: 10,
+                                embeddingParameters: {
+                                    model: "voyage-3-large",
+                                    outputDimension: 256,
+                                },
+                            },
+                        },
+                        {
+                            $project: {
+                                description_embedding: 0,
+                            },
+                        },
+                    ],
+                },
+            });
+
+            const responseContent = getResponseContent(response);
+            expect(responseContent).toContain(
+                `Error running aggregate: Could not find an index with name "non_existing" in namespace "${integration.randomDbName()}.databases".`
+            );
+        });
+
         for (const [dataType, embedding] of Object.entries(DOCUMENT_EMBEDDINGS)) {
             for (const similarity of ["euclidean", "cosine", "dotProduct"]) {
                 describe.skipIf(!process.env.TEST_MDB_MCP_VOYAGE_API_KEY)(
@@ -417,6 +458,7 @@ describeWithMongoDB(
                                 .mongoClient()
                                 .db(integration.randomDbName())
                                 .collection("databases");
+
                             await collection.insertOne({ name: "mongodb", description_embedding: embedding });
 
                             await createVectorSearchIndexAndWait(
@@ -686,6 +728,7 @@ describeWithMongoDB(
             previewFeatures: ["vectorSearch"],
             maxDocumentsPerQuery: -1,
             maxBytesPerQuery: -1,
+            indexCheck: true,
         }),
         downloadOptions: { search: true },
     }