Merge staging into feature/redshift

Author: aws-toolkit-automation

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)