chore: add orgId and projectid (#159)

blva · web-flow · commit aaf96d1d1e12 · 2025-04-29T16:47:25.000+01:00
diff --git 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
@@ -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
@@ -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,45 +16,22 @@ export class Telemetry {
 
     constructor(
         private readonly session: Session,
+        private readonly userConfig: UserConfig,
         private readonly eventCache: EventCache = EventCache.getInstance()
     ) {
         this.commonProperties = {
             ...MACHINE_METADATA,
         };
     }
 
-    /**
-     * 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<void> {
         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
@@ -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<ToolCallback<typeof this.argsShape>>
+    ): 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
@@ -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<void> {
         return this.session.connectToMongoDB(connectionString, this.config.connectOptions);
     }
+
+    protected resolveTelemetryMetadata(
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        args: ToolArgs<typeof this.argsShape>
+    ): 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
@@ -11,6 +11,10 @@ export type ToolArgs<Args extends ZodRawShape> = z.objectOutputType<Args, ZodNev
 
 export type OperationType = "metadata" | "read" | "create" | "delete" | "update";
 export type ToolCategory = "mongodb" | "atlas";
+export type TelemetryToolMetadata = {
+    projectId?: string;
+    orgId?: string;
+};
 
 export abstract class ToolBase {
     protected abstract name: string;
@@ -31,28 +35,6 @@ export abstract class ToolBase {
         protected readonly telemetry: Telemetry
     ) {}
 
-    /**
-     * Creates and emits a tool telemetry event
-     * @param startTime - Start time in milliseconds
-     * @param result - Whether the command succeeded or failed
-     * @param error - Optional error if the command failed
-     */
-    private async emitToolEvent(startTime: number, result: CallToolResult): Promise<void> {
-        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<typeof this.argsShape>);
-                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<ToolCallback<typeof this.argsShape>>
+    ): 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<ToolCallback<typeof this.argsShape>>
+    ): Promise<void> {
+        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
@@ -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
@@ -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,15 +15,15 @@ describe("Server integration test", () => {
             });
         },
         () => ({
-            ...config,
+            ...defaultTestConfig,
             apiClientId: undefined,
             apiClientSecret: undefined,
         })
     );
 
     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
@@ -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
diff --git a/tests/unit/telemetry.test.ts b/tests/unit/telemetry.test.ts

Original file line number	Diff line number	Diff line change
`@@ -28,7 +28,7 @@ export class Server {`
`28`	`28`	`constructor({ session, mcpServer, userConfig }: ServerOptions) {`
`29`	`29`	`this.startTime = Date.now();`
`30`	`30`	`this.session = session;`
`31`		`- this.telemetry = new Telemetry(session);`
	`31`	`+ this.telemetry = new Telemetry(session, userConfig);`
`32`	`32`	`this.mcpServer = mcpServer;`
`33`	`33`	`this.userConfig = userConfig;`
`34`	`34`	`}`