From 0b60d8abe4cb4d62a59e8486fd2cec1788417dea Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 11 Dec 2025 16:35:46 +0100 Subject: [PATCH 1/3] chore: add telemetry properties for vector search --- src/server.ts | 2 ++ src/telemetry/types.ts | 2 ++ 2 files changed, 4 insertions(+) diff --git a/src/server.ts b/src/server.ts index daec7a0a2..b4cda3320 100644 --- 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 index 46b91f3bd..164823bdc 100644 --- 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; From 80f9e7e424909b2568e02f57513cc57379537fd0 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 11 Dec 2025 23:22:05 +0100 Subject: [PATCH 2/3] chore: add mcp auto-embed usage property --- src/telemetry/types.ts | 18 +++++-- src/tools/mongodb/create/insertMany.ts | 15 ++++++ src/tools/mongodb/read/aggregate.ts | 20 ++++++++ .../tools/mongodb/create/insertMany.test.ts | 40 ++++++++++++++- .../tools/mongodb/read/aggregate.test.ts | 50 +++++++++++++++++++ 5 files changed, 137 insertions(+), 6 deletions(-) diff --git a/src/telemetry/types.ts b/src/telemetry/types.ts index 164823bdc..1e68a522f 100644 --- a/src/telemetry/types.ts +++ b/src/telemetry/types.ts @@ -143,15 +143,19 @@ 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 & @@ -159,6 +163,10 @@ export type ConnectionMetadata = AtlasMetadata & connection_auth_type?: string; }; -type AtlasLocalToolMetadata = { - atlas_local_deployment_id?: string; +export type PerfAdvisorToolMetadata = AtlasMetadata & { + operations: string[]; +}; + +export type AutoEmbeddingsUsageMetadata = ConnectionMetadata & { + embeddingsGeneratedBy: "mcp" | "mongot"; }; diff --git a/src/tools/mongodb/create/insertMany.ts b/src/tools/mongodb/create/insertMany.ts index 6e1041e8c..f365c0fae 100644 --- 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, + { 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 index 8c87dbc7d..a04b8f80e 100644 --- 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,23 @@ The aggregation resulted in ${aggResultsCount === undefined ? "indeterminable nu Returning ${documents.length} documents${appliedLimitText ? ` ${appliedLimitText}` : "."}\ `; } + + protected resolveTelemetryMetadata( + args: ToolArgs, + { result }: { result: CallToolResult } + ): ConnectionMetadata | AutoEmbeddingsUsageMetadata { + const [maybeVectorStage] = args.pipeline; + if ( + maybeVectorStage && + (maybeVectorStage as z.infer)?.["$vectorSearch"]?.embeddingParameters && + 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 index 2170efb95..f9b2d9281 100644 --- 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( @@ -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"); + }); }); }, { diff --git a/tests/integration/tools/mongodb/read/aggregate.test.ts b/tests/integration/tools/mongodb/read/aggregate.test.ts index 8a0c0adda..5ae2fed43 100644 --- 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(() => { @@ -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}`, () => { From de74005adefbbd939815378cc578a361a0539231 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Thu, 11 Dec 2025 23:27:20 +0100 Subject: [PATCH 3/3] Update src/telemetry/types.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/telemetry/types.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/telemetry/types.ts b/src/telemetry/types.ts index 1e68a522f..2bfded7f6 100644 --- a/src/telemetry/types.ts +++ b/src/telemetry/types.ts @@ -168,5 +168,10 @@ export type PerfAdvisorToolMetadata = AtlasMetadata & { }; 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"; };