chore: add connection metadata to telemetry (#716)

Author: blva

src/common/session.ts

@@ -11,6 +11,7 @@ import type {
     ConnectionSettings,
     ConnectionStateConnected,
     ConnectionStateErrored,
+    ConnectionStringAuthType,
 } from "./connectionManager.js";
 import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { ErrorCodes, MongoDBError } from "./errors.js";
@@ -182,4 +183,8 @@ export class Session extends EventEmitter<SessionEvents> {
     get connectedAtlasCluster(): AtlasClusterConnectionInfo | undefined {
         return this.connectionManager.currentConnectionState.connectedAtlasCluster;
     }
+
+    get connectionStringAuthType(): ConnectionStringAuthType | undefined {
+        return this.connectionManager.currentConnectionState.connectionStringAuthType;
+    }
 }

src/telemetry/types.ts

@@ -141,17 +141,22 @@ 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 = AtlasLocalToolMetadata | AtlasToolMetadata | PerfAdvisorToolMetadata;
+export type TelemetryToolMetadata = AtlasMetadata | PerfAdvisorToolMetadata | ConnectionMetadata;
 
-export type AtlasLocalToolMetadata = {
-    atlas_local_deployment_id?: string;
-};
-
-export type AtlasToolMetadata = {
+export type AtlasMetadata = {
     project_id?: string;
     org_id?: string;
 };
 
-export type PerfAdvisorToolMetadata = AtlasToolMetadata & {
+export type PerfAdvisorToolMetadata = AtlasMetadata & {
     operations: string[];
 };
+
+export type ConnectionMetadata = AtlasMetadata &
+    AtlasLocalToolMetadata & {
+        connection_auth_type?: string;
+    };
+
+type AtlasLocalToolMetadata = {
+    atlas_local_deployment_id?: string;
+};

src/tools/atlas/atlasTool.ts

@@ -1,5 +1,5 @@
 import type { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
-import type { AtlasToolMetadata } from "../../telemetry/types.js";
+import type { AtlasMetadata } from "../../telemetry/types.js";
 import { ToolBase, type ToolArgs, type ToolCategory } from "../tool.js";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { LogId } from "../../common/logger.js";
@@ -85,8 +85,8 @@ For more information on Atlas API access roles, visit: https://www.mongodb.com/d
     protected resolveTelemetryMetadata(
         result: CallToolResult,
         ...args: Parameters<ToolCallback<typeof this.argsShape>>
-    ): AtlasToolMetadata {
-        const toolMetadata: AtlasToolMetadata = {};
+    ): AtlasMetadata {
+        const toolMetadata: AtlasMetadata = {};
         if (!args.length) {
             return toolMetadata;
         }

src/tools/atlas/connect/connectCluster.ts

@@ -8,6 +8,8 @@ import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUti
 import type { AtlasClusterConnectionInfo } from "../../../common/connectionManager.js";
 import { getDefaultRoleFromConfig } from "../../../common/atlas/roles.js";
 import { AtlasArgs } from "../../args.js";
+import type { ConnectionMetadata } from "../../../telemetry/types.js";
+import type { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
 
 const addedIpAccessListMessage =
     "Note: Your current IP address has been added to the Atlas project's IP access list to enable secure connection.";
@@ -317,4 +319,17 @@ export class ConnectClusterTool extends AtlasToolBase {
 
         return { content };
     }
+
+    protected override resolveTelemetryMetadata(
+        result: CallToolResult,
+        ...args: Parameters<ToolCallback<typeof this.argsShape>>
+    ): ConnectionMetadata {
+        const parentMetadata = super.resolveTelemetryMetadata(result, ...args);
+        const connectionMetadata = this.getConnectionInfoMetadata();
+        if (connectionMetadata && connectionMetadata.project_id !== undefined) {
+            // delete the project_id from the parent metadata to avoid duplication
+            delete parentMetadata.project_id;
+        }
+        return { ...parentMetadata, ...connectionMetadata };
+    }
 }

src/tools/atlasLocal/atlasLocalTool.ts

@@ -4,7 +4,7 @@ import { ToolBase } from "../tool.js";
 import type { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
 import type { Client } from "@mongodb-js/atlas-local";
 import { LogId } from "../../common/logger.js";
-import type { AtlasLocalToolMetadata } from "../../telemetry/types.js";
+import type { ConnectionMetadata } from "../../telemetry/types.js";
 
 export const AtlasLocalToolMetadataDeploymentIdKey = "deploymentId";
 
@@ -119,8 +119,8 @@ please log a ticket here: https://github.com/mongodb-js/mongodb-mcp-server/issue
         return super.handleError(error, args);
     }
 
-    protected resolveTelemetryMetadata(result: CallToolResult): AtlasLocalToolMetadata {
-        const toolMetadata: AtlasLocalToolMetadata = {};
+    protected resolveTelemetryMetadata(result: CallToolResult): ConnectionMetadata {
+        const toolMetadata: ConnectionMetadata = {};
 
         // Atlas Local tools set the deployment ID in the result metadata for telemetry
         // If the deployment ID is set, we use it for telemetry

src/tools/atlasLocal/connect/connectDeployment.ts

@@ -3,6 +3,7 @@ import { AtlasLocalToolBase } from "../atlasLocalTool.js";
 import type { OperationType, ToolArgs } from "../../tool.js";
 import type { Client } from "@mongodb-js/atlas-local";
 import { CommonArgs } from "../../args.js";
+import type { ConnectionMetadata } from "../../../telemetry/types.js";
 
 export class ConnectDeploymentTool extends AtlasLocalToolBase {
     public name = "atlas-local-connect-deployment";
@@ -34,4 +35,8 @@ export class ConnectDeploymentTool extends AtlasLocalToolBase {
             },
         };
     }
+
+    protected override resolveTelemetryMetadata(result: CallToolResult): ConnectionMetadata {
+        return { ...super.resolveTelemetryMetadata(result), ...this.getConnectionInfoMetadata() };
+    }
 }

src/tools/mongodb/mongodbTool.ts

@@ -6,7 +6,8 @@ import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { ErrorCodes, MongoDBError } from "../../common/errors.js";
 import { LogId } from "../../common/logger.js";
 import type { Server } from "../../server.js";
-import type { AtlasToolMetadata } from "../../telemetry/types.js";
+import type { ConnectionMetadata } from "../../telemetry/types.js";
+import type { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
 
 export const DbOperationArgs = {
     database: z.string().describe("Database name"),
@@ -111,19 +112,21 @@ export abstract class MongoDBToolBase extends ToolBase {
         return this.session.connectToMongoDB({ connectionString });
     }
 
+    /**
+     * Resolves the tool metadata from the arguments passed to the mongoDB tools.
+     *
+     * Since MongoDB tools are executed against a MongoDB instance, the tool calls will always have the connection information.
+     *
+     * @param result - The result of the tool call.
+     * @param args - The arguments passed to the tool
+     * @returns The tool metadata
+     */
     protected resolveTelemetryMetadata(
         // eslint-disable-next-line @typescript-eslint/no-unused-vars
-        result: CallToolResult,
+        _result: CallToolResult,
         // eslint-disable-next-line @typescript-eslint/no-unused-vars
-        args: ToolArgs<typeof this.argsShape>
-    ): AtlasToolMetadata {
-        const metadata: AtlasToolMetadata = {};
-
-        // Add projectId to the metadata if running a MongoDB operation to an Atlas cluster
-        if (this.session.connectedAtlasCluster?.projectId) {
-            metadata.project_id = this.session.connectedAtlasCluster.projectId;
-        }
-
-        return metadata;
+        _args: Parameters<ToolCallback<typeof this.argsShape>>
+    ): ConnectionMetadata {
+        return this.getConnectionInfoMetadata();
     }
 }

src/tools/tool.ts

@@ -5,7 +5,7 @@ import type { CallToolResult, ToolAnnotations } from "@modelcontextprotocol/sdk/
 import type { Session } from "../common/session.js";
 import { LogId } from "../common/logger.js";
 import type { Telemetry } from "../telemetry/telemetry.js";
-import type { TelemetryToolMetadata, ToolEvent } from "../telemetry/types.js";
+import type { ConnectionMetadata, TelemetryToolMetadata, ToolEvent } from "../telemetry/types.js";
 import type { UserConfig } from "../common/config.js";
 import type { Server } from "../server.js";
 import type { Elicitation } from "../elicitation.js";
@@ -303,6 +303,20 @@ export abstract class ToolBase {
     protected isFeatureEnabled(feature: PreviewFeature): boolean {
         return this.config.previewFeatures.includes(feature);
     }
+
+    protected getConnectionInfoMetadata(): ConnectionMetadata {
+        const metadata: ConnectionMetadata = {};
+        if (this.session.connectedAtlasCluster?.projectId) {
+            metadata.project_id = this.session.connectedAtlasCluster.projectId;
+        }
+
+        const connectionStringAuthType = this.session.connectionStringAuthType;
+        if (connectionStringAuthType !== undefined) {
+            metadata.connection_auth_type = connectionStringAuthType;
+        }
+
+        return metadata;
+    }
 }
 
 /**

tests/integration/tools/mongodb/mongodbTool.test.ts

@@ -14,7 +14,7 @@ import { InMemoryTransport } from "../../inMemoryTransport.js";
 import { Telemetry } from "../../../../src/telemetry/telemetry.js";
 import { Server } from "../../../../src/server.js";
 import { type ConnectionErrorHandler, connectionErrorHandler } from "../../../../src/common/connectionErrorHandler.js";
-import { defaultTestConfig } from "../../helpers.js";
+import { defaultTestConfig, expectDefined } from "../../helpers.js";
 import { setupMongoDBIntegrationTest } from "./mongodbHelpers.js";
 import { ErrorCodes } from "../../../../src/common/errors.js";
 import { Keychain } from "../../../../src/common/keychain.js";
@@ -325,4 +325,42 @@ describe("MongoDBTool implementations", () => {
             expect(tools?.tools.find((tool) => tool.name === "UnusableVoyageTool")).toBeUndefined();
         });
     });
+
+    describe("resolveTelemetryMetadata", () => {
+        it("should return empty metadata when not connected", async () => {
+            await cleanupAndStartServer();
+            const tool = mcpServer?.tools.find((t) => t.name === "Random");
+            expectDefined(tool);
+            const randomTool = tool as RandomTool;
+
+            const result: CallToolResult = { content: [{ type: "text", text: "test" }] };
+            const metadata = randomTool["resolveTelemetryMetadata"](result, {} as never);
+
+            expect(metadata).toEqual({});
+            expect(metadata).not.toHaveProperty("project_id");
+            expect(metadata).not.toHaveProperty("connection_auth_type");
+        });
+
+        it("should return metadata with connection_auth_type when connected via connection string", async () => {
+            await cleanupAndStartServer({ connectionString: mdbIntegration.connectionString() });
+            // Connect to MongoDB to set the connection state
+            await mcpClient?.callTool({
+                name: "Random",
+                arguments: {},
+            });
+
+            const tool = mcpServer?.tools.find((t) => t.name === "Random");
+            expectDefined(tool);
+            const randomTool = tool as RandomTool;
+
+            const result: CallToolResult = { content: [{ type: "text", text: "test" }] };
+            const metadata = randomTool["resolveTelemetryMetadata"](result, {} as never);
+
+            // When connected via connection string, connection_auth_type should be set
+            // The actual value depends on the connection string, but it should be present
+            expect(metadata).toHaveProperty("connection_auth_type");
+            expect(typeof metadata.connection_auth_type).toBe("string");
+            expect(metadata.connection_auth_type).toBe("scram");
+        });
+    });
 });

tests/unit/toolBase.test.ts

@@ -175,6 +175,90 @@ describe("ToolBase", () => {
             expect(event.properties).toHaveProperty("test_param2", "three");
         });
     });
+
+    describe("getConnectionInfoMetadata", () => {
+        it("should return empty metadata when neither connectedAtlasCluster nor connectionStringAuthType are set", () => {
+            (mockSession as { connectedAtlasCluster?: unknown }).connectedAtlasCluster = undefined;
+            (mockSession as { connectionStringAuthType?: unknown }).connectionStringAuthType = undefined;
+
+            const metadata = testTool["getConnectionInfoMetadata"]();
+
+            expect(metadata).toEqual({});
+            expect(metadata).not.toHaveProperty("project_id");
+            expect(metadata).not.toHaveProperty("connection_auth_type");
+        });
+
+        it("should return metadata with project_id when connectedAtlasCluster.projectId is set", () => {
+            (mockSession as { connectedAtlasCluster?: unknown }).connectedAtlasCluster = {
+                projectId: "test-project-id",
+                username: "test-user",
+                clusterName: "test-cluster",
+                expiryDate: new Date(),
+            };
+            (mockSession as { connectionStringAuthType?: unknown }).connectionStringAuthType = undefined;
+
+            const metadata = testTool["getConnectionInfoMetadata"]();
+
+            expect(metadata).toEqual({
+                project_id: "test-project-id",
+            });
+            expect(metadata).not.toHaveProperty("connection_auth_type");
+        });
+
+        it("should return empty metadata when connectedAtlasCluster exists but projectId is falsy", () => {
+            (mockSession as { connectedAtlasCluster?: unknown }).connectedAtlasCluster = {
+                projectId: "",
+                username: "test-user",
+                clusterName: "test-cluster",
+                expiryDate: new Date(),
+            };
+            (mockSession as { connectionStringAuthType?: unknown }).connectionStringAuthType = undefined;
+
+            const metadata = testTool["getConnectionInfoMetadata"]();
+
+            expect(metadata).toEqual({});
+            expect(metadata).not.toHaveProperty("project_id");
+        });
+
+        it("should return metadata with connection_auth_type when connectionStringAuthType is set", () => {
+            (mockSession as { connectedAtlasCluster?: unknown }).connectedAtlasCluster = undefined;
+            (mockSession as { connectionStringAuthType?: unknown }).connectionStringAuthType = "scram";
+
+            const metadata = testTool["getConnectionInfoMetadata"]();
+
+            expect(metadata).toEqual({
+                connection_auth_type: "scram",
+            });
+            expect(metadata).not.toHaveProperty("project_id");
+        });
+
+        it("should return metadata with both project_id and connection_auth_type when both are set", () => {
+            (mockSession as { connectedAtlasCluster?: unknown }).connectedAtlasCluster = {
+                projectId: "test-project-id",
+                username: "test-user",
+                clusterName: "test-cluster",
+                expiryDate: new Date(),
+            };
+            (mockSession as { connectionStringAuthType?: unknown }).connectionStringAuthType = "oidc-auth-flow";
+
+            const metadata = testTool["getConnectionInfoMetadata"]();
+
+            expect(metadata).toEqual({
+                project_id: "test-project-id",
+                connection_auth_type: "oidc-auth-flow",
+            });
+        });
+
+        it("should handle different connectionStringAuthType values", () => {
+            const authTypes = ["scram", "ldap", "kerberos", "oidc-auth-flow", "oidc-device-flow", "x.509"] as const;
+
+            for (const authType of authTypes) {
+                (mockSession as { connectionStringAuthType?: unknown }).connectionStringAuthType = authType;
+                const metadata = testTool["getConnectionInfoMetadata"]();
+                expect(metadata.connection_auth_type).toBe(authType);
+            }
+        });
+    });
 });
 
 class TestTool extends ToolBase {