Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
25 changes: 20 additions & 5 deletions src/telemetry/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@ export type ServerEventProperties = {
read_only_mode?: TelemetryBoolSet;
disabled_tools?: string[];
confirmation_required_tools?: string[];
previewFeatures?: string[];
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

note: while we can add arrays in telemetry, this really makes stuff like amplitude difficult as the values appear as one text with spaces normally.

ordering of strings is also a problem as you can't queries based on keys or position.

since this pattern is already stablished let's proceed, but ideally we should be removing arrays from segment events.

embeddingProviderConfigured?: boolean;
};

export type ServerEvent = TelemetryEvent<ServerEventProperties>;
Expand Down Expand Up @@ -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";
};
15 changes: 15 additions & 0 deletions src/tools/mongodb/create/insertMany.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 });
}
}
}
22 changes: 22 additions & 0 deletions src/tools/mongodb/read/aggregate.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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 });
}
}
}
63 changes: 62 additions & 1 deletion tests/integration/tools/mongodb/create/insertMany.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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: {
Expand Down Expand Up @@ -237,6 +261,7 @@ describeWithMongoDB(

afterEach(async () => {
await collection.drop();
vi.clearAllMocks();
});

it("generates embeddings for a single document with one field", async () => {
Expand Down Expand Up @@ -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");
});
});
},
{
Expand Down
80 changes: 80 additions & 0 deletions tests/integration/tools/mongodb/read/aggregate.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import {
validateThrowsForInvalidArguments,
getResponseContent,
defaultTestConfig,
expectDefined,
} from "../../../helpers.js";
import { beforeEach, describe, expect, it, vi, afterEach } from "vitest";
import {
Expand All @@ -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(() => {
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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,
{
Expand Down Expand Up @@ -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}`, () => {
Expand Down