Skip to content
Open
Show file tree
Hide file tree
Changes from all 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
22 changes: 22 additions & 0 deletions src/tools/atlas/connect/connectCluster.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.";
Expand Down Expand Up @@ -303,4 +305,24 @@ 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();
// 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 }),
};
Comment on lines +313 to +326
Copy link
Collaborator

Choose a reason for hiding this comment

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

This bit is duplicated, maybe we can extract this to a small function and re-use?

}
}
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
18 changes: 18 additions & 0 deletions src/tools/atlasLocal/connect/connectDeployment.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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 }),
};
}
}
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
40 changes: 39 additions & 1 deletion tests/integration/tools/mongodb/mongodbTool.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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");
});
});
});
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