chore: add telemetry properties for vector search and mcp auto-embed usage (#796)

himanshusinghs · web-flow · commit fd0e4b6b029c · 2025-12-12T13:51:01.000+01:00
diff --git a/src/server.ts b/src/server.ts
@@ -236,6 +236,8 @@ export class Server {
             event.properties.read_only_mode = this.userConfig.readOnly ? "true" : "false";
             event.properties.disabled_tools = this.userConfig.disabledTools || [];
             event.properties.confirmation_required_tools = this.userConfig.confirmationRequiredTools || [];
+            event.properties.previewFeatures = this.userConfig.previewFeatures;
+            event.properties.embeddingProviderConfigured = !!this.userConfig.voyageApiKey;
         }
         if (command === "stop") {
             event.properties.runtime_duration_ms = Date.now() - this.startTime;
diff --git a/src/telemetry/types.ts b/src/telemetry/types.ts
@@ -45,6 +45,8 @@ export type ServerEventProperties = {
     read_only_mode?: TelemetryBoolSet;
     disabled_tools?: string[];
     confirmation_required_tools?: string[];
+    previewFeatures?: string[];
+    embeddingProviderConfigured?: boolean;
 };
 
 export type ServerEvent = TelemetryEvent<ServerEventProperties>;
@@ -141,22 +143,35 @@ 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 & {
+    /**
+     * Indicates which component generated the embeddings.
+     * "mcp" is used when embeddings are generated by the MCP server.
+     * "mongot" is reserved for future use, when embeddings may be generated by MongoDB's mongot process.
+     */
+    embeddingsGeneratedBy: "mcp" | "mongot";
 };
diff --git a/src/tools/mongodb/create/insertMany.ts b/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 });
+        }
+    }
 }
diff --git a/src/tools/mongodb/read/aggregate.ts b/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,25 @@ 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 !== null &&
+            maybeVectorStage instanceof Object &&
+            "$vectorSearch" in maybeVectorStage &&
+            "embeddingParameters" in maybeVectorStage["$vectorSearch"] &&
+            this.config.voyageApiKey
+        ) {
+            return {
+                ...super.resolveTelemetryMetadata(args, { result }),
+                embeddingsGeneratedBy: "mcp",
+            };
+        } else {
+            return super.resolveTelemetryMetadata(args, { result });
+        }
+    }
 }
diff --git a/tests/integration/tools/mongodb/create/insertMany.test.ts b/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(
@@ -101,6 +102,29 @@ describeWithMongoDB("insertMany tool when search is disabled", (integration) =>
         expect(content).toContain(insertedIds[0]?.toString());
     });
 
+    it("should emit tool event without auto-embedding usage metadata", async () => {
+        const mockEmitEvents = vi.spyOn(integration.mcpServer()["telemetry"], "emitEvents");
+        vi.spyOn(integration.mcpServer()["telemetry"], "isTelemetryEnabled").mockReturnValue(true);
+        await integration.connectMcpClient();
+
+        const response = await integration.mcpClient().callTool({
+            name: "insert-many",
+            arguments: {
+                database: integration.randomDbName(),
+                collection: "test",
+                documents: [{ title: "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).toBeUndefined();
+    });
+
     validateAutoConnectBehavior(integration, "insert-many", () => {
         return {
             args: {
@@ -237,6 +261,7 @@ describeWithMongoDB(
 
             afterEach(async () => {
                 await collection.drop();
+                vi.clearAllMocks();
             });
 
             it("generates embeddings for a single document with one field", async () => {
@@ -627,6 +652,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");
+            });
         });
     },
     {
diff --git a/tests/integration/tools/mongodb/read/aggregate.test.ts b/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(() => {
@@ -151,6 +153,36 @@ describeWithMongoDB("aggregate tool", (integration) => {
         );
     });
 
+    it("should emit tool event without auto-embedding usage metadata", async () => {
+        const mockEmitEvents = vi.spyOn(integration.mcpServer()["telemetry"], "emitEvents");
+        vi.spyOn(integration.mcpServer()["telemetry"], "isTelemetryEnabled").mockReturnValue(true);
+
+        const mongoClient = integration.mongoClient();
+        await mongoClient
+            .db(integration.randomDbName())
+            .collection("people")
+            .insertMany([
+                { name: "Peter", age: 5 },
+                { name: "Laura", age: 10 },
+                { name: "Søren", age: 15 },
+            ]);
+
+        await integration.connectMcpClient();
+        await integration.mcpClient().callTool({
+            name: "aggregate",
+            arguments: {
+                database: integration.randomDbName(),
+                collection: "people",
+                pipeline: [{ $match: { age: { $gt: 8 } } }, { $sort: { name: -1 } }],
+            },
+        });
+
+        expect(mockEmitEvents).toHaveBeenCalled();
+        const emittedEvent = mockEmitEvents.mock.lastCall?.[0][0] as ToolEvent;
+        expectDefined(emittedEvent);
+        expect(emittedEvent.properties.embeddingsGeneratedBy).toBeUndefined();
+    });
+
     for (const disabledOpType of ["create", "update", "delete"] as const) {
         it(`can not run $out stages when ${disabledOpType} operation is disabled`, async () => {
             await integration.connectMcpClient();
@@ -393,6 +425,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 +498,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}`, () => {

Original file line number	Diff line number	Diff line change
`@@ -236,6 +236,8 @@ export class Server {`
`236`	`236`	`event.properties.read_only_mode = this.userConfig.readOnly ? "true" : "false";`
`237`	`237`	`event.properties.disabled_tools = this.userConfig.disabledTools \|\| [];`
`238`	`238`	`event.properties.confirmation_required_tools = this.userConfig.confirmationRequiredTools \|\| [];`
	`239`	`+ event.properties.previewFeatures = this.userConfig.previewFeatures;`
	`240`	`+ event.properties.embeddingProviderConfigured = !!this.userConfig.voyageApiKey;`
`239`	`241`	`}`
`240`	`242`	`if (command === "stop") {`
`241`	`243`	`event.properties.runtime_duration_ms = Date.now() - this.startTime;`