Skip to content
Open
Show file tree
Hide file tree
Changes from 7 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
5 changes: 5 additions & 0 deletions src/common/session.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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;
}
}
19 changes: 12 additions & 7 deletions src/telemetry/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

planning to add host info here - starting with connection auth type as an example of something we can capture

};

type AtlasLocalToolMetadata = {
atlas_local_deployment_id?: string;
};
6 changes: 3 additions & 3 deletions src/tools/atlas/atlasTool.ts
Original file line number Diff line number Diff line change
@@ -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";
Expand Down Expand Up @@ -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;
}
Expand Down
16 changes: 16 additions & 0 deletions src/tools/atlas/connect/connectCluster.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ 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 { ConnectionMetadata } from "../../../telemetry/types.js";

const addedIpAccessListMessage =
"Note: Your current IP address has been added to the Atlas project's IP access list to enable secure connection.";
Expand Down Expand Up @@ -303,4 +304,19 @@ export class ConnectClusterTool extends AtlasToolBase {

return { content };
}

// eslint-disable-next-line @typescript-eslint/no-unused-vars
protected override resolveTelemetryMetadata(_args: ToolArgs<typeof this.argsShape>): ConnectionMetadata {
const metadata: ConnectionMetadata = {};
const connectionStringAuthType = this.session.connectionManager.currentConnectionState.connectionStringAuthType;
if (connectionStringAuthType) {
metadata.connection_auth_type = connectionStringAuthType;
}

if (this.session.connectedAtlasCluster?.projectId) {
metadata.project_id = this.session.connectedAtlasCluster.projectId;
}

return metadata;
}
}
6 changes: 3 additions & 3 deletions src/tools/atlasLocal/atlasLocalTool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions src/tools/mongodb/connect/connect.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import type { ToolArgs, OperationType, ToolConstructorParams } from "../../tool.
import assert from "assert";
import type { Server } from "../../../server.js";
import { LogId } from "../../../common/logger.js";
import { ConnectionMetadata } from "../../../telemetry/types.js";

const disconnectedSchema = z
.object({
Expand Down
27 changes: 15 additions & 12 deletions src/tools/mongodb/mongodbTool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
Expand Down Expand Up @@ -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();
}
}
16 changes: 15 additions & 1 deletion src/tools/tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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;
}
}

/**
Expand Down
84 changes: 84 additions & 0 deletions tests/unit/toolBase.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Loading