diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index c34ae30292d..2217904c2f2 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import vscode, { env, version } from 'vscode' +import vscode, { version } from 'vscode' import * as nls from 'vscode-nls' import { LanguageClient, LanguageClientOptions, RequestType, State } from 'vscode-languageclient' import { InlineCompletionManager } from '../app/inline/completion' @@ -38,6 +38,7 @@ import { getOptOutPreference, isAmazonLinux2, getClientId, + getClientName, extensionVersion, isSageMaker, DevSettings, @@ -163,7 +164,7 @@ export async function startLanguageServer( initializationOptions: { aws: { clientInfo: { - name: env.appName, + name: getClientName(), version: version, extension: { name: 'AmazonQ-For-VSCode', diff --git a/packages/core/src/shared/extensionUtilities.ts b/packages/core/src/shared/extensionUtilities.ts index dc6faeaf1dc..80bedf1e0f6 100644 --- a/packages/core/src/shared/extensionUtilities.ts +++ b/packages/core/src/shared/extensionUtilities.ts @@ -150,7 +150,11 @@ function createCloud9Properties(company: string): IdeProperties { } } -function isSageMakerUnifiedStudio(): boolean { +/** + * export method - for testing purposes only + * @internal + */ +export function isSageMakerUnifiedStudio(): boolean { if (serviceName === notInitialized) { serviceName = process.env.SERVICE_NAME ?? '' isSMUS = serviceName === sageMakerUnifiedStudio @@ -158,6 +162,15 @@ function isSageMakerUnifiedStudio(): boolean { return isSMUS } +/** + * Reset cached SageMaker state - for testing purposes only + * @internal + */ +export function resetSageMakerState(): void { + serviceName = notInitialized + isSMUS = false +} + /** * Decides if the current system is (the specified flavor of) Cloud9. */ @@ -177,17 +190,17 @@ export function isCloud9(flavor: 'classic' | 'codecatalyst' | 'any' = 'any'): bo */ export function isSageMaker(appName: 'SMAI' | 'SMUS' = 'SMAI'): boolean { // Check for SageMaker-specific environment variables first + let hasSMEnvVars: boolean = false if (hasSageMakerEnvVars()) { getLogger().debug('SageMaker environment detected via environment variables') - return true + hasSMEnvVars = true } - // Fall back to app name checks switch (appName) { case 'SMAI': - return vscode.env.appName === sageMakerAppname + return vscode.env.appName === sageMakerAppname && hasSMEnvVars case 'SMUS': - return vscode.env.appName === sageMakerAppname && isSageMakerUnifiedStudio() + return vscode.env.appName === sageMakerAppname && isSageMakerUnifiedStudio() && hasSMEnvVars default: return false } diff --git a/packages/core/src/shared/index.ts b/packages/core/src/shared/index.ts index c89360a01dd..8b62fd3c5dc 100644 --- a/packages/core/src/shared/index.ts +++ b/packages/core/src/shared/index.ts @@ -27,7 +27,7 @@ export { Prompter } from './ui/prompter' export { VirtualFileSystem } from './virtualFilesystem' export { VirtualMemoryFile } from './virtualMemoryFile' export { AmazonqCreateUpload, Metric } from './telemetry/telemetry' -export { getClientId, getOperatingSystem, getOptOutPreference } from './telemetry/util' +export { getClientId, getClientName, getOperatingSystem, getOptOutPreference } from './telemetry/util' export { extensionVersion } from './vscode/env' export { cast } from './utilities/typeConstructors' export * as workspaceUtils from './utilities/workspaceUtils' diff --git a/packages/core/src/shared/telemetry/util.ts b/packages/core/src/shared/telemetry/util.ts index 310c36b82d6..dc57148393b 100644 --- a/packages/core/src/shared/telemetry/util.ts +++ b/packages/core/src/shared/telemetry/util.ts @@ -481,3 +481,15 @@ export function withTelemetryContext(opts: TelemetryContextArgs) { }) } } + +/** + * Used to identify the q client info and send the respective origin parameter from LSP to invoke Maestro service at CW API level + * + * Returns default value of vscode appName or AmazonQ-For-SMUS-CE in case of a sagemaker unified studio environment + */ +export function getClientName(): string { + if (isSageMaker('SMUS')) { + return 'AmazonQ-For-SMUS-CE' + } + return env.appName +} diff --git a/packages/core/src/test/shared/extensionUtilities.test.ts b/packages/core/src/test/shared/extensionUtilities.test.ts index 621b31d6603..61238394126 100644 --- a/packages/core/src/test/shared/extensionUtilities.test.ts +++ b/packages/core/src/test/shared/extensionUtilities.test.ts @@ -18,6 +18,8 @@ import globals from '../../shared/extensionGlobals' import { maybeShowMinVscodeWarning } from '../../shared/extensionStartup' import { getTestWindow } from './vscode/window' import { assertTelemetry } from '../testUtil' +import { isSageMaker } from '../../shared/extensionUtilities' +import { hasSageMakerEnvVars } from '../../shared/vscode/env' describe('extensionUtilities', function () { it('maybeShowMinVscodeWarning', async () => { @@ -361,3 +363,146 @@ describe('UserActivity', function () { return event.event } }) + +describe('isSageMaker', function () { + let sandbox: sinon.SinonSandbox + const env = require('../../shared/vscode/env') + const utils = require('../../shared/extensionUtilities') + + beforeEach(function () { + sandbox = sinon.createSandbox() + utils.resetSageMakerState() + }) + + afterEach(function () { + sandbox.restore() + delete process.env.SERVICE_NAME + }) + + describe('SMAI detection', function () { + it('returns true when both app name and env vars match', function () { + sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor') + sandbox.stub(env, 'hasSageMakerEnvVars').returns(true) + + assert.strictEqual(isSageMaker('SMAI'), true) + }) + + it('returns false when app name is different', function () { + sandbox.stub(vscode.env, 'appName').value('Visual Studio Code') + sandbox.stub(env, 'hasSageMakerEnvVars').returns(true) + + assert.strictEqual(isSageMaker('SMAI'), false) + }) + + it('returns false when env vars are missing', function () { + sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor') + sandbox.stub(env, 'hasSageMakerEnvVars').returns(false) + + assert.strictEqual(isSageMaker('SMAI'), false) + }) + + it('defaults to SMAI when no parameter provided', function () { + sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor') + sandbox.stub(env, 'hasSageMakerEnvVars').returns(true) + + assert.strictEqual(isSageMaker(), true) + }) + }) + + describe('SMUS detection', function () { + it('returns true when all conditions are met', function () { + sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor') + sandbox.stub(env, 'hasSageMakerEnvVars').returns(true) + process.env.SERVICE_NAME = 'SageMakerUnifiedStudio' + + assert.strictEqual(isSageMaker('SMUS'), true) + }) + + it('returns false when unified studio is missing', function () { + sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor') + sandbox.stub(env, 'hasSageMakerEnvVars').returns(true) + process.env.SERVICE_NAME = 'SomeOtherService' + + assert.strictEqual(isSageMaker('SMUS'), false) + }) + + it('returns false when env vars are missing', function () { + sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor') + sandbox.stub(env, 'hasSageMakerEnvVars').returns(false) + process.env.SERVICE_NAME = 'SageMakerUnifiedStudio' + + assert.strictEqual(isSageMaker('SMUS'), false) + }) + + it('returns false when app name is different', function () { + sandbox.stub(vscode.env, 'appName').value('Visual Studio Code') + sandbox.stub(env, 'hasSageMakerEnvVars').returns(true) + process.env.SERVICE_NAME = 'SageMakerUnifiedStudio' + + assert.strictEqual(isSageMaker('SMUS'), false) + }) + }) + + it('returns false for invalid appName parameter', function () { + sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor') + sandbox.stub(env, 'hasSageMakerEnvVars').returns(true) + + // @ts-ignore - Testing invalid input + assert.strictEqual(isSageMaker('INVALID'), false) + }) +}) + +describe('hasSageMakerEnvVars', function () { + let originalEnv: NodeJS.ProcessEnv + + beforeEach(function () { + originalEnv = { ...process.env } + // Clear all SageMaker-related env vars + delete process.env.SAGEMAKER_APP_TYPE + delete process.env.SAGEMAKER_INTERNAL_IMAGE_URI + delete process.env.STUDIO_LOGGING_DIR + delete process.env.SM_APP_TYPE + delete process.env.SM_INTERNAL_IMAGE_URI + delete process.env.SERVICE_NAME + }) + + afterEach(function () { + process.env = originalEnv + }) + + const testCases = [ + { env: 'SAGEMAKER_APP_TYPE', value: 'JupyterServer', expected: true }, + { env: 'SAGEMAKER_INTERNAL_IMAGE_URI', value: 'some-uri', expected: true }, + { env: 'STUDIO_LOGGING_DIR', value: '/var/log/studio/app.log', expected: true }, + { env: 'STUDIO_LOGGING_DIR', value: '/var/log/other/app.log', expected: false }, + { env: 'SM_APP_TYPE', value: 'JupyterServer', expected: true }, + { env: 'SM_INTERNAL_IMAGE_URI', value: 'some-uri', expected: true }, + { env: 'SERVICE_NAME', value: 'SageMakerUnifiedStudio', expected: true }, + { env: 'SERVICE_NAME', value: 'SomeOtherService', expected: false }, + ] + + for (const { env, value, expected } of testCases) { + it(`returns ${expected} when ${env} is set to "${value}"`, function () { + process.env[env] = value + + const result = hasSageMakerEnvVars() + + assert.strictEqual(result, expected) + }) + } + + it('returns true when multiple SageMaker env vars are set', function () { + process.env.SAGEMAKER_APP_TYPE = 'JupyterServer' + process.env.SM_APP_TYPE = 'CodeEditor' + + const result = hasSageMakerEnvVars() + + assert.strictEqual(result, true) + }) + + it('returns false when no SageMaker env vars are set', function () { + const result = hasSageMakerEnvVars() + + assert.strictEqual(result, false) + }) +}) diff --git a/packages/core/src/test/shared/telemetry/util.test.ts b/packages/core/src/test/shared/telemetry/util.test.ts index 8d6f3ddc53f..aa4957eea52 100644 --- a/packages/core/src/test/shared/telemetry/util.test.ts +++ b/packages/core/src/test/shared/telemetry/util.test.ts @@ -24,6 +24,10 @@ import { randomUUID } from 'crypto' import { isUuid } from '../../../shared/crypto' import { MetricDatum } from '../../../shared/telemetry/clienttelemetry' import { assertLogsContain } from '../../globalSetup.test' +import { getClientName } from '../../../shared/telemetry/util' +import * as extensionUtilities from '../../../shared/extensionUtilities' +import * as sinon from 'sinon' +import * as vscode from 'vscode' describe('TelemetryConfig', function () { const settingKey = 'aws.telemetry' @@ -391,3 +395,56 @@ describe('validateMetricEvent', function () { assertLogsContain('invalid Metric', false, 'warn') }) }) + +describe('getClientName', function () { + let sandbox: sinon.SinonSandbox + let isSageMakerStub: sinon.SinonStub + + beforeEach(function () { + sandbox = sinon.createSandbox() + isSageMakerStub = sandbox.stub(extensionUtilities, 'isSageMaker') + }) + + afterEach(function () { + sandbox.restore() + }) + + it('returns "AmazonQ-For-SMUS-CE" when in SMUS environment', function () { + isSageMakerStub.withArgs('SMUS').returns(true) + sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor') + + const result = getClientName() + + assert.strictEqual(result, 'AmazonQ-For-SMUS-CE') + assert.ok(isSageMakerStub.calledOnceWith('SMUS')) + }) + + it('returns vscode app name when not in SMUS environment', function () { + const mockAppName = 'Visual Studio Code' + isSageMakerStub.withArgs('SMUS').returns(false) + sandbox.stub(vscode.env, 'appName').value(mockAppName) + + const result = getClientName() + + assert.strictEqual(result, mockAppName) + assert.ok(isSageMakerStub.calledOnceWith('SMUS')) + }) + + it('handles undefined app name gracefully', function () { + isSageMakerStub.withArgs('SMUS').returns(false) + sandbox.stub(vscode.env, 'appName').value(undefined) + + const result = getClientName() + + assert.strictEqual(result, undefined) + }) + + it('prioritizes SMUS detection over app name', function () { + isSageMakerStub.withArgs('SMUS').returns(true) + sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor') + + const result = getClientName() + + assert.strictEqual(result, 'AmazonQ-For-SMUS-CE') + }) +})