fix(amazonq): add regionalization support to .NET Transform server (#952)

awschristou · web-flow · commit 7571ffdb8766 · 2025-04-14T16:09:31.000-07:00
This change updates the .NET Transform server to use AmazonQTokenServiceManager when producing the service client. This way, the service calls are able to provide the developer profile arn or region details for IdC based connections. This change is based on #859, which made a similar change for security scanning. I worked on this change with @shruti0085 , and we tested it by using the Visual Studio Toolkit to perform a .NET Transform using a local build of the server. We were able to see the profileArn field being passed along with all of the .NET Transform related service calls. License By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
diff --git a/server/aws-lsp-codewhisperer/src/language-server/netTransform/netTransformServer.ts b/server/aws-lsp-codewhisperer/src/language-server/netTransform/netTransformServer.ts
@@ -6,7 +6,6 @@ import {
     Server,
     Workspace,
 } from '@aws/language-server-runtimes/server-interface'
-import { StreamingClient } from '../../client/streamingClient/codewhispererStreamingClient'
 import { CodeWhispererServiceToken } from '../../shared/codeWhispererService'
 import {
     emitTransformationJobArtifactsDownloadedFailure,
@@ -32,8 +31,6 @@ import {
     StartTransformRequest,
 } from './models'
 import { TransformHandler } from './transformHandler'
-import { CodeWhispererStreamingClientConfig } from '@amzn/codewhisperer-streaming'
-import { getUserAgent } from '../../shared/telemetryUtils'
 
 export const validStatesForGettingPlan = ['COMPLETED', 'PARTIALLY_COMPLETED', 'PLANNED', 'TRANSFORMING', 'TRANSFORMED']
 export const validStatesForComplete = ['COMPLETED']
@@ -45,13 +42,8 @@ const PollTransformForPlanCommand = 'aws/qNetTransform/pollTransformForPlan'
 const GetTransformPlanCommand = 'aws/qNetTransform/getTransformPlan'
 const CancelTransformCommand = 'aws/qNetTransform/cancelTransform'
 const DownloadArtifactsCommand = 'aws/qNetTransform/downloadArtifacts'
-import {
-    DEFAULT_AWS_Q_REGION,
-    DEFAULT_AWS_Q_ENDPOINT_URL,
-    AWS_Q_REGION_ENV_VAR,
-    AWS_Q_ENDPOINT_URL_ENV_VAR,
-} from '../../shared/constants'
 import { SDKInitializator } from '@aws/language-server-runtimes/server-interface'
+import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager'
 
 /**
  *
@@ -69,14 +61,8 @@ export const QNetTransformServerToken =
         ) => CodeWhispererServiceToken
     ): Server =>
     ({ credentialsProvider, workspace, logging, lsp, telemetry, runtime, sdkInitializator }) => {
-        const codewhispererclient = service(
-            credentialsProvider,
-            workspace,
-            runtime.getConfiguration(AWS_Q_REGION_ENV_VAR) ?? DEFAULT_AWS_Q_REGION,
-            runtime.getConfiguration(AWS_Q_ENDPOINT_URL_ENV_VAR) ?? DEFAULT_AWS_Q_ENDPOINT_URL,
-            sdkInitializator
-        )
-        const transformHandler = new TransformHandler(codewhispererclient, workspace, logging, runtime)
+        let amazonQServiceManager: AmazonQTokenServiceManager
+        let transformHandler: TransformHandler
         const runTransformCommand = async (params: ExecuteCommandParams, _token: CancellationToken) => {
             try {
                 switch (params.command) {
@@ -128,17 +114,8 @@ export const QNetTransformServerToken =
                     }
                     case DownloadArtifactsCommand: {
                         const request = params as DownloadArtifactsRequest
-                        const cwStreamingClientInstance = new StreamingClient()
-                        const cwStreamingClient = await cwStreamingClientInstance.getStreamingClient(
-                            credentialsProvider,
-                            runtime.getConfiguration(AWS_Q_REGION_ENV_VAR) ?? DEFAULT_AWS_Q_REGION,
-                            runtime.getConfiguration(AWS_Q_ENDPOINT_URL_ENV_VAR) ?? DEFAULT_AWS_Q_ENDPOINT_URL,
-                            sdkInitializator,
-                            customCWClientConfig
-                        )
 
                         const response = await transformHandler.downloadExportResultArchive(
-                            cwStreamingClient,
                             request.TransformationJobId,
                             request.SolutionRootPath
                         )
@@ -202,15 +179,25 @@ export const QNetTransformServerToken =
             return runTransformCommand(params, _token)
         }
 
-        const customCWClientConfig: CodeWhispererStreamingClientConfig = {}
-        const onInitializeHandler = (params: InitializeParams) => {
-            // Cache user agent to reuse between commands calls
-            customCWClientConfig.customUserAgent = getUserAgent(params, runtime.serverInfo)
-
-            codewhispererclient.updateClientConfig({
-                customUserAgent: customCWClientConfig.customUserAgent,
+        const onInitializeHandler = async (params: InitializeParams) => {
+            amazonQServiceManager = AmazonQTokenServiceManager.getInstance({
+                lsp,
+                logging,
+                runtime,
+                credentialsProvider,
+                sdkInitializator,
+                workspace,
             })
 
+            transformHandler = new TransformHandler(amazonQServiceManager, workspace, logging, runtime)
+
+            /* 
+                Calling handleDidChangeConfiguration once to ensure we get configuration atleast once at start up
+                
+                TODO: TODO: consider refactoring such responsibilities to common service manager config/initialisation server
+            */
+            await amazonQServiceManager.handleDidChangeConfiguration()
+
             return {
                 capabilities: {
                     executeCommandProvider: {
diff --git a/server/aws-lsp-codewhisperer/src/language-server/netTransform/tests/transformHandler.test.ts b/server/aws-lsp-codewhisperer/src/language-server/netTransform/tests/transformHandler.test.ts
@@ -32,6 +32,7 @@ import { Readable } from 'stream'
 import { ArtifactManager } from '../artifactManager'
 import path = require('path')
 import { IZipEntry } from 'adm-zip'
+import { AmazonQTokenServiceManager } from '../../../shared/amazonQServiceManager/AmazonQTokenServiceManager'
 
 const mocked$Response = {
     $response: {
@@ -62,7 +63,12 @@ describe('Test Transform handler ', () => {
         client = stubInterface<CodeWhispererServiceToken>()
         workspace = stubInterface<Workspace>()
         runtime = stubInterface<Runtime>()
-        transformHandler = new TransformHandler(client, workspace, mockedLogging, runtime)
+
+        const serviceManager = stubInterface<AmazonQTokenServiceManager>()
+        client = stubInterface<CodeWhispererServiceToken>()
+        serviceManager.getCodewhispererService.returns(client)
+
+        transformHandler = new TransformHandler(serviceManager, workspace, mockedLogging, runtime)
     })
 
     describe('test upload artifact', () => {
diff --git a/server/aws-lsp-codewhisperer/src/language-server/netTransform/transformHandler.ts b/server/aws-lsp-codewhisperer/src/language-server/netTransform/transformHandler.ts
@@ -1,4 +1,4 @@
-import { CodeWhispererStreaming, ExportIntent } from '@amzn/codewhisperer-streaming'
+import { ExportIntent } from '@amzn/codewhisperer-streaming'
 import { Logging, Runtime, Workspace } from '@aws/language-server-runtimes/server-interface'
 import * as fs from 'fs'
 import got from 'got'
@@ -9,7 +9,6 @@ import {
     StopTransformationRequest,
     TransformationJob,
 } from '../../client/token/codewhispererbearertokenclient'
-import { CodeWhispererServiceToken } from '../../shared/codeWhispererService'
 import { ArtifactManager } from './artifactManager'
 import { getCWStartTransformRequest, getCWStartTransformResponse } from './converter'
 import {
@@ -29,21 +28,17 @@ import {
 import * as validation from './validation'
 import path = require('path')
 import AdmZip = require('adm-zip')
-import { Console } from 'console'
-import { supportedProjects, unsupportedViewComponents } from './resources/SupportedProjects'
-import { String } from 'aws-sdk/clients/codebuild'
-import { ProjectMetadata } from 'aws-sdk/clients/lookoutvision'
-import { httpstatus } from 'aws-sdk/clients/glacier'
+import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager'
 
 const workspaceFolderName = 'artifactWorkspace'
 
 export class TransformHandler {
-    private client: CodeWhispererServiceToken
+    private serviceManager: AmazonQTokenServiceManager
     private workspace: Workspace
     private logging: Logging
     private runtime: Runtime
-    constructor(client: CodeWhispererServiceToken, workspace: Workspace, logging: Logging, runtime: Runtime) {
-        this.client = client
+    constructor(serviceManager: AmazonQTokenServiceManager, workspace: Workspace, logging: Logging, runtime: Runtime) {
+        this.serviceManager = serviceManager
         this.workspace = workspace
         this.logging = logging
         this.runtime = runtime
@@ -82,7 +77,9 @@ export class TransformHandler {
             const uploadId = await this.preTransformationUploadCode(payloadFilePath)
             const request = getCWStartTransformRequest(userInputrequest, uploadId, this.logging)
             this.logging.log('Sending request to start transform api: ' + JSON.stringify(request))
-            const response = await this.client.codeModernizerStartCodeTransformation(request)
+            const response = await this.serviceManager
+                .getCodewhispererService()
+                .codeModernizerStartCodeTransformation(request)
             this.logging.log('Received transformation job Id: ' + response?.transformationJobId)
             return getCWStartTransformResponse(
                 response,
@@ -122,7 +119,7 @@ export class TransformHandler {
         const sha256 = await ArtifactManager.getSha256Async(payloadFileName)
         let response: CreateUploadUrlResponse
         try {
-            response = await this.client.codeModernizerCreateUploadUrl({
+            response = await this.serviceManager.getCodewhispererService().codeModernizerCreateUploadUrl({
                 contentChecksum: sha256,
                 contentChecksumType: 'SHA_256',
                 uploadIntent: 'TRANSFORMATION',
@@ -192,7 +189,9 @@ export class TransformHandler {
             const getCodeTransformationRequest = {
                 transformationJobId: request.TransformationJobId,
             } as GetTransformationRequest
-            const response = await this.client.codeModernizerGetCodeTransformation(getCodeTransformationRequest)
+            const response = await this.serviceManager
+                .getCodewhispererService()
+                .codeModernizerGetCodeTransformation(getCodeTransformationRequest)
             this.logging.log('Transformation status: ' + response.transformationJob?.status)
             return {
                 TransformationJob: response.transformationJob,
@@ -211,9 +210,9 @@ export class TransformHandler {
                 const getCodeTransformationPlanRequest = {
                     transformationJobId: request.TransformationJobId,
                 } as GetTransformationRequest
-                const response = await this.client.codeModernizerGetCodeTransformationPlan(
-                    getCodeTransformationPlanRequest
-                )
+                const response = await this.serviceManager
+                    .getCodewhispererService()
+                    .codeModernizerGetCodeTransformationPlan(getCodeTransformationPlanRequest)
                 return {
                     TransformationPlan: response.transformationPlan,
                 } as GetTransformPlanResponse
@@ -247,7 +246,9 @@ export class TransformHandler {
                 this.logging.log(
                     'Sending CancelTransformRequest with job Id: ' + stopCodeTransformationRequest.transformationJobId
                 )
-                const response = await this.client.codeModernizerStopCodeTransformation(stopCodeTransformationRequest)
+                const response = await this.serviceManager
+                    .getCodewhispererService()
+                    .codeModernizerStopCodeTransformation(stopCodeTransformationRequest)
                 this.logging.log('Transformation status: ' + response.transformationStatus)
                 let status: CancellationJobStatus
                 switch (response.transformationStatus) {
@@ -293,7 +294,9 @@ export class TransformHandler {
         const getCodeTransformationRequest = {
             transformationJobId: request.TransformationJobId,
         } as GetTransformationRequest
-        let response = await this.client.codeModernizerGetCodeTransformation(getCodeTransformationRequest)
+        let response = await this.serviceManager
+            .getCodewhispererService()
+            .codeModernizerGetCodeTransformation(getCodeTransformationRequest)
         this.logging.log('Start polling for transformation plan.')
         this.logging.log('The valid status to exit polling are: ' + validExitStatus)
         this.logging.log('The failure status are: ' + failureStates)
@@ -308,7 +311,9 @@ export class TransformHandler {
                 const getCodeTransformationRequest = {
                     transformationJobId: request.TransformationJobId,
                 } as GetTransformationRequest
-                response = await this.client.codeModernizerGetCodeTransformation(getCodeTransformationRequest)
+                response = await this.serviceManager
+                    .getCodewhispererService()
+                    .codeModernizerGetCodeTransformation(getCodeTransformationRequest)
                 this.logging.log('Transformation status: ' + response.transformationJob?.status)
 
                 if (validExitStatus.includes(status)) {
@@ -350,10 +355,10 @@ export class TransformHandler {
         } as GetTransformResponse
     }
 
-    async downloadExportResultArchive(cwStreamingClient: CodeWhispererStreaming, exportId: string, saveToDir: string) {
+    async downloadExportResultArchive(exportId: string, saveToDir: string) {
         let result
         try {
-            result = await cwStreamingClient.exportResultArchive({
+            result = await this.serviceManager.getStreamingClient().client.exportResultArchive({
                 exportId,
                 exportIntent: ExportIntent.TRANSFORMATION,
             })
