feat: add support for SMUS Q CodeEditor client to send MD IDE origin (aws#2032)

* feat: add support for SMUS Q CodeEditor client to send MD IDE origin

* refactor: move smus service name to constants file and typecase argument getClientName util

Author: parameja1

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

@@ -124,6 +124,7 @@ import {
     isUsageLimitError,
     isNullish,
     getOriginFromClientInfo,
+    getClientName,
     sanitizeInput,
     sanitizeRequestInput,
 } from '../../shared/utils'
@@ -366,7 +367,7 @@ export class AgenticChatController implements ChatHandlers {
             this.#features.lsp
         )
         this.#mcpEventHandler = new McpEventHandler(features, telemetryService)
-        this.#origin = getOriginFromClientInfo(this.#features.lsp.getClientInitializeParams()?.clientInfo?.name)
+        this.#origin = getOriginFromClientInfo(getClientName(this.#features.lsp.getClientInitializeParams()))
         this.#activeUserTracker = ActiveUserTracker.getInstance(this.#features)
     }
 

server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts

@@ -27,7 +27,7 @@ import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/Ba
 import { loggingUtils } from '@aws/lsp-core'
 import { Logging } from '@aws/language-server-runtimes/server-interface'
 import { Features } from '../types'
-import { getOriginFromClientInfo, getRequestID, isUsageLimitError } from '../../shared/utils'
+import { getOriginFromClientInfo, getClientName, getRequestID, isUsageLimitError } from '../../shared/utils'
 import { enabledModelSelection } from '../../shared/utils'
 
 export type ChatSessionServiceConfig = CodeWhispererStreamingClientConfig
@@ -138,7 +138,7 @@ export class ChatSessionService {
         this.#serviceManager = serviceManager
         this.#lsp = lsp
         this.#logging = logging
-        this.#origin = getOriginFromClientInfo(this.#lsp?.getClientInitializeParams()?.clientInfo?.name)
+        this.#origin = getOriginFromClientInfo(getClientName(this.#lsp?.getClientInitializeParams()))
     }
 
     public async sendMessage(request: SendMessageCommandInput): Promise<SendMessageCommandOutput> {

server/aws-lsp-codewhisperer/src/shared/constants.ts

@@ -18,6 +18,8 @@ export const AWS_Q_ENDPOINT_URL_ENV_VAR = 'AWS_Q_ENDPOINT_URL'
 export const Q_CONFIGURATION_SECTION = 'aws.q'
 export const CODE_WHISPERER_CONFIGURATION_SECTION = 'aws.codeWhisperer'
 
+export const SAGEMAKER_UNIFIED_STUDIO_SERVICE = 'SageMakerUnifiedStudio'
+
 /**
  * Names of directories relevant to the crash reporting functionality.
  *

server/aws-lsp-codewhisperer/src/shared/utils.test.ts

@@ -3,14 +3,14 @@ import {
     ThrottlingException,
     ThrottlingExceptionReason,
 } from '@amzn/codewhisperer-streaming'
-import { CredentialsProvider, Position } from '@aws/language-server-runtimes/server-interface'
+import { CredentialsProvider, Position, InitializeParams } from '@aws/language-server-runtimes/server-interface'
 import * as assert from 'assert'
 import { AWSError } from 'aws-sdk'
 import { expect } from 'chai'
 import * as sinon from 'sinon'
 import * as os from 'os'
 import * as path from 'path'
-import { BUILDER_ID_START_URL } from './constants'
+import { BUILDER_ID_START_URL, SAGEMAKER_UNIFIED_STUDIO_SERVICE } from './constants'
 import {
     getBearerTokenFromProvider,
     getEndPositionForAcceptedSuggestion,
@@ -24,6 +24,7 @@ import {
     getFileExtensionName,
     listFilesWithGitignore,
     getOriginFromClientInfo,
+    getClientName,
     sanitizeInput,
     sanitizeRequestInput,
 } from './utils'
@@ -73,12 +74,86 @@ describe('getBearerTokenFromProvider', () => {
     })
 })
 
+describe('getClientName', () => {
+    let originalEnv: string | undefined
+
+    beforeEach(() => {
+        originalEnv = process.env.SERVICE_NAME
+    })
+
+    afterEach(() => {
+        if (originalEnv !== undefined) {
+            process.env.SERVICE_NAME = originalEnv
+        } else {
+            delete process.env.SERVICE_NAME
+        }
+    })
+
+    it('returns client name from initializationOptions path when SERVICE_NAME is SageMakerUnifiedStudio', () => {
+        process.env.SERVICE_NAME = SAGEMAKER_UNIFIED_STUDIO_SERVICE
+        const lspParams = {
+            initializationOptions: {
+                aws: {
+                    clientInfo: {
+                        name: 'AmazonQ-For-SMUS-CE-1.0.0',
+                    },
+                },
+            },
+            clientInfo: {
+                name: 'VSCode-Extension',
+            },
+        } as InitializeParams
+
+        const result = getClientName(lspParams)
+        assert.strictEqual(result, 'AmazonQ-For-SMUS-CE-1.0.0')
+    })
+
+    it('returns client name from clientInfo path when SERVICE_NAME is not SageMakerUnifiedStudio', () => {
+        process.env.SERVICE_NAME = 'SomeOtherService'
+        const lspParams = {
+            initializationOptions: {
+                aws: {
+                    clientInfo: {
+                        name: 'AmazonQ-For-SMUS-CE-1.0.0',
+                    },
+                },
+            },
+            clientInfo: {
+                name: 'VSCode-Extension',
+            },
+        } as InitializeParams
+
+        const result = getClientName(lspParams)
+        assert.strictEqual(result, 'VSCode-Extension')
+    })
+
+    it('returns undefined when lspParams is undefined', () => {
+        const result = getClientName(undefined)
+        assert.strictEqual(result, undefined)
+    })
+})
+
 describe('getOriginFromClientInfo', () => {
-    it('returns MD_IDE for SMUS client name', () => {
+    it('returns MD_IDE for SMUS-IDE client name', () => {
         const result = getOriginFromClientInfo('AmazonQ-For-SMUS-IDE-1.0.0')
         assert.strictEqual(result, 'MD_IDE')
     })
 
+    it('returns MD_IDE for SMUS-CE client name', () => {
+        const result = getOriginFromClientInfo('AmazonQ-For-SMUS-CE-1.0.0')
+        assert.strictEqual(result, 'MD_IDE')
+    })
+
+    it('returns MD_IDE for client names starting with SMUS-IDE prefix', () => {
+        const result = getOriginFromClientInfo('AmazonQ-For-SMUS-IDE')
+        assert.strictEqual(result, 'MD_IDE')
+    })
+
+    it('returns MD_IDE for client names starting with SMUS-CE prefix', () => {
+        const result = getOriginFromClientInfo('AmazonQ-For-SMUS-CE')
+        assert.strictEqual(result, 'MD_IDE')
+    })
+
     it('returns IDE for non-SMUS client name', () => {
         const result = getOriginFromClientInfo('VSCode-Extension')
         assert.strictEqual(result, 'IDE')
@@ -93,6 +168,11 @@ describe('getOriginFromClientInfo', () => {
         const result = getOriginFromClientInfo('')
         assert.strictEqual(result, 'IDE')
     })
+
+    it('returns IDE for client names that do not match SMUS patterns', () => {
+        const result = getOriginFromClientInfo('AmazonQ-For-Other-IDE')
+        assert.strictEqual(result, 'IDE')
+    })
 })
 
 describe('getSsoConnectionType', () => {

server/aws-lsp-codewhisperer/src/shared/utils.ts

@@ -14,6 +14,7 @@ import {
     crashMonitoringDirName,
     driveLetterRegex,
     MISSING_BEARER_TOKEN_ERROR,
+    SAGEMAKER_UNIFIED_STUDIO_SERVICE,
 } from './constants'
 import {
     CodeWhispererStreamingServiceException,
@@ -373,8 +374,14 @@ export function getBearerTokenFromProvider(credentialsProvider: CredentialsProvi
     return credentials.token
 }
 
+export function getClientName(lspParams: InitializeParams | undefined): string | undefined {
+    return process.env.SERVICE_NAME === SAGEMAKER_UNIFIED_STUDIO_SERVICE
+        ? lspParams?.initializationOptions?.aws?.clientInfo?.name
+        : lspParams?.clientInfo?.name
+}
+
 export function getOriginFromClientInfo(clientName: string | undefined): Origin {
-    if (clientName?.startsWith('AmazonQ-For-SMUS-IDE')) {
+    if (clientName?.startsWith('AmazonQ-For-SMUS-IDE') || clientName?.startsWith('AmazonQ-For-SMUS-CE')) {
         return 'MD_IDE'
     }
     return 'IDE'