chore: addresses a few common concerns

himanshusinghs · himanshusinghs · commit 766dd57f3b66 · 2025-10-15T11:34:10.000+02:00
1. Wraps the toolcall result in untrusted data wrapper
2. Add additional test for confirming that we don't drop when
   confirmation is not provided
diff --git a/src/tools/mongodb/delete/dropSearchIndex.ts b/src/tools/mongodb/delete/dropSearchIndex.ts
@@ -1,7 +1,7 @@
 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 { formatUntrustedData, type OperationType, type ToolArgs } from "../../tool.js";
 import { ListSearchIndexesTool } from "../search/listSearchIndexes.js";
 
 export class DropSearchIndexTool extends MongoDBToolBase {
@@ -30,24 +30,23 @@ export class DropSearchIndexTool extends MongoDBToolBase {
         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}".`,
-                    },
-                ],
+                content: formatUntrustedData(
+                    "Index does not exist in the provided namespace.",
+                    JSON.stringify({ indexName, 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}".`,
-                },
-            ],
+            content: formatUntrustedData(
+                "Successfully dropped the index from the provided namespace.",
+                JSON.stringify({
+                    indexName,
+                    namespace: `${database}.${collection}`,
+                })
+            ),
         };
     }
 
diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts
@@ -393,57 +393,90 @@ export function sleep(ms: number): Promise<void> {
     return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
-export async function waitUntilSearchManagementServiceIsReady(
-    collection: Collection,
-    abortSignal?: AbortSignal
+export async function waitFor(
+    condition: () => boolean | Promise<boolean>,
+    {
+        retries,
+        retryTimeout,
+        abortSignal,
+        shouldRetryOnError,
+    }: {
+        retries: number;
+        retryTimeout: number;
+        abortSignal?: AbortSignal;
+        shouldRetryOnError?: (error: unknown) => boolean | Promise<boolean>;
+    } = {
+        retries: 100,
+        retryTimeout: 100,
+    }
 ): Promise<void> {
-    for (let i = 0; i < SEARCH_READY_CHECK_RETRIES && !abortSignal?.aborted; i++) {
+    for (let i = 0; i < retries && !abortSignal?.aborted; i++) {
         try {
-            await collection.listSearchIndexes({}).toArray();
-            return;
+            if (await condition()) {
+                return;
+            }
+
+            await sleep(retryTimeout);
         } catch (error) {
-            if (
-                error instanceof Error &&
-                error.message.includes("Error connecting to Search Index Management service")
-            ) {
-                await sleep(100);
+            if (shouldRetryOnError && (await shouldRetryOnError(error))) {
+                await sleep(retryTimeout);
                 continue;
-            } else {
-                throw error;
             }
+            throw error;
         }
     }
 }
 
+export async function waitUntilSearchManagementServiceIsReady(
+    collection: Collection,
+    abortSignal?: AbortSignal
+): Promise<void> {
+    await waitFor(
+        async (): Promise<boolean> => {
+            await collection.listSearchIndexes({}).toArray();
+            return true;
+        },
+        {
+            retries: SEARCH_READY_CHECK_RETRIES,
+            retryTimeout: 100,
+            abortSignal,
+            shouldRetryOnError: (error: unknown) => {
+                return (
+                    error instanceof Error &&
+                    error.message.includes("Error connecting to Search Index Management service")
+                );
+            },
+        }
+    );
+}
+
 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 {
+    await waitFor(
+        async (): Promise<boolean> => {
             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;
-            }
+            return searchIndexes.some((index) => indexValidator(index));
+        },
+        {
+            retries: SEARCH_INDEX_STATUS_CHECK_RETRIES,
+            retryTimeout: 100,
+            abortSignal,
+            shouldRetryOnError: (error: unknown) => {
+                return (
+                    error instanceof Error &&
+                    error.message.includes("Error connecting to Search Index Management service")
+                );
+            },
         }
-    }
+    );
 }
 
 export async function waitUntilSearchIndexIsListed(
diff --git a/tests/integration/tools/mongodb/delete/dropSearchIndex.test.ts b/tests/integration/tools/mongodb/delete/dropSearchIndex.test.ts
@@ -1,4 +1,4 @@
-import { beforeEach, afterEach, describe, expect, it } from "vitest";
+import { beforeEach, afterEach, describe, expect, it, vi, type MockInstance } from "vitest";
 import {
     databaseCollectionInvalidArgs,
     databaseCollectionParameters,
@@ -10,8 +10,9 @@ import {
     validateToolMetadata,
     waitUntilSearchManagementServiceIsReady,
     waitUntilSearchIndexIsListed,
+    getDataFromUntrustedContent,
 } from "../../../helpers.js";
-import { describeWithMongoDB } from "../mongodbHelpers.js";
+import { describeWithMongoDB, setupMongoDBIntegrationTest } from "../mongodbHelpers.js";
 import type { Collection } from "mongodb";
 import { createMockElicitInput } from "../../../../utils/elicitationMocks.js";
 import { Elicitation } from "../../../../../src/elicitation.js";
@@ -77,9 +78,10 @@ describeWithMongoDB(
                 });
                 expect(response.isError).toBe(true);
                 const content = getResponseContent(response.content);
-                expect(content).toEqual(
-                    'Index with name "non-existent" does not exist in the provided namespace "any.foo".'
-                );
+                expect(content).toContain("Index does not exist in the provided namespace.");
+
+                const data = getDataFromUntrustedContent(content);
+                expect(JSON.parse(data)).toMatchObject({ indexName: "non-existent", namespace: "any.foo" });
             });
         });
 
@@ -91,7 +93,7 @@ describeWithMongoDB(
                 await moviesCollection.insertMany([
                     {
                         name: "Movie1",
-                        plot: "This is a horrible movie about a database called BongoDB and how it tried to copy the original MangoDB.",
+                        plot: "This is a horrible movie about a database called BongoDB and how it tried to copy the OG MangoDB.",
                     },
                 ]);
                 await waitUntilSearchManagementServiceIsReady(moviesCollection, signal);
@@ -103,13 +105,7 @@ describeWithMongoDB(
             });
 
             afterEach(async () => {
-                try {
-                    await moviesCollection.dropSearchIndex("searchIdx");
-                } catch (error) {
-                    if (error instanceof Error && !error.message.includes("not found in")) {
-                        throw error;
-                    }
-                }
+                // dropping collection also drops the associated search indexes
                 await moviesCollection.drop();
             });
 
@@ -119,9 +115,10 @@ describeWithMongoDB(
                     arguments: { database: "mflix", collection: "movies", indexName: "searchIdx" },
                 });
                 const content = getResponseContent(response.content);
-                expect(content).toEqual(
-                    'Successfully dropped the index with name "searchIdx" from the provided namespace "mflix.movies".'
-                );
+                expect(content).toContain("Successfully dropped the index from the provided namespace.");
+
+                const data = getDataFromUntrustedContent(content);
+                expect(JSON.parse(data)).toMatchObject({ indexName: "searchIdx", namespace: "mflix.movies" });
             });
         });
     },
@@ -132,23 +129,82 @@ describeWithMongoDB(
 
 describe("drop-search-index tool - when invoked via an elicitation enabled client", () => {
     const mockElicitInput = createMockElicitInput();
+    const mdbIntegration = setupMongoDBIntegrationTest({ search: true });
     const integration = setupIntegrationTest(
         () => defaultTestConfig,
         () => defaultDriverOptions,
         { elicitInput: mockElicitInput }
     );
 
+    let moviesCollection: Collection;
+    let dropSearchIndexSpy: MockInstance;
+
+    beforeEach(async ({ signal }) => {
+        const mongoClient = mdbIntegration.mongoClient();
+        moviesCollection = mongoClient.db("mflix").collection("movies");
+        await moviesCollection.insertMany([
+            {
+                name: "Movie1",
+                plot: "This is a horrible movie about a database called BongoDB and how it tried to copy the OG MangoDB.",
+            },
+        ]);
+        await waitUntilSearchManagementServiceIsReady(moviesCollection, signal);
+        await moviesCollection.createSearchIndex({
+            name: "searchIdx",
+            definition: { mappings: { dynamic: true } },
+        });
+        await waitUntilSearchIndexIsListed(moviesCollection, "searchIdx", signal);
+
+        await integration.mcpClient().callTool({
+            name: "connect",
+            arguments: {
+                connectionString: mdbIntegration.connectionString(),
+            },
+        });
+
+        // Note: Unlike drop-index tool test, we don't test the final state of
+        // indexes because of possible longer wait periods for changes to
+        // reflect, at-times taking >30 seconds.
+        dropSearchIndexSpy = vi.spyOn(integration.mcpServer().session.serviceProvider, "dropSearchIndex");
+    });
+
+    afterEach(async () => {
+        // dropping collection also drops the associated search indexes
+        await moviesCollection.drop();
+    });
+
     it("should ask for confirmation before proceeding with tool call", async () => {
         mockElicitInput.confirmYes();
         await integration.mcpClient().callTool({
             name: "drop-search-index",
-            arguments: { database: "any", collection: "foo", indexName: "default" },
+            arguments: { database: "mflix", collection: "movies", indexName: "searchIdx" },
+        });
+        expect(mockElicitInput.mock).toHaveBeenCalledTimes(1);
+        expect(mockElicitInput.mock).toHaveBeenCalledWith({
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+            message: expect.stringContaining(
+                "You are about to drop the `searchIdx` index from the `mflix.movies` namespace"
+            ),
+            requestedSchema: Elicitation.CONFIRMATION_SCHEMA,
+        });
+
+        expect(dropSearchIndexSpy).toHaveBeenCalledExactlyOnceWith("mflix", "movies", "searchIdx");
+    });
+
+    it("should not drop the index if the confirmation was not provided", async () => {
+        mockElicitInput.confirmNo();
+        await integration.mcpClient().callTool({
+            name: "drop-search-index",
+            arguments: { database: "mflix", collection: "movies", indexName: "searchIdx" },
         });
         expect(mockElicitInput.mock).toHaveBeenCalledTimes(1);
         expect(mockElicitInput.mock).toHaveBeenCalledWith({
             // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-            message: expect.stringContaining("You are about to drop the `default` index from the `any.foo` namespace"),
+            message: expect.stringContaining(
+                "You are about to drop the `searchIdx` index from the `mflix.movies` namespace"
+            ),
             requestedSchema: Elicitation.CONFIRMATION_SCHEMA,
         });
+        expect(dropSearchIndexSpy).not.toHaveBeenCalled();
     });
 });
