Skip to content

Commit b0972bd

Browse files
blvablva
committed
address comment: inject deviceId
1 parent 0bc1a9f commit b0972bd

File tree

13 files changed

+75
-123
lines changed

13 files changed

+75
-123
lines changed

‎src/common/connectionManager.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
7979
private userConfig: UserConfig,
8080
private driverOptions: DriverOptions,
8181
private logger: CompositeLogger,
82+
deviceId: DeviceId,
8283
bus?: EventEmitter
8384
) {
8485
super();
@@ -89,7 +90,7 @@ export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
8990
this.bus.on("mongodb-oidc-plugin:auth-failed", this.onOidcAuthFailed.bind(this));
9091
this.bus.on("mongodb-oidc-plugin:auth-succeeded", this.onOidcAuthSucceeded.bind(this));
9192

92-
this.deviceId = DeviceId.create(this.logger);
93+
this.deviceId = deviceId;
9394
this.clientName = "unknown";
9495
}
9596

‎src/helpers/deviceId.ts‎

Lines changed: 21 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ export class DeviceId {
2121

2222
public static create(logger: LoggerBase, timeout?: number): DeviceId {
2323
if (this.instance) {
24-
return this.instance;
24+
throw new Error("DeviceId instance already exists, use get() to retrieve the device ID");
2525
}
2626

2727
const instance = new DeviceId(logger, timeout ?? DEVICE_ID_TIMEOUT);
@@ -32,14 +32,8 @@ export class DeviceId {
3232
return instance;
3333
}
3434

35-
/**
36-
* Sets up the device ID calculation promise and abort controller.
37-
*/
3835
private setup(): void {
39-
this.abortController = new AbortController();
4036
this.deviceIdPromise = this.calculateDeviceId();
41-
// start the promise upon setup
42-
void this.deviceIdPromise;
4337
}
4438

4539
/**
@@ -50,6 +44,7 @@ export class DeviceId {
5044
this.abortController.abort();
5145
this.abortController = undefined;
5246
}
47+
5348
this.deviceId = undefined;
5449
this.deviceIdPromise = undefined;
5550
DeviceId.instance = undefined;
@@ -59,60 +54,41 @@ export class DeviceId {
5954
* Gets the device ID, waiting for the calculation to complete if necessary.
6055
* @returns Promise that resolves to the device ID string
6156
*/
62-
public async get(): Promise<string> {
63-
if (this.deviceId !== undefined) {
64-
return this.deviceId;
57+
public get(): Promise<string> {
58+
if (this.deviceId) {
59+
return Promise.resolve(this.deviceId);
6560
}
6661

67-
if (!this.deviceIdPromise) {
68-
this.deviceIdPromise = this.calculateDeviceId();
62+
if (this.deviceIdPromise) {
63+
return this.deviceIdPromise;
6964
}
7065

71-
return this.deviceIdPromise;
66+
return this.calculateDeviceId();
7267
}
7368

7469
/**
7570
* Internal method that performs the actual device ID calculation.
7671
*/
7772
private async calculateDeviceId(): Promise<string> {
7873
if (!this.abortController) {
79-
throw new Error("Device ID calculation not started");
74+
this.abortController = new AbortController();
8075
}
8176

82-
try {
83-
const deviceId = await getDeviceId({
84-
getMachineId: this.getMachineId,
85-
onError: (reason, error) => {
86-
this.handleDeviceIdError(reason, String(error));
87-
},
88-
timeout: this.timeout,
89-
abortSignal: this.abortController.signal,
90-
});
91-
92-
// Cache the result
93-
this.deviceId = deviceId;
94-
return deviceId;
95-
} catch (error) {
96-
// Check if this was an abort error
97-
if (error instanceof Error && error.name === "AbortError") {
98-
throw error; // Re-throw abort errors
99-
}
100-
101-
this.logger.debug({
102-
id: LogId.deviceIdResolutionError,
103-
context: "deviceId",
104-
message: `Failed to get device ID: ${String(error)}`,
105-
});
106-
107-
// Cache the fallback value
108-
this.deviceId = "unknown";
109-
return "unknown";
110-
} finally {
111-
this.abortController = undefined;
112-
}
77+
this.deviceIdPromise = getDeviceId({
78+
getMachineId: this.getMachineId,
79+
onError: (reason, error) => {
80+
this.handleDeviceIdError(reason, String(error));
81+
},
82+
timeout: this.timeout,
83+
abortSignal: this.abortController.signal,
84+
});
85+
86+
return this.deviceIdPromise;
11387
}
11488

11589
private handleDeviceIdError(reason: string, error: string): void {
90+
this.deviceIdPromise = Promise.resolve("unknown");
91+
11692
switch (reason) {
11793
case "resolutionError":
11894
this.logger.debug({

‎src/index.ts‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,6 @@ 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 { DeviceId } from "./helpers/deviceId.js";
4645

4746
async function main(): Promise<void> {
4847
systemCA().catch(() => undefined); // load system CA asynchronously as in mongosh
@@ -51,7 +50,6 @@ async function main(): Promise<void> {
5150
assertVersionMode();
5251

5352
const transportRunner = config.transport === "stdio" ? new StdioRunner(config) : new StreamableHttpRunner(config);
54-
const deviceId = DeviceId.create(transportRunner.logger);
5553
const shutdown = (): void => {
5654
transportRunner.logger.info({
5755
id: LogId.serverCloseRequested,
@@ -62,7 +60,6 @@ async function main(): Promise<void> {
6260
transportRunner
6361
.close()
6462
.then(() => {
65-
deviceId.close();
6663
transportRunner.logger.info({
6764
id: LogId.serverClosed,
6865
context: "server",
@@ -71,7 +68,6 @@ async function main(): Promise<void> {
7168
process.exit(0);
7269
})
7370
.catch((error: unknown) => {
74-
deviceId.close();
7571
transportRunner.logger.error({
7672
id: LogId.serverCloseFailure,
7773
context: "server",

‎src/telemetry/telemetry.ts‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,13 +33,12 @@ export class Telemetry {
3333
static create(
3434
session: Session,
3535
userConfig: UserConfig,
36+
deviceId: DeviceId,
3637
{
3738
commonProperties = { ...MACHINE_METADATA },
3839
eventCache = EventCache.getInstance(),
39-
deviceId = DeviceId.create(session.logger),
4040
}: {
4141
eventCache?: EventCache;
42-
deviceId?: DeviceId;
4342
commonProperties?: CommonProperties;
4443
} = {}
4544
): Telemetry {

‎src/transports/base.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,11 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
77
import { CompositeLogger, ConsoleLogger, DiskLogger, LoggerBase, McpLogger } from "../common/logger.js";
88
import { ExportsManager } from "../common/exportsManager.js";
99
import { ConnectionManager } from "../common/connectionManager.js";
10+
import { DeviceId } from "../helpers/deviceId.js";
1011

1112
export abstract class TransportRunnerBase {
1213
public logger: LoggerBase;
14+
public deviceId: DeviceId;
1315

1416
protected constructor(protected readonly userConfig: UserConfig) {
1517
const loggers: LoggerBase[] = [];
@@ -28,6 +30,7 @@ export abstract class TransportRunnerBase {
2830
}
2931

3032
this.logger = new CompositeLogger(...loggers);
33+
this.deviceId = DeviceId.create(this.logger);
3134
}
3235

3336
protected setupServer(userConfig: UserConfig): Server {
@@ -43,7 +46,7 @@ export abstract class TransportRunnerBase {
4346

4447
const logger = new CompositeLogger(...loggers);
4548
const exportsManager = ExportsManager.init(userConfig, logger);
46-
const connectionManager = new ConnectionManager(userConfig, driverOptions, logger);
49+
const connectionManager = new ConnectionManager(userConfig, driverOptions, logger, this.deviceId);
4750

4851
const session = new Session({
4952
apiBaseUrl: userConfig.apiBaseUrl,
@@ -54,7 +57,7 @@ export abstract class TransportRunnerBase {
5457
connectionManager,
5558
});
5659

57-
const telemetry = Telemetry.create(session, userConfig);
60+
const telemetry = Telemetry.create(session, userConfig, this.deviceId);
5861

5962
return new Server({
6063
mcpServer,

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

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,19 +5,20 @@ import nodeMachineId from "node-machine-id";
55

66
describe("Device ID", () => {
77
let testLogger: CompositeLogger;
8+
let deviceId: DeviceId;
89

910
beforeEach(() => {
1011
testLogger = new CompositeLogger();
1112
testLogger.debug = vi.fn();
1213
});
1314

1415
afterEach(() => {
15-
DeviceId.create(testLogger).close();
16+
deviceId?.close();
1617
});
1718

1819
describe("when resolving device ID", () => {
1920
it("should successfully resolve device ID in real environment", async () => {
20-
const deviceId = DeviceId.create(testLogger);
21+
deviceId = DeviceId.create(testLogger);
2122
const result = await deviceId.get();
2223

2324
expect(result).not.toBe("unknown");
@@ -29,7 +30,7 @@ describe("Device ID", () => {
2930
it("should cache device ID after first resolution", async () => {
3031
// spy on machineId
3132
const machineIdSpy = vi.spyOn(nodeMachineId, "machineId");
32-
const deviceId = DeviceId.create(testLogger);
33+
deviceId = DeviceId.create(testLogger);
3334

3435
// First call
3536
const result1 = await deviceId.get();
@@ -43,7 +44,7 @@ describe("Device ID", () => {
4344
});
4445

4546
it("should handle concurrent device ID requests correctly", async () => {
46-
const deviceId = DeviceId.create(testLogger);
47+
deviceId = DeviceId.create(testLogger);
4748

4849
const promises = Array.from({ length: 5 }, () => deviceId.get());
4950

@@ -79,7 +80,7 @@ describe("Device ID", () => {
7980
reject(new Error("Machine ID failed"));
8081
});
8182
});
82-
const deviceId = DeviceId.create(testLogger);
83+
deviceId = DeviceId.create(testLogger);
8384
const handleDeviceIdErrorSpy = vi.spyOn(deviceId, "handleDeviceIdError" as keyof DeviceId);
8485

8586
const result = await deviceId.get();
@@ -99,7 +100,7 @@ describe("Device ID", () => {
99100
});
100101
});
101102

102-
const deviceId = DeviceId.create(testLogger, 100); // Short timeout
103+
deviceId = DeviceId.create(testLogger, 100); // Short timeout
103104
const handleDeviceIdErrorSpy = vi.spyOn(deviceId, "handleDeviceIdError" as keyof DeviceId);
104105

105106
const result = await deviceId.get();

‎tests/integration/helpers.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { McpError, ResourceUpdatedNotificationSchema } from "@modelcontextprotoc
1111
import { config, driverOptions } from "../../src/common/config.js";
1212
import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest";
1313
import { ConnectionManager, ConnectionState } from "../../src/common/connectionManager.js";
14+
import { DeviceId } from "../../src/helpers/deviceId.js";
1415

1516
interface ParameterInfo {
1617
name: string;
@@ -68,7 +69,8 @@ export function setupIntegrationTest(
6869

6970
const exportsManager = ExportsManager.init(userConfig, logger);
7071

71-
const connectionManager = new ConnectionManager(userConfig, driverOptions, logger);
72+
const deviceId = DeviceId.create(logger);
73+
const connectionManager = new ConnectionManager(userConfig, driverOptions, logger, deviceId);
7274

7375
const session = new Session({
7476
apiBaseUrl: userConfig.apiBaseUrl,
@@ -87,7 +89,7 @@ export function setupIntegrationTest(
8789

8890
userConfig.telemetry = "disabled";
8991

90-
const telemetry = Telemetry.create(session, userConfig);
92+
const telemetry = Telemetry.create(session, userConfig, deviceId);
9193

9294
mcpServer = new Server({
9395
session,

‎tests/integration/telemetry.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,10 @@ describe("Telemetry", () => {
1919
apiBaseUrl: "",
2020
logger,
2121
exportsManager: ExportsManager.init(config, logger),
22-
connectionManager: new ConnectionManager(config, driverOptions, logger),
22+
connectionManager: new ConnectionManager(config, driverOptions, logger, deviceId),
2323
}),
2424
config,
25-
{ deviceId }
25+
deviceId
2626
);
2727

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

‎tests/unit/common/session.test.ts‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,25 +10,27 @@ import { DeviceId } from "../../../src/helpers/deviceId.js";
1010
vi.mock("@mongosh/service-provider-node-driver");
1111

1212
const MockNodeDriverServiceProvider = vi.mocked(NodeDriverServiceProvider);
13+
const MockDeviceId = vi.mocked(DeviceId.create(new CompositeLogger()));
1314

1415
describe("Session", () => {
1516
let session: Session;
1617
let mockDeviceId: Mocked<DeviceId>;
1718

1819
beforeEach(() => {
1920
const logger = new CompositeLogger();
20-
mockDeviceId = vi.mocked(DeviceId.create(logger));
21-
mockDeviceId.get = vi.fn().mockResolvedValue("test-device-id");
21+
22+
mockDeviceId = MockDeviceId;
2223

2324
session = new Session({
2425
apiClientId: "test-client-id",
2526
apiBaseUrl: "https://api.test.com",
2627
logger,
2728
exportsManager: ExportsManager.init(config, logger),
28-
connectionManager: new ConnectionManager(config, driverOptions, logger),
29+
connectionManager: new ConnectionManager(config, driverOptions, logger, mockDeviceId),
2930
});
3031

3132
MockNodeDriverServiceProvider.connect = vi.fn().mockResolvedValue({} as unknown as NodeDriverServiceProvider);
33+
MockDeviceId.get = vi.fn().mockResolvedValue("test-device-id");
3234
});
3335

3436
describe("connectToMongoDB", () => {

‎tests/unit/helpers/connectionOptions.test.ts‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,17 @@
1-
import { beforeEach, describe, expect, it, Mocked, vi } from "vitest";
1+
import { beforeEach, describe, expect, it, vi } from "vitest";
22
import { setAppNameParamIfMissing } from "../../../src/helpers/connectionOptions.js";
33
import { DeviceId } from "../../../src/helpers/deviceId.js";
44
import { CompositeLogger } from "../../../src/common/logger.js";
55

6+
const MockDeviceId = vi.mocked(DeviceId.create(new CompositeLogger()));
7+
68
describe("Connection Options", () => {
79
let testLogger: CompositeLogger;
8-
let mockDeviceId: Mocked<DeviceId>;
910

1011
beforeEach(() => {
1112
testLogger = new CompositeLogger();
1213
testLogger.debug = vi.fn();
13-
14-
mockDeviceId = vi.mocked(DeviceId.create(testLogger));
15-
mockDeviceId.get = vi.fn().mockResolvedValue("test-device-id");
14+
MockDeviceId.get = vi.fn().mockResolvedValue("test-device-id");
1615
});
1716

1817
describe("setAppNameParamIfMissing", () => {
@@ -23,7 +22,7 @@ describe("Connection Options", () => {
2322
components: {
2423
appName: "TestApp",
2524
clientName: "TestClient",
26-
deviceId: mockDeviceId.get(),
25+
deviceId: MockDeviceId.get(),
2726
},
2827
});
2928

@@ -65,7 +64,7 @@ describe("Connection Options", () => {
6564
connectionString,
6665
components: {
6766
appName: "TestApp",
68-
deviceId: mockDeviceId.get(),
67+
deviceId: MockDeviceId.get(),
6968
},
7069
});
7170

@@ -92,7 +91,7 @@ describe("Connection Options", () => {
9291
components: {
9392
appName: "TestApp",
9493
clientName: "TestClient",
95-
deviceId: mockDeviceId.get(),
94+
deviceId: MockDeviceId.get(),
9695
},
9796
});
9897

0 commit comments

Comments
 (0)