Skip to content

feat: core telemetry functionality #87

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 52 commits into from
Apr 23, 2025
Merged
Show file tree
Hide file tree
Changes from 51 commits
Commits
Show all changes
52 commits
Select commit Hold shift + click to select a range
9cf0447
wip
blva Apr 16, 2025
1b8ea02
WIP
blva Apr 16, 2025
af714b4
wip
blva Apr 16, 2025
40aaef0
wip: simple approach
blva Apr 16, 2025
8adc20d
add cache and client logic
blva Apr 17, 2025
cff01e0
Merge remote-tracking branch 'origin/main' into telemetry
blva Apr 22, 2025
5336e4a
update deps
blva Apr 22, 2025
76e2a6f
update
blva Apr 22, 2025
1850676
update logs
blva Apr 22, 2025
3a5d4c1
reformat
blva Apr 22, 2025
23fff05
address comment: server main
blva Apr 22, 2025
754ce8d
address comment: make dynamic config part of session
blva Apr 22, 2025
8301e43
address comment: remove redundant method
blva Apr 22, 2025
994b698
address comment: move tool emission to tool.ts
blva Apr 22, 2025
3bfc9f2
address comment: update agentClient to agentRunner
blva Apr 22, 2025
6d92021
address comment: update agentClient to agentRunner
blva Apr 22, 2025
7997c54
address comment: isolate machine data
blva Apr 22, 2025
f18b916
address comment: inject config to constructor
blva Apr 22, 2025
a9cd835
feat: use machineid library
blva Apr 22, 2025
6813518
chore: reformat
blva Apr 22, 2025
65fad1d
chore: reformat
blva Apr 22, 2025
7e98c33
address comment: fix misleading comment
blva Apr 22, 2025
5209073
address npm check
blva Apr 22, 2025
2e33584
address comment: support do_not_track
blva Apr 22, 2025
d92adf1
Merge branch 'main' into telemetry
blva Apr 22, 2025
8ec7d9c
address comment: use in-memory cache
blva Apr 23, 2025
6ad2a19
address comments: use const
blva Apr 23, 2025
807109d
Merge branch 'main' into telemetry
blva Apr 23, 2025
f9a46f9
chore: fix lint and tests
blva Apr 23, 2025
599201d
address comment: remove unused const
blva Apr 23, 2025
7a889b4
address comment: do not use object freeze
blva Apr 23, 2025
eb360e2
address comment: as const
blva Apr 23, 2025
710b131
Update src/common/atlas/apiClient.ts
blva Apr 23, 2025
a15eeb6
address comment: make common props readonly
blva Apr 23, 2025
90caa25
Merge remote-tracking branch 'origin/main' into telemetry
blva Apr 23, 2025
89113a5
address comment: do not use enum
blva Apr 23, 2025
143f898
fix
blva Apr 23, 2025
f751a30
address comment: inject and use eventcache
blva Apr 23, 2025
0927d28
minor fixeS
blva Apr 23, 2025
fb0b8af
chore: gen sessionId
blva Apr 23, 2025
9e625aa
add lru cache
blva Apr 23, 2025
a61d9b4
chore: disable telemetry in tests
blva Apr 23, 2025
2bd00cf
Merge remote-tracking branch 'origin/main' into telemetry
blva Apr 23, 2025
3636fde
chore: reformat
blva Apr 23, 2025
0df870c
clean up
blva Apr 23, 2025
4b7563d
fix check
blva Apr 23, 2025
6dd1f79
address comment for category and update emit tool to get tool result …
blva Apr 23, 2025
807b2ca
address comment: fix telemetry init
blva Apr 23, 2025
139c3ee
Merge remote-tracking branch 'origin/main' into telemetry
blva Apr 23, 2025
3a05d31
format
blva Apr 23, 2025
8505d91
fix check
blva Apr 23, 2025
024a3d1
address comment
blva Apr 23, 2025
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
33 changes: 30 additions & 3 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@
"jest-environment-node": "^29.7.0",
"jest-extended": "^4.0.2",
"mongodb-runner": "^5.8.2",
"native-machine-id": "^0.0.8",
"openapi-types": "^12.1.3",
"openapi-typescript": "^7.6.1",
"prettier": "^3.5.3",
Expand All @@ -61,6 +62,7 @@
"@mongodb-js/devtools-connect": "^3.7.2",
"@mongosh/service-provider-node-driver": "^3.6.0",
"bson": "^6.10.3",
"lru-cache": "^11.1.0",
"mongodb": "^6.15.0",
"mongodb-log-writer": "^2.4.1",
"mongodb-redact": "^1.1.6",
Expand Down
41 changes: 40 additions & 1 deletion src/common/atlas/apiClient.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,11 @@
import createClient, { Client, FetchOptions, Middleware } from "openapi-fetch";
import createClient, { Client, Middleware } from "openapi-fetch";
import type { FetchOptions } from "openapi-fetch";
import { AccessToken, ClientCredentials } from "simple-oauth2";
import { ApiClientError } from "./apiClientError.js";
import { paths, operations } from "./openapi.js";
import { BaseEvent } from "../../telemetry/types.js";
import { mongoLogId } from "mongodb-log-writer";
import logger from "../../logger.js";
import { packageInfo } from "../../packageInfo.js";

const ATLAS_API_VERSION = "2025-03-12";
Expand Down Expand Up @@ -93,6 +97,15 @@ export class ApiClient {
this.client.use(this.errorMiddleware);
}

public hasCredentials(): boolean {
logger.info(
mongoLogId(1_000_000),
"api-client",
`Checking if API client has credentials: ${!!(this.oauth2Client && this.accessToken)}`
);
return !!(this.oauth2Client && this.accessToken);
}

public async getIpInfo(): Promise<{
currentIpv4Address: string;
}> {
Expand All @@ -118,6 +131,32 @@ export class ApiClient {
}>;
}

async sendEvents(events: BaseEvent[]): Promise<void> {
let endpoint = "api/private/unauth/telemetry/events";
const headers: Record<string, string> = {
Accept: "application/json",
"Content-Type": "application/json",
"User-Agent": this.options.userAgent,
};

const accessToken = await this.getAccessToken();
if (accessToken) {
endpoint = "api/private/v1.0/telemetry/events";
headers["Authorization"] = `Bearer ${accessToken}`;
}

const url = new URL(endpoint, this.options.baseUrl);
const response = await fetch(url, {
method: "POST",
headers,
body: JSON.stringify(events),
});

if (!response.ok) {
throw await ApiClientError.fromResponse(response);
}
}

// DO NOT EDIT. This is auto-generated code.
async listClustersForAllProjects(options?: FetchOptions<operations["listClustersForAllProjects"]>) {
const { data } = await this.client.GET("/api/atlas/v2/clusters", options);
Expand Down
1 change: 1 addition & 0 deletions src/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ export interface UserConfig {
apiBaseUrl?: string;
apiClientId?: string;
apiClientSecret?: string;
telemetry?: "enabled" | "disabled";
logPath: string;
connectionString?: string;
connectOptions: {
Expand Down
1 change: 0 additions & 1 deletion src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,5 @@ try {
await server.connect(transport);
} catch (error: unknown) {
logger.emergency(mongoLogId(1_000_004), "server", `Fatal error running server: ${error as string}`);

process.exit(1);
}
18 changes: 15 additions & 3 deletions src/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ import { AtlasTools } from "./tools/atlas/tools.js";
import { MongoDbTools } from "./tools/mongodb/tools.js";
import logger, { initializeLogger } from "./logger.js";
import { mongoLogId } from "mongodb-log-writer";
import { ObjectId } from "mongodb";
import { Telemetry } from "./telemetry/telemetry.js";
import { UserConfig } from "./config.js";

export interface ServerOptions {
Expand All @@ -16,25 +18,35 @@ export interface ServerOptions {
export class Server {
public readonly session: Session;
private readonly mcpServer: McpServer;
private readonly telemetry: Telemetry;
private readonly userConfig: UserConfig;

constructor({ session, mcpServer, userConfig }: ServerOptions) {
this.session = session;
this.telemetry = new Telemetry(session);
this.mcpServer = mcpServer;
this.userConfig = userConfig;
}

async connect(transport: Transport) {
this.mcpServer.server.registerCapabilities({ logging: {} });

this.registerTools();
this.registerResources();

await initializeLogger(this.mcpServer, this.userConfig.logPath);

await this.mcpServer.connect(transport);

logger.info(mongoLogId(1_000_004), "server", `Server started with transport ${transport.constructor.name}`);
this.mcpServer.server.oninitialized = () => {
this.session.setAgentRunner(this.mcpServer.server.getClientVersion());
this.session.sessionId = new ObjectId().toString();

logger.info(
mongoLogId(1_000_004),
"server",
`Server started with transport ${transport.constructor.name} and agent runner ${this.session.agentRunner?.name}`
);
};
}

async close(): Promise<void> {
Expand All @@ -44,7 +56,7 @@ export class Server {

private registerTools() {
for (const tool of [...AtlasTools, ...MongoDbTools]) {
new tool(this.session, this.userConfig).register(this.mcpServer);
new tool(this.session, this.userConfig, this.telemetry).register(this.mcpServer);
}
}

Expand Down
15 changes: 15 additions & 0 deletions src/session.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
import { ApiClient, ApiClientCredentials } from "./common/atlas/apiClient.js";
import { Implementation } from "@modelcontextprotocol/sdk/types.js";

export interface SessionOptions {
apiBaseUrl?: string;
Expand All @@ -8,8 +9,13 @@ export interface SessionOptions {
}

export class Session {
sessionId?: string;
serviceProvider?: NodeDriverServiceProvider;
apiClient: ApiClient;
agentRunner?: {
name: string;
version: string;
};

constructor({ apiBaseUrl, apiClientId, apiClientSecret }: SessionOptions = {}) {
const credentials: ApiClientCredentials | undefined =
Expand All @@ -26,6 +32,15 @@ export class Session {
});
}

setAgentRunner(agentRunner: Implementation | undefined) {
if (agentRunner?.name && agentRunner?.version) {
this.agentRunner = {
name: agentRunner.name,
version: agentRunner.version,
};
}
}

async close(): Promise<void> {
if (this.serviceProvider) {
try {
Expand Down
15 changes: 15 additions & 0 deletions src/telemetry/constants.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
import { getMachineIdSync } from "native-machine-id";
import { packageInfo } from "../packageInfo.js";

/**
* Machine-specific metadata formatted for telemetry
*/
export const MACHINE_METADATA = {
device_id: getMachineIdSync(),
mcp_server_version: packageInfo.version,
mcp_server_name: packageInfo.mcpServerName,
platform: process.platform,
arch: process.arch,
os_type: process.platform,
os_version: process.version,
} as const;
62 changes: 62 additions & 0 deletions src/telemetry/eventCache.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
import { BaseEvent } from "./types.js";
import { LRUCache } from "lru-cache";

/**
* Singleton class for in-memory telemetry event caching
* Provides a central storage for telemetry events that couldn't be sent
* Uses LRU cache to automatically drop oldest events when limit is exceeded
*/
export class EventCache {
private static instance: EventCache;
private static readonly MAX_EVENTS = 1000;

private cache: LRUCache<number, BaseEvent>;
private nextId = 0;

private constructor() {
this.cache = new LRUCache({
max: EventCache.MAX_EVENTS,
// Using FIFO eviction strategy for events
allowStale: false,
updateAgeOnGet: false,
});
}

/**
* Gets the singleton instance of EventCache
* @returns The EventCache instance
*/
public static getInstance(): EventCache {
if (!EventCache.instance) {
EventCache.instance = new EventCache();
}
return EventCache.instance;
}

/**
* Gets a copy of the currently cached events
* @returns Array of cached BaseEvent objects
*/
public getEvents(): BaseEvent[] {
return Array.from(this.cache.values());
}

/**
* Appends new events to the cached events
* LRU cache automatically handles dropping oldest events when limit is exceeded
* @param events - The events to append
*/
public appendEvents(events: BaseEvent[]): void {
for (const event of events) {
this.cache.set(this.nextId++, event);
}
}

/**
* Clears all cached events
*/
public clearEvents(): void {
this.cache.clear();
this.nextId = 0;
}
}
Loading
Loading