Skip to content

Commit af714b4

Browse files
blvablva
committed
wip
1 parent 1b8ea02 commit af714b4

File tree

9 files changed

+211
-23
lines changed

9 files changed

+211
-23
lines changed

‎package-lock.json‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@
6464
"mongodb-log-writer": "^2.4.1",
6565
"mongodb-redact": "^1.1.6",
6666
"mongodb-schema": "^12.6.2",
67+
"node-machine-id": "^1.1.12",
6768
"openapi-fetch": "^0.13.5",
6869
"simple-oauth2": "^5.1.0",
6970
"yargs-parser": "^21.1.1",

‎src/common/atlas/apiClient.ts‎

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import type { FetchOptions } from "openapi-fetch";
44
import { AccessToken, ClientCredentials } from "simple-oauth2";
55
import { ApiClientError } from "./apiClientError.js";
66
import { paths, operations } from "./openapi.js";
7+
import { BaseEvent } from "../../telemetry/types.js";
78

89
const ATLAS_API_VERSION = "2025-03-12";
910

@@ -63,10 +64,10 @@ export class ApiClient {
6364
constructor(options?: ApiClientOptions) {
6465
this.options = {
6566
...options,
66-
baseUrl: options?.baseUrl || "https://cloud.mongodb.com/",
67+
baseUrl: options?.baseUrl || "https://cloud-dev.mongodb.com/",
6768
userAgent:
6869
options?.userAgent ||
69-
`${config.mcp_server_name}/${config.version} (${process.platform}; ${process.arch}; ${process.env.HOSTNAME || "unknown"})`,
70+
`${config.mcpServerName}/${config.version} (${process.platform}; ${process.arch}; ${process.env.HOSTNAME || "unknown"})`,
7071
};
7172

7273
this.client = createClient<paths>({
@@ -117,6 +118,32 @@ export class ApiClient {
117118
}>;
118119
}
119120

121+
async sendEvents(events: BaseEvent[]): Promise<void> {
122+
let endpoint = "api/private/unauth/telemetry/events";
123+
const headers: Record<string, string> = {
124+
Accept: "application/json",
125+
"Content-Type": "application/json",
126+
"User-Agent": this.options.userAgent,
127+
};
128+
129+
const accessToken = await this.getAccessToken();
130+
if (accessToken) {
131+
endpoint = "api/private/v1.0/telemetry/events";
132+
headers["Authorization"] = `Bearer ${accessToken}`;
133+
}
134+
135+
const url = new URL(endpoint, this.options.baseUrl);
136+
const response = await fetch(url, {
137+
method: "POST",
138+
headers,
139+
body: JSON.stringify(events),
140+
});
141+
142+
if (!response.ok) {
143+
throw await ApiClientError.fromResponse(response);
144+
}
145+
}
146+
120147
// DO NOT EDIT. This is auto-generated code.
121148
async listClustersForAllProjects(options?: FetchOptions<operations["listClustersForAllProjects"]>) {
122149
const { data } = await this.client.GET("/api/atlas/v2/clusters", options);

‎src/config.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ interface UserConfig {
1111
apiBaseUrl?: string;
1212
apiClientId?: string;
1313
apiClientSecret?: string;
14+
telemetry?: 'enabled' | 'disabled';
1415
logPath: string;
1516
connectionString?: string;
1617
connectOptions: {
@@ -52,7 +53,8 @@ const config = {
5253
...mergedUserConfig,
5354
...machineMetadata,
5455
version: packageJson.version,
55-
mcp_server_name: "MdbMcpServer"
56+
mcpServerName: "MdbMcpServer",
57+
isTelemetryEnabled: true
5658
};
5759

5860
export default config;

‎src/server.ts‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,11 +25,8 @@ export class Server {
2525
await this.mcpServer.connect(transport);
2626

2727
this.mcpServer.server.oninitialized = () => {
28-
const client = this.mcpServer.server.getClientVersion();
29-
this.session.clientName = client?.name;
30-
this.session.clientVersion = client?.version;
31-
32-
logger.info(mongoLogId(1_000_004), "server", `Server started with transport ${transport.constructor.name}`);
28+
this.session.setAgentClientData(this.mcpServer.server.getClientVersion());
29+
logger.info(mongoLogId(1_000_004), "server", `Server started with transport ${transport.constructor.name}`);
3330
};
3431
}
3532

‎src/session.ts‎

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,32 @@ import { ApiClient } from "./common/atlas/apiClient.js";
33
import config from "./config.js";
44
import logger from "./logger.js";
55
import { mongoLogId } from "mongodb-log-writer";
6+
import { Implementation } from "@modelcontextprotocol/sdk/types.js";
67

78
export class Session {
89
sessionId?: string;
910
serviceProvider?: NodeDriverServiceProvider;
1011
apiClient?: ApiClient;
11-
clientName?: string;
12-
clientVersion?: string;
12+
agentClientName?: string;
13+
agentClientVersion?: string;
14+
15+
constructor() {
16+
// configure api client if credentials are set
17+
if (config.apiClientId && config.apiClientSecret) {
18+
this.apiClient = new ApiClient({
19+
baseUrl: config.apiBaseUrl,
20+
credentials: {
21+
clientId: config.apiClientId,
22+
clientSecret: config.apiClientSecret,
23+
},
24+
});
25+
}
26+
}
27+
28+
setAgentClientData(agentClient: Implementation | undefined) {
29+
this.agentClientName = agentClient?.name;
30+
this.agentClientVersion = agentClient?.version;
31+
}
1332

1433
ensureAuthenticated(): asserts this is { apiClient: ApiClient } {
1534
if (!this.apiClient) {
@@ -39,12 +58,4 @@ export class Session {
3958
this.serviceProvider = undefined;
4059
}
4160
}
42-
43-
async emitTelemetry(todo: unknown): Promise<void> {
44-
logger.info(
45-
mongoLogId(1_000_001),
46-
"telemetry",
47-
`Telemetry event: ${JSON.stringify(todo)}`
48-
);
49-
}
5061
}

‎src/telemetry/telemetry.ts‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
import { Session } from '../session.js';
2+
import { BaseEvent, type ToolEvent } from './types.js';
3+
import pkg from '../../package.json' with { type: 'json' };
4+
import config from '../config.js';
5+
import logger from '../logger.js';
6+
import { mongoLogId } from 'mongodb-log-writer';
7+
import { ApiClient } from '../common/atlas/apiClient.js';
8+
import { ApiClientError } from '../common/atlas/apiClientError.js';
9+
10+
export class Telemetry {
11+
constructor(private readonly session: Session) {}
12+
13+
private readonly commonProperties = {
14+
mcp_server_version: pkg.version,
15+
mcp_server_name: config.mcpServerName,
16+
mcp_client_version: this.session.agentClientVersion,
17+
mcp_client_name: this.session.agentClientName,
18+
session_id: this.session.sessionId,
19+
device_id: config.device_id,
20+
platform: config.platform,
21+
arch: config.arch,
22+
os_type: config.os_type,
23+
os_version: config.os_version,
24+
};
25+
26+
private readonly isTelemetryEnabled = config.telemetry === 'enabled';
27+
28+
async emitToolEvent(command: string, category: string, startTime: number, result: 'success' | 'failure', error?: Error): Promise<void> {
29+
if (!this.isTelemetryEnabled) {
30+
logger.debug(mongoLogId(1_000_000), "telemetry", `Telemetry is disabled, skipping event.`);
31+
return;
32+
}
33+
34+
const duration = Date.now() - startTime;
35+
36+
const event: ToolEvent = {
37+
timestamp: new Date().toISOString(),
38+
source: 'mdbmcp',
39+
properties: {
40+
...this.commonProperties,
41+
command: command,
42+
category: category,
43+
duration_ms: duration,
44+
result: result
45+
}
46+
};
47+
48+
if (result === 'failure') {
49+
event.properties.error_type = error?.name;
50+
event.properties.error_code = error?.message;
51+
}
52+
53+
await this.emit(event);
54+
}
55+
56+
private async emit(event: BaseEvent): Promise<void> {
57+
try {
58+
if (this.session.apiClient) {
59+
await this.session.apiClient.sendEvents([event]);
60+
}
61+
} catch (error) {
62+
logger.warning(mongoLogId(1_000_000), "telemetry", `Error sending event to authenticated client: ${error}`);
63+
}
64+
65+
// if it is unauthenticated, send to temp client
66+
try {
67+
const tempApiClient = new ApiClient({
68+
baseUrl: config.apiBaseUrl,
69+
});
70+
await tempApiClient.sendEvents([event]);
71+
} catch (error) {
72+
logger.warning(mongoLogId(1_000_000), "telemetry", `Error sending event to unauthenticated client: ${error}`);
73+
}
74+
}
75+
}

‎src/telemetry/types.ts‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
/**
2+
* Base interface for all events
3+
*/
4+
export interface Event {
5+
timestamp: string;
6+
source: 'mdbmcp';
7+
properties: Record<string, unknown>;
8+
}
9+
10+
export interface BaseEvent extends Event {
11+
properties: {
12+
device_id: string;
13+
mcp_server_version: string;
14+
mcp_server_name: string;
15+
mcp_client_version?: string;
16+
mcp_client_name?: string;
17+
platform: string;
18+
arch: string;
19+
os_type: string;
20+
os_version?: string;
21+
session_id?: string;
22+
} & Event['properties'];
23+
}
24+
25+
/**
26+
* Interface for tool events
27+
*/
28+
export interface ToolEvent extends BaseEvent {
29+
properties: {
30+
command: string;
31+
category: string;
32+
duration_ms: number;
33+
result: 'success' | 'failure';
34+
error_code?: string;
35+
error_type?: string;
36+
project_id?: string;
37+
org_id?: string;
38+
cluster_name?: string;
39+
is_atlas?: boolean;
40+
} & BaseEvent['properties'];
41+
}

‎src/tools/tool.ts‎

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,22 @@
1-
import { McpServer, ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
2-
import { z, ZodNever, ZodRawShape } from "zod";
3-
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
1+
import { z, type ZodRawShape, type ZodNever } from "zod";
2+
import type { McpServer, ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
3+
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
44
import { Session } from "../session.js";
55
import logger from "../logger.js";
66
import { mongoLogId } from "mongodb-log-writer";
7+
<<<<<<< HEAD
78
import config from "../config.js";
9+
=======
10+
import { Telemetry } from "../telemetry/telemetry.js";
11+
>>>>>>> 61e295a (wip)
812

913
export type ToolArgs<Args extends ZodRawShape> = z.objectOutputType<Args, ZodNever>;
1014

1115
export type OperationType = "metadata" | "read" | "create" | "delete" | "update" | "cluster";
1216
export type ToolCategory = "mongodb" | "atlas";
1317

1418
export abstract class ToolBase {
19+
<<<<<<< HEAD
1520
protected abstract name: string;
1621

1722
protected abstract category: ToolCategory;
@@ -21,37 +26,57 @@ export abstract class ToolBase {
2126
protected abstract description: string;
2227

2328
protected abstract argsShape: ZodRawShape;
29+
=======
30+
protected abstract readonly name: string;
31+
protected abstract readonly description: string;
32+
protected abstract readonly argsShape: ZodRawShape;
33+
>>>>>>> 61e295a (wip)
2434

2535
protected abstract category: string;
36+
private readonly telemetry: Telemetry;
2637

2738
protected abstract execute(...args: Parameters<ToolCallback<typeof this.argsShape>>): Promise<CallToolResult>;
2839

29-
protected constructor(protected session: Session) {}
40+
protected constructor(protected session: Session) {
41+
this.telemetry = new Telemetry(session);
42+
}
3043

3144
public register(server: McpServer): void {
3245
if (!this.verifyAllowed()) {
3346
return;
3447
}
3548

3649
const callback: ToolCallback<typeof this.argsShape> = async (...args) => {
50+
const startTime = Date.now();
3751
try {
3852
logger.debug(
3953
mongoLogId(1_000_006),
4054
"tool",
4155
`Executing ${this.name} with args: ${JSON.stringify(args)}`
4256
);
4357

44-
return await this.execute(...args);
58+
const result = await this.execute(...args);
59+
await this.telemetry.emitToolEvent(this.name, this.category, startTime, "success");
60+
return result;
4561
} catch (error: unknown) {
4662
logger.error(mongoLogId(1_000_000), "tool", `Error executing ${this.name}: ${error as string}`);
4763

64+
await this.telemetry.emitToolEvent(
65+
this.name,
66+
this.category,
67+
startTime,
68+
"failure",
69+
error instanceof Error ? error : new Error(String(error))
70+
);
71+
4872
return await this.handleError(error);
4973
}
5074
};
5175

5276
server.tool(this.name, this.description, this.argsShape, callback);
5377
}
5478

79+
<<<<<<< HEAD
5580
// Checks if a tool is allowed to run based on the config
5681
private verifyAllowed(): boolean {
5782
let errorClarification: string | undefined;
@@ -76,6 +101,8 @@ export abstract class ToolBase {
76101
return true;
77102
}
78103

104+
=======
105+
>>>>>>> 61e295a (wip)
79106
// This method is intended to be overridden by subclasses to handle errors
80107
protected handleError(error: unknown): Promise<CallToolResult> | CallToolResult {
81108
return {

0 commit comments

Comments
 (0)