fix: update client references inside unified service manager

Author: liramon1

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

@@ -29,7 +29,7 @@ import {
     setTokenCredentialsForAmazonQServiceManagerFactory,
     setIamCredentialsForAmazonQServiceManagerFactory,
 } from '../testUtils'
-import { StreamingClientService } from '../streamingClientService'
+import { StreamingClientServiceToken, StreamingClientServiceIAM } from '../streamingClientService'
 import { generateSingletonInitializationTests } from './testUtils'
 import * as utils from '../utils'
 
@@ -61,8 +61,8 @@ const TEST_ENDPOINT_US_EAST_1 = 'http://amazon-q-in-us-east-1-endpoint'
 const TEST_ENDPOINT_EU_CENTRAL_1 = 'http://amazon-q-in-eu-central-1-endpoint'
 
 describe('Token', () => {
-    let codewhispererServiceStub: StubbedInstance<CodeWhispererService>
-    let codewhispererStubFactory: sinon.SinonStub<any[], StubbedInstance<CodeWhispererService>>
+    let codewhispererServiceStub: StubbedInstance<CodeWhispererServiceToken>
+    let codewhispererStubFactory: sinon.SinonStub<any[], StubbedInstance<CodeWhispererServiceToken>>
     let sdkInitializatorSpy: sinon.SinonSpy
     let getListAllAvailableProfilesHandlerStub: sinon.SinonStub
 
@@ -94,7 +94,7 @@ describe('Token', () => {
             v2: sinon.spy(features.sdkInitializator.v2),
         })
 
-        codewhispererServiceStub = stubInterface<CodeWhispererService>()
+        codewhispererServiceStub = stubInterface<CodeWhispererServiceToken>()
         // @ts-ignore
         codewhispererServiceStub.client = sinon.stub()
         codewhispererServiceStub.customizationArn = undefined
@@ -141,7 +141,7 @@ describe('Token', () => {
 
     const setupServiceManagerWithProfile = async (
         profileArn = 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ'
-    ): Promise<CodeWhispererService> => {
+    ): Promise<CodeWhispererServiceToken> => {
         setupServiceManager(true)
         assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION')
 
@@ -161,7 +161,7 @@ describe('Token', () => {
         assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED')
         assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter')
 
-        return service
+        return service as CodeWhispererServiceToken
     }
 
     describe('Initialization process', () => {
@@ -245,7 +245,7 @@ describe('Token', () => {
             assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED')
             assert.strictEqual(amazonQServiceManager.getConnectionType(), 'builderId')
 
-            assert(streamingClient instanceof StreamingClientService)
+            assert(streamingClient instanceof StreamingClientServiceToken)
             assert(codewhispererServiceStub.generateSuggestions.calledOnce)
         })
 
@@ -318,7 +318,7 @@ describe('Token', () => {
                 assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter')
                 assert(codewhispererServiceStub.generateSuggestions.calledOnce)
 
-                assert(streamingClient instanceof StreamingClientService)
+                assert(streamingClient instanceof StreamingClientServiceToken)
             })
         })
 
@@ -379,7 +379,7 @@ describe('Token', () => {
                 assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter')
                 assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1))
 
-                assert(streamingClient instanceof StreamingClientService)
+                assert(streamingClient instanceof StreamingClientServiceToken)
                 assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1')
             })
 
@@ -456,7 +456,7 @@ describe('Token', () => {
                 )
 
                 assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1))
-                assert(streamingClient1 instanceof StreamingClientService)
+                assert(streamingClient1 instanceof StreamingClientServiceToken)
                 assert.strictEqual(await streamingClient1.client.config.region(), 'us-east-1')
 
                 // Profile change
@@ -483,7 +483,7 @@ describe('Token', () => {
                 // CodeWhisperer Service was not recreated
                 assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1))
 
-                assert(streamingClient2 instanceof StreamingClientService)
+                assert(streamingClient2 instanceof StreamingClientServiceToken)
                 assert.strictEqual(streamingClient1, streamingClient2)
                 assert.strictEqual(await streamingClient2.client.config.region(), 'us-east-1')
             })
@@ -516,7 +516,7 @@ describe('Token', () => {
                 )
                 assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1))
 
-                assert(streamingClient1 instanceof StreamingClientService)
+                assert(streamingClient1 instanceof StreamingClientServiceToken)
                 assert.strictEqual(await streamingClient1.client.config.region(), 'us-east-1')
 
                 // Profile change
@@ -548,7 +548,7 @@ describe('Token', () => {
                 ])
 
                 // Streaming Client was recreated
-                assert(streamingClient2 instanceof StreamingClientService)
+                assert(streamingClient2 instanceof StreamingClientServiceToken)
                 assert.notStrictEqual(streamingClient1, streamingClient2)
                 assert.strictEqual(await streamingClient2.client.config.region(), 'eu-central-1')
             })
@@ -582,7 +582,7 @@ describe('Token', () => {
                 )
                 assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1))
 
-                assert(streamingClient instanceof StreamingClientService)
+                assert(streamingClient instanceof StreamingClientServiceToken)
                 assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1')
 
                 // Profile change to invalid profile
@@ -691,7 +691,7 @@ describe('Token', () => {
                     TEST_ENDPOINT_EU_CENTRAL_1,
                 ])
 
-                assert(streamingClient instanceof StreamingClientService)
+                assert(streamingClient instanceof StreamingClientServiceToken)
                 assert.strictEqual(await streamingClient.client.config.region(), 'eu-central-1')
             })
 
@@ -726,7 +726,7 @@ describe('Token', () => {
                 )
                 assert.deepStrictEqual(codewhispererStubFactory.lastCall.args, ['us-east-1', TEST_ENDPOINT_US_EAST_1])
 
-                assert(streamingClient instanceof StreamingClientService)
+                assert(streamingClient instanceof StreamingClientServiceToken)
                 assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1')
 
                 // Updaing profile
@@ -907,7 +907,7 @@ describe('Token', () => {
                 assert.strictEqual(amazonQServiceManager.getConnectionType(), 'builderId')
                 assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined)
 
-                assert(streamingClient instanceof StreamingClientService)
+                assert(streamingClient instanceof StreamingClientServiceToken)
                 assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1')
 
                 setCredentials('identityCenter')
@@ -921,7 +921,7 @@ describe('Token', () => {
                 assert(codewhispererStubFactory.calledTwice)
                 assert(codewhispererStubFactory.calledWithExactly(DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL))
 
-                assert(streamingClient2 instanceof StreamingClientService)
+                assert(streamingClient2 instanceof StreamingClientServiceToken)
                 assert.strictEqual(await streamingClient2.client.config.region(), DEFAULT_AWS_Q_REGION)
             })
 
@@ -937,7 +937,7 @@ describe('Token', () => {
                 assert.strictEqual(amazonQServiceManager.getConnectionType(), 'builderId')
                 assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined)
 
-                assert(streamingClient instanceof StreamingClientService)
+                assert(streamingClient instanceof StreamingClientServiceToken)
                 assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1')
 
                 setCredentials('identityCenter')
@@ -967,7 +967,7 @@ describe('Token', () => {
                 assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter')
                 assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined)
 
-                assert(streamingClient instanceof StreamingClientService)
+                assert(streamingClient instanceof StreamingClientServiceToken)
                 assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1')
 
                 setCredentials('builderId')
@@ -981,7 +981,7 @@ describe('Token', () => {
                 assert(codewhispererStubFactory.calledTwice)
                 assert(codewhispererStubFactory.calledWithExactly(DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL))
 
-                assert(streamingClient2 instanceof StreamingClientService)
+                assert(streamingClient2 instanceof StreamingClientServiceToken)
                 assert.strictEqual(await streamingClient2.client.config.region(), 'us-east-1')
             })
         })
@@ -1081,19 +1081,12 @@ describe('IAM', () => {
         })
 
         it('should initialize the streaming client only once', () => {
+            // Mock the credentials provider to return credentials when requested
             setCredentials()
-            // Mock getIAMCredentialsFromProvider to return dummy credentials
-            const getIAMCredentialsStub = sinon.stub(utils, 'getIAMCredentialsFromProvider').returns({
-                accessKeyId: 'dummy-access-key',
-                secretAccessKey: 'dummy-secret-key',
-                sessionToken: 'dummy-session-token',
-            })
-
             const streamingClient = serviceManager.getStreamingClient()
 
+            // Verify that getting the client again returns the same instance
             assert.deepStrictEqual(serviceManager.getStreamingClient(), streamingClient)
-
-            getIAMCredentialsStub.restore()
         })
     })
 })

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

@@ -32,8 +32,13 @@ import { AmazonQDeveloperProfile, signalsAWSQDeveloperProfilesEnabled } from './
 import { isStringOrNull } from '../utils'
 import { getAmazonQRegionAndEndpoint } from './configurationUtils'
 import { getUserAgent } from '../telemetryUtils'
-import { StreamingClientServiceToken, StreamingClientServiceIAM } from '../streamingClientService'
+import {
+    StreamingClientServiceToken,
+    StreamingClientServiceIAM,
+    StreamingClientServiceBase,
+} from '../streamingClientService'
 import { parse } from '@aws-sdk/util-arn-parser'
+import { CodeWhispererServiceBase } from '../codeWhispererService/codeWhispererServiceBase'
 
 /**
  * AmazonQServiceManager manages state and provides centralized access to
@@ -60,7 +65,10 @@ import { parse } from '@aws-sdk/util-arn-parser'
  * the LSP hand shake.
  *
  */
-export class AmazonQServiceManager extends BaseAmazonQServiceManager<CodeWhispererService, StreamingClientService> {
+export class AmazonQServiceManager extends BaseAmazonQServiceManager<
+    CodeWhispererServiceBase,
+    StreamingClientServiceBase
+> {
     private static instance: AmazonQServiceManager | null = null
     private enableDeveloperProfileSupport?: boolean
     private activeIdcProfile?: AmazonQDeveloperProfile
@@ -185,33 +193,6 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<CodeWhisper
         }
     }
 
-    private handleConnectionChange() {
-        if (this.features.credentialsProvider.hasCredentials('iam')) {
-            if (!this.cachedCodewhispererService) {
-                const amazonQRegionAndEndpoint = getAmazonQRegionAndEndpoint(
-                    this.features.runtime,
-                    this.features.logging
-                )
-                this.region = amazonQRegionAndEndpoint.region
-                this.endpoint = amazonQRegionAndEndpoint.endpoint
-                this.cachedCodewhispererService = new CodeWhispererService(
-                    this.features.credentialsProvider,
-                    this.features.workspace,
-                    this.features.logging,
-                    this.region,
-                    this.endpoint,
-                    this.features.sdkInitializator
-                )
-                this.updateCachedServiceConfig()
-            }
-            this.state = 'INITIALIZED'
-            return
-        } else {
-            this.handleSsoConnectionChange()
-            return
-        }
-    }
-
     /**
      * Validate if Bearer Token Connection type has changed mid-session.
      * When connection type change is detected: reinitialize CodeWhispererService class with current connection type.
@@ -435,13 +416,34 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<CodeWhisper
         return
     }
 
-    public getCodewhispererService(): CodeWhispererService {
+    public getCodewhispererService(): CodeWhispererServiceBase {
         // Prevent initiating requests while profile change is in progress.
         if (this.state === 'PENDING_Q_PROFILE_UPDATE') {
             throw new AmazonQServicePendingProfileUpdateError()
         }
 
-        this.handleConnectionChange()
+        if (this.features.credentialsProvider.hasCredentials('iam')) {
+            if (!this.cachedCodewhispererService) {
+                const amazonQRegionAndEndpoint = getAmazonQRegionAndEndpoint(
+                    this.features.runtime,
+                    this.features.logging
+                )
+                this.region = amazonQRegionAndEndpoint.region
+                this.endpoint = amazonQRegionAndEndpoint.endpoint
+                this.cachedCodewhispererService = new CodeWhispererServiceIAM(
+                    this.features.credentialsProvider,
+                    this.features.workspace,
+                    this.features.logging,
+                    this.region,
+                    this.endpoint,
+                    this.features.sdkInitializator
+                )
+                this.updateCachedServiceConfig()
+            }
+            this.state = 'INITIALIZED'
+        } else {
+            this.handleSsoConnectionChange()
+        }
 
         if (this.state === 'INITIALIZED' && this.cachedCodewhispererService) {
             return this.cachedCodewhispererService
@@ -458,7 +460,7 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<CodeWhisper
         throw new AmazonQServiceNotInitializedError()
     }
 
-    public getStreamingClient() {
+    public getStreamingClient(): StreamingClientServiceBase {
         this.log('Getting instance of CodeWhispererStreaming client')
 
         // Trigger checks in token service
@@ -517,8 +519,8 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<CodeWhisper
         return getUserAgent(initializeParams as InitializeParams, this.features.runtime.serverInfo)
     }
 
-    private serviceFactory(region: string, endpoint: string): CodeWhispererService {
-        const service = new CodeWhispererService(
+    private serviceFactory(region: string, endpoint: string): CodeWhispererServiceToken {
+        const service = new CodeWhispererServiceToken(
             this.features.credentialsProvider,
             this.features.workspace,
             this.features.logging,
@@ -545,17 +547,25 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<CodeWhisper
         return service
     }
 
-    private streamingClientFactory(region: string, endpoint: string): StreamingClientService {
-        const streamingClient = new StreamingClientService(
-            this.features.credentialsProvider,
-            this.features.sdkInitializator,
-            this.features.logging,
-            region,
-            endpoint,
-            this.getCustomUserAgent()
-        )
-
-        if (this.features.credentialsProvider.hasCredentials('bearer')) {
+    private streamingClientFactory(region: string, endpoint: string): StreamingClientServiceBase {
+        let streamingClient: StreamingClientServiceBase
+        if (this.features.credentialsProvider.hasCredentials('iam')) {
+            streamingClient = new StreamingClientServiceIAM(
+                this.features.credentialsProvider,
+                this.features.sdkInitializator,
+                this.features.logging,
+                region,
+                endpoint
+            )
+        } else {
+            streamingClient = new StreamingClientServiceToken(
+                this.features.credentialsProvider,
+                this.features.sdkInitializator,
+                this.features.logging,
+                region,
+                endpoint,
+                this.getCustomUserAgent()
+            )
             streamingClient.profileArn = this.activeIdcProfile?.arn
         }
 
@@ -598,7 +608,7 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager<CodeWhisper
         return this.activeIdcProfile?.arn
     }
 
-    public setServiceFactory(factory: (region: string, endpoint: string) => CodeWhispererService) {
+    public setServiceFactory(factory: (region: string, endpoint: string) => CodeWhispererServiceToken) {
         this.serviceFactory = factory.bind(this)
     }