Skip to content

Commit dae4f06

Browse files
blvablva
committed
simplify device id
1 parent 658cdc8 commit dae4f06

File tree

14 files changed

+179
-222
lines changed

14 files changed

+179
-222
lines changed

‎src/common/connectionManager.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import { packageInfo } from "./packageInfo.js";
66
import ConnectionString from "mongodb-connection-string-url";
77
import { MongoClientOptions } from "mongodb";
88
import { ErrorCodes, MongoDBError } from "./errors.js";
9-
import { DeviceIdService } from "../helpers/deviceId.js";
9+
import { DeviceId } from "../helpers/deviceId.js";
1010
import { AppNameComponents } from "../helpers/connectionOptions.js";
1111
import { CompositeLogger, LogId } from "./logger.js";
1212
import { ConnectionInfo, generateConnectionInfoFromCliArgs } from "@mongosh/arg-parser";
@@ -71,7 +71,7 @@ export interface ConnectionManagerEvents {
7171

7272
export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
7373
private state: AnyConnectionState;
74-
private deviceId: DeviceIdService;
74+
private deviceId: DeviceId;
7575
private clientName: string;
7676
private bus: EventEmitter;
7777

@@ -88,8 +88,8 @@ export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
8888

8989
this.bus.on("mongodb-oidc-plugin:auth-failed", this.onOidcAuthFailed.bind(this));
9090
this.bus.on("mongodb-oidc-plugin:auth-succeeded", this.onOidcAuthSucceeded.bind(this));
91-
92-
this.deviceId = DeviceIdService.getInstance();
91+
92+
this.deviceId = DeviceId.create(this.logger);
9393
this.clientName = "unknown";
9494
}
9595

@@ -111,7 +111,7 @@ export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
111111
settings = { ...settings };
112112
const appNameComponents: AppNameComponents = {
113113
appName: `${packageInfo.mcpServerName} ${packageInfo.version}`,
114-
deviceId: this.deviceId.getDeviceId(),
114+
deviceId: this.deviceId.get(),
115115
clientName: this.clientName,
116116
};
117117

‎src/helpers/deviceId.ts‎

Lines changed: 24 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4,62 +4,68 @@ import { LogId, LoggerBase } from "../common/logger.js";
44

55
export const DEVICE_ID_TIMEOUT = 3000;
66

7-
export class DeviceIdService {
8-
private static instance: DeviceIdService | undefined = undefined;
7+
export class DeviceId {
98
private deviceId: string | undefined = undefined;
109
private deviceIdPromise: Promise<string> | undefined = undefined;
1110
private abortController: AbortController | undefined = undefined;
1211
private logger: LoggerBase;
1312
private readonly getMachineId: () => Promise<string>;
1413
private timeout: number;
14+
private static instance: DeviceId | undefined = undefined;
1515

16-
private constructor(logger: LoggerBase, timeout: number) {
16+
private constructor(logger: LoggerBase, timeout: number = DEVICE_ID_TIMEOUT) {
1717
this.logger = logger;
1818
this.timeout = timeout;
1919
this.getMachineId = (): Promise<string> => nodeMachineId.machineId(true);
2020
}
2121

22-
static create(
23-
logger: LoggerBase,
24-
{ timeout = DEVICE_ID_TIMEOUT }: { timeout?: number } = {},
25-
): DeviceIdService {
26-
const instance = new DeviceIdService(logger, timeout);
22+
public static create(logger: LoggerBase, timeout?: number): DeviceId {
23+
if (this.instance) {
24+
return this.instance;
25+
}
2726

27+
const instance = new DeviceId(logger, timeout ?? DEVICE_ID_TIMEOUT);
2828
void instance.setup();
29+
30+
this.instance = instance;
31+
2932
return instance;
3033
}
3134

35+
/**
36+
* Sets up the device ID calculation promise and abort controller.
37+
*/
3238
private async setup(): Promise<void> {
3339
this.abortController = new AbortController();
34-
3540
this.deviceIdPromise = this.calculateDeviceId();
36-
const deviceId = await this.deviceIdPromise;
37-
this.deviceId = deviceId;
41+
// start the promise upon setup
42+
void this.deviceIdPromise;
3843
}
3944

40-
41-
public close(): void {
45+
/**
46+
* Closes the device ID calculation promise and abort controller.
47+
*/
48+
public close(): void {
4249
if (this.abortController) {
4350
this.abortController.abort();
4451
this.abortController = undefined;
4552
}
4653
this.deviceId = undefined;
4754
this.deviceIdPromise = undefined;
48-
DeviceIdService.instance = undefined;
55+
DeviceId.instance = undefined;
4956
}
50-
5157

5258
/**
5359
* Gets the device ID, waiting for the calculation to complete if necessary.
5460
* @returns Promise that resolves to the device ID string
5561
*/
56-
public async getDeviceId(): Promise<string> {
62+
public async get(): Promise<string> {
5763
if (this.deviceId !== undefined) {
5864
return this.deviceId;
5965
}
6066

6167
if (!this.deviceIdPromise) {
62-
throw new Error("DeviceIdService calculation not started");
68+
this.deviceIdPromise = this.calculateDeviceId();
6369
}
6470

6571
return this.deviceIdPromise;
@@ -124,7 +130,7 @@ export class DeviceIdService {
124130
});
125131
break;
126132
case "abort":
127-
// No need to log in the case of aborts
133+
// No need to log in the case of 'abort' errors
128134
break;
129135
}
130136
}

‎src/index.ts‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ import { packageInfo } from "./common/packageInfo.js";
4242
import { StdioRunner } from "./transports/stdio.js";
4343
import { StreamableHttpRunner } from "./transports/streamableHttp.js";
4444
import { systemCA } from "@mongodb-js/devtools-proxy-support";
45-
import { DeviceIdService } from "./helpers/deviceId.js";
45+
import { DeviceId } from "./helpers/deviceId.js";
4646

4747
async function main(): Promise<void> {
4848
systemCA().catch(() => undefined); // load system CA asynchronously as in mongosh
@@ -51,8 +51,7 @@ async function main(): Promise<void> {
5151
assertVersionMode();
5252

5353
const transportRunner = config.transport === "stdio" ? new StdioRunner(config) : new StreamableHttpRunner(config);
54-
const deviceId = DeviceIdService.init(transportRunner.logger);
55-
54+
const deviceId = DeviceId.create(transportRunner.logger);
5655
const shutdown = (): void => {
5756
transportRunner.logger.info({
5857
id: LogId.serverCloseRequested,

‎src/telemetry/telemetry.ts‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import { ApiClient } from "../common/atlas/apiClient.js";
66
import { MACHINE_METADATA } from "./constants.js";
77
import { EventCache } from "./eventCache.js";
88
import { detectContainerEnv } from "../helpers/container.js";
9-
import { DeviceIdService } from "../helpers/deviceId.js";
9+
import { DeviceId } from "../helpers/deviceId.js";
1010

1111
type EventResult = {
1212
success: boolean;
@@ -18,13 +18,13 @@ export class Telemetry {
1818
/** Resolves when the setup is complete or a timeout occurs */
1919
public setupPromise: Promise<[string, boolean]> | undefined;
2020
private eventCache: EventCache;
21-
private deviceId: DeviceIdService;
21+
private deviceId: DeviceId;
2222

2323
private constructor(
2424
private readonly session: Session,
2525
private readonly userConfig: UserConfig,
2626
private readonly commonProperties: CommonProperties,
27-
{ eventCache, deviceId }: { eventCache: EventCache; deviceId: DeviceIdService }
27+
{ eventCache, deviceId }: { eventCache: EventCache; deviceId: DeviceId }
2828
) {
2929
this.eventCache = eventCache;
3030
this.deviceId = deviceId;
@@ -36,10 +36,10 @@ export class Telemetry {
3636
{
3737
commonProperties = { ...MACHINE_METADATA },
3838
eventCache = EventCache.getInstance(),
39-
deviceId = DeviceIdService.getInstance(),
39+
deviceId = DeviceId.create(session.logger),
4040
}: {
4141
eventCache?: EventCache;
42-
deviceId?: DeviceIdService;
42+
deviceId?: DeviceId;
4343
commonProperties?: CommonProperties;
4444
} = {}
4545
): Telemetry {
@@ -54,7 +54,7 @@ export class Telemetry {
5454
return;
5555
}
5656

57-
this.setupPromise = Promise.all([this.deviceId.getDeviceId(), detectContainerEnv()]);
57+
this.setupPromise = Promise.all([this.deviceId.get(), detectContainerEnv()]);
5858
const [deviceIdValue, containerEnv] = await this.setupPromise;
5959

6060
this.commonProperties.device_id = deviceIdValue;

‎tests/integration/common/deviceId.test.ts‎

Lines changed: 17 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { describe, expect, it, vi, beforeEach, afterEach } from "vitest";
2-
import { DeviceIdService } from "../../../src/helpers/deviceId.js";
2+
import { DeviceId } from "../../../src/helpers/deviceId.js";
33
import { CompositeLogger } from "../../../src/common/logger.js";
44
import nodeMachineId from "node-machine-id";
55

@@ -12,15 +12,13 @@ describe("Device ID", () => {
1212
});
1313

1414
afterEach(() => {
15-
if (DeviceIdService.isInitialized()) {
16-
DeviceIdService.getInstance().close();
17-
}
15+
DeviceId.create(testLogger).close();
1816
});
1917

2018
describe("when resolving device ID", () => {
2119
it("should successfully resolve device ID in real environment", async () => {
22-
const deviceId = DeviceIdService.init(testLogger);
23-
const result = await deviceId.getDeviceId();
20+
const deviceId = DeviceId.create(testLogger);
21+
const result = await deviceId.get();
2422

2523
expect(result).not.toBe("unknown");
2624
expect(result).toBeTruthy();
@@ -31,23 +29,23 @@ describe("Device ID", () => {
3129
it("should cache device ID after first resolution", async () => {
3230
// spy on machineId
3331
const machineIdSpy = vi.spyOn(nodeMachineId, "machineId");
34-
const deviceId = DeviceIdService.init(testLogger);
32+
const deviceId = DeviceId.create(testLogger);
3533

3634
// First call
37-
const result1 = await deviceId.getDeviceId();
35+
const result1 = await deviceId.get();
3836
expect(result1).not.toBe("unknown");
3937

4038
// Second call should be cached
41-
const result2 = await deviceId.getDeviceId();
39+
const result2 = await deviceId.get();
4240
expect(result2).toBe(result1);
4341
// check that machineId was called only once
4442
expect(machineIdSpy).toHaveBeenCalledOnce();
4543
});
4644

4745
it("should handle concurrent device ID requests correctly", async () => {
48-
const deviceId = DeviceIdService.init(testLogger);
46+
const deviceId = DeviceId.create(testLogger);
4947

50-
const promises = Array.from({ length: 5 }, () => deviceId.getDeviceId());
48+
const promises = Array.from({ length: 5 }, () => deviceId.get());
5149

5250
// All should resolve to the same value
5351
const results = await Promise.all(promises);
@@ -81,10 +79,10 @@ describe("Device ID", () => {
8179
reject(new Error("Machine ID failed"));
8280
});
8381
});
84-
const deviceId = DeviceIdService.init(testLogger);
85-
const handleDeviceIdErrorSpy = vi.spyOn(deviceId, "handleDeviceIdError" as keyof DeviceIdService);
82+
const deviceId = DeviceId.create(testLogger);
83+
const handleDeviceIdErrorSpy = vi.spyOn(deviceId, "handleDeviceIdError" as keyof DeviceId);
8684

87-
const result = await deviceId.getDeviceId();
85+
const result = await deviceId.get();
8886

8987
expect(result).toBe("unknown");
9088
expect(handleDeviceIdErrorSpy).toHaveBeenCalledWith(
@@ -101,38 +99,13 @@ describe("Device ID", () => {
10199
});
102100
});
103101

104-
const deviceId = DeviceIdService.init(testLogger);
105-
const handleDeviceIdErrorSpy = vi.spyOn(deviceId, "handleDeviceIdError" as keyof DeviceIdService);
102+
const deviceId = DeviceId.create(testLogger, 100); // Short timeout
103+
const handleDeviceIdErrorSpy = vi.spyOn(deviceId, "handleDeviceIdError" as keyof DeviceId);
106104

107-
deviceId.close();
108-
109-
// expect the deviceId service to throw an error
110-
await expect(deviceId.getDeviceId()).rejects.toThrow(Error);
111-
// test that the private function handleDeviceIdError was called with reason "abort"
112-
expect(handleDeviceIdErrorSpy).toHaveBeenCalledWith(
113-
"abort",
114-
expect.stringContaining("Aborted by abort signal")
115-
);
116-
117-
// check that the deviceId service is not initialized anymore
118-
expect(() => DeviceIdService.getInstance()).toThrow(Error);
119-
});
120-
121-
it("should handle timeout scenarios gracefully", async () => {
122-
nodeMachineId.machineId = vi.fn().mockImplementation(() => {
123-
return new Promise<string>((resolve) => {
124-
setTimeout(() => resolve("delayed-id"), 200);
125-
});
126-
});
127-
128-
// override the timeout to 100ms
129-
const deviceId = DeviceIdService.init(testLogger, 100);
130-
const handleDeviceIdErrorSpy = vi.spyOn(deviceId, "handleDeviceIdError" as keyof DeviceIdService);
131-
132-
const result = await deviceId.getDeviceId();
105+
const result = await deviceId.get();
133106

134107
expect(result).toBe("unknown");
135-
expect(handleDeviceIdErrorSpy).toHaveBeenCalledWith("timeout", expect.stringContaining("Timeout"));
136-
}, 5000);
108+
expect(handleDeviceIdErrorSpy).toHaveBeenCalledWith("timeout", expect.any(String));
109+
});
137110
});
138111
});

‎tests/integration/helpers.ts‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
77
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
88
import { InMemoryTransport } from "./inMemoryTransport.js";
99
import { UserConfig, DriverOptions } from "../../src/common/config.js";
10-
import { DeviceIdService } from "../../src/helpers/deviceId.js";
10+
import { DeviceId } from "../../src/helpers/deviceId.js";
1111
import { McpError, ResourceUpdatedNotificationSchema } from "@modelcontextprotocol/sdk/types.js";
1212
import { config, driverOptions } from "../../src/common/config.js";
1313
import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest";
@@ -69,9 +69,6 @@ export function setupIntegrationTest(
6969

7070
const exportsManager = ExportsManager.init(userConfig, logger);
7171

72-
// Initialize DeviceIdService for tests
73-
DeviceIdService.init(logger);
74-
7572
const connectionManager = new ConnectionManager(userConfig, driverOptions, logger);
7673

7774
const session = new Session({

‎tests/integration/telemetry.test.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { Telemetry } from "../../src/telemetry/telemetry.js";
22
import { Session } from "../../src/common/session.js";
33
import { config, driverOptions } from "../../src/common/config.js";
4-
import { DeviceIdService } from "../../src/helpers/deviceId.js";
4+
import { DeviceId } from "../../src/helpers/deviceId.js";
55
import { describe, expect, it } from "vitest";
66
import { CompositeLogger } from "../../src/common/logger.js";
77
import { ConnectionManager } from "../../src/common/connectionManager.js";
@@ -11,9 +11,8 @@ describe("Telemetry", () => {
1111
it("should resolve the actual device ID", async () => {
1212
const logger = new CompositeLogger();
1313

14-
// Initialize DeviceIdService like the main application does
15-
DeviceIdService.init(logger);
16-
const actualDeviceId = await DeviceIdService.getInstance().getDeviceId();
14+
const deviceId = DeviceId.create(logger);
15+
const actualDeviceId = await deviceId.get();
1716

1817
const telemetry = Telemetry.create(
1918
new Session({
@@ -22,7 +21,8 @@ describe("Telemetry", () => {
2221
exportsManager: ExportsManager.init(config, logger),
2322
connectionManager: new ConnectionManager(config, driverOptions, logger),
2423
}),
25-
config
24+
config,
25+
{ deviceId }
2626
);
2727

2828
expect(telemetry.getCommonProperties().device_id).toBe(undefined);

‎tests/integration/transports/stdio.test.ts‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { describe, expect, it, beforeAll, afterAll } from "vitest";
22
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
33
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
44
import { describeWithMongoDB } from "../tools/mongodb/mongodbHelpers.js";
5-
import { DeviceIdService } from "../../../src/helpers/deviceId.js";
5+
import { DeviceId } from "../../../src/helpers/deviceId.js";
66
import { CompositeLogger } from "../../../src/common/logger.js";
77

88
describeWithMongoDB("StdioRunner", (integration) => {
@@ -22,7 +22,6 @@ describeWithMongoDB("StdioRunner", (integration) => {
2222
name: "test",
2323
version: "0.0.0",
2424
});
25-
DeviceIdService.init(new CompositeLogger());
2625
await client.connect(transport);
2726
});
2827

‎tests/integration/transports/streamableHttp.test.ts‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js";
33
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
44
import { describe, expect, it, beforeAll, afterAll } from "vitest";
55
import { config } from "../../../src/common/config.js";
6-
import { DeviceIdService } from "../../../src/helpers/deviceId.js";
6+
import { DeviceId } from "../../../src/helpers/deviceId.js";
77

88
describe("StreamableHttpRunner", () => {
99
let runner: StreamableHttpRunner;
@@ -16,7 +16,6 @@ describe("StreamableHttpRunner", () => {
1616
config.telemetry = "disabled";
1717
config.loggers = ["stderr"];
1818
runner = new StreamableHttpRunner(config);
19-
DeviceIdService.init(runner.logger);
2019
await runner.start();
2120
});
2221

0 commit comments

Comments
 (0)