Skip to content

chore: add orgId and projectid #159

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 19 commits into from
Apr 29, 2025
Merged
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
1 change: 1 addition & 0 deletions src/logger.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
2 changes: 1 addition & 1 deletion src/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
56 changes: 27 additions & 29 deletions src/telemetry/telemetry.ts
Original file line number Diff line number Diff line change
@@ -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";
Expand All @@ -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;
}

Expand All @@ -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
Expand All @@ -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;
}

Expand Down
47 changes: 46 additions & 1 deletion src/tools/atlas/atlasTool.ts
Original file line number Diff line number Diff line change
@@ -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";
Expand All @@ -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;
}
}
16 changes: 15 additions & 1 deletion src/tools/mongodb/mongodbTool.ts
Original file line number Diff line number Diff line change
@@ -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";
Expand Down Expand Up @@ -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;
}
}
73 changes: 49 additions & 24 deletions src/tools/tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
}
};
Expand Down Expand Up @@ -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]);
}
}
14 changes: 6 additions & 8 deletions tests/integration/helpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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();
Expand Down
9 changes: 4 additions & 5 deletions tests/integration/server.test.ts
Original file line number Diff line number Diff line change
@@ -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", () => {
Expand All @@ -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",
}));
Expand Down Expand Up @@ -59,7 +58,7 @@ describe("Server integration test", () => {

describe("with read-only mode", () => {
const integration = setupIntegrationTest(() => ({
...config,
...defaultTestConfig,
readOnly: true,
apiClientId: "test",
apiClientSecret: "test",
Expand Down
5 changes: 2 additions & 3 deletions tests/integration/tools/atlas/atlasHelpers.ts
Original file line number Diff line number Diff line change
@@ -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,
}));
Expand Down
Loading
Loading