Skip to content

Commit ca68195

Browse files
blvablva
authored
feat: update connectionString appName param - [MCP-68] (#406)
1 parent 4314bc3 commit ca68195

File tree

20 files changed

+708
-156
lines changed

20 files changed

+708
-156
lines changed

‎src/common/connectionManager.ts

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ 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 { DeviceId } from "../helpers/deviceId.js";
10+
import { AppNameComponents } from "../helpers/connectionOptions.js";
911
import { CompositeLogger, LogId } from "./logger.js";
1012
import { ConnectionInfo, generateConnectionInfoFromCliArgs } from "@mongosh/arg-parser";
1113

@@ -69,12 +71,15 @@ export interface ConnectionManagerEvents {
6971

7072
export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
7173
private state: AnyConnectionState;
74+
private deviceId: DeviceId;
75+
private clientName: string;
7276
private bus: EventEmitter;
7377

7478
constructor(
7579
private userConfig: UserConfig,
7680
private driverOptions: DriverOptions,
7781
private logger: CompositeLogger,
82+
deviceId: DeviceId,
7883
bus?: EventEmitter
7984
) {
8085
super();
@@ -84,6 +89,13 @@ export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
8489

8590
this.bus.on("mongodb-oidc-plugin:auth-failed", this.onOidcAuthFailed.bind(this));
8691
this.bus.on("mongodb-oidc-plugin:auth-succeeded", this.onOidcAuthSucceeded.bind(this));
92+
93+
this.deviceId = deviceId;
94+
this.clientName = "unknown";
95+
}
96+
97+
setClientName(clientName: string): void {
98+
this.clientName = clientName;
8799
}
88100

89101
async connect(settings: ConnectionSettings): Promise<AnyConnectionState> {
@@ -98,9 +110,15 @@ export class ConnectionManager extends EventEmitter<ConnectionManagerEvents> {
98110

99111
try {
100112
settings = { ...settings };
101-
settings.connectionString = setAppNameParamIfMissing({
113+
const appNameComponents: AppNameComponents = {
114+
appName: `${packageInfo.mcpServerName} ${packageInfo.version}`,
115+
deviceId: this.deviceId.get(),
116+
clientName: this.clientName,
117+
};
118+
119+
settings.connectionString = await setAppNameParamIfMissing({
102120
connectionString: settings.connectionString,
103-
defaultAppName: `${packageInfo.mcpServerName} ${packageInfo.version}`,
121+
components: appNameComponents,
104122
});
105123

106124
connectionInfo = generateConnectionInfoFromCliArgs({

‎src/common/logger.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ export const LogId = {
1414
serverClosed: mongoLogId(1_000_004),
1515
serverCloseFailure: mongoLogId(1_000_005),
1616
serverDuplicateLoggers: mongoLogId(1_000_006),
17+
serverMcpClientSet: mongoLogId(1_000_007),
1718

1819
atlasCheckCredentials: mongoLogId(1_001_001),
1920
atlasDeleteDatabaseUserFailure: mongoLogId(1_001_002),
@@ -30,8 +31,8 @@ export const LogId = {
3031
telemetryEmitStart: mongoLogId(1_002_003),
3132
telemetryEmitSuccess: mongoLogId(1_002_004),
3233
telemetryMetadataError: mongoLogId(1_002_005),
33-
telemetryDeviceIdFailure: mongoLogId(1_002_006),
34-
telemetryDeviceIdTimeout: mongoLogId(1_002_007),
34+
deviceIdResolutionError: mongoLogId(1_002_006),
35+
deviceIdTimeout: mongoLogId(1_002_007),
3536

3637
toolExecute: mongoLogId(1_003_001),
3738
toolExecuteFailure: mongoLogId(1_003_002),

‎src/common/session.ts

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,10 @@ export class Session extends EventEmitter<SessionEvents> {
3434
readonly exportsManager: ExportsManager;
3535
readonly connectionManager: ConnectionManager;
3636
readonly apiClient: ApiClient;
37-
agentRunner?: {
38-
name: string;
39-
version: string;
37+
mcpClient?: {
38+
name?: string;
39+
version?: string;
40+
title?: string;
4041
};
4142

4243
public logger: CompositeLogger;
@@ -69,13 +70,24 @@ export class Session extends EventEmitter<SessionEvents> {
6970
this.connectionManager.on("connection-errored", (error) => this.emit("connection-error", error.errorReason));
7071
}
7172

72-
setAgentRunner(agentRunner: Implementation | undefined): void {
73-
if (agentRunner?.name && agentRunner?.version) {
74-
this.agentRunner = {
75-
name: agentRunner.name,
76-
version: agentRunner.version,
77-
};
73+
setMcpClient(mcpClient: Implementation | undefined): void {
74+
if (!mcpClient) {
75+
this.connectionManager.setClientName("unknown");
76+
this.logger.debug({
77+
id: LogId.serverMcpClientSet,
78+
context: "session",
79+
message: "MCP client info not found",
80+
});
7881
}
82+
83+
this.mcpClient = {
84+
name: mcpClient?.name || "unknown",
85+
version: mcpClient?.version || "unknown",
86+
title: mcpClient?.title || "unknown",
87+
};
88+
89+
// Set the client name on the connection manager for appName generation
90+
this.connectionManager.setClientName(this.mcpClient.name || "unknown");
7991
}
8092

8193
async disconnect(): Promise<void> {

‎src/helpers/connectionOptions.ts

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,59 @@
11
import { MongoClientOptions } from "mongodb";
22
import ConnectionString from "mongodb-connection-string-url";
33

4-
export function setAppNameParamIfMissing({
4+
export interface AppNameComponents {
5+
appName: string;
6+
deviceId?: Promise<string>;
7+
clientName?: string;
8+
}
9+
10+
/**
11+
* Sets the appName parameter with the extended format: appName--deviceId--clientName
12+
* Only sets the appName if it's not already present in the connection string
13+
* @param connectionString - The connection string to modify
14+
* @param components - The components to build the appName from
15+
* @returns The modified connection string
16+
*/
17+
export async function setAppNameParamIfMissing({
518
connectionString,
6-
defaultAppName,
19+
components,
720
}: {
821
connectionString: string;
9-
defaultAppName?: string;
10-
}): string {
22+
components: AppNameComponents;
23+
}): Promise<string> {
1124
const connectionStringUrl = new ConnectionString(connectionString);
12-
1325
const searchParams = connectionStringUrl.typedSearchParams<MongoClientOptions>();
1426

15-
if (!searchParams.has("appName") && defaultAppName !== undefined) {
16-
searchParams.set("appName", defaultAppName);
27+
// Only set appName if it's not already present
28+
if (searchParams.has("appName")) {
29+
return connectionStringUrl.toString();
1730
}
1831

32+
const appName = components.appName || "unknown";
33+
const deviceId = components.deviceId ? await components.deviceId : "unknown";
34+
const clientName = components.clientName || "unknown";
35+
36+
// Build the extended appName format: appName--deviceId--clientName
37+
const extendedAppName = `${appName}--${deviceId}--${clientName}`;
38+
39+
searchParams.set("appName", extendedAppName);
40+
1941
return connectionStringUrl.toString();
2042
}
43+
44+
/**
45+
* Validates the connection string
46+
* @param connectionString - The connection string to validate
47+
* @param looseValidation - Whether to allow loose validation
48+
* @returns void
49+
* @throws Error if the connection string is invalid
50+
*/
51+
export function validateConnectionString(connectionString: string, looseValidation: boolean): void {
52+
try {
53+
new ConnectionString(connectionString, { looseValidation });
54+
} catch (error) {
55+
throw new Error(
56+
`Invalid connection string with error: ${error instanceof Error ? error.message : String(error)}`
57+
);
58+
}
59+
}

‎src/helpers/deviceId.ts

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
import { getDeviceId } from "@mongodb-js/device-id";
2+
import nodeMachineId from "node-machine-id";
3+
import { LogId, LoggerBase } from "../common/logger.js";
4+
5+
export const DEVICE_ID_TIMEOUT = 3000;
6+
7+
export class DeviceId {
8+
private deviceId: string | undefined = undefined;
9+
private deviceIdPromise: Promise<string> | undefined = undefined;
10+
private abortController: AbortController | undefined = undefined;
11+
private logger: LoggerBase;
12+
private readonly getMachineId: () => Promise<string>;
13+
private timeout: number;
14+
private static instance: DeviceId | undefined = undefined;
15+
16+
private constructor(logger: LoggerBase, timeout: number = DEVICE_ID_TIMEOUT) {
17+
this.logger = logger;
18+
this.timeout = timeout;
19+
this.getMachineId = (): Promise<string> => nodeMachineId.machineId(true);
20+
}
21+
22+
public static create(logger: LoggerBase, timeout?: number): DeviceId {
23+
if (this.instance) {
24+
throw new Error("DeviceId instance already exists, use get() to retrieve the device ID");
25+
}
26+
27+
const instance = new DeviceId(logger, timeout ?? DEVICE_ID_TIMEOUT);
28+
instance.setup();
29+
30+
this.instance = instance;
31+
32+
return instance;
33+
}
34+
35+
private setup(): void {
36+
this.deviceIdPromise = this.calculateDeviceId();
37+
}
38+
39+
/**
40+
* Closes the device ID calculation promise and abort controller.
41+
*/
42+
public close(): void {
43+
if (this.abortController) {
44+
this.abortController.abort();
45+
this.abortController = undefined;
46+
}
47+
48+
this.deviceId = undefined;
49+
this.deviceIdPromise = undefined;
50+
DeviceId.instance = undefined;
51+
}
52+
53+
/**
54+
* Gets the device ID, waiting for the calculation to complete if necessary.
55+
* @returns Promise that resolves to the device ID string
56+
*/
57+
public get(): Promise<string> {
58+
if (this.deviceId) {
59+
return Promise.resolve(this.deviceId);
60+
}
61+
62+
if (this.deviceIdPromise) {
63+
return this.deviceIdPromise;
64+
}
65+
66+
return this.calculateDeviceId();
67+
}
68+
69+
/**
70+
* Internal method that performs the actual device ID calculation.
71+
*/
72+
private async calculateDeviceId(): Promise<string> {
73+
if (!this.abortController) {
74+
this.abortController = new AbortController();
75+
}
76+
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;
87+
}
88+
89+
private handleDeviceIdError(reason: string, error: string): void {
90+
this.deviceIdPromise = Promise.resolve("unknown");
91+
92+
switch (reason) {
93+
case "resolutionError":
94+
this.logger.debug({
95+
id: LogId.deviceIdResolutionError,
96+
context: "deviceId",
97+
message: `Resolution error: ${String(error)}`,
98+
});
99+
break;
100+
case "timeout":
101+
this.logger.debug({
102+
id: LogId.deviceIdTimeout,
103+
context: "deviceId",
104+
message: "Device ID retrieval timed out",
105+
noRedaction: true,
106+
});
107+
break;
108+
case "abort":
109+
// No need to log in the case of 'abort' errors
110+
break;
111+
}
112+
}
113+
}

‎src/index.ts

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,6 @@ async function main(): Promise<void> {
5050
assertVersionMode();
5151

5252
const transportRunner = config.transport === "stdio" ? new StdioRunner(config) : new StreamableHttpRunner(config);
53-
5453
const shutdown = (): void => {
5554
transportRunner.logger.info({
5655
id: LogId.serverCloseRequested,

‎src/server.ts

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import {
1717
} from "@modelcontextprotocol/sdk/types.js";
1818
import assert from "assert";
1919
import { ToolBase } from "./tools/tool.js";
20+
import { validateConnectionString } from "./helpers/connectionOptions.js";
2021

2122
export interface ServerOptions {
2223
session: Session;
@@ -97,12 +98,14 @@ export class Server {
9798
});
9899

99100
this.mcpServer.server.oninitialized = (): void => {
100-
this.session.setAgentRunner(this.mcpServer.server.getClientVersion());
101+
this.session.setMcpClient(this.mcpServer.server.getClientVersion());
102+
// Placed here to start the connection to the config connection string as soon as the server is initialized.
103+
void this.connectToConfigConnectionString();
101104

102105
this.session.logger.info({
103106
id: LogId.serverInitialized,
104107
context: "server",
105-
message: `Server started with transport ${transport.constructor.name} and agent runner ${this.session.agentRunner?.name}`,
108+
message: `Server started with transport ${transport.constructor.name} and agent runner ${this.session.mcpClient?.name}`,
106109
});
107110

108111
this.emitServerEvent("start", Date.now() - this.startTime);
@@ -188,20 +191,20 @@ export class Server {
188191
}
189192

190193
private async validateConfig(): Promise<void> {
194+
// Validate connection string
191195
if (this.userConfig.connectionString) {
192196
try {
193-
await this.session.connectToMongoDB({
194-
connectionString: this.userConfig.connectionString,
195-
});
197+
validateConnectionString(this.userConfig.connectionString, false);
196198
} catch (error) {
197-
console.error(
198-
"Failed to connect to MongoDB instance using the connection string from the config: ",
199-
error
199+
console.error("Connection string validation failed with error: ", error);
200+
throw new Error(
201+
"Connection string validation failed with error: " +
202+
(error instanceof Error ? error.message : String(error))
200203
);
201-
throw new Error("Failed to connect to MongoDB instance using the connection string from the config");
202204
}
203205
}
204206

207+
// Validate API client credentials
205208
if (this.userConfig.apiClientId && this.userConfig.apiClientSecret) {
206209
try {
207210
await this.session.apiClient.validateAccessToken();
@@ -219,4 +222,20 @@ export class Server {
219222
}
220223
}
221224
}
225+
226+
private async connectToConfigConnectionString(): Promise<void> {
227+
if (this.userConfig.connectionString) {
228+
try {
229+
await this.session.connectToMongoDB({
230+
connectionString: this.userConfig.connectionString,
231+
});
232+
} catch (error) {
233+
console.error(
234+
"Failed to connect to MongoDB instance using the connection string from the config: ",
235+
error
236+
);
237+
throw new Error("Failed to connect to MongoDB instance using the connection string from the config");
238+
}
239+
}
240+
}
222241
}

0 commit comments

Comments
 (0)