From c90d0d085a99ab8a89b077f22dd25e8ec92edff0 Mon Sep 17 00:00:00 2001 From: Jayakrishna Parameswaramangalam Date: Tue, 29 Jul 2025 15:15:03 -0700 Subject: [PATCH 1/4] fix(amazonq): update LSP client info name for sagemaker unified studio --- packages/amazonq/src/lsp/client.ts | 5 +++-- packages/core/src/shared/extensionUtilities.ts | 2 +- packages/core/src/shared/index.ts | 2 +- packages/core/src/shared/telemetry/util.ts | 14 +++++++++++++- 4 files changed, 18 insertions(+), 5 deletions(-) diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index d335dae40ef..fd7c858ba34 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' @@ -37,6 +37,7 @@ import { getOptOutPreference, isAmazonLinux2, getClientId, + getClientName, extensionVersion, isSageMaker, DevSettings, @@ -161,7 +162,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..caac78436c8 100644 --- a/packages/core/src/shared/extensionUtilities.ts +++ b/packages/core/src/shared/extensionUtilities.ts @@ -150,7 +150,7 @@ function createCloud9Properties(company: string): IdeProperties { } } -function isSageMakerUnifiedStudio(): boolean { +export function isSageMakerUnifiedStudio(): boolean { if (serviceName === notInitialized) { serviceName = process.env.SERVICE_NAME ?? '' isSMUS = serviceName === sageMakerUnifiedStudio 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..b21d7a087aa 100644 --- a/packages/core/src/shared/telemetry/util.ts +++ b/packages/core/src/shared/telemetry/util.ts @@ -23,7 +23,7 @@ import { mapMetadata, MetadataObj } from './telemetryLogger' import { Result } from './telemetry.gen' import { MetricDatum } from './clienttelemetry' import { isValidationExemptMetric } from './exemptMetrics' -import { isAmazonQ, isCloud9, isSageMaker } from '../../shared/extensionUtilities' +import { isAmazonQ, isCloud9, isSageMaker, isSageMakerUnifiedStudio } from '../../shared/extensionUtilities' import { isUuid, randomUUID } from '../crypto' import { ClassToInterfaceType } from '../utilities/tsUtils' import { asStringifiedStack, FunctionEntry } from './spans' @@ -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() { + if (isSageMaker() && isSageMakerUnifiedStudio()) { + return 'AmazonQ-For-SMUS-CE' + } + return env.appName +} From 8838b61a3b72bc150d613086f0be55b80aa5117c Mon Sep 17 00:00:00 2001 From: Jayakrishna Parameswaramangalam Date: Tue, 29 Jul 2025 15:18:12 -0700 Subject: [PATCH 2/4] update return type for getClientName --- packages/core/src/shared/telemetry/util.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/core/src/shared/telemetry/util.ts b/packages/core/src/shared/telemetry/util.ts index b21d7a087aa..d9e6ff840e6 100644 --- a/packages/core/src/shared/telemetry/util.ts +++ b/packages/core/src/shared/telemetry/util.ts @@ -487,7 +487,7 @@ export function withTelemetryContext(opts: TelemetryContextArgs) { * * Returns default value of vscode appName or AmazonQ-For-SMUS-CE in case of a sagemaker unified studio environment */ -export function getClientName() { +export function getClientName(): string { if (isSageMaker() && isSageMakerUnifiedStudio()) { return 'AmazonQ-For-SMUS-CE' } From cee9380e81d6954158e38291af99d4b82d3d082b Mon Sep 17 00:00:00 2001 From: Jayakrishna Parameswaramangalam Date: Wed, 30 Jul 2025 18:17:13 -0700 Subject: [PATCH 3/4] update isSageMaker logic to use env vars properly and add unit tests --- .../core/src/shared/extensionUtilities.ts | 17 +- packages/core/src/shared/telemetry/util.ts | 4 +- .../test/shared/extensionUtilities.test.ts | 145 ++++++++++++++++++ .../src/test/shared/telemetry/util.test.ts | 57 +++++++ 4 files changed, 217 insertions(+), 6 deletions(-) diff --git a/packages/core/src/shared/extensionUtilities.ts b/packages/core/src/shared/extensionUtilities.ts index caac78436c8..4720572a5d2 100644 --- a/packages/core/src/shared/extensionUtilities.ts +++ b/packages/core/src/shared/extensionUtilities.ts @@ -158,6 +158,15 @@ export 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 +186,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/telemetry/util.ts b/packages/core/src/shared/telemetry/util.ts index d9e6ff840e6..dc57148393b 100644 --- a/packages/core/src/shared/telemetry/util.ts +++ b/packages/core/src/shared/telemetry/util.ts @@ -23,7 +23,7 @@ import { mapMetadata, MetadataObj } from './telemetryLogger' import { Result } from './telemetry.gen' import { MetricDatum } from './clienttelemetry' import { isValidationExemptMetric } from './exemptMetrics' -import { isAmazonQ, isCloud9, isSageMaker, isSageMakerUnifiedStudio } from '../../shared/extensionUtilities' +import { isAmazonQ, isCloud9, isSageMaker } from '../../shared/extensionUtilities' import { isUuid, randomUUID } from '../crypto' import { ClassToInterfaceType } from '../utilities/tsUtils' import { asStringifiedStack, FunctionEntry } from './spans' @@ -488,7 +488,7 @@ export function withTelemetryContext(opts: TelemetryContextArgs) { * 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() && isSageMakerUnifiedStudio()) { + 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') + }) +}) From 8fe741e85752c9398d741d723097e2854167e59f Mon Sep 17 00:00:00 2001 From: Jayakrishna Parameswaramangalam Date: Fri, 1 Aug 2025 18:21:55 -0700 Subject: [PATCH 4/4] fix(amazonq): add comment to smus util method --- packages/core/src/shared/extensionUtilities.ts | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/core/src/shared/extensionUtilities.ts b/packages/core/src/shared/extensionUtilities.ts index 4720572a5d2..80bedf1e0f6 100644 --- a/packages/core/src/shared/extensionUtilities.ts +++ b/packages/core/src/shared/extensionUtilities.ts @@ -150,6 +150,10 @@ function createCloud9Properties(company: string): IdeProperties { } } +/** + * export method - for testing purposes only + * @internal + */ export function isSageMakerUnifiedStudio(): boolean { if (serviceName === notInitialized) { serviceName = process.env.SERVICE_NAME ?? ''