amazonQ: only add missing scopes intentionally (#4152)

* amazonQ: only add missing scopes intentionally

Problem:

Before we would always add missing scopes whenever an
existing CW user was missing the new scopes for chat.
This would be done during the reauthentication process
automatically.

The issue is that we did not want this and only wanted
to add missing Amazon Q/chat scopes when the user intentionally
wanted them to be added.

Solution:
- Don't always add missing scopes during reauthenticate(). Instead
  use an arg to determine if we should
- We will add missing scopes in the Chat "reauthenticate" prompt as well
  as the Amazon Q node.
- Everywhere else will simply reauthenticate without adding new scopes

Signed-off-by: nkomonen <[email protected]>

* add type to signature

Signed-off-by: nkomonen <[email protected]>

* add check for invalid arg

There is a possibility that addMissingScopes is incorrectly set by vscode,
so check and set a default if necessary

Signed-off-by: nkomonen <[email protected]>

---------

Signed-off-by: nkomonen <[email protected]>

Author: nkomonen-amazon

src/amazonq/auth/controller.ts

@@ -30,6 +30,6 @@ export class AuthController {
     }
 
     private handleReAuth() {
-        reconnect.execute(placeholder, amazonQChatSource)
+        reconnect.execute(placeholder, amazonQChatSource, true)
     }
 }

src/amazonq/explorer/amazonQChildrenNodes.ts

@@ -50,7 +50,7 @@ export const switchToAmazonQNode = () =>
  */
 export const enableAmazonQNode = () =>
     // Simply trigger re-auth to obtain the proper scopes- same functionality as if requested in the chat window.
-    reconnect.build(placeholder, cwTreeNodeSource).asTreeNode({
+    reconnect.build(placeholder, cwTreeNodeSource, true).asTreeNode({
         label: localize('AWS.amazonq.enable', 'Enable'),
         iconPath: getIcon('vscode-debug-start'),
         contextValue: 'awsEnableAmazonQ',

src/codewhisperer/commands/basicCommands.ts

@@ -151,9 +151,13 @@ export const selectCustomizationPrompt = Commands.declare(
 
 export const reconnect = Commands.declare(
     { id: 'aws.codewhisperer.reconnect', compositeKey: { 1: 'source' } },
-    () => async (_: VsCodeCommandArg, source: CodeWhispererSource) => {
-        await AuthUtil.instance.reauthenticate()
-    }
+    () =>
+        async (_: VsCodeCommandArg, source: CodeWhispererSource, addMissingScopes: boolean = false) => {
+            if (typeof addMissingScopes !== 'boolean') {
+                addMissingScopes = false
+            }
+            await AuthUtil.instance.reauthenticate(addMissingScopes)
+        }
 )
 
 /** Opens the Add Connections webview with CW highlighted */

src/codewhisperer/util/authUtil.ts

@@ -36,8 +36,6 @@ export const codeWhispererCoreScopes = [...scopesSsoAccountAccess, ...scopesCode
 export const codeWhispererChatScopes = [...codeWhispererCoreScopes, ...scopesCodeWhispererChat]
 export const amazonQScopes = [...codeWhispererChatScopes, ...scopesGumby, ...scopesFeatureDev]
 
-export const awsBuilderIdSsoProfile = createBuilderIdProfile(codeWhispererChatScopes)
-
 /**
  * "Core" are the CW scopes that existed before the addition of new scopes
  * for Amazon Q.
@@ -211,7 +209,7 @@ export class AuthUtil {
         let conn = (await this.auth.listConnections()).find(isBuilderIdConnection)
 
         if (!conn) {
-            conn = await this.auth.createConnection(awsBuilderIdSsoProfile)
+            conn = await this.auth.createConnection(createBuilderIdProfile(codeWhispererChatScopes))
         } else if (!isValidCodeWhispererChatConnection(conn)) {
             conn = await this.secondaryAuth.addScopes(conn, codeWhispererChatScopes)
         }
@@ -304,23 +302,27 @@ export class AuthUtil {
         return connectionExpired
     }
 
-    public async reauthenticate() {
+    public async reauthenticate(addMissingScopes: boolean = false) {
         try {
             if (this.conn?.type !== 'sso') {
                 return
             }
 
             // Edge Case: With the addition of Amazon Q/Chat scopes we may need to add
-            // the new scopes to existing connections.
-            if (isBuilderIdConnection(this.conn) && !isValidCodeWhispererChatConnection(this.conn)) {
-                const conn = await this.secondaryAuth.addScopes(this.conn, codeWhispererChatScopes)
-                this.secondaryAuth.useNewConnection(conn)
-            } else if (isIdcSsoConnection(this.conn) && !isValidAmazonQConnection(this.conn)) {
-                const conn = await this.secondaryAuth.addScopes(this.conn, amazonQScopes)
-                this.secondaryAuth.useNewConnection(conn)
-            } else {
-                await this.auth.reauthenticate(this.conn)
+            // 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)) {
+                    const conn = await this.secondaryAuth.addScopes(this.conn, amazonQScopes)
+                    await this.secondaryAuth.useNewConnection(conn)
+                    return
+                }
             }
+
+            await this.auth.reauthenticate(this.conn)
         } catch (err) {
             throw ToolkitError.chain(err, 'Unable to authenticate connection')
         } finally {

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

@@ -145,23 +145,35 @@ describe('AuthUtil', async function () {
         assert.strictEqual(auth.getConnectionState(conn), 'valid')
     })
 
-    it('reauthenticate adds missing CodeWhisperer Chat Builder ID scopes', async function () {
+    it('reauthenticate does NOT add missing CodeWhisperer scopes if not required to', async function () {
         const conn = await auth.createConnection(createBuilderIdProfile({ scopes: codeWhispererCoreScopes }))
         await auth.useConnection(conn)
 
         await authUtil.reauthenticate()
 
+        assert.strictEqual(authUtil.conn?.type, 'sso')
+        assert.deepStrictEqual(authUtil.conn?.scopes, codeWhispererCoreScopes)
+    })
+
+    it('reauthenticate adds missing CodeWhisperer Chat 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)
     })
 
-    it('reauthenticate adds missing Amazon Q IdC scopes', async function () {
+    it('reauthenticate adds missing Amazon Q IdC scopes when explicitly required', async function () {
         const conn = await auth.createConnection(
             createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: codeWhispererCoreScopes })
         )
         await auth.useConnection(conn)
 
-        await authUtil.reauthenticate()
+        // method under test
+        await authUtil.reauthenticate(true)
 
         assert.strictEqual(authUtil.conn?.type, 'sso')
         assert.deepStrictEqual(authUtil.conn?.scopes, amazonQScopes)