feat: introduces drop-search-index tool

- Tool is currently only visible to tests
- Tool has been added to the list of tools requiring elicitation
- Tests have been added for basic tool functionality and elicitation
- listSearchIndexes tool and its tests have been modified to re-use some
  of the existing code.

Author: himanshusinghs

src/common/config.ts

@@ -204,6 +204,7 @@ export const defaultUserConfig: UserConfig = {
         "drop-collection",
         "delete-many",
         "drop-index",
+        "drop-search-index",
     ],
     transport: "stdio",
     httpPort: 3000,

src/tools/mongodb/delete/dropSearchIndex.ts

@@ -0,0 +1,79 @@
+import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { CommonArgs } from "../../args.js";
+import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
+import { type OperationType, type ToolArgs } from "../../tool.js";
+import { ListSearchIndexesTool } from "../search/listSearchIndexes.js";
+
+export class DropSearchIndexTool extends MongoDBToolBase {
+    public name = "drop-search-index";
+    protected description = "Drop a search index or vector search index for the provided database and collection.";
+    protected argsShape = {
+        ...DbOperationArgs,
+        indexName: CommonArgs.string()
+            .nonempty()
+            .describe("The name of the search or vector search index to be dropped."),
+    };
+    public operationType: OperationType = "delete";
+
+    protected override verifyAllowed(): boolean {
+        // Only enable this on tests for now.
+        return process.env.VITEST === "true";
+    }
+
+    protected override async execute({
+        database,
+        collection,
+        indexName,
+    }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
+        const provider = await this.ensureConnected();
+        const searchIndexes = await ListSearchIndexesTool.getSearchIndexes(provider, database, collection);
+        const indexDoesNotExist = !searchIndexes.find((index) => index.name === indexName);
+        if (indexDoesNotExist) {
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: `Index with name "${indexName}" does not exist in the provided namespace "${database}.${collection}".`,
+                    },
+                ],
+                isError: true,
+            };
+        }
+
+        await provider.dropSearchIndex(database, collection, indexName);
+        return {
+            content: [
+                {
+                    type: "text",
+                    text: `Successfully dropped the index with name "${indexName}" from the provided namespace "${database}.${collection}".`,
+                },
+            ],
+        };
+    }
+
+    protected getConfirmationMessage({ database, collection, indexName }: ToolArgs<typeof this.argsShape>): string {
+        return (
+            `You are about to drop the \`${indexName}\` index from the \`${database}.${collection}\` namespace:\n\n` +
+            "This operation will permanently remove the index and might affect the performance of queries relying on this index.\n\n" +
+            "**Do you confirm the execution of the action?**"
+        );
+    }
+
+    protected handleError(
+        error: unknown,
+        args: ToolArgs<typeof DbOperationArgs>
+    ): Promise<CallToolResult> | CallToolResult {
+        if (error instanceof Error && "codeName" in error && error.codeName === "SearchNotEnabled") {
+            return {
+                content: [
+                    {
+                        text: "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.",
+                        type: "text",
+                        isError: true,
+                    },
+                ],
+            };
+        }
+        return super.handleError(error, args);
+    }
+}

src/tools/mongodb/search/listSearchIndexes.ts

@@ -1,10 +1,11 @@
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import type { ToolArgs, OperationType } from "../../tool.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { formatUntrustedData } from "../../tool.js";
 import { EJSON } from "bson";
 
-export type SearchIndexStatus = {
+export type SearchIndexWithStatus = {
     name: string;
     type: string;
     status: string;
@@ -20,14 +21,13 @@ export class ListSearchIndexesTool extends MongoDBToolBase {
 
     protected async execute({ database, collection }: ToolArgs<typeof DbOperationArgs>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
-        const indexes = await provider.getSearchIndexes(database, collection);
-        const trimmedIndexDefinitions = this.pickRelevantInformation(indexes);
+        const searchIndexes = await ListSearchIndexesTool.getSearchIndexes(provider, database, collection);
 
-        if (trimmedIndexDefinitions.length > 0) {
+        if (searchIndexes.length > 0) {
             return {
                 content: formatUntrustedData(
-                    `Found ${trimmedIndexDefinitions.length} search and vector search indexes in ${database}.${collection}`,
-                    trimmedIndexDefinitions.map((index) => EJSON.stringify(index)).join("\n")
+                    `Found ${searchIndexes.length} search and vector search indexes in ${database}.${collection}`,
+                    searchIndexes.map((index) => EJSON.stringify(index)).join("\n")
                 ),
             };
         } else {
@@ -45,22 +45,6 @@ export class ListSearchIndexesTool extends MongoDBToolBase {
         return process.env.VITEST === "true";
     }
 
-    /**
-     * Atlas Search index status contains a lot of information that is not relevant for the agent at this stage.
-     * Like for example, the status on each of the dedicated nodes. We only care about the main status, if it's
-     * queryable and the index name. We are also picking the index definition as it can be used by the agent to
-     * understand which fields are available for searching.
-     **/
-    protected pickRelevantInformation(indexes: Record<string, unknown>[]): SearchIndexStatus[] {
-        return indexes.map((index) => ({
-            name: (index["name"] ?? "default") as string,
-            type: (index["type"] ?? "UNKNOWN") as string,
-            status: (index["status"] ?? "UNKNOWN") as string,
-            queryable: (index["queryable"] ?? false) as boolean,
-            latestDefinition: index["latestDefinition"] as Document,
-        }));
-    }
-
     protected handleError(
         error: unknown,
         args: ToolArgs<typeof DbOperationArgs>
@@ -78,4 +62,25 @@ export class ListSearchIndexesTool extends MongoDBToolBase {
         }
         return super.handleError(error, args);
     }
+
+    static async getSearchIndexes(
+        provider: NodeDriverServiceProvider,
+        database: string,
+        collection: string
+    ): Promise<SearchIndexWithStatus[]> {
+        const searchIndexes = await provider.getSearchIndexes(database, collection);
+        /**
+         * Atlas Search index status contains a lot of information that is not relevant for the agent at this stage.
+         * Like for example, the status on each of the dedicated nodes. We only care about the main status, if it's
+         * queryable and the index name. We are also picking the index definition as it can be used by the agent to
+         * understand which fields are available for searching.
+         **/
+        return searchIndexes.map<SearchIndexWithStatus>((index) => ({
+            name: (index["name"] ?? "default") as string,
+            type: (index["type"] ?? "UNKNOWN") as string,
+            status: (index["status"] ?? "UNKNOWN") as string,
+            queryable: (index["queryable"] ?? false) as boolean,
+            latestDefinition: index["latestDefinition"] as Document,
+        }));
+    }
 }

src/tools/mongodb/tools.ts

@@ -21,6 +21,7 @@ import { LogsTool } from "./metadata/logs.js";
 import { ExportTool } from "./read/export.js";
 import { ListSearchIndexesTool } from "./search/listSearchIndexes.js";
 import { DropIndexTool } from "./delete/dropIndex.js";
+import { DropSearchIndexTool } from "./delete/dropSearchIndex.js";
 
 export const MongoDbTools = [
     ConnectTool,
@@ -46,4 +47,5 @@ export const MongoDbTools = [
     LogsTool,
     ExportTool,
     ListSearchIndexesTool,
+    DropSearchIndexTool,
 ];

tests/integration/helpers.ts

@@ -1,3 +1,4 @@
+import type { Collection, Document } from "mongodb";
 import { CompositeLogger } from "../../src/common/logger.js";
 import { ExportsManager } from "../../src/common/exportsManager.js";
 import { Session } from "../../src/common/session.js";
@@ -22,6 +23,9 @@ import { Keychain } from "../../src/common/keychain.js";
 import { Elicitation } from "../../src/elicitation.js";
 import type { MockClientCapabilities, createMockElicitInput } from "../utils/elicitationMocks.js";
 
+const SEARCH_READY_CHECK_RETRIES = 200;
+const SEARCH_INDEX_STATUS_CHECK_RETRIES = 100;
+
 export const driverOptions = setupDriverConfig({
     config,
     defaults: defaultDriverOptionsFromConfig,
@@ -417,3 +421,77 @@ export function getDataFromUntrustedContent(content: string): string {
 export function sleep(ms: number): Promise<void> {
     return new Promise((resolve) => setTimeout(resolve, ms));
 }
+
+export async function waitUntilSearchManagementServiceIsReady(
+    collection: Collection,
+    abortSignal?: AbortSignal
+): Promise<void> {
+    for (let i = 0; i < SEARCH_READY_CHECK_RETRIES && !abortSignal?.aborted; i++) {
+        try {
+            await collection.listSearchIndexes({}).toArray();
+            return;
+        } catch (error) {
+            if (
+                error instanceof Error &&
+                error.message.includes("Error connecting to Search Index Management service")
+            ) {
+                await sleep(100);
+                continue;
+            } else {
+                throw error;
+            }
+        }
+    }
+}
+
+async function waitUntilSearchIndexIs(
+    collection: Collection,
+    searchIndex: string,
+    indexValidator: (index: { name: string; queryable: boolean }) => boolean,
+    abortSignal?: AbortSignal
+): Promise<void> {
+    for (let i = 0; i < SEARCH_INDEX_STATUS_CHECK_RETRIES && !abortSignal?.aborted; i++) {
+        try {
+            const searchIndexes = (await collection.listSearchIndexes(searchIndex).toArray()) as {
+                name: string;
+                queryable: boolean;
+            }[];
+
+            if (searchIndexes.some((index) => indexValidator(index))) {
+                return;
+            }
+            await sleep(100);
+        } catch (error) {
+            if (
+                error instanceof Error &&
+                error.message.includes("Error connecting to Search Index Management service")
+            ) {
+                await sleep(100);
+                continue;
+            } else {
+                throw error;
+            }
+        }
+    }
+}
+
+export async function waitUntilSearchIndexIsListed(
+    collection: Collection,
+    searchIndex: string,
+    abortSignal?: AbortSignal
+): Promise<void> {
+    return waitUntilSearchIndexIs(collection, searchIndex, (index) => index.name === searchIndex, abortSignal);
+}
+
+export async function waitUntilSearchIndexIsQueryable(
+    collection: Collection,
+    searchIndex: string,
+    abortSignal?: AbortSignal
+): Promise<void> {
+    return waitUntilSearchIndexIs(
+        collection,
+        searchIndex,
+        (index) => index.name === searchIndex && index.queryable,
+        abortSignal
+    );
+}