chore: simplify, assume search indexes are available just by listing them

kmruiz · kmruiz · commit 19a333c2f62d · 2025-10-16T10:26:26.000+02:00
diff --git a/src/common/connectionManager.ts b/src/common/connectionManager.ts
@@ -25,7 +25,6 @@ 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;
@@ -34,38 +33,24 @@ export interface ConnectionState {
 }
 
 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;
-
-    public async getSearchAvailability(): Promise<SearchAvailability> {
-        if ((await this.isSearchSupported()) === true) {
-            if ((await this.isSearchAvailable()) === true) {
-                return "available";
-            }
-
-            return "not-available-yet";
-        }
 
-        return false;
-    }
-
-    private async isSearchSupported(): Promise<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
+                // 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 {
@@ -75,57 +60,6 @@ export class ConnectionStateConnected implements ConnectionState {
 
         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;
-            }
-        }
-
-        resolve(false);
-    }
 }
 
 export interface ConnectionStateConnecting extends ConnectionState {
diff --git a/src/common/search/vectorSearchEmbeddings.ts b/src/common/search/vectorSearchEmbeddings.ts
@@ -93,7 +93,7 @@ export class VectorSearchEmbeddings {
     private async assertAtlasSearchIsAvailable(): Promise<NodeDriverServiceProvider | null> {
         const connectionState = this.connectionManager.currentConnectionState;
         if (connectionState.tag === "connected") {
-            if ((await connectionState.getSearchAvailability()) === "available") {
+            if (await connectionState.isSearchSupported()) {
                 return connectionState.serviceProvider;
             }
         }
diff --git a/src/common/session.ts b/src/common/session.ts
@@ -11,7 +11,6 @@ import type {
     ConnectionSettings,
     ConnectionStateConnected,
     ConnectionStateErrored,
-    SearchAvailability,
 } from "./connectionManager.js";
 import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { ErrorCodes, MongoDBError } from "./errors.js";
@@ -147,31 +146,24 @@ export class Session extends EventEmitter<SessionEvents> {
         return this.connectionManager.currentConnectionState.tag === "connected";
     }
 
-    async isSearchAvailable(): Promise<SearchAvailability> {
+    async isSearchSupported(): Promise<boolean> {
         const state = this.connectionManager.currentConnectionState;
         if (state.tag === "connected") {
-            return await state.getSearchAvailability();
+            return await state.isSearchSupported();
         }
 
         return false;
     }
 
-    async assertSearchAvailable(): Promise<void> {
-        const availability = await this.isSearchAvailable();
+    async assertSearchSupported(): Promise<void> {
+        const availability = await this.isSearchSupported();
         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;
     }
 
diff --git a/src/resources/common/debug.ts b/src/resources/common/debug.ts
@@ -61,8 +61,7 @@ export class DebugResource extends ReactiveResource<
 
         switch (this.current.tag) {
             case "connected": {
-                const searchAvailability = await this.session.isSearchAvailable();
-                const searchIndexesSupported = searchAvailability !== false;
+                const searchIndexesSupported = await this.session.isSearchSupported();
                 result += `The user is connected to the MongoDB cluster${searchIndexesSupported ? " with support for search indexes" : " without any support for search indexes"}.`;
                 break;
             }
diff --git a/src/tools/mongodb/create/createIndex.ts b/src/tools/mongodb/create/createIndex.ts
@@ -112,7 +112,7 @@ export class CreateIndexTool extends MongoDBToolBase {
                 break;
             case "vectorSearch":
                 {
-                    await this.ensureSearchIsAvailable();
+                    await this.ensureSearchIsSupported();
                     indexes = await provider.createSearchIndexes(database, collection, [
                         {
                             name,
diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts
@@ -46,8 +46,8 @@ export abstract class MongoDBToolBase extends ToolBase {
         return this.session.serviceProvider;
     }
 
-    protected async ensureSearchIsAvailable(): Promise<void> {
-        return await this.session.assertSearchAvailable();
+    protected async ensureSearchIsSupported(): Promise<void> {
+        return await this.session.assertSearchSupported();
     }
 
     public register(server: Server): boolean {
diff --git a/src/tools/mongodb/search/listSearchIndexes.ts b/src/tools/mongodb/search/listSearchIndexes.ts
@@ -20,7 +20,7 @@ export class ListSearchIndexesTool extends MongoDBToolBase {
 
     protected async execute({ database, collection }: ToolArgs<typeof DbOperationArgs>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
-        await this.session.assertSearchAvailable();
+        await this.ensureSearchIsSupported();
 
         const indexes = await provider.getSearchIndexes(database, collection);
         const trimmedIndexDefinitions = this.pickRelevantInformation(indexes);
diff --git a/tests/unit/common/session.test.ts b/tests/unit/common/session.test.ts
@@ -151,29 +151,7 @@ describe("Session", () => {
                 connectionString: "mongodb://localhost:27017",
             });
 
-            expect(await session.isSearchAvailable()).toEqual("available");
-        });
-
-        it("should return 'available' if listing search indexes succeed and we don't have write permissions", async () => {
-            getSearchIndexesMock.mockResolvedValue([]);
-            insertOneMock.mockRejectedValue(new Error("Read only mode"));
-            createSearchIndexesMock.mockResolvedValue([]);
-
-            await session.connectToMongoDB({
-                connectionString: "mongodb://localhost:27017",
-            });
-
-            expect(await session.isSearchAvailable()).toEqual("available");
-        });
-
-        it("should return 'not-available-yet' if listing search indexes work but can not create an index", async () => {
-            getSearchIndexesMock.mockResolvedValue([]);
-            insertOneMock.mockResolvedValue([]);
-            createSearchIndexesMock.mockRejectedValue(new Error("SearchNotAvailable"));
-            await session.connectToMongoDB({
-                connectionString: "mongodb://localhost:27017",
-            });
-            expect(await session.isSearchAvailable()).toEqual("not-available-yet");
+            expect(await session.isSearchSupported()).toBeTruthy();
         });
 
         it("should return false if listing search indexes fail with search error", async () => {
@@ -182,51 +160,29 @@ describe("Session", () => {
             await session.connectToMongoDB({
                 connectionString: "mongodb://localhost:27017",
             });
-            expect(await session.isSearchAvailable()).toEqual(false);
+            expect(await session.isSearchSupported()).toEqual(false);
         });
     });
 
-    describe("assertSearchAvailable", () => {
+    describe("assertSearchSupported", () => {
         let getSearchIndexesMock: MockedFunction<() => unknown>;
-        let createSearchIndexesMock: MockedFunction<() => unknown>;
 
         beforeEach(() => {
             getSearchIndexesMock = vi.fn();
-            createSearchIndexesMock = vi.fn();
 
             MockNodeDriverServiceProvider.connect = vi.fn().mockResolvedValue({
                 getSearchIndexes: getSearchIndexesMock,
-                createSearchIndexes: createSearchIndexesMock,
-                insertOne: vi.fn().mockResolvedValue({}),
-                dropDatabase: vi.fn().mockResolvedValue({}),
             } as unknown as NodeDriverServiceProvider);
         });
 
         it("should not throw if it is available", async () => {
             getSearchIndexesMock.mockResolvedValue([]);
-            createSearchIndexesMock.mockResolvedValue([]);
-
-            await session.connectToMongoDB({
-                connectionString: "mongodb://localhost:27017",
-            });
-
-            await expect(session.assertSearchAvailable()).resolves.not.toThrowError();
-        });
-
-        it("should throw if it is supported but not available", async () => {
-            getSearchIndexesMock.mockResolvedValue([]);
-            createSearchIndexesMock.mockRejectedValue(new Error("Not ready yet"));
 
             await session.connectToMongoDB({
                 connectionString: "mongodb://localhost:27017",
             });
 
-            await expect(session.assertSearchAvailable()).rejects.toThrowError(
-                new MongoDBError(
-                    ErrorCodes.AtlasSearchNotAvailable,
-                    "Atlas Search is supported in the current cluster but not available yet."
-                )
-            );
+            await expect(session.assertSearchSupported()).resolves.not.toThrowError();
         });
 
         it("should throw if it is not supported", async () => {
@@ -236,7 +192,7 @@ describe("Session", () => {
                 connectionString: "mongodb://localhost:27017",
             });
 
-            await expect(session.assertSearchAvailable()).rejects.toThrowError(
+            await expect(session.assertSearchSupported()).rejects.toThrowError(
                 new MongoDBError(
                     ErrorCodes.AtlasSearchNotSupported,
                     "Atlas Search is not supported in the current cluster."
diff --git a/tests/unit/resources/common/debug.test.ts b/tests/unit/resources/common/debug.test.ts
@@ -108,7 +108,7 @@ describe("debug resource", () => {
     });
 
     it("should notify if a cluster supports search indexes", async () => {
-        vi.spyOn(session, "isSearchAvailable").mockImplementation(() => Promise.resolve("available"));
+        vi.spyOn(session, "isSearchSupported").mockImplementation(() => Promise.resolve(true));
         debugResource.reduceApply("connect", undefined);
         const output = await debugResource.toOutput();
 

Original file line number	Diff line number	Diff line change
`@@ -93,7 +93,7 @@ export class VectorSearchEmbeddings {`
`93`	`93`	`private async assertAtlasSearchIsAvailable(): Promise<NodeDriverServiceProvider \| null> {`
`94`	`94`	`const connectionState = this.connectionManager.currentConnectionState;`
`95`	`95`	`if (connectionState.tag === "connected") {`
`96`		`- if ((await connectionState.getSearchAvailability()) === "available") {`
	`96`	`+ if (await connectionState.isSearchSupported()) {`
`97`	`97`	`return connectionState.serviceProvider;`
`98`	`98`	`}`
`99`	`99`	`}`
Original file line number	Diff line number	Diff line change
`@@ -61,8 +61,7 @@ export class DebugResource extends ReactiveResource<`
`61`	`61`
`62`	`62`	`switch (this.current.tag) {`
`63`	`63`	`case "connected": {`
`64`		`- const searchAvailability = await this.session.isSearchAvailable();`
`65`		`- const searchIndexesSupported = searchAvailability !== false;`
	`64`	`+ const searchIndexesSupported = await this.session.isSearchSupported();`
`66`	`65`	result += `The user is connected to the MongoDB cluster${searchIndexesSupported ? " with support for search indexes" : " without any support for search indexes"}.`;
`67`	`66`	`break;`
`68`	`67`	`}`
Original file line number	Diff line number	Diff line change
`@@ -112,7 +112,7 @@ export class CreateIndexTool extends MongoDBToolBase {`
`112`	`112`	`break;`
`113`	`113`	`case "vectorSearch":`
`114`	`114`	`{`
`115`		`- await this.ensureSearchIsAvailable();`
	`115`	`+ await this.ensureSearchIsSupported();`
`116`	`116`	`indexes = await provider.createSearchIndexes(database, collection, [`
`117`	`117`	`{`
`118`	`118`	`name,`
Original file line number	Diff line number	Diff line change
`@@ -46,8 +46,8 @@ export abstract class MongoDBToolBase extends ToolBase {`
`46`	`46`	`return this.session.serviceProvider;`
`47`	`47`	`}`
`48`	`48`
`49`		`- protected async ensureSearchIsAvailable(): Promise<void> {`
`50`		`- return await this.session.assertSearchAvailable();`
	`49`	`+ protected async ensureSearchIsSupported(): Promise<void> {`
	`50`	`+ return await this.session.assertSearchSupported();`
`51`	`51`	`}`
`52`	`52`
`53`	`53`	`public register(server: Server): boolean {`