diff --git a/src/common/session.ts b/src/common/session.ts index e692a7a4..4a8a0213 100644 --- a/src/common/session.ts +++ b/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 { get connectedAtlasCluster(): AtlasClusterConnectionInfo | undefined { return this.connectionManager.currentConnectionState.connectedAtlasCluster; } + + get connectionStringAuthType(): ConnectionStringAuthType | undefined { + return this.connectionManager.currentConnectionState.connectionStringAuthType; + } } diff --git a/src/telemetry/types.ts b/src/telemetry/types.ts index 268c7400..46b91f3b 100644 --- a/src/telemetry/types.ts +++ b/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; +}; diff --git a/src/tools/atlas/atlasTool.ts b/src/tools/atlas/atlasTool.ts index 594d832a..f1d6ee4e 100644 --- a/src/tools/atlas/atlasTool.ts +++ b/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> - ): AtlasToolMetadata { - const toolMetadata: AtlasToolMetadata = {}; + ): AtlasMetadata { + const toolMetadata: AtlasMetadata = {}; if (!args.length) { return toolMetadata; } diff --git a/src/tools/atlas/connect/connectCluster.ts b/src/tools/atlas/connect/connectCluster.ts index 3ba519fc..c12f1709 100644 --- a/src/tools/atlas/connect/connectCluster.ts +++ b/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."; @@ -303,4 +305,24 @@ export class ConnectClusterTool extends AtlasToolBase { return { content }; } + + protected override resolveTelemetryMetadata( + result: CallToolResult, + args: Parameters> + ): ConnectionMetadata { + const parentMetadata = super.resolveTelemetryMetadata(result, ...args); + const connectionMetadata = this.getConnectionInfoMetadata(); + // Explicitly merge, preferring parentMetadata for known overlapping keys (project_id, org_id) + // since parent has more complete information from tool arguments + const { project_id, org_id, ...restConnectionMetadata } = connectionMetadata; + const finalProjectId = parentMetadata.project_id ?? project_id; + const finalOrgId = parentMetadata.org_id ?? org_id; + return { + ...parentMetadata, + ...restConnectionMetadata, + // Only include project_id and org_id if they are defined + ...(finalProjectId !== undefined && { project_id: finalProjectId }), + ...(finalOrgId !== undefined && { org_id: finalOrgId }), + }; + } } diff --git a/src/tools/atlasLocal/atlasLocalTool.ts b/src/tools/atlasLocal/atlasLocalTool.ts index dc548e20..67b66872 100644 --- a/src/tools/atlasLocal/atlasLocalTool.ts +++ b/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 diff --git a/src/tools/atlasLocal/connect/connectDeployment.ts b/src/tools/atlasLocal/connect/connectDeployment.ts index c8523bb1..f183a556 100644 --- a/src/tools/atlasLocal/connect/connectDeployment.ts +++ b/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,21 @@ export class ConnectDeploymentTool extends AtlasLocalToolBase { }, }; } + + protected override resolveTelemetryMetadata(result: CallToolResult): ConnectionMetadata { + const parentMetadata = super.resolveTelemetryMetadata(result); + const connectionMetadata = this.getConnectionInfoMetadata(); + // Explicitly merge, preferring parentMetadata for known overlapping keys (project_id, org_id) + // since parent has deployment-specific information + const { project_id, org_id, ...restConnectionMetadata } = connectionMetadata; + const finalProjectId = parentMetadata.project_id ?? project_id; + const finalOrgId = parentMetadata.org_id ?? org_id; + return { + ...parentMetadata, + ...restConnectionMetadata, + // Only include project_id and org_id if they are defined + ...(finalProjectId !== undefined && { project_id: finalProjectId }), + ...(finalOrgId !== undefined && { org_id: finalOrgId }), + }; + } } diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts index 735579a7..e28338b3 100644 --- a/src/tools/mongodb/mongodbTool.ts +++ b/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 - ): 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> + ): ConnectionMetadata { + return this.getConnectionInfoMetadata(); } } diff --git a/src/tools/tool.ts b/src/tools/tool.ts index 60cc9a04..ac650215 100644 --- a/src/tools/tool.ts +++ b/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; + } } /** diff --git a/tests/integration/tools/mongodb/mongodbTool.test.ts b/tests/integration/tools/mongodb/mongodbTool.test.ts index ca3bc423..6692c6ef 100644 --- a/tests/integration/tools/mongodb/mongodbTool.test.ts +++ b/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"; @@ -327,4 +327,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"); + }); + }); }); diff --git a/tests/unit/toolBase.test.ts b/tests/unit/toolBase.test.ts index a71c483b..e6d7c58e 100644 --- a/tests/unit/toolBase.test.ts +++ b/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 {