feat(amazonq): enable builder id auth for amazon q (aws#4608)

* feat: enable builder id auth for amazon q

* address comment to update if condition

* Fix verification of amazonQ connections and adding correct scopes

---------

Co-authored-by: David Hasani <[email protected]>
Co-authored-by: Joao Salles <[email protected]>

Author: 3 people

.changes/next-release/Feature-51070f02-21b0-447b-9743-e7ce616337fe.json

@@ -0,0 +1,4 @@
+{
+    "type": "Feature",
+    "description": "Enable Amazon Q feature development and Amazon Q transform capabilities (/dev and /transform) for AWS Builder ID users."
+}

packages/core/src/amazonq/explorer/amazonQTreeNode.ts

@@ -9,7 +9,7 @@ import { ResourceTreeDataProvider, TreeNode } from '../../shared/treeview/resour
 import { AuthUtil, amazonQScopes, codeWhispererChatScopes } from '../../codewhisperer/util/authUtil'
 import { createLearnMoreNode, enableAmazonQNode, switchToAmazonQNode } from './amazonQChildrenNodes'
 import { Command, Commands } from '../../shared/vscode/commands2'
-import { hasScopes, isSsoConnection } from '../../auth/connection'
+import { hasScopes, isBuilderIdConnection, isSsoConnection } from '../../auth/connection'
 import { listCodeWhispererCommands } from '../../codewhisperer/ui/statusBarMenu'
 import { getIcon } from '../../shared/icons'
 import { vsCodeState } from '../../codewhisperer/models/model'
@@ -77,6 +77,16 @@ export class AmazonQNode implements TreeNode {
             }
         }
 
+        if (isBuilderIdConnection(AuthUtil.instance.conn)) {
+            const missingScopes =
+                (AuthUtil.instance.isBuilderIdInUse() && !hasScopes(AuthUtil.instance.conn, amazonQScopes)) ||
+                !hasScopes(AuthUtil.instance.conn, codeWhispererChatScopes)
+
+            if (missingScopes) {
+                return [enableAmazonQNode(), createLearnMoreNode()]
+            }
+        }
+
         return [
             vsCodeState.isFreeTierLimitReached ? createFreeTierLimitMet('tree') : switchToAmazonQNode('tree'),
             createNewMenuButton(),

packages/core/src/codewhisperer/models/constants.ts

@@ -310,9 +310,6 @@ export const transformByQPartiallyCompletedMessage =
 export const noPomXmlFoundMessage =
     'None of your open Java projects are supported by Amazon Q Code Transformation. Currently, Amazon Q can only upgrade Java projects built on Maven. A pom.xml must be present in the root of your project to upgrade it. For more information, see the [Amazon Q documentation](LINK_HERE).'
 
-export const noActiveIdCMessage =
-    'Amazon Q Code Transformation requires an active IAM Identity Center connection. For more information, see the [Code Transformation documentation](LINK_HERE).'
-
 export const noOngoingJobMessage = 'No job is in-progress at the moment'
 
 export const jobInProgressMessage = 'Job is already in-progress'

packages/core/src/codewhisperer/util/authUtil.ts

@@ -57,20 +57,15 @@ export const isValidCodeWhispererCoreConnection = (conn?: Connection): conn is C
         (isSsoConnection(conn) && hasScopes(conn, codeWhispererCoreScopes))
     )
 }
-/** For Builder ID only, if using IdC then use {@link isValidAmazonQConnection} */
-export const isValidCodeWhispererChatConnection = (conn?: Connection): conn is Connection => {
+/** Superset that includes all of CodeWhisperer + Amazon Q */
+export const isValidAmazonQConnection = (conn?: Connection): conn is Connection => {
     return (
-        isBuilderIdConnection(conn) &&
+        (isSsoConnection(conn) || isBuilderIdConnection(conn)) &&
         isValidCodeWhispererCoreConnection(conn) &&
-        hasScopes(conn, codeWhispererChatScopes)
+        hasScopes(conn, amazonQScopes)
     )
 }
 
-/** Superset that includes all of CodeWhisperer + Amazon Q */
-export const isValidAmazonQConnection = (conn?: Connection): conn is Connection => {
-    return isSsoConnection(conn) && isValidCodeWhispererCoreConnection(conn) && hasScopes(conn, amazonQScopes)
-}
-
 interface HasAlreadySeenQWelcome {
     local?: boolean
     devEnv?: boolean
@@ -217,9 +212,9 @@ export class AuthUtil {
         let conn = (await this.auth.listConnections()).find(isBuilderIdConnection)
 
         if (!conn) {
-            conn = await this.auth.createConnection(createBuilderIdProfile(codeWhispererChatScopes))
-        } else if (!isValidCodeWhispererChatConnection(conn)) {
-            conn = await this.secondaryAuth.addScopes(conn, codeWhispererChatScopes)
+            conn = await this.auth.createConnection(createBuilderIdProfile(amazonQScopes))
+        } else if (!isValidAmazonQConnection(conn)) {
+            conn = await this.secondaryAuth.addScopes(conn, amazonQScopes)
         }
 
         if (this.auth.getConnectionState(conn) === 'invalid') {
@@ -333,11 +328,10 @@ export class AuthUtil {
             // Edge Case: With the addition of Amazon Q/Chat scopes we may need to add
             // the new scopes to existing pre-chat connections.
             if (addMissingScopes) {
-                if (isBuilderIdConnection(this.conn) && !isValidCodeWhispererChatConnection(this.conn)) {
-                    const conn = await this.secondaryAuth.addScopes(this.conn, codeWhispererChatScopes)
-                    await this.secondaryAuth.useNewConnection(conn)
-                    return
-                } else if (isIdcSsoConnection(this.conn) && !isValidAmazonQConnection(this.conn)) {
+                if (
+                    (isBuilderIdConnection(this.conn) || isIdcSsoConnection(this.conn)) &&
+                    !isValidAmazonQConnection(this.conn)
+                ) {
                     const conn = await this.secondaryAuth.addScopes(this.conn, amazonQScopes)
                     await this.secondaryAuth.useNewConnection(conn)
                     return
@@ -386,7 +380,7 @@ export class AuthUtil {
     }
 
     public isValidCodeTransformationAuthUser(): boolean {
-        return this.isEnterpriseSsoInUse() && this.isConnectionValid()
+        return (this.isEnterpriseSsoInUse() || this.isBuilderIdInUse()) && this.isConnectionValid()
     }
 }
 
@@ -411,23 +405,11 @@ export async function getChatAuthState(cwAuth = AuthUtil.instance): Promise<Feat
     // default to expired to indicate reauth is needed if unmodified
     const state: FeatureAuthState = buildFeatureAuthState(AuthStates.expired)
 
-    if (isBuilderIdConnection(currentConnection)) {
-        // Regardless, if using Builder ID, Amazon Q is unsupported
-        state[Features.amazonQ] = AuthStates.unsupported
-    }
-
     if (cwAuth.isConnectionExpired()) {
         return state
     }
 
-    if (isBuilderIdConnection(currentConnection)) {
-        if (isValidCodeWhispererCoreConnection(currentConnection)) {
-            state[Features.codewhispererCore] = AuthStates.connected
-        }
-        if (isValidCodeWhispererChatConnection(currentConnection)) {
-            state[Features.codewhispererChat] = AuthStates.connected
-        }
-    } else if (isIdcSsoConnection(currentConnection)) {
+    if (isBuilderIdConnection(currentConnection) || isIdcSsoConnection(currentConnection)) {
         if (isValidCodeWhispererCoreConnection(currentConnection)) {
             state[Features.codewhispererCore] = AuthStates.connected
         }

packages/core/src/test/codewhisperer/util/authUtil.test.ts

@@ -44,7 +44,7 @@ describe('AuthUtil', async function () {
         const conn = authUtil.conn
         assert.strictEqual(conn?.type, 'sso')
         assert.strictEqual(conn.label, 'AWS Builder ID')
-        assert.deepStrictEqual(conn.scopes, codeWhispererChatScopes)
+        assert.deepStrictEqual(conn.scopes, amazonQScopes)
     })
 
     it('if there IS an existing AwsBuilderID conn, it will upgrade the scopes and use it', async function () {
@@ -61,7 +61,7 @@ describe('AuthUtil', async function () {
         const conn = authUtil.conn
         assert.strictEqual(conn?.type, 'sso')
         assert.strictEqual(conn.id, existingBuilderId.id)
-        assert.deepStrictEqual(conn.scopes, codeWhispererChatScopes)
+        assert.deepStrictEqual(conn.scopes, amazonQScopes)
     })
 
     it('if there is no valid enterprise SSO conn, will create and use one', async function () {
@@ -155,15 +155,15 @@ describe('AuthUtil', async function () {
         assert.deepStrictEqual(authUtil.conn?.scopes, codeWhispererCoreScopes)
     })
 
-    it('reauthenticate adds missing CodeWhisperer Chat Builder ID scopes when explicitly required', async function () {
+    it('reauthenticate adds missing Builder ID scopes when explicitly required', async function () {
         const conn = await auth.createConnection(createBuilderIdProfile({ scopes: codeWhispererCoreScopes }))
         await auth.useConnection(conn)
 
         // method under test
         await authUtil.reauthenticate(true)
 
         assert.strictEqual(authUtil.conn?.type, 'sso')
-        assert.deepStrictEqual(authUtil.conn?.scopes, codeWhispererChatScopes)
+        assert.deepStrictEqual(authUtil.conn?.scopes, amazonQScopes)
     })
 
     it('reauthenticate adds missing Amazon Q IdC scopes when explicitly required', async function () {
@@ -222,7 +222,7 @@ describe('AuthUtil', async function () {
         assert.strictEqual(authUtil.conn?.id, upgradeableConn.id)
         assert.strictEqual(authUtil.conn.startUrl, upgradeableConn.startUrl)
         assert.strictEqual(authUtil.conn.ssoRegion, upgradeableConn.ssoRegion)
-        assert.deepStrictEqual(authUtil.conn.scopes, codeWhispererChatScopes)
+        assert.deepStrictEqual(authUtil.conn.scopes, amazonQScopes)
         assert.strictEqual((await auth.listConnections()).filter(isAnySsoConnection).length, 1)
     })
 
@@ -277,20 +277,20 @@ describe('getChatAuthState()', function () {
             assert.deepStrictEqual(result, {
                 codewhispererCore: AuthStates.connected,
                 codewhispererChat: AuthStates.expired,
-                amazonQ: AuthStates.unsupported,
+                amazonQ: AuthStates.expired,
             })
         })
 
         it('indicates all SUPPORTED features connected when all scopes are set', async function () {
-            const conn = await auth.createConnection(createBuilderIdProfile({ scopes: codeWhispererChatScopes }))
+            const conn = await auth.createConnection(createBuilderIdProfile({ scopes: amazonQScopes }))
             createToken(conn)
             await auth.useConnection(conn)
 
             const result = await getChatAuthState(authUtil)
             assert.deepStrictEqual(result, {
                 codewhispererCore: AuthStates.connected,
                 codewhispererChat: AuthStates.connected,
-                amazonQ: AuthStates.unsupported,
+                amazonQ: AuthStates.connected,
             })
         })
 
@@ -304,7 +304,7 @@ describe('getChatAuthState()', function () {
             assert.deepStrictEqual(result, {
                 codewhispererCore: AuthStates.expired,
                 codewhispererChat: AuthStates.expired,
-                amazonQ: AuthStates.unsupported,
+                amazonQ: AuthStates.expired,
             })
         })
     })

packages/core/src/testE2E/util/codewhispererUtil.ts

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { isValidCodeWhispererChatConnection } from '../../codewhisperer/util/authUtil'
+import { isValidAmazonQConnection } from '../../codewhisperer/util/authUtil'
 import { Auth } from '../../auth/auth'
 
 /*
@@ -17,7 +17,7 @@ If user has an expired connection they must reauthenticate prior to running test
 */
 
 async function getValidConnection() {
-    return (await Auth.instance.listConnections()).find(isValidCodeWhispererChatConnection)
+    return (await Auth.instance.listConnections()).find(isValidAmazonQConnection)
 }
 
 //Returns true if a valid connection is found and set, false if not