auth: remove multiple auth prompts (#3887)

* auth: remove yes/no prompt when using multiple auths

Problem:

If a user has an auth already set up, eg credentials, then
I add a new auth, eg builder id, they will be shown a confusing
message. The message asks if they want to explicitly switch to the new
connection (no), or use both at the same time (yes).

Solution:

Remove this yes/no prompt and simply choose the (yes) option.

Additional:

Does some refactoring of secondaryAuth#registerAuthListener() since
it was clear what it was intended to be doing.

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

* auth: dont prompt to upgrade builder id scopes

Problem:

If a CodeCatalyst Builder ID was setup and then
a CodeCatalyst Builder ID setup was requested, the
user would be asked if they want to 'switch' their
connection.

Solution:

Don't prompt the user to switch their connect and
just upgrade their CodeWhisperer Builder ID scopes
to also include CodeCatalyst scopes.

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

* authNode: remove from developer tools menu

Removes the auth node from the developer tools menu
since the actual services in the developer tools menu
may not be using the same connection as the auth node
and it is confusing to users.

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

* codewhisperer: fix unit tests

Now that we have removed the prompts when adding/switching
sso connections, the tests need to be updated since they user
message that was being expected is not shown anymore.

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

---------

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

Author: nkomonen-amazon

src/auth/secondaryAuth.ts

@@ -7,79 +7,30 @@ import globals from '../shared/extensionGlobals'
 
 import * as vscode from 'vscode'
 import { getLogger } from '../shared/logger'
-import { showQuickPick } from '../shared/ui/pickerPrompter'
 import { cast, Optional } from '../shared/utilities/typeConstructors'
 import { Auth } from './auth'
 import { once } from '../shared/utilities/functionUtils'
-import { telemetry } from '../shared/telemetry/telemetry'
-import { createExitButton, createHelpButton } from '../shared/ui/buttons'
 import { isNonNullable } from '../shared/utilities/tsUtils'
 import { CancellationError } from '../shared/utilities/timeoutUtils'
 import { Connection, SsoConnection, StatefulConnection } from './connection'
 
-async function promptUseNewConnection(newConn: Connection, oldConn: Connection, tools: string[], swapNo: boolean) {
-    // Multi-select picker would be better ?
-    const saveConnectionItem = {
-        label: `Yes, keep using ${newConn.label} with ${tools.join(', ')} while using ${
-            oldConn.label
-        } with other services.`,
-        detail: `To remove later, select "Remove Connection from Tool" from the tool's context (right-click) menu.`,
-        data: 'yes',
-    } as const
-
-    const useConnectionItem = {
-        label: `No, switch everything to authenticate with ${(swapNo ? newConn : oldConn).label}.`,
-        detail: 'This will not log you out; you can reconnect at any time by switching connections.',
-        data: 'no',
-    } as const
-
-    const helpButton = createHelpButton()
-    const openLink = helpButton.onClick.bind(helpButton)
-    helpButton.onClick = () => {
-        telemetry.ui_click.emit({ elementId: 'connection_multiple_auths_help' })
-        openLink()
-    }
-
-    const resp = await showQuickPick([saveConnectionItem, useConnectionItem], {
-        title: `Some tools you've been using don't work with ${oldConn.label}. Keep using ${newConn.label} in the background while using ${oldConn.label}?`,
-        placeholder: 'Confirm choice',
-        buttons: [helpButton, createExitButton()],
-    })
-
-    switch (resp) {
-        case 'yes':
-            telemetry.ui_click.emit({ elementId: 'connection_multiple_auths_yes' })
-            break
-        case 'no':
-            telemetry.ui_click.emit({ elementId: 'connection_multiple_auths_no' })
-            break
-        default:
-            telemetry.ui_click.emit({ elementId: 'connection_multiple_auths_exit' })
-    }
-
-    return resp
-}
-
-let oldConn: Auth['activeConnection']
+let currentConn: Auth['activeConnection']
 const auths = new Map<string, SecondaryAuth>()
 const multiConnectionListeners = new WeakMap<Auth, vscode.Disposable>()
 const registerAuthListener = (auth: Auth) => {
-    return auth.onDidChangeActiveConnection(async conn => {
-        const potentialConn = oldConn
-        if (conn !== undefined && potentialConn?.state === 'valid') {
+    return auth.onDidChangeActiveConnection(async newConn => {
+        // When we change the active connection, there may be
+        // secondary auths that were dependent on the previous active connection.
+        // To ensure secondary auths still work, when we change to a new active connection,
+        // the following will "save" the oldConn with the secondary auths that are using it.
+        const oldConn = currentConn
+        if (newConn && oldConn?.state === 'valid') {
             const saveableAuths = Array.from(auths.values()).filter(
-                a => !a.isUsingSavedConnection && a.isUsable(potentialConn) && !a.isUsable(conn)
+                a => !a.hasSavedConnection && a.isUsable(oldConn) && !a.isUsable(newConn)
             )
-            const toolNames = saveableAuths.map(a => a.toolLabel)
-            if (
-                saveableAuths.length > 0 &&
-                (await promptUseNewConnection(potentialConn, conn, toolNames, false)) === 'yes'
-            ) {
-                await Promise.all(saveableAuths.map(a => a.saveConnection(potentialConn)))
-            }
+            await Promise.all(saveableAuths.map(a => a.saveConnection(oldConn)))
         }
-
-        oldConn = conn
+        currentConn = newConn
     })
 }
 
@@ -104,15 +55,13 @@ export function getSecondaryAuth<T extends Connection>(
  * Gets all {@link SecondaryAuth} instances that have saved the connection
  */
 export function getDependentAuths(conn: Connection): SecondaryAuth[] {
-    return Array.from(auths.values()).filter(
-        auth => auth.isUsingSavedConnection && auth.activeConnection?.id === conn.id
-    )
+    return Array.from(auths.values()).filter(auth => auth.hasSavedConnection && auth.activeConnection?.id === conn.id)
 }
 
 export function getAllConnectionsInUse(auth: Auth): StatefulConnection[] {
     const connMap = new Map<Connection['id'], StatefulConnection>()
     const toolConns = Array.from(auths.values())
-        .filter(a => a.isUsingSavedConnection)
+        .filter(a => a.hasSavedConnection)
         .map(a => a.activeConnection)
 
     for (const conn of [auth.activeConnection, ...toolConns].filter(isNonNullable)) {
@@ -174,13 +123,18 @@ export class SecondaryAuth<T extends Connection = Connection> {
     }
 
     public get activeConnection(): T | undefined {
-        return (
-            this.#savedConnection ??
-            (this.#activeConnection && this.isUsable(this.#activeConnection) ? this.#activeConnection : undefined)
-        )
+        if (this.#savedConnection) {
+            return this.#savedConnection
+        }
+
+        if (this.#activeConnection && this.isUsable(this.#activeConnection)) {
+            return this.#activeConnection
+        }
+
+        return undefined
     }
 
-    public get isUsingSavedConnection() {
+    public get hasSavedConnection() {
         return this.#savedConnection !== undefined
     }
 
@@ -202,11 +156,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
 
     public async useNewConnection(conn: T) {
         if (this.auth.activeConnection !== undefined && !this.isUsable(this.auth.activeConnection)) {
-            if ((await promptUseNewConnection(conn, this.auth.activeConnection, [this.toolLabel], true)) === 'yes') {
-                await this.saveConnection(conn)
-            } else {
-                await this.auth.useConnection(conn)
-            }
+            await this.saveConnection(conn)
         } else {
             await this.auth.useConnection(conn)
         }

src/awsexplorer/activation.ts

@@ -27,11 +27,9 @@ import { telemetry } from '../shared/telemetry/telemetry'
 import { cdkNode, CdkRootNode } from '../cdk/explorer/rootNode'
 import { CodeWhispererNode, codewhispererNode } from '../codewhisperer/explorer/codewhispererNode'
 import { once } from '../shared/utilities/functionUtils'
-import { Auth } from '../auth/auth'
 import { CodeCatalystRootNode } from '../codecatalyst/explorer'
 import { CodeCatalystAuthenticationProvider } from '../codecatalyst/auth'
 import { S3FolderNode } from '../s3/explorer/s3FolderNode'
-import { AuthNode } from '../auth/utils'
 import { TreeNode } from '../shared/treeview/resourceTreeDataProvider'
 
 /**
@@ -78,9 +76,8 @@ export async function activate(args: {
     )
 
     const authProvider = CodeCatalystAuthenticationProvider.fromContext(args.context.extensionContext)
-    const authNode = new AuthNode(Auth.instance)
     const codecatalystNode = isCloud9('classic') ? [] : [new CodeCatalystRootNode(authProvider)]
-    const nodes = [authNode, ...codecatalystNode, cdkNode, codewhispererNode]
+    const nodes = [...codecatalystNode, cdkNode, codewhispererNode]
     const developerTools = createLocalExplorerView(nodes)
     args.context.extensionContext.subscriptions.push(developerTools)
 
@@ -99,7 +96,6 @@ export async function activate(args: {
     })
 
     registerDeveloperToolsCommands(args.context.extensionContext, developerTools, {
-        auth: authNode,
         codeCatalyst: codecatalystNode ? codecatalystNode[0] : undefined,
         codeWhisperer: codewhispererNode,
     })
@@ -175,7 +171,6 @@ async function registerDeveloperToolsCommands(
     ctx: vscode.ExtensionContext,
     developerTools: vscode.TreeView<TreeNode>,
     nodes: {
-        auth: AuthNode
         codeWhisperer: CodeWhispererNode
         codeCatalyst: CodeCatalystRootNode | undefined
     }
@@ -203,10 +198,9 @@ async function registerDeveloperToolsCommands(
         )
     }
 
-    registerShowDeveloperToolsNode('CodeWhisperer', codewhispererNode)
+    registerShowDeveloperToolsNode('CodeWhisperer', nodes.codeWhisperer)
 
-    const codeCatalystNode = nodes.codeCatalyst
-    if (codeCatalystNode) {
-        registerShowDeveloperToolsNode('CodeCatalyst', codeCatalystNode)
+    if (nodes.codeCatalyst) {
+        registerShowDeveloperToolsNode('CodeCatalyst', nodes.codeCatalyst)
     }
 }

src/codecatalyst/auth.ts

@@ -8,7 +8,6 @@ import { CodeCatalystClient, createClient } from '../shared/clients/codecatalyst
 import { Auth } from '../auth/auth'
 import { getSecondaryAuth } from '../auth/secondaryAuth'
 import { getLogger } from '../shared/logger'
-import * as localizedText from '../shared/localizedText'
 import { ToolkitError, isAwsError } from '../shared/errors'
 import { MetricName, MetricShapes, telemetry } from '../shared/telemetry/telemetry'
 import { openUrl } from '../shared/utilities/vsCodeUtils'
@@ -66,7 +65,7 @@ export class CodeCatalystAuthenticationProvider {
     }
 
     public get isUsingSavedConnection() {
-        return this.secondaryAuth.isUsingSavedConnection
+        return this.secondaryAuth.hasSavedConnection
     }
 
     public isConnectionValid(): boolean {
@@ -123,14 +122,22 @@ export class CodeCatalystAuthenticationProvider {
         throw new ToolkitError('Not onboarded with CodeCatalyst', { code: 'NotOnboarded', cancelled: true })
     }
 
-    public async promptNotConnected(): Promise<SsoConnection> {
+    /**
+     * Return a Builder ID connection that works with CodeCatalyst.
+     *
+     * This cannot create a Builder ID, but will return an existing Builder ID,
+     * upgrading the scopes if necessary.
+     */
+    public async tryGetBuilderIdConnection(): Promise<SsoConnection> {
+        if (this.activeConnection && isBuilderIdConnection(this.activeConnection)) {
+            return this.activeConnection
+        }
+
         type ConnectionFlowEvent = Partial<MetricShapes[MetricName]> & {
             readonly codecatalyst_connectionFlow: 'Create' | 'Switch' | 'Upgrade' // eslint-disable-line @typescript-eslint/naming-convention
         }
 
         const conn = (await this.auth.listConnections()).find(isBuilderIdConnection)
-        const continueItem: vscode.MessageItem = { title: localizedText.continueText }
-        const cancelItem: vscode.MessageItem = { title: localizedText.cancel, isCloseAffordance: true }
 
         if (conn === undefined) {
             telemetry.record({
@@ -153,21 +160,11 @@ export class CodeCatalystAuthenticationProvider {
             return this.secondaryAuth.addScopes(conn, defaultScopes)
         }
 
-        if (isBuilderIdConnection(conn) && this.auth.activeConnection?.id !== conn.id) {
+        if (this.auth.activeConnection?.id !== conn.id) {
             telemetry.record({
                 codecatalyst_connectionFlow: 'Switch',
             } satisfies ConnectionFlowEvent as MetricShapes[MetricName])
 
-            const resp = await vscode.window.showInformationMessage(
-                'CodeCatalyst requires an AWS Builder ID connection.\n\n Switch to it now?',
-                { modal: true },
-                continueItem,
-                cancelItem
-            )
-            if (resp !== continueItem) {
-                throw new ToolkitError('Not connected to CodeCatalyst', { code: 'NoConnection', cancelled: true })
-            }
-
             if (isUpgradeableConnection(conn)) {
                 await upgrade()
             }

src/codecatalyst/commands.ts

@@ -136,7 +136,7 @@ function createClientInjector(authProvider: CodeCatalystAuthenticationProvider):
         telemetry.record({ userId: AccountStatus.NotSet })
 
         await authProvider.restore()
-        const conn = authProvider.activeConnection ?? (await authProvider.promptNotConnected())
+        const conn = await authProvider.tryGetBuilderIdConnection()
         const validatedConn = await validateConnection(conn, authProvider.auth)
         const client = await createClient(validatedConn)
         telemetry.record({ userId: client.identity.id })

src/codecatalyst/utils.ts

@@ -69,7 +69,7 @@ export function isInDevEnv(): boolean {
 export const getStartedCommand = Commands.register('aws.codecatalyst.getStarted', setupCodeCatalystBuilderId)
 
 export async function setupCodeCatalystBuilderId(authProvider: CodeCatalystAuthenticationProvider) {
-    let conn = authProvider.activeConnection ?? (await authProvider.promptNotConnected())
+    let conn = await authProvider.tryGetBuilderIdConnection()
 
     if (authProvider.auth.getConnectionState(conn) === 'invalid') {
         conn = await authProvider.auth.reauthenticate(conn)

src/codewhisperer/util/authUtil.ts

@@ -104,7 +104,7 @@ export class AuthUtil {
     }
 
     public get isUsingSavedConnection() {
-        return this.conn !== undefined && this.secondaryAuth.isUsingSavedConnection
+        return this.conn !== undefined && this.secondaryAuth.hasSavedConnection
     }
 
     public isConnected(): boolean {

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

@@ -22,6 +22,10 @@ describe('AuthUtil', async function () {
         authUtil = new AuthUtil(auth)
     })
 
+    afterEach(async function () {
+        await auth.logout()
+    })
+
     it('if there is no valid AwsBuilderID conn, it will create one and use it', async function () {
         getTestWindow().onDidShowQuickPick(async picker => {
             await picker.untilReady()
@@ -83,22 +87,20 @@ describe('AuthUtil', async function () {
         )
     })
 
-    it('prompts to attach connection to CodeWhisperer when switching to an unsupported connection', async function () {
+    it('CodeWhisperer uses fallback connection when switching to an unsupported connection', async function () {
         const supportedConn = await auth.createConnection(createBuilderIdProfile({ scopes: codewhispererScopes }))
         const unsupportedConn = await auth.createConnection(createSsoProfile())
 
-        getTestWindow().onDidShowQuickPick(picker => {
-            assert.ok(picker.title?.startsWith(`Some tools you've been using don't work with ${unsupportedConn.label}`))
-            const keepUsing = picker.findItemOrThrow(new RegExp(`keep using ${supportedConn.label}`))
-            picker.acceptItem(keepUsing)
-        })
-
         await auth.useConnection(supportedConn)
         assert.ok(authUtil.isConnected())
         assert.strictEqual(auth.activeConnection?.id, authUtil.conn?.id)
 
+        // Switch to unsupported connection
+        const cwAuthUpdatedConnection = captureEventOnce(authUtil.secondaryAuth.onDidChangeActiveConnection)
         await auth.useConnection(unsupportedConn)
-        await captureEventOnce(authUtil.secondaryAuth.onDidChangeActiveConnection)
+        await cwAuthUpdatedConnection
+
+        // Is using the fallback connection
         assert.ok(authUtil.isConnected())
         assert.ok(authUtil.isUsingSavedConnection)
         assert.notStrictEqual(auth.activeConnection?.id, authUtil.conn?.id)
@@ -114,13 +116,7 @@ describe('AuthUtil', async function () {
         assert.strictEqual((await auth.listConnections()).filter(isSsoConnection).length, 1)
     })
 
-    it('prompts to upgrade connections if they do not have the required scopes', async function () {
-        getTestWindow().onDidShowMessage(message => {
-            assert.ok(message.modal)
-            message.assertMessage(/CodeWhisperer requires access to your/)
-            message.selectItem('Proceed')
-        })
-
+    it('automatically upgrades connections if they do not have the required scopes', async function () {
         const upgradeableConn = await auth.createConnection(createBuilderIdProfile())
         await auth.useConnection(upgradeableConn)
         assert.strictEqual(authUtil.isConnected(), false)