chore: identify if autoembeddings are supported or not

Author: himanshusinghs

src/common/connectionManager.ts

@@ -1,5 +1,5 @@
 import { EventEmitter } from "events";
-import type { MongoClientOptions } from "mongodb";
+import type { MongoClientOptions, MongoServerError } from "mongodb";
 import { ConnectionString } from "mongodb-connection-string-url";
 import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { generateConnectionInfoFromCliArgs, type ConnectionInfo } from "@mongosh/arg-parser";
@@ -33,6 +33,8 @@ export interface ConnectionState {
 }
 
 const MCP_TEST_DATABASE = "#mongodb-mcp";
+const MCP_TEST_VECTOR_SEARCH_INDEX_FIELD = "__mongodb-mcp-field";
+const MCP_TEST_VECTOR_SEARCH_INDEX_NAME = "mongodb-mcp-vector-search-index";
 
 export const defaultDriverOptions: ConnectionInfo["driverOptions"] = {
     readConcern: {
@@ -56,23 +58,72 @@ export class ConnectionStateConnected implements ConnectionState {
         public connectedAtlasCluster?: AtlasClusterConnectionInfo
     ) {}
 
-    private _isSearchSupported?: boolean;
+    private connectionMetadata?: { searchSupported: boolean; autoEmbeddingIndexSupported: boolean };
 
-    public async isSearchSupported(): Promise<boolean> {
-        if (this._isSearchSupported === undefined) {
-            try {
-                // If a cluster supports search indexes, the call below will succeed
-                // with a cursor otherwise will throw an Error.
-                // the Search Index Management Service might not be ready yet, but
-                // we assume that the agent can retry in that situation.
-                await this.serviceProvider.getSearchIndexes(MCP_TEST_DATABASE, "test");
-                this._isSearchSupported = true;
-            } catch {
-                this._isSearchSupported = false;
+    private async populateConnectionMetadata(): Promise<void> {
+        try {
+            await this.serviceProvider.createCollection(MCP_TEST_DATABASE, "test");
+            // If a cluster supports search indexes, the call below will succeed
+            // with a cursor otherwise will throw an Error. The Search Index
+            // Management Service might not be ready yet, but we assume that the
+            // agent can retry in that situation.
+            await this.serviceProvider.createSearchIndexes(MCP_TEST_DATABASE, "test", [
+                {
+                    name: MCP_TEST_VECTOR_SEARCH_INDEX_NAME,
+                    type: "vectorSearch",
+                    definition: {
+                        fields: [
+                            {
+                                // TODO: Before public preview this needs to be
+                                // changed to either "autoEmbedText" or "autoEmbed"
+                                // depending on which syntax is finalized.
+                                type: "text",
+                                path: MCP_TEST_VECTOR_SEARCH_INDEX_FIELD,
+                                model: "voyage-3-large",
+                            },
+                        ],
+                    },
+                },
+            ]);
+            this.connectionMetadata = {
+                searchSupported: true,
+                autoEmbeddingIndexSupported: true,
+            };
+        } catch (error) {
+            if ((error as MongoServerError).codeName === "SearchNotEnabled") {
+                this.connectionMetadata = {
+                    searchSupported: false,
+                    autoEmbeddingIndexSupported: false,
+                };
             }
+            // If the error if because we tried creating an index with autoEmbed
+            // field definition then we can safely assume that search is
+            // supported but auto-embeddings are not.
+            else if ((error as Error).message.includes('"userCommand.indexes[0].fields[0].type" must be one of')) {
+                this.connectionMetadata = {
+                    searchSupported: true,
+                    autoEmbeddingIndexSupported: false,
+                };
+            }
+        } finally {
+            await this.serviceProvider.dropDatabase(MCP_TEST_DATABASE);
+        }
+    }
+
+    public async isSearchSupported(): Promise<boolean> {
+        if (this.connectionMetadata === undefined) {
+            await this.populateConnectionMetadata();
+        }
+
+        return this.connectionMetadata?.searchSupported ?? false;
+    }
+
+    public async isAutoEmbeddingIndexSupported(): Promise<boolean> {
+        if (this.connectionMetadata === undefined) {
+            await this.populateConnectionMetadata();
         }
 
-        return this._isSearchSupported;
+        return this.connectionMetadata?.autoEmbeddingIndexSupported ?? false;
     }
 }
 

src/common/errors.ts

@@ -8,6 +8,7 @@ export enum ErrorCodes {
     AtlasVectorSearchIndexNotFound = 1_000_006,
     AtlasVectorSearchInvalidQuery = 1_000_007,
     Unexpected = 1_000_008,
+    AutoEmbeddingIndexNotSupported = 1_000_009,
 }
 
 export class MongoDBError<ErrorCode extends ErrorCodes = ErrorCodes> extends Error {

src/common/session.ts

@@ -179,6 +179,25 @@ export class Session extends EventEmitter<SessionEvents> {
         }
     }
 
+    async isAutoEmbeddingIndexSupported(): Promise<boolean> {
+        const state = this.connectionManager.currentConnectionState;
+        if (state.tag === "connected") {
+            return await state.isAutoEmbeddingIndexSupported();
+        }
+
+        return false;
+    }
+
+    async assertAutoEmbeddingIndexSupported(): Promise<void> {
+        const isSearchSupported = await this.isSearchSupported();
+        if (!isSearchSupported) {
+            throw new MongoDBError(
+                ErrorCodes.AutoEmbeddingIndexNotSupported,
+                "The connected search management service does not support creating auto-embedding indexes."
+            );
+        }
+    }
+
     get serviceProvider(): NodeDriverServiceProvider {
         if (this.isConnectedToMongoDB) {
             const state = this.connectionManager.currentConnectionState as ConnectionStateConnected;

src/tools/mongodb/create/createIndex.ts

@@ -183,7 +183,7 @@ export class CreateIndexTool extends MongoDBToolBase {
                 break;
             case "vectorSearch":
                 {
-                    await this.ensureSearchIsSupported();
+                    await this.session.assertSearchSupported();
                     indexes = await provider.createSearchIndexes(database, collection, [
                         {
                             name,
@@ -204,7 +204,7 @@ export class CreateIndexTool extends MongoDBToolBase {
                 break;
             case "search":
                 {
-                    await this.ensureSearchIsSupported();
+                    await this.session.assertSearchSupported();
                     indexes = await provider.createSearchIndexes(database, collection, [
                         {
                             name,

src/tools/mongodb/delete/dropIndex.ts

@@ -58,7 +58,7 @@ export class DropIndexTool extends MongoDBToolBase {
         provider: NodeDriverServiceProvider,
         { database, collection, indexName }: ToolArgs<typeof this.argsShape>
     ): Promise<CallToolResult> {
-        await this.ensureSearchIsSupported();
+        await this.session.assertSearchSupported();
         const indexes = await provider.getSearchIndexes(database, collection, indexName);
         if (indexes.length === 0) {
             return {

src/tools/mongodb/mongodbTool.ts

@@ -47,10 +47,6 @@ export abstract class MongoDBToolBase extends ToolBase {
         return this.session.serviceProvider;
     }
 
-    protected ensureSearchIsSupported(): Promise<void> {
-        return this.session.assertSearchSupported();
-    }
-
     public register(server: Server): boolean {
         this.server = server;
         return super.register(server);