chore: Merge reliably search permission detection

Author: kmruiz

src/common/connectionManager.ts

@@ -25,39 +25,106 @@ export interface ConnectionSettings {
 type ConnectionTag = "connected" | "connecting" | "disconnected" | "errored";
 type OIDCConnectionAuthType = "oidc-auth-flow" | "oidc-device-flow";
 export type ConnectionStringAuthType = "scram" | "ldap" | "kerberos" | OIDCConnectionAuthType | "x.509";
+export type SearchAvailability = false | "not-available-yet" | "available";
 
 export interface ConnectionState {
     tag: ConnectionTag;
     connectionStringAuthType?: ConnectionStringAuthType;
     connectedAtlasCluster?: AtlasClusterConnectionInfo;
 }
 
+const MCP_TEST_DATABASE = "#mongodb-mcp";
+const SEARCH_AVAILABILITY_CHECK_TIMEOUT_MS = 500;
 export class ConnectionStateConnected implements ConnectionState {
     public tag = "connected" as const;
 
     constructor(
         public serviceProvider: NodeDriverServiceProvider,
         public connectionStringAuthType?: ConnectionStringAuthType,
         public connectedAtlasCluster?: AtlasClusterConnectionInfo
-    ) {}
+    ) {
+        this.#isSearchAvailable = false;
+    }
+
+    #isSearchSupported?: boolean;
+    #isSearchAvailable: boolean;
 
-    private _isSearchSupported?: boolean;
+    public async getSearchAvailability(): Promise<SearchAvailability> {
+        if ((await this.isSearchSupported()) === true) {
+            if ((await this.isSearchAvailable()) === true) {
+                return "available";
+            }
+
+            return "not-available-yet";
+        }
 
-    public async isSearchSupported(): Promise<boolean> {
-        if (this._isSearchSupported === undefined) {
+        return false;
+    }
+
+    private async isSearchSupported(): Promise<boolean> {
+        if (this.#isSearchSupported === undefined) {
             try {
-                const dummyDatabase = "test";
-                const dummyCollection = "test";
                 // If a cluster supports search indexes, the call below will succeed
                 // with a cursor otherwise will throw an Error
-                await this.serviceProvider.getSearchIndexes(dummyDatabase, dummyCollection);
-                this._isSearchSupported = true;
+                await this.serviceProvider.getSearchIndexes(MCP_TEST_DATABASE, "test");
+                this.#isSearchSupported = true;
             } catch {
-                this._isSearchSupported = false;
+                this.#isSearchSupported = false;
+            }
+        }
+
+        return this.#isSearchSupported;
+    }
+
+    private async isSearchAvailable(): Promise<boolean> {
+        if (this.#isSearchAvailable === true) {
+            return true;
+        }
+
+        const timeoutPromise = new Promise<boolean>((_resolve, reject) =>
+            setTimeout(
+                () =>
+                    reject(
+                        new MongoDBError(
+                            ErrorCodes.AtlasSearchNotAvailable,
+                            "Atlas Search is supported in your environment but is not available yet. Retry again later."
+                        )
+                    ),
+                SEARCH_AVAILABILITY_CHECK_TIMEOUT_MS
+            )
+        );
+
+        const checkPromise = new Promise<boolean>((resolve) => {
+            void this.doCheckSearchIndexIsAvailable(resolve);
+        });
+
+        return await Promise.race([checkPromise, timeoutPromise]);
+    }
+
+    private async doCheckSearchIndexIsAvailable(resolve: (result: boolean) => void): Promise<void> {
+        for (let i = 0; i < 100; i++) {
+            try {
+                try {
+                    await this.serviceProvider.insertOne(MCP_TEST_DATABASE, "test", { search: "search is available" });
+                } catch (err) {
+                    // if inserting one document fails, it means we are in readOnly mode. We can't verify reliably if
+                    // Search is available, so assume it is.
+                    void err;
+                    resolve(true);
+                    return;
+                }
+                await this.serviceProvider.createSearchIndexes(MCP_TEST_DATABASE, "test", [
+                    { definition: { mappings: { dynamic: true } } },
+                ]);
+                await this.serviceProvider.dropDatabase(MCP_TEST_DATABASE);
+                resolve(true);
+                return;
+            } catch (err) {
+                void err;
             }
         }
 
-        return this._isSearchSupported;
+        resolve(false);
     }
 }
 

src/common/errors.ts

@@ -3,7 +3,8 @@ export enum ErrorCodes {
     MisconfiguredConnectionString = 1_000_001,
     ForbiddenCollscan = 1_000_002,
     ForbiddenWriteOperation = 1_000_003,
-    AtlasSearchNotAvailable = 1_000_004,
+    AtlasSearchNotSupported = 1_000_004,
+    AtlasSearchNotAvailable = 1_000_005,
 }
 
 export class MongoDBError<ErrorCode extends ErrorCodes = ErrorCodes> extends Error {

src/common/search/vectorSearchEmbeddings.ts

@@ -1,6 +1,7 @@
 import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { BSON, type Document } from "bson";
 import type { UserConfig } from "../config.js";
+import type { ConnectionManager } from "../connectionManager.js";
 
 export type VectorFieldIndexDefinition = {
     type: "vector";
@@ -14,8 +15,8 @@ export type EmbeddingNamespace = `${string}.${string}`;
 export class VectorSearchEmbeddings {
     constructor(
         private readonly config: UserConfig,
-        private readonly embeddings: Map<EmbeddingNamespace, VectorFieldIndexDefinition[]> = new Map(),
-        private readonly atlasSearchStatus: Map<string, boolean> = new Map()
+        private readonly connectionManager: ConnectionManager,
+        private readonly embeddings: Map<EmbeddingNamespace, VectorFieldIndexDefinition[]> = new Map()
     ) {}
 
     cleanupEmbeddingsForNamespace({ database, collection }: { database: string; collection: string }): void {
@@ -26,13 +27,12 @@ export class VectorSearchEmbeddings {
     async embeddingsForNamespace({
         database,
         collection,
-        provider,
     }: {
         database: string;
         collection: string;
-        provider: NodeDriverServiceProvider;
     }): Promise<VectorFieldIndexDefinition[]> {
-        if (!(await this.isAtlasSearchAvailable(provider))) {
+        const provider = await this.assertAtlasSearchIsAvailable();
+        if (!provider) {
             return [];
         }
 
@@ -64,15 +64,14 @@ export class VectorSearchEmbeddings {
         {
             database,
             collection,
-            provider,
         }: {
             database: string;
             collection: string;
-            provider: NodeDriverServiceProvider;
         },
         document: Document
     ): Promise<VectorFieldIndexDefinition[]> {
-        if (!(await this.isAtlasSearchAvailable(provider))) {
+        const provider = await this.assertAtlasSearchIsAvailable();
+        if (!provider) {
             return [];
         }
 
@@ -83,25 +82,19 @@ export class VectorSearchEmbeddings {
             return [];
         }
 
-        const embeddings = await this.embeddingsForNamespace({ database, collection, provider });
+        const embeddings = await this.embeddingsForNamespace({ database, collection });
         return embeddings.filter((emb) => !this.documentPassesEmbeddingValidation(emb, document));
     }
 
-    async isAtlasSearchAvailable(provider: NodeDriverServiceProvider): Promise<boolean> {
-        const providerUri = provider.getURI();
-        if (!providerUri) {
-            // no URI? can't be cached
-            return await this.canListAtlasSearchIndexes(provider);
-        }
-
-        if (this.atlasSearchStatus.has(providerUri)) {
-            // has should ensure that get is always defined
-            return this.atlasSearchStatus.get(providerUri) ?? false;
+    private async assertAtlasSearchIsAvailable(): Promise<NodeDriverServiceProvider | null> {
+        const connectionState = this.connectionManager.currentConnectionState;
+        if (connectionState.tag === "connected") {
+            if ((await connectionState.getSearchAvailability()) === "available") {
+                return connectionState.serviceProvider;
+            }
         }
 
-        const availability = await this.canListAtlasSearchIndexes(provider);
-        this.atlasSearchStatus.set(providerUri, availability);
-        return availability;
+        return null;
     }
 
     private isVectorFieldIndexDefinition(doc: Document): doc is VectorFieldIndexDefinition {
@@ -160,15 +153,6 @@ export class VectorSearchEmbeddings {
         return true;
     }
 
-    private async canListAtlasSearchIndexes(provider: NodeDriverServiceProvider): Promise<boolean> {
-        try {
-            await provider.getSearchIndexes("test", "test");
-            return true;
-        } catch {
-            return false;
-        }
-    }
-
     private isANumber(value: unknown): boolean {
         if (typeof value === "number") {
             return true;

src/common/session.ts

@@ -11,6 +11,7 @@ import type {
     ConnectionSettings,
     ConnectionStateConnected,
     ConnectionStateErrored,
+    SearchAvailability,
 } from "./connectionManager.js";
 import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { ErrorCodes, MongoDBError } from "./errors.js";
@@ -146,13 +147,32 @@ export class Session extends EventEmitter<SessionEvents> {
         return this.connectionManager.currentConnectionState.tag === "connected";
     }
 
-    isSearchSupported(): Promise<boolean> {
+    async isSearchAvailable(): Promise<SearchAvailability> {
         const state = this.connectionManager.currentConnectionState;
         if (state.tag === "connected") {
-            return state.isSearchSupported();
+            return await state.getSearchAvailability();
         }
 
-        return Promise.resolve(false);
+        return false;
+    }
+
+    async assertSearchAvailable(): Promise<void> {
+        const availability = await this.isSearchAvailable();
+        if (!availability) {
+            throw new MongoDBError(
+                ErrorCodes.AtlasSearchNotSupported,
+                "Atlas Search is not supported in the current cluster."
+            );
+        }
+
+        if (availability === "not-available-yet") {
+            throw new MongoDBError(
+                ErrorCodes.AtlasSearchNotAvailable,
+                "Atlas Search is supported in the current cluster but not available yet."
+            );
+        }
+
+        return;
     }
 
     get serviceProvider(): NodeDriverServiceProvider {

src/resources/common/debug.ts

@@ -61,7 +61,8 @@ export class DebugResource extends ReactiveResource<
 
         switch (this.current.tag) {
             case "connected": {
-                const searchIndexesSupported = await this.session.isSearchSupported();
+                const searchAvailability = await this.session.isSearchAvailable();
+                const searchIndexesSupported = searchAvailability !== false;
                 result += `The user is connected to the MongoDB cluster${searchIndexesSupported ? " with support for search indexes" : " without any support for search indexes"}.`;
                 break;
             }

src/tools/mongodb/create/createIndex.ts

@@ -1,7 +1,6 @@
 import { z } from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
-import type { ToolCategory } from "../../tool.js";
 import { type ToolArgs, type OperationType, FeatureFlags } from "../../tool.js";
 import type { IndexDirection } from "mongodb";
 
@@ -113,25 +112,7 @@ export class CreateIndexTool extends MongoDBToolBase {
                 break;
             case "vectorSearch":
                 {
-                    const isVectorSearchSupported = await this.session.isSearchSupported();
-                    if (!isVectorSearchSupported) {
-                        // TODO: remove hacky casts once we merge the local dev tools
-                        const isLocalAtlasAvailable =
-                            (this.server?.tools.filter((t) => t.category === ("atlas-local" as unknown as ToolCategory))
-                                .length ?? 0) > 0;
-
-                        const CTA = isLocalAtlasAvailable ? "`atlas-local` tools" : "Atlas CLI";
-                        return {
-                            content: [
-                                {
-                                    text: `The connected MongoDB deployment does not support vector search indexes. Either connect to a MongoDB Atlas cluster or use the ${CTA} to create and manage a local Atlas deployment.`,
-                                    type: "text",
-                                },
-                            ],
-                            isError: true,
-                        };
-                    }
-
+                    await this.ensureSearchIsAvailable();
                     indexes = await provider.createSearchIndexes(database, collection, [
                         {
                             name,

src/tools/mongodb/create/insertMany.ts

@@ -28,7 +28,7 @@ export class InsertManyTool extends MongoDBToolBase {
             ...(await Promise.all(
                 documents.flatMap((document) =>
                     this.session.vectorSearchEmbeddings.findFieldsWithWrongEmbeddings(
-                        { database, collection, provider },
+                        { database, collection },
                         document
                     )
                 )

src/tools/mongodb/mongodbTool.ts

@@ -46,16 +46,8 @@ export abstract class MongoDBToolBase extends ToolBase {
         return this.session.serviceProvider;
     }
 
-    protected async ensureSearchAvailable(): Promise<NodeDriverServiceProvider> {
-        const provider = await this.ensureConnected();
-        if (!(await this.session.vectorSearchEmbeddings.isAtlasSearchAvailable(provider))) {
-            throw new MongoDBError(
-                ErrorCodes.AtlasSearchNotAvailable,
-                "This MongoDB cluster does not support Search Indexes. Make sure you are using an Atlas Cluster, either remotely in Atlas or using the Atlas Local image, or your cluster supports MongoDB Search."
-            );
-        }
-
-        return provider;
+    protected async ensureSearchIsAvailable(): Promise<void> {
+        return await this.session.assertSearchAvailable();
     }
 
     public register(server: Server): boolean {
@@ -94,6 +86,30 @@ export abstract class MongoDBToolBase extends ToolBase {
                         ],
                         isError: true,
                     };
+                case ErrorCodes.AtlasSearchNotSupported: {
+                    const CTA = this.isToolCategoryAvailable("atlas-local" as unknown as ToolCategory)
+                        ? "`atlas-local` tools"
+                        : "Atlas CLI";
+                    return {
+                        content: [
+                            {
+                                text: `The connected MongoDB deployment does not support vector search indexes. Either connect to a MongoDB Atlas cluster or use the ${CTA} to create and manage a local Atlas deployment.`,
+                                type: "text",
+                            },
+                        ],
+                        isError: true,
+                    };
+                }
+                case ErrorCodes.AtlasSearchNotAvailable:
+                    return {
+                        content: [
+                            {
+                                text: `The connected MongoDB deployment does support vector search indexes but they are not ready yet. Try again later.`,
+                                type: "text",
+                            },
+                        ],
+                        isError: true,
+                    };
             }
         }
 
@@ -117,4 +133,8 @@ export abstract class MongoDBToolBase extends ToolBase {
 
         return metadata;
     }
+
+    protected isToolCategoryAvailable(name: ToolCategory): boolean {
+        return (this.server?.tools.filter((t) => t.category === name).length ?? 0) > 0;
+    }
 }

src/tools/mongodb/search/listSearchIndexes.ts

@@ -19,7 +19,9 @@ export class ListSearchIndexesTool extends MongoDBToolBase {
     public operationType: OperationType = "metadata";
 
     protected async execute({ database, collection }: ToolArgs<typeof DbOperationArgs>): Promise<CallToolResult> {
-        const provider = await this.ensureSearchAvailable();
+        const provider = await this.ensureConnected();
+        await this.session.assertSearchAvailable();
+
         const indexes = await provider.getSearchIndexes(database, collection);
         const trimmedIndexDefinitions = this.pickRelevantInformation(indexes);