diff --git a/src/connectionController.ts b/src/connectionController.ts index 03248ea93..e0f7a58ba 100644 --- a/src/connectionController.ts +++ b/src/connectionController.ts @@ -79,6 +79,8 @@ interface NewConnectionParams { reuseExisting?: boolean; } +export const DEFAULT_TELEMETRY_APP_NAME = `${packageJSON.name} ${packageJSON.version}`; + function isOIDCAuth(connectionString: string): boolean { const authMechanismString = ( new ConnectionString(connectionString).searchParams.get('authMechanism') || @@ -481,7 +483,7 @@ export default class ConnectionController { const connectionOptions = adjustConnectionOptionsBeforeConnect({ connectionOptions: connectionInfo.connectionOptions, connectionId, - defaultAppName: `${packageJSON.name} ${packageJSON.version}`, + defaultAppName: DEFAULT_TELEMETRY_APP_NAME, notifyDeviceFlow, preferences: { forceConnectionOptions: [], diff --git a/src/mcp/mcpConnectionManager.ts b/src/mcp/mcpConnectionManager.ts index a1e9edd61..1020e88b8 100644 --- a/src/mcp/mcpConnectionManager.ts +++ b/src/mcp/mcpConnectionManager.ts @@ -9,8 +9,10 @@ import { type DevtoolsConnectOptions, } from '@mongosh/service-provider-node-driver'; import type { ServiceProvider } from '@mongosh/service-provider-core'; -import { isAtlasStream } from 'mongodb-build-info'; +import { isAtlas, isAtlasStream } from 'mongodb-build-info'; import { MCPLogIds } from './mcpLogIds'; +import ConnectionString from 'mongodb-connection-string-url'; +import { DEFAULT_TELEMETRY_APP_NAME } from '../connectionController'; export interface MCPConnectParams { connectionId: string; @@ -18,14 +20,25 @@ export interface MCPConnectParams { connectOptions: DevtoolsConnectOptions; } +export const MCP_SERVER_TELEMETRY_APP_NAME_SUFFIX = 'MongoDB MCP Server'; + +type MCPConnectionManagerConfig = { + logger: LoggerBase; + getTelemetryAnonymousId: () => string; +}; + export class MCPConnectionManager extends ConnectionManager { + private logger: LoggerBase; + private getTelemetryAnonymousId: () => string; private activeConnection: { id: string; provider: ServiceProvider; } | null = null; - constructor(private readonly logger: LoggerBase) { + constructor({ logger, getTelemetryAnonymousId }: MCPConnectionManagerConfig) { super(); + this.logger = logger; + this.getTelemetryAnonymousId = getTelemetryAnonymousId; } connect(): Promise { @@ -42,13 +55,15 @@ To connect, choose a connection from MongoDB VSCode extensions's sidepanel - htt connectParams: MCPConnectParams, ): Promise { try { + const { connectionId, connectOptions, connectionString } = + this.overridePresetAppName(connectParams); const serviceProvider = await NodeDriverServiceProvider.connect( - connectParams.connectionString, - connectParams.connectOptions, + connectionString, + connectOptions, ); await serviceProvider.runCommand('admin', { hello: 1 }); this.activeConnection = { - id: connectParams.connectionId, + id: connectionId, provider: serviceProvider, }; return this.changeState('connection-success', { @@ -114,4 +129,38 @@ To connect, choose a connection from MongoDB VSCode extensions's sidepanel - htt await this.connectToVSCodeConnection(connectParams); } + + overridePresetAppName(connectParams: MCPConnectParams): MCPConnectParams { + const connectionURL = new ConnectionString(connectParams.connectionString); + const connectOptions: DevtoolsConnectOptions = { + ...connectParams.connectOptions, + }; + const searchParams = + connectionURL.typedSearchParams(); + const appName = searchParams.get('appName'); + + if ( + !appName || + (appName.startsWith(DEFAULT_TELEMETRY_APP_NAME) && + !appName.includes(MCP_SERVER_TELEMETRY_APP_NAME_SUFFIX)) + ) { + const defaultAppName = `${DEFAULT_TELEMETRY_APP_NAME} ${MCP_SERVER_TELEMETRY_APP_NAME_SUFFIX}`; + const telemetryAnonymousId = this.getTelemetryAnonymousId(); + const connectionId = connectParams.connectionId; + const newAppName = isAtlas(connectParams.connectionString) + ? `${defaultAppName}${ + telemetryAnonymousId ? `--${telemetryAnonymousId}` : '' + }--${connectionId}` + : defaultAppName; + + searchParams.set('appName', newAppName); + connectOptions.appName = newAppName; + } + + return { + connectionId: connectParams.connectionId, + connectionString: connectionURL.toString(), + connectOptions, + }; + } } diff --git a/src/mcp/mcpController.ts b/src/mcp/mcpController.ts index 733992bba..0fd34d7e4 100644 --- a/src/mcp/mcpController.ts +++ b/src/mcp/mcpController.ts @@ -39,15 +39,29 @@ export type MCPServerInfo = { headers: Record; }; +type MCPControllerConfig = { + context: vscode.ExtensionContext; + connectionController: ConnectionController; + getTelemetryAnonymousId: () => string; +}; + export class MCPController { + private context: vscode.ExtensionContext; + private connectionController: ConnectionController; + private getTelemetryAnonymousId: () => string; + private didChangeEmitter = new vscode.EventEmitter(); private server?: MCPServerInfo; private mcpConnectionManager?: MCPConnectionManager; - constructor( - private readonly context: vscode.ExtensionContext, - private readonly connectionController: ConnectionController, - ) { + constructor({ + context, + connectionController, + getTelemetryAnonymousId, + }: MCPControllerConfig) { + this.context = context; + this.connectionController = connectionController; + this.getTelemetryAnonymousId = getTelemetryAnonymousId; this.context.subscriptions.push( vscode.lm.registerMcpServerDefinitionProvider('mongodb', { onDidChangeMcpServerDefinitions: this.didChangeEmitter.event, @@ -93,7 +107,10 @@ export class MCPController { logger, }) => { const connectionManager = (this.mcpConnectionManager = - new MCPConnectionManager(logger)); + new MCPConnectionManager({ + logger, + getTelemetryAnonymousId: this.getTelemetryAnonymousId, + })); await this.switchConnectionManagerToCurrentConnection(); return connectionManager; }; diff --git a/src/mdbExtensionController.ts b/src/mdbExtensionController.ts index 329935105..d633f7eb6 100644 --- a/src/mdbExtensionController.ts +++ b/src/mdbExtensionController.ts @@ -167,10 +167,12 @@ export default class MDBExtensionController implements vscode.Disposable { telemetryService: this._telemetryService, }); this._editorsController.registerProviders(); - this._mcpController = new MCPController( + this._mcpController = new MCPController({ context, - this._connectionController, - ); + connectionController: this._connectionController, + getTelemetryAnonymousId: (): string => + this._connectionStorage.getUserAnonymousId(), + }); } subscribeToConfigurationChanges(): void { diff --git a/src/test/suite/mcp/mcpConnectionManager.test.ts b/src/test/suite/mcp/mcpConnectionManager.test.ts index b977a67c6..e753d7cc4 100644 --- a/src/test/suite/mcp/mcpConnectionManager.test.ts +++ b/src/test/suite/mcp/mcpConnectionManager.test.ts @@ -2,11 +2,17 @@ import sinon from 'sinon'; import { afterEach, beforeEach } from 'mocha'; import chai, { expect } from 'chai'; import chaiAsPromised from 'chai-as-promised'; +import ConnectionString from 'mongodb-connection-string-url'; import type { LoggerBase } from '@himanshusinghs/mongodb-mcp-server'; import type { ConnectionStateErrored } from '@himanshusinghs/mongodb-mcp-server'; import type { DevtoolsConnectOptions } from '@mongosh/service-provider-node-driver'; import { NodeDriverServiceProvider } from '@mongosh/service-provider-node-driver'; -import { MCPConnectionManager } from '../../../mcp/mcpConnectionManager'; +import type { MCPConnectParams } from '../../../mcp/mcpConnectionManager'; +import { + MCP_SERVER_TELEMETRY_APP_NAME_SUFFIX, + MCPConnectionManager, +} from '../../../mcp/mcpConnectionManager'; +import { DEFAULT_TELEMETRY_APP_NAME } from '../../../connectionController'; chai.use(chaiAsPromised); @@ -17,9 +23,9 @@ suite('MCPConnectionManager Test Suite', function () { beforeEach(() => { mcpConnectionManager = new MCPConnectionManager({ - error: () => {}, - warning: () => {}, - } as unknown as LoggerBase); + logger: { error: () => {}, warning: () => {} } as unknown as LoggerBase, + getTelemetryAnonymousId: (): string => '1FOO', + }); fakeServiceProvider = { runCommand: (() => Promise.resolve({})) as NodeDriverServiceProvider['runCommand'], @@ -287,4 +293,146 @@ suite('MCPConnectionManager Test Suite', function () { }); }); }); + + suite('#overrideAppNameIfContainsVSCode', function () { + let localConnectionURL: ConnectionString; + let atlasConnectionURL: ConnectionString; + beforeEach(() => { + localConnectionURL = new ConnectionString( + `mongodb://localhost:27017/?appName=${DEFAULT_TELEMETRY_APP_NAME}`, + ); + atlasConnectionURL = new ConnectionString( + 'mongodb://cat-data-sets.cats.mongodb.net/admin', + ); + }); + + for (const { + suiteName, + getConnectionURL, + getConnectionManager, + expectedAppName, + expectedString, + } of [ + { + suiteName: 'when connection string is not atlas', + getConnectionURL: (): ConnectionString => localConnectionURL.clone(), + getConnectionManager: (): MCPConnectionManager => mcpConnectionManager, + expectedAppName: `${DEFAULT_TELEMETRY_APP_NAME} ${MCP_SERVER_TELEMETRY_APP_NAME_SUFFIX}`, + expectedString: (): string => { + const url = localConnectionURL.clone(); + const expectedAppName = `${DEFAULT_TELEMETRY_APP_NAME} ${MCP_SERVER_TELEMETRY_APP_NAME_SUFFIX}`; + url.searchParams.set('appName', expectedAppName); + return url.toString(); + }, + }, + { + suiteName: 'when connection string is atlas', + getConnectionURL: (): ConnectionString => atlasConnectionURL.clone(), + getConnectionManager: (): MCPConnectionManager => mcpConnectionManager, + expectedAppName: `${DEFAULT_TELEMETRY_APP_NAME} ${MCP_SERVER_TELEMETRY_APP_NAME_SUFFIX}--1FOO--1`, + expectedString: (): string => { + const url = atlasConnectionURL.clone(); + const expectedAppName = `${DEFAULT_TELEMETRY_APP_NAME} ${MCP_SERVER_TELEMETRY_APP_NAME_SUFFIX}--1FOO--1`; + url.searchParams.set('appName', expectedAppName); + return url.toString(); + }, + }, + ]) { + suite(suiteName, function () { + suite('and appName is not set', function () { + test('should set appName attribute both in connection string and connection options', function () { + const url = getConnectionURL(); + url.searchParams.delete('appName'); + const connectParams: MCPConnectParams = { + connectionId: '1', + connectionString: url.toString(), + connectOptions: { + productName: 'VSCode', + productDocsLink: 'https://mongodb.com', + appName: DEFAULT_TELEMETRY_APP_NAME, + }, + }; + + expect( + getConnectionManager().overridePresetAppName(connectParams), + ).to.deep.equal({ + connectionId: '1', + connectionString: expectedString(), + connectOptions: { + productName: 'VSCode', + productDocsLink: 'https://mongodb.com', + appName: expectedAppName, + }, + }); + }); + }); + + suite('if appName is set to default vscode app name', function () { + test('should set appName attribute both in connection string and connection options', function () { + const url = getConnectionURL(); + const connectParams: MCPConnectParams = { + connectionId: '1', + connectionString: url.toString(), + connectOptions: { + productName: 'VSCode', + productDocsLink: 'https://mongodb.com', + appName: DEFAULT_TELEMETRY_APP_NAME, + }, + }; + + expect( + getConnectionManager().overridePresetAppName(connectParams), + ).to.deep.equal({ + connectionId: '1', + connectionString: expectedString(), + connectOptions: { + productName: 'VSCode', + productDocsLink: 'https://mongodb.com', + appName: expectedAppName, + }, + }); + }); + }); + + suite('if appName is set to something else', function () { + test('should not override appName attribute both in connection string and connection options', function () { + const url = getConnectionURL(); + url.searchParams.set('appName', 'MongoDB MCP Server 0.0.0'); + const connectParams: MCPConnectParams = { + connectionId: '1', + connectionString: url.toString(), + connectOptions: { + productName: 'VSCode', + productDocsLink: 'https://mongodb.com', + appName: DEFAULT_TELEMETRY_APP_NAME, + }, + }; + + expect( + getConnectionManager().overridePresetAppName(connectParams), + ).to.deep.equal(connectParams); + + // Now for the case when appName is already set to expected MCP server appname + url.searchParams.set( + 'appName', + `${DEFAULT_TELEMETRY_APP_NAME} ${MCP_SERVER_TELEMETRY_APP_NAME_SUFFIX}`, + ); + const nextConnectParams: MCPConnectParams = { + connectionId: '1', + connectionString: url.toString(), + connectOptions: { + productName: 'VSCode', + productDocsLink: 'https://mongodb.com', + appName: DEFAULT_TELEMETRY_APP_NAME, + }, + }; + + expect( + getConnectionManager().overridePresetAppName(nextConnectParams), + ).to.deep.equal(nextConnectParams); + }); + }); + }); + } + }); }); diff --git a/src/test/suite/mcp/mcpController.test.ts b/src/test/suite/mcp/mcpController.test.ts index c91439173..ac55a9f4b 100644 --- a/src/test/suite/mcp/mcpController.test.ts +++ b/src/test/suite/mcp/mcpController.test.ts @@ -33,7 +33,11 @@ suite('MCPController test suite', function () { telemetryService: testTelemetryService, }); - mcpController = new MCPController(extensionContext, connectionController); + mcpController = new MCPController({ + context: extensionContext, + connectionController: connectionController, + getTelemetryAnonymousId: (): string => '1FOO', + }); }); afterEach(async () => {