chore: defer machine ID resolution (#161)

gagik · web-flow · commit 645196ebc20f · 2025-05-01T14:33:27.000+02:00
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -69,7 +69,7 @@
     "mongodb-log-writer": "^2.4.1",
     "mongodb-redact": "^1.1.6",
     "mongodb-schema": "^12.6.2",
-    "node-machine-id": "^1.1.12",
+    "node-machine-id": "1.1.12",
     "openapi-fetch": "^0.13.5",
     "simple-oauth2": "^5.1.0",
     "yargs-parser": "^21.1.1",
diff --git a/src/deferred-promise.ts b/src/deferred-promise.ts
@@ -0,0 +1,58 @@
+type DeferredPromiseOptions<T> = {
+    timeout?: number;
+    onTimeout?: (resolve: (value: T) => void, reject: (reason: Error) => void) => void;
+};
+
+/** Creates a promise and exposes its resolve and reject methods, with an optional timeout. */
+export class DeferredPromise<T> extends Promise<T> {
+    resolve: (value: T) => void;
+    reject: (reason: unknown) => void;
+    private timeoutId?: NodeJS.Timeout;
+
+    constructor(
+        executor: (resolve: (value: T) => void, reject: (reason: Error) => void) => void,
+        { timeout, onTimeout }: DeferredPromiseOptions<T> = {}
+    ) {
+        let resolveFn: (value: T) => void;
+        let rejectFn: (reason?: unknown) => void;
+
+        super((resolve, reject) => {
+            resolveFn = resolve;
+            rejectFn = reject;
+        });
+
+        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+        this.resolve = resolveFn!;
+        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+        this.reject = rejectFn!;
+
+        if (timeout !== undefined && onTimeout) {
+            this.timeoutId = setTimeout(() => {
+                onTimeout(this.resolve, this.reject);
+            }, timeout);
+        }
+
+        executor(
+            (value: T) => {
+                if (this.timeoutId) clearTimeout(this.timeoutId);
+                this.resolve(value);
+            },
+            (reason: Error) => {
+                if (this.timeoutId) clearTimeout(this.timeoutId);
+                this.reject(reason);
+            }
+        );
+    }
+
+    static fromPromise<T>(promise: Promise<T>, options: DeferredPromiseOptions<T> = {}): DeferredPromise<T> {
+        return new DeferredPromise<T>((resolve, reject) => {
+            promise
+                .then((value) => {
+                    resolve(value);
+                })
+                .catch((reason) => {
+                    reject(reason as Error);
+                });
+        }, options);
+    }
+}
diff --git a/src/index.ts b/src/index.ts
@@ -7,6 +7,7 @@ import { config } from "./config.js";
 import { Session } from "./session.js";
 import { Server } from "./server.js";
 import { packageInfo } from "./packageInfo.js";
+import { Telemetry } from "./telemetry/telemetry.js";
 
 try {
     const session = new Session({
@@ -19,9 +20,12 @@ try {
         version: packageInfo.version,
     });
 
+    const telemetry = Telemetry.create(session, config);
+
     const server = new Server({
         mcpServer,
         session,
+        telemetry,
         userConfig: config,
     });
 
diff --git a/src/logger.ts b/src/logger.ts
@@ -20,6 +20,7 @@ export const LogId = {
     telemetryEmitSuccess: mongoLogId(1_002_004),
     telemetryMetadataError: mongoLogId(1_002_005),
     telemetryDeviceIdFailure: mongoLogId(1_002_006),
+    telemetryDeviceIdTimeout: mongoLogId(1_002_007),
 
     toolExecute: mongoLogId(1_003_001),
     toolExecuteFailure: mongoLogId(1_003_002),
diff --git a/src/server.ts b/src/server.ts
@@ -16,6 +16,7 @@ export interface ServerOptions {
     session: Session;
     userConfig: UserConfig;
     mcpServer: McpServer;
+    telemetry: Telemetry;
 }
 
 export class Server {
@@ -25,10 +26,10 @@ export class Server {
     public readonly userConfig: UserConfig;
     private readonly startTime: number;
 
-    constructor({ session, mcpServer, userConfig }: ServerOptions) {
+    constructor({ session, mcpServer, userConfig, telemetry }: ServerOptions) {
         this.startTime = Date.now();
         this.session = session;
-        this.telemetry = new Telemetry(session, userConfig);
+        this.telemetry = telemetry;
         this.mcpServer = mcpServer;
         this.userConfig = userConfig;
     }
@@ -93,6 +94,7 @@ export class Server {
     }
 
     async close(): Promise<void> {
+        await this.telemetry.close();
         await this.session.close();
         await this.mcpServer.close();
     }
diff --git a/src/telemetry/constants.ts b/src/telemetry/constants.ts
@@ -1,11 +1,10 @@
 import { packageInfo } from "../packageInfo.js";
 import { type CommonStaticProperties } from "./types.js";
-import { getDeviceId } from "./device-id.js";
+
 /**
  * Machine-specific metadata formatted for telemetry
  */
 export const MACHINE_METADATA: CommonStaticProperties = {
-    device_id: getDeviceId(),
     mcp_server_version: packageInfo.version,
     mcp_server_name: packageInfo.mcpServerName,
     platform: process.platform,
diff --git a/src/telemetry/device-id.ts b/src/telemetry/device-id.ts
diff --git a/src/telemetry/eventCache.ts b/src/telemetry/eventCache.ts
@@ -1,5 +1,5 @@
-import { BaseEvent } from "./types.js";
 import { LRUCache } from "lru-cache";
+import { BaseEvent } from "./types.js";
 
 /**
  * Singleton class for in-memory telemetry event caching
diff --git a/src/telemetry/telemetry.ts b/src/telemetry/telemetry.ts
@@ -5,23 +5,101 @@ import logger, { LogId } from "../logger.js";
 import { ApiClient } from "../common/atlas/apiClient.js";
 import { MACHINE_METADATA } from "./constants.js";
 import { EventCache } from "./eventCache.js";
+import { createHmac } from "crypto";
+import nodeMachineId from "node-machine-id";
+import { DeferredPromise } from "../deferred-promise.js";
 
 type EventResult = {
     success: boolean;
     error?: Error;
 };
 
+export const DEVICE_ID_TIMEOUT = 3000;
+
 export class Telemetry {
-    private readonly commonProperties: CommonProperties;
+    private isBufferingEvents: boolean = true;
+    /** Resolves when the device ID is retrieved or timeout occurs */
+    public deviceIdPromise: DeferredPromise<string> | undefined;
+    private eventCache: EventCache;
+    private getRawMachineId: () => Promise<string>;
 
-    constructor(
+    private constructor(
         private readonly session: Session,
         private readonly userConfig: UserConfig,
-        private readonly eventCache: EventCache = EventCache.getInstance()
+        private readonly commonProperties: CommonProperties,
+        { eventCache, getRawMachineId }: { eventCache: EventCache; getRawMachineId: () => Promise<string> }
     ) {
-        this.commonProperties = {
-            ...MACHINE_METADATA,
-        };
+        this.eventCache = eventCache;
+        this.getRawMachineId = getRawMachineId;
+    }
+
+    static create(
+        session: Session,
+        userConfig: UserConfig,
+        {
+            commonProperties = { ...MACHINE_METADATA },
+            eventCache = EventCache.getInstance(),
+
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-return, @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-member-access
+            getRawMachineId = () => nodeMachineId.machineId(true),
+        }: {
+            eventCache?: EventCache;
+            getRawMachineId?: () => Promise<string>;
+            commonProperties?: CommonProperties;
+        } = {}
+    ): Telemetry {
+        const instance = new Telemetry(session, userConfig, commonProperties, { eventCache, getRawMachineId });
+
+        void instance.start();
+        return instance;
+    }
+
+    private async start(): Promise<void> {
+        if (!this.isTelemetryEnabled()) {
+            return;
+        }
+        this.deviceIdPromise = DeferredPromise.fromPromise(this.getDeviceId(), {
+            timeout: DEVICE_ID_TIMEOUT,
+            onTimeout: (resolve) => {
+                resolve("unknown");
+                logger.debug(LogId.telemetryDeviceIdTimeout, "telemetry", "Device ID retrieval timed out");
+            },
+        });
+        this.commonProperties.device_id = await this.deviceIdPromise;
+
+        this.isBufferingEvents = false;
+    }
+
+    public async close(): Promise<void> {
+        this.deviceIdPromise?.resolve("unknown");
+        this.isBufferingEvents = false;
+        await this.emitEvents(this.eventCache.getEvents());
+    }
+
+    /**
+     * @returns A hashed, unique identifier for the running device or `"unknown"` if not known.
+     */
+    private async getDeviceId(): Promise<string> {
+        try {
+            if (this.commonProperties.device_id) {
+                return this.commonProperties.device_id;
+            }
+
+            const originalId: string = await this.getRawMachineId();
+
+            // Create a hashed format from the all uppercase version of the machine ID
+            // to match it exactly with the denisbrodbeck/machineid library that Atlas CLI uses.
+            const hmac = createHmac("sha256", originalId.toUpperCase());
+
+            /** This matches the message used to create the hashes in Atlas CLI */
+            const DEVICE_ID_HASH_MESSAGE = "atlascli";
+
+            hmac.update(DEVICE_ID_HASH_MESSAGE);
+            return hmac.digest("hex");
+        } catch (error) {
+            logger.debug(LogId.telemetryDeviceIdFailure, "telemetry", String(error));
+            return "unknown";
+        }
     }
 
     /**
@@ -78,6 +156,11 @@ export class Telemetry {
      * Falls back to caching if both attempts fail
      */
     private async emit(events: BaseEvent[]): Promise<void> {
+        if (this.isBufferingEvents) {
+            this.eventCache.appendEvents(events);
+            return;
+        }
+
         const cachedEvents = this.eventCache.getEvents();
         const allEvents = [...cachedEvents, ...events];
 
diff --git a/src/telemetry/types.ts b/src/telemetry/types.ts
@@ -53,7 +53,6 @@ export type ServerEvent = TelemetryEvent<ServerEventProperties>;
  * Interface for static properties, they can be fetched once and reused.
  */
 export type CommonStaticProperties = {
-    device_id?: string;
     mcp_server_version: string;
     mcp_server_name: string;
     platform: string;
diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts
@@ -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 { Telemetry } from "../../src/telemetry/telemetry.js";
 import { config } from "../../src/config.js";
 
 interface ParameterInfo {
@@ -56,9 +57,14 @@ export function setupIntegrationTest(getUserConfig: () => UserConfig): Integrati
             apiClientSecret: userConfig.apiClientSecret,
         });
 
+        userConfig.telemetry = "disabled";
+
+        const telemetry = Telemetry.create(session, userConfig);
+
         mcpServer = new Server({
             session,
             userConfig,
+            telemetry,
             mcpServer: new McpServer({
                 name: "test-server",
                 version: "5.2.3",
diff --git a/tests/integration/telemetry.test.ts b/tests/integration/telemetry.test.ts
@@ -0,0 +1,29 @@
+import { createHmac } from "crypto";
+import { Telemetry } from "../../src/telemetry/telemetry.js";
+import { Session } from "../../src/session.js";
+import { config } from "../../src/config.js";
+import nodeMachineId from "node-machine-id";
+
+describe("Telemetry", () => {
+    it("should resolve the actual machine ID", async () => {
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-member-access
+        const actualId: string = await nodeMachineId.machineId(true);
+
+        const actualHashedId = createHmac("sha256", actualId.toUpperCase()).update("atlascli").digest("hex");
+
+        const telemetry = Telemetry.create(
+            new Session({
+                apiBaseUrl: "",
+            }),
+            config
+        );
+
+        expect(telemetry.getCommonProperties().device_id).toBe(undefined);
+        expect(telemetry["isBufferingEvents"]).toBe(true);
+
+        await telemetry.deviceIdPromise;
+
+        expect(telemetry.getCommonProperties().device_id).toBe(actualHashedId);
+        expect(telemetry["isBufferingEvents"]).toBe(false);
+    });
+});
diff --git a/tests/integration/tools/mongodb/mongodbHelpers.ts b/tests/integration/tools/mongodb/mongodbHelpers.ts
@@ -76,8 +76,6 @@ export function setupMongoDBIntegrationTest(): MongoDBIntegrationTest {
         let dbsDir = path.join(tmpDir, "mongodb-runner", "dbs");
         for (let i = 0; i < 10; i++) {
             try {
-                // TODO: Fix this type once mongodb-runner is updated.
-                // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call
                 mongoCluster = await MongoCluster.start({
                     tmpDir: dbsDir,
                     logDir: path.join(tmpDir, "mongodb-runner", "logs"),
diff --git a/tests/unit/deferred-promise.test.ts b/tests/unit/deferred-promise.test.ts
diff --git a/tests/unit/telemetry.test.ts b/tests/unit/telemetry.test.ts
diff --git a/tsconfig.build.json b/tsconfig.build.json

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@ export interface ServerOptions {`
`16`	`16`	`session: Session;`
`17`	`17`	`userConfig: UserConfig;`
`18`	`18`	`mcpServer: McpServer;`
	`19`	`+ telemetry: Telemetry;`
`19`	`20`	`}`
`20`	`21`
`21`	`22`	`export class Server {`
`@@ -25,10 +26,10 @@ export class Server {`
`25`	`26`	`public readonly userConfig: UserConfig;`
`26`	`27`	`private readonly startTime: number;`
`27`	`28`
`28`		`- constructor({ session, mcpServer, userConfig }: ServerOptions) {`
	`29`	`+ constructor({ session, mcpServer, userConfig, telemetry }: ServerOptions) {`
`29`	`30`	`this.startTime = Date.now();`
`30`	`31`	`this.session = session;`
`31`		`- this.telemetry = new Telemetry(session, userConfig);`
	`32`	`+ this.telemetry = telemetry;`
`32`	`33`	`this.mcpServer = mcpServer;`
`33`	`34`	`this.userConfig = userConfig;`
`34`	`35`	`}`
`@@ -93,6 +94,7 @@ export class Server {`
`93`	`94`	`}`
`94`	`95`
`95`	`96`	`async close(): Promise<void> {`
	`97`	`+ await this.telemetry.close();`
`96`	`98`	`await this.session.close();`
`97`	`99`	`await this.mcpServer.close();`
`98`	`100`	`}`