diff --git a/src/logger.ts b/src/logger.ts index 534bfb80..19a311c2 100644 --- a/src/logger.ts +++ b/src/logger.ts @@ -17,6 +17,7 @@ export const LogId = { telemetryEmitFailure: mongoLogId(1_002_002), telemetryEmitStart: mongoLogId(1_002_003), telemetryEmitSuccess: mongoLogId(1_002_004), + telemetryMetadataError: mongoLogId(1_002_005), toolExecute: mongoLogId(1_003_001), toolExecuteFailure: mongoLogId(1_003_002), diff --git a/src/server.ts b/src/server.ts index b11ba31d..effdee63 100644 --- a/src/server.ts +++ b/src/server.ts @@ -28,7 +28,7 @@ export class Server { constructor({ session, mcpServer, userConfig }: ServerOptions) { this.startTime = Date.now(); this.session = session; - this.telemetry = new Telemetry(session); + this.telemetry = new Telemetry(session, userConfig); this.mcpServer = mcpServer; this.userConfig = userConfig; } diff --git a/src/telemetry/telemetry.ts b/src/telemetry/telemetry.ts index 53431232..31760ff4 100644 --- a/src/telemetry/telemetry.ts +++ b/src/telemetry/telemetry.ts @@ -1,6 +1,6 @@ import { Session } from "../session.js"; import { BaseEvent, CommonProperties } from "./types.js"; -import { config } from "../config.js"; +import { UserConfig } from "../config.js"; import logger, { LogId } from "../logger.js"; import { ApiClient } from "../common/atlas/apiClient.js"; import { MACHINE_METADATA } from "./constants.js"; @@ -16,6 +16,7 @@ export class Telemetry { constructor( private readonly session: Session, + private readonly userConfig: UserConfig, private readonly eventCache: EventCache = EventCache.getInstance() ) { this.commonProperties = { @@ -23,38 +24,14 @@ export class Telemetry { }; } - /** - * Checks if telemetry is currently enabled - * This is a method rather than a constant to capture runtime config changes - * - * Follows the Console Do Not Track standard (https://consoledonottrack.com/) - * by respecting the DO_NOT_TRACK environment variable - */ - private static isTelemetryEnabled(): boolean { - // Check if telemetry is explicitly disabled in config - if (config.telemetry === "disabled") { - return false; - } - - const doNotTrack = process.env.DO_NOT_TRACK; - if (doNotTrack) { - const value = doNotTrack.toLowerCase(); - // Telemetry should be disabled if DO_NOT_TRACK is "1", "true", or "yes" - if (value === "1" || value === "true" || value === "yes") { - return false; - } - } - - return true; - } - /** * Emits events through the telemetry pipeline * @param events - The events to emit */ public async emitEvents(events: BaseEvent[]): Promise { try { - if (!Telemetry.isTelemetryEnabled()) { + if (!this.isTelemetryEnabled()) { + logger.info(LogId.telemetryEmitFailure, "telemetry", `Telemetry is disabled.`); return; } @@ -75,10 +52,27 @@ export class Telemetry { mcp_client_name: this.session.agentRunner?.name, session_id: this.session.sessionId, config_atlas_auth: this.session.apiClient.hasCredentials() ? "true" : "false", - config_connection_string: config.connectionString ? "true" : "false", + config_connection_string: this.userConfig.connectionString ? "true" : "false", }; } + /** + * Checks if telemetry is currently enabled + * This is a method rather than a constant to capture runtime config changes + * + * Follows the Console Do Not Track standard (https://consoledonottrack.com/) + * by respecting the DO_NOT_TRACK environment variable + */ + public isTelemetryEnabled(): boolean { + // Check if telemetry is explicitly disabled in config + if (this.userConfig.telemetry === "disabled") { + return false; + } + + const doNotTrack = "DO_NOT_TRACK" in process.env; + return !doNotTrack; + } + /** * Attempts to emit events through authenticated and unauthenticated clients * Falls back to caching if both attempts fail @@ -96,7 +90,11 @@ export class Telemetry { const result = await this.sendEvents(this.session.apiClient, allEvents); if (result.success) { this.eventCache.clearEvents(); - logger.debug(LogId.telemetryEmitSuccess, "telemetry", `Sent ${allEvents.length} events successfully`); + logger.debug( + LogId.telemetryEmitSuccess, + "telemetry", + `Sent ${allEvents.length} events successfully: ${JSON.stringify(allEvents, null, 2)}` + ); return; } diff --git a/src/tools/atlas/atlasTool.ts b/src/tools/atlas/atlasTool.ts index 6ca5282d..6c74bb88 100644 --- a/src/tools/atlas/atlasTool.ts +++ b/src/tools/atlas/atlasTool.ts @@ -1,4 +1,7 @@ -import { ToolBase, ToolCategory } from "../tool.js"; +import { ToolBase, ToolCategory, TelemetryToolMetadata } from "../tool.js"; +import { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js"; +import logger, { LogId } from "../../logger.js"; +import { z } from "zod"; export abstract class AtlasToolBase extends ToolBase { protected category: ToolCategory = "atlas"; @@ -9,4 +12,46 @@ export abstract class AtlasToolBase extends ToolBase { } return super.verifyAllowed(); } + + /** + * + * Resolves the tool metadata from the arguments passed to the tool + * + * @param args - The arguments passed to the tool + * @returns The tool metadata + */ + protected resolveTelemetryMetadata( + ...args: Parameters> + ): TelemetryToolMetadata { + const toolMetadata: TelemetryToolMetadata = {}; + if (!args.length) { + return toolMetadata; + } + + // Create a typed parser for the exact shape we expect + const argsShape = z.object(this.argsShape); + const parsedResult = argsShape.safeParse(args[0]); + + if (!parsedResult.success) { + logger.debug( + LogId.telemetryMetadataError, + "tool", + `Error parsing tool arguments: ${parsedResult.error.message}` + ); + return toolMetadata; + } + + const data = parsedResult.data; + + // Extract projectId using type guard + if ("projectId" in data && typeof data.projectId === "string" && data.projectId.trim() !== "") { + toolMetadata.projectId = data.projectId; + } + + // Extract orgId using type guard + if ("orgId" in data && typeof data.orgId === "string" && data.orgId.trim() !== "") { + toolMetadata.orgId = data.orgId; + } + return toolMetadata; + } } diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts index d0e59b8b..2ef1aee0 100644 --- a/src/tools/mongodb/mongodbTool.ts +++ b/src/tools/mongodb/mongodbTool.ts @@ -1,5 +1,5 @@ import { z } from "zod"; -import { ToolArgs, ToolBase, ToolCategory } from "../tool.js"; +import { ToolArgs, ToolBase, ToolCategory, TelemetryToolMetadata } from "../tool.js"; import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { ErrorCodes, MongoDBError } from "../../errors.js"; @@ -73,4 +73,18 @@ export abstract class MongoDBToolBase extends ToolBase { protected connectToMongoDB(connectionString: string): Promise { return this.session.connectToMongoDB(connectionString, this.config.connectOptions); } + + protected resolveTelemetryMetadata( + // eslint-disable-next-line @typescript-eslint/no-unused-vars + args: ToolArgs + ): TelemetryToolMetadata { + const metadata: TelemetryToolMetadata = {}; + + // Add projectId to the metadata if running a MongoDB operation to an Atlas cluster + if (this.session.connectedAtlasCluster?.projectId) { + metadata.projectId = this.session.connectedAtlasCluster.projectId; + } + + return metadata; + } } diff --git a/src/tools/tool.ts b/src/tools/tool.ts index d7ea909e..b6b00eda 100644 --- a/src/tools/tool.ts +++ b/src/tools/tool.ts @@ -11,6 +11,10 @@ export type ToolArgs = z.objectOutputType { - const duration = Date.now() - startTime; - const event: ToolEvent = { - timestamp: new Date().toISOString(), - source: "mdbmcp", - properties: { - command: this.name, - category: this.category, - component: "tool", - duration_ms: duration, - result: result.isError ? "failure" : "success", - }, - }; - await this.telemetry.emitEvents([event]); - } - public register(server: McpServer): void { if (!this.verifyAllowed()) { return; @@ -64,12 +46,12 @@ export abstract class ToolBase { logger.debug(LogId.toolExecute, "tool", `Executing ${this.name} with args: ${JSON.stringify(args)}`); const result = await this.execute(...args); - await this.emitToolEvent(startTime, result); + await this.emitToolEvent(startTime, result, ...args).catch(() => {}); return result; } catch (error: unknown) { logger.error(LogId.toolExecuteFailure, "tool", `Error executing ${this.name}: ${error as string}`); const toolResult = await this.handleError(error, args[0] as ToolArgs); - await this.emitToolEvent(startTime, toolResult).catch(() => {}); + await this.emitToolEvent(startTime, toolResult, ...args).catch(() => {}); return toolResult; } }; @@ -149,4 +131,47 @@ export abstract class ToolBase { ], }; } + + protected abstract resolveTelemetryMetadata( + ...args: Parameters> + ): TelemetryToolMetadata; + + /** + * Creates and emits a tool telemetry event + * @param startTime - Start time in milliseconds + * @param result - Whether the command succeeded or failed + * @param args - The arguments passed to the tool + */ + private async emitToolEvent( + startTime: number, + result: CallToolResult, + ...args: Parameters> + ): Promise { + if (!this.telemetry.isTelemetryEnabled()) { + return; + } + const duration = Date.now() - startTime; + const metadata = this.resolveTelemetryMetadata(...args); + const event: ToolEvent = { + timestamp: new Date().toISOString(), + source: "mdbmcp", + properties: { + command: this.name, + category: this.category, + component: "tool", + duration_ms: duration, + result: result.isError ? "failure" : "success", + }, + }; + + if (metadata?.orgId) { + event.properties.org_id = metadata.orgId; + } + + if (metadata?.projectId) { + event.properties.project_id = metadata.projectId; + } + + await this.telemetry.emitEvents([event]); + } } diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts index c57deda8..bacc89b9 100644 --- a/tests/integration/helpers.ts +++ b/tests/integration/helpers.ts @@ -5,6 +5,7 @@ import { UserConfig } from "../../src/config.js"; import { McpError } from "@modelcontextprotocol/sdk/types.js"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { Session } from "../../src/session.js"; +import { config } from "../../src/config.js"; interface ParameterInfo { name: string; @@ -19,6 +20,10 @@ export interface IntegrationTest { mcpClient: () => Client; mcpServer: () => Server; } +export const defaultTestConfig: UserConfig = { + ...config, + telemetry: "disabled", +}; export function setupIntegrationTest(getUserConfig: () => UserConfig): IntegrationTest { let mcpClient: Client | undefined; @@ -51,25 +56,18 @@ export function setupIntegrationTest(getUserConfig: () => UserConfig): Integrati apiClientSecret: userConfig.apiClientSecret, }); - userConfig.telemetry = "disabled"; mcpServer = new Server({ session, userConfig, mcpServer: new McpServer({ name: "test-server", - version: "1.2.3", + version: "5.2.3", }), }); await mcpServer.connect(serverTransport); await mcpClient.connect(clientTransport); }); - beforeEach(() => { - if (mcpServer) { - mcpServer.userConfig.telemetry = "disabled"; - } - }); - afterEach(async () => { if (mcpServer) { await mcpServer.session.close(); diff --git a/tests/integration/server.test.ts b/tests/integration/server.test.ts index aec15add..3b4c1858 100644 --- a/tests/integration/server.test.ts +++ b/tests/integration/server.test.ts @@ -1,5 +1,4 @@ -import { expectDefined, setupIntegrationTest } from "./helpers.js"; -import { config } from "../../src/config.js"; +import { defaultTestConfig, expectDefined, setupIntegrationTest } from "./helpers.js"; import { describeWithMongoDB } from "./tools/mongodb/mongodbHelpers.js"; describe("Server integration test", () => { @@ -16,7 +15,7 @@ describe("Server integration test", () => { }); }, () => ({ - ...config, + ...defaultTestConfig, apiClientId: undefined, apiClientSecret: undefined, }) @@ -24,7 +23,7 @@ describe("Server integration test", () => { describe("with atlas", () => { const integration = setupIntegrationTest(() => ({ - ...config, + ...defaultTestConfig, apiClientId: "test", apiClientSecret: "test", })); @@ -59,7 +58,7 @@ describe("Server integration test", () => { describe("with read-only mode", () => { const integration = setupIntegrationTest(() => ({ - ...config, + ...defaultTestConfig, readOnly: true, apiClientId: "test", apiClientSecret: "test", diff --git a/tests/integration/tools/atlas/atlasHelpers.ts b/tests/integration/tools/atlas/atlasHelpers.ts index d66a4041..aecf0479 100644 --- a/tests/integration/tools/atlas/atlasHelpers.ts +++ b/tests/integration/tools/atlas/atlasHelpers.ts @@ -1,15 +1,14 @@ import { ObjectId } from "mongodb"; import { Group } from "../../../../src/common/atlas/openapi.js"; import { ApiClient } from "../../../../src/common/atlas/apiClient.js"; -import { setupIntegrationTest, IntegrationTest } from "../../helpers.js"; -import { config } from "../../../../src/config.js"; +import { setupIntegrationTest, IntegrationTest, defaultTestConfig } from "../../helpers.js"; export type IntegrationTestFunction = (integration: IntegrationTest) => void; export function describeWithAtlas(name: string, fn: IntegrationTestFunction) { const testDefinition = () => { const integration = setupIntegrationTest(() => ({ - ...config, + ...defaultTestConfig, apiClientId: process.env.MDB_MCP_API_CLIENT_ID, apiClientSecret: process.env.MDB_MCP_API_CLIENT_SECRET, })); diff --git a/tests/integration/tools/mongodb/mongodbHelpers.ts b/tests/integration/tools/mongodb/mongodbHelpers.ts index 2b4ea6a0..087d675c 100644 --- a/tests/integration/tools/mongodb/mongodbHelpers.ts +++ b/tests/integration/tools/mongodb/mongodbHelpers.ts @@ -2,8 +2,8 @@ import { MongoCluster } from "mongodb-runner"; import path from "path"; import fs from "fs/promises"; import { MongoClient, ObjectId } from "mongodb"; -import { getResponseContent, IntegrationTest, setupIntegrationTest } from "../../helpers.js"; -import { config, UserConfig } from "../../../../src/config.js"; +import { getResponseContent, IntegrationTest, setupIntegrationTest, defaultTestConfig } from "../../helpers.js"; +import { UserConfig } from "../../../../src/config.js"; interface MongoDBIntegrationTest { mongoClient: () => MongoClient; @@ -14,7 +14,7 @@ interface MongoDBIntegrationTest { export function describeWithMongoDB( name: string, fn: (integration: IntegrationTest & MongoDBIntegrationTest & { connectMcpClient: () => Promise }) => void, - getUserConfig: (mdbIntegration: MongoDBIntegrationTest) => UserConfig = () => config, + getUserConfig: (mdbIntegration: MongoDBIntegrationTest) => UserConfig = () => defaultTestConfig, describeFn = describe ) { describeFn(name, () => { diff --git a/tests/unit/telemetry.test.ts b/tests/unit/telemetry.test.ts index 5b37da8e..4165b60c 100644 --- a/tests/unit/telemetry.test.ts +++ b/tests/unit/telemetry.test.ts @@ -113,8 +113,7 @@ describe("Telemetry", () => { } as unknown as Session; // Create the telemetry instance with mocked dependencies - telemetry = new Telemetry(session, mockEventCache); - + telemetry = new Telemetry(session, config, mockEventCache); config.telemetry = "enabled"; });