chore: add mcp auto-embed usage property

Author: himanshusinghs

src/telemetry/types.ts

@@ -143,22 +143,30 @@ export type CommonProperties = {
  * For MongoDB tools, this is typically empty, while for Atlas tools, this should include
  * the project and organization IDs if available.
  */
-export type TelemetryToolMetadata = AtlasMetadata | PerfAdvisorToolMetadata | ConnectionMetadata;
+export type TelemetryToolMetadata =
+    | AtlasMetadata
+    | ConnectionMetadata
+    | PerfAdvisorToolMetadata
+    | AutoEmbeddingsUsageMetadata;
 
 export type AtlasMetadata = {
     project_id?: string;
     org_id?: string;
 };
 
-export type PerfAdvisorToolMetadata = AtlasMetadata & {
-    operations: string[];
+type AtlasLocalToolMetadata = {
+    atlas_local_deployment_id?: string;
 };
 
 export type ConnectionMetadata = AtlasMetadata &
     AtlasLocalToolMetadata & {
         connection_auth_type?: string;
     };
 
-type AtlasLocalToolMetadata = {
-    atlas_local_deployment_id?: string;
+export type PerfAdvisorToolMetadata = AtlasMetadata & {
+    operations: string[];
+};
+
+export type AutoEmbeddingsUsageMetadata = ConnectionMetadata & {
+    embeddingsGeneratedBy: "mcp" | "mongot";
 };

src/tools/mongodb/create/insertMany.ts

@@ -6,6 +6,7 @@ import { zEJSON } from "../../args.js";
 import { type Document } from "bson";
 import { zSupportedEmbeddingParameters } from "../mongodbSchemas.js";
 import { ErrorCodes, MongoDBError } from "../../../common/errors.js";
+import type { ConnectionMetadata, AutoEmbeddingsUsageMetadata } from "../../../telemetry/types.js";
 
 const zSupportedEmbeddingParametersWithInput = zSupportedEmbeddingParameters.extend({
     input: z
@@ -155,4 +156,18 @@ export class InsertManyTool extends MongoDBToolBase {
             }
         }
     }
+
+    protected resolveTelemetryMetadata(
+        args: ToolArgs<typeof this.argsShape>,
+        { result }: { result: CallToolResult }
+    ): ConnectionMetadata | AutoEmbeddingsUsageMetadata {
+        if ("embeddingParameters" in args && this.config.voyageApiKey) {
+            return {
+                ...super.resolveTelemetryMetadata(args, { result }),
+                embeddingsGeneratedBy: "mcp",
+            };
+        } else {
+            return super.resolveTelemetryMetadata(args, { result });
+        }
+    }
 }

src/tools/mongodb/read/aggregate.ts

@@ -17,6 +17,7 @@ import {
     assertVectorSearchFilterFieldsAreIndexed,
     type SearchIndex,
 } from "../../../helpers/assertVectorSearchFilterFieldsAreIndexed.js";
+import type { AutoEmbeddingsUsageMetadata, ConnectionMetadata } from "../../../telemetry/types.js";
 
 const pipelineDescriptionWithVectorSearch = `\
 An array of aggregation stages to execute.
@@ -344,4 +345,23 @@ The aggregation resulted in ${aggResultsCount === undefined ? "indeterminable nu
 Returning ${documents.length} documents${appliedLimitText ? ` ${appliedLimitText}` : "."}\
 `;
     }
+
+    protected resolveTelemetryMetadata(
+        args: ToolArgs<typeof this.argsShape>,
+        { result }: { result: CallToolResult }
+    ): ConnectionMetadata | AutoEmbeddingsUsageMetadata {
+        const [maybeVectorStage] = args.pipeline;
+        if (
+            maybeVectorStage &&
+            (maybeVectorStage as z.infer<typeof VectorSearchStage>)?.["$vectorSearch"]?.embeddingParameters &&
+            this.config.voyageApiKey
+        ) {
+            return {
+                ...super.resolveTelemetryMetadata(args, { result }),
+                embeddingsGeneratedBy: "mcp",
+            };
+        } else {
+            return super.resolveTelemetryMetadata(args, { result });
+        }
+    }
 }

tests/integration/tools/mongodb/create/insertMany.test.ts

@@ -14,9 +14,10 @@ import {
     getDataFromUntrustedContent,
     defaultTestConfig,
 } from "../../../helpers.js";
-import { beforeEach, afterEach, expect, it, describe } from "vitest";
+import { beforeEach, afterEach, expect, it, describe, vi } from "vitest";
 import { ObjectId } from "bson";
 import type { Collection } from "mongodb";
+import type { ToolEvent } from "../../../../../src/telemetry/types.js";
 
 describeWithMongoDB("insertMany tool when search is disabled", (integration) => {
     validateToolMetadata(
@@ -237,6 +238,7 @@ describeWithMongoDB(
 
             afterEach(async () => {
                 await collection.drop();
+                vi.clearAllMocks();
             });
 
             it("generates embeddings for a single document with one field", async () => {
@@ -627,6 +629,42 @@ describeWithMongoDB(
                 // Verify embeddings are different for different text
                 expect(doc?.titleEmbeddings).not.toEqual(doc?.plotEmbeddings);
             });
+
+            it("should emit tool event with auto-embedding usage metadata", async () => {
+                const mockEmitEvents = vi.spyOn(integration.mcpServer()["telemetry"], "emitEvents");
+                vi.spyOn(integration.mcpServer()["telemetry"], "isTelemetryEnabled").mockReturnValue(true);
+
+                await createVectorSearchIndexAndWait(integration.mongoClient(), database, "test", [
+                    {
+                        type: "vector",
+                        path: "titleEmbeddings",
+                        numDimensions: 1024,
+                        similarity: "cosine",
+                        quantization: "scalar",
+                    },
+                ]);
+
+                const response = await integration.mcpClient().callTool({
+                    name: "insert-many",
+                    arguments: {
+                        database,
+                        collection: "test",
+                        documents: [{ title: "The Matrix" }],
+                        embeddingParameters: {
+                            model: "voyage-3.5-lite",
+                            input: [{ titleEmbeddings: "The Matrix" }],
+                        },
+                    },
+                });
+
+                const content = getResponseContent(response.content);
+                expect(content).toContain("Documents were inserted successfully.");
+
+                expect(mockEmitEvents).toHaveBeenCalled();
+                const emittedEvent = mockEmitEvents.mock.lastCall?.[0][0] as ToolEvent;
+                expectDefined(emittedEvent);
+                expect(emittedEvent.properties.embeddingsGeneratedBy).toBe("mcp");
+            });
         });
     },
     {

tests/integration/tools/mongodb/read/aggregate.test.ts

@@ -4,6 +4,7 @@ import {
     validateThrowsForInvalidArguments,
     getResponseContent,
     defaultTestConfig,
+    expectDefined,
 } from "../../../helpers.js";
 import { beforeEach, describe, expect, it, vi, afterEach } from "vitest";
 import {
@@ -17,6 +18,7 @@ import * as constants from "../../../../../src/helpers/constants.js";
 import { freshInsertDocuments } from "./find.test.js";
 import { BSON } from "bson";
 import { DOCUMENT_EMBEDDINGS } from "./vyai/embeddings.js";
+import type { ToolEvent } from "../../../../../src/telemetry/types.js";
 
 describeWithMongoDB("aggregate tool", (integration) => {
     afterEach(() => {
@@ -393,6 +395,10 @@ describeWithMongoDB(
             await integration.mongoClient().db(integration.randomDbName()).collection("databases").drop();
         });
 
+        afterEach(() => {
+            vi.clearAllMocks();
+        });
+
         validateToolMetadata(integration, "aggregate", "Run an aggregation against a MongoDB collection", "read", [
             ...databaseCollectionParameters,
             {
@@ -462,6 +468,50 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi
             );
         });
 
+        it("should emit tool event with auto-embedding usage metadata", async () => {
+            const mockEmitEvents = vi.spyOn(integration.mcpServer()["telemetry"], "emitEvents");
+            vi.spyOn(integration.mcpServer()["telemetry"], "isTelemetryEnabled").mockReturnValue(true);
+
+            await waitUntilSearchIsReady(integration.mongoClient());
+            await createVectorSearchIndexAndWait(integration.mongoClient(), integration.randomDbName(), "databases", [
+                {
+                    type: "vector",
+                    path: "description_embedding",
+                    numDimensions: 256,
+                    similarity: "cosine",
+                    quantization: "none",
+                },
+            ]);
+
+            await integration.mcpClient().callTool({
+                name: "aggregate",
+                arguments: {
+                    database: integration.randomDbName(),
+                    collection: "databases",
+                    pipeline: [
+                        {
+                            $vectorSearch: {
+                                index: "default",
+                                path: "description_embedding",
+                                queryVector: "some data",
+                                numCandidates: 10,
+                                limit: 10,
+                                embeddingParameters: {
+                                    model: "voyage-3-large",
+                                    outputDimension: "256",
+                                },
+                            },
+                        },
+                    ],
+                },
+            });
+
+            expect(mockEmitEvents).toHaveBeenCalled();
+            const emittedEvent = mockEmitEvents.mock.lastCall?.[0][0] as ToolEvent;
+            expectDefined(emittedEvent);
+            expect(emittedEvent.properties.embeddingsGeneratedBy).toBe("mcp");
+        });
+
         for (const [dataType, embedding] of Object.entries(DOCUMENT_EMBEDDINGS)) {
             for (const similarity of ["euclidean", "cosine", "dotProduct"]) {
                 describe(`querying with dataType ${dataType} and similarity ${similarity}`, () => {