chore: defer machine ID resolution (#161)

Author: gagik

package-lock.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",

package.json

@@ -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);
+    }
+}

src/deferred-promise.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,
     });
 

src/index.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),

src/logger.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();
     }

src/server.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,

src/telemetry/constants.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

src/telemetry/device-id.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];