fix(amazonq): profile is not set after re-auth (#1690)

Author: ctlai95

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

@@ -838,27 +838,45 @@ describe('AmazonQTokenServiceManager', () => {
     })
 
     describe('Connection types with no Developer Profiles support', () => {
-        it('returns error when profile update is requested and connection type is none', async () => {
+        it('handles reauthentication scenario when connection type is none but profile ARN is provided', async () => {
             setupServiceManager(true)
             clearCredentials()
 
             assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION')
+            assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'none')
 
-            await assert.rejects(
-                amazonQTokenServiceManager.handleOnUpdateConfiguration(
-                    {
-                        section: 'aws.q',
-                        settings: {
-                            profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ',
-                        },
+            await amazonQTokenServiceManager.handleOnUpdateConfiguration(
+                {
+                    section: 'aws.q',
+                    settings: {
+                        profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ',
                     },
-                    {} as CancellationToken
-                ),
-                new ResponseError(LSPErrorCodes.RequestFailed, 'Amazon Q service is not signed in', {
-                    awsErrorCode: 'E_AMAZON_Q_PENDING_CONNECTION',
-                })
+                },
+                {} as CancellationToken
             )
 
+            assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter')
+            assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED')
+        })
+
+        it('ignores null profile when connection type is none', async () => {
+            setupServiceManager(true)
+            clearCredentials()
+
+            assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION')
+            assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'none')
+
+            await amazonQTokenServiceManager.handleOnUpdateConfiguration(
+                {
+                    section: 'aws.q',
+                    settings: {
+                        profileArn: null,
+                    },
+                },
+                {} as CancellationToken
+            )
+
+            assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'none')
             assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION')
         })
 

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

@@ -290,11 +290,14 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
 
         if (this.connectionType === 'none') {
             if (newProfileArn !== null) {
-                throw new AmazonQServicePendingSigninError()
+                // During reauthentication, connection might be temporarily 'none' but user is providing a profile
+                // Set connection type to identityCenter to proceed with profile setting
+                this.connectionType = 'identityCenter'
+                this.state = 'PENDING_Q_PROFILE_UPDATE'
+            } else {
+                this.logServiceState('Received null profile while not connected, ignoring request')
+                return
             }
-
-            this.logServiceState('Received null profile while not connected, ignoring request')
-            return
         }
 
         if (this.connectionType !== 'identityCenter') {