chore: sync AmazonQServiceManager with AmazonQTokenServiceManager

Author: liramon1

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts

@@ -97,6 +97,11 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<
         this.state = state
     }
 
+    endpointOverride(): string | undefined {
+        return this.features.lsp.getClientInitializeParams()?.initializationOptions?.aws?.awsClientCapabilities
+            ?.textDocument?.inlineCompletionWithReferences?.endpointOverride
+    }
+
     public static initInstance(features: QServiceManagerFeatures): AmazonQServiceManager {
         if (!AmazonQServiceManager.instance) {
             AmazonQServiceManager.instance = new AmazonQServiceManager(features)
@@ -224,6 +229,10 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<
             return
         }
 
+        const endpointOverride =
+            this.features.lsp.getClientInitializeParams()?.initializationOptions?.aws?.awsClientCapabilities
+                ?.textDocument?.inlineCompletionWithReferences?.endpointOverride
+
         // Connection type changed to 'builderId'
 
         if (newConnectionType === 'builderId') {
@@ -234,7 +243,11 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<
             // region set by client -> runtime region -> default region
             const clientParams = this.features.lsp.getClientInitializeParams()
 
-            this.createCodewhispererServiceInstances('builderId', clientParams?.initializationOptions?.aws?.region)
+            this.createCodewhispererServiceInstances(
+                'builderId',
+                clientParams?.initializationOptions?.aws?.region,
+                endpointOverride
+            )
             this.state = 'INITIALIZED'
             this.log('Initialized Amazon Q service with builderId connection')
 
@@ -256,7 +269,7 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<
                 return
             }
 
-            this.createCodewhispererServiceInstances('identityCenter')
+            this.createCodewhispererServiceInstances('identityCenter', undefined, endpointOverride)
             this.state = 'INITIALIZED'
             this.log('Initialized Amazon Q service with identityCenter connection')
 
@@ -358,7 +371,11 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<
 
         if (!this.activeIdcProfile) {
             this.activeIdcProfile = newProfile
-            this.createCodewhispererServiceInstances('identityCenter', newProfile.identityDetails.region)
+            this.createCodewhispererServiceInstances(
+                'identityCenter',
+                newProfile.identityDetails.region,
+                this.endpointOverride()
+            )
             this.state = 'INITIALIZED'
             this.log(
                 `Initialized identityCenter connection to region ${newProfile.identityDetails.region} for profile ${newProfile.arn}`
@@ -410,7 +427,11 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<
 
         this.activeIdcProfile = newProfile
 
-        this.createCodewhispererServiceInstances('identityCenter', newProfile.identityDetails.region)
+        this.createCodewhispererServiceInstances(
+            'identityCenter',
+            newProfile.identityDetails.region,
+            this.endpointOverride()
+        )
         this.state = 'INITIALIZED'
 
         return
@@ -489,7 +510,8 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<
 
     private createCodewhispererServiceInstances(
         connectionType: 'builderId' | 'identityCenter',
-        clientOrProfileRegion?: string
+        clientOrProfileRegion: string | undefined,
+        endpointOverride: string | undefined
     ) {
         this.logServiceState('Initializing CodewhispererService')
 
@@ -504,10 +526,14 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<
         this.region = region
         this.endpoint = endpoint
 
-        this.cachedCodewhispererService = this.serviceFactory(region, endpoint)
+        if (endpointOverride) {
+            this.endpoint = endpointOverride
+        }
+
+        this.cachedCodewhispererService = this.serviceFactory(region, this.endpoint)
         this.log(`CodeWhispererToken service for connection type ${connectionType} was initialized, region=${region}`)
 
-        this.cachedStreamingClient = this.streamingClientFactory(region, endpoint)
+        this.cachedStreamingClient = this.streamingClientFactory(region, this.endpoint)
         this.log(`StreamingClient service for connection type ${connectionType} was initialized, region=${region}`)
 
         this.logServiceState('CodewhispererService and StreamingClient Initialization finished')