Merge pull request #7707 from rhamilt/ssoQuickPickUpdate

feat(auth): Display second quickpick for SSO linked IAM profiles

Author: bywang56

packages/core/src/auth/utils.ts

@@ -40,6 +40,7 @@ import {
     hasScopes,
     scopesSsoAccountAccess,
     isSsoConnection,
+    IamConnection,
 } from './connection'
 import { Commands, placeholder } from '../shared/vscode/commands2'
 import { Auth } from './auth'
@@ -79,6 +80,18 @@ export async function promptForConnection(auth: Auth, type?: 'iam' | 'iam-only'
         return globals.awsContextCommands.onCommandEditCredentials()
     }
 
+    // If selected connection is SSO connection and has linked IAM profiles, show second quick pick with the linked IAM profiles
+    if (isSsoConnection(resp)) {
+        const linkedProfiles = await getLinkedIamProfiles(auth, resp)
+
+        if (linkedProfiles.length > 0) {
+            const linkedResp = await showLinkedProfilePicker(linkedProfiles, resp)
+            if (linkedResp) {
+                return linkedResp
+            }
+        }
+    }
+
     return resp
 }
 
@@ -340,6 +353,36 @@ export const createDeleteConnectionButton: () => vscode.QuickInputButton = () =>
     return { tooltip: deleteConnection, iconPath: getIcon('vscode-trash') }
 }
 
+async function getLinkedIamProfiles(auth: Auth, ssoConnection: SsoConnection): Promise<IamConnection[]> {
+    const allConnections = await auth.listAndTraverseConnections().promise()
+
+    return allConnections.filter(
+        (conn) => isIamConnection(conn) && conn.id.startsWith(`sso:${ssoConnection.id}#`)
+    ) as IamConnection[]
+}
+
+/**
+ * Shows a quick pick with linked IAM profiles for a selected SSO connection
+ */
+async function showLinkedProfilePicker(
+    linkedProfiles: IamConnection[],
+    ssoConnection: SsoConnection
+): Promise<IamConnection | undefined> {
+    const title = `Select an IAM Role for ${ssoConnection.label}`
+
+    const items: DataQuickPickItem<IamConnection>[] = linkedProfiles.map((profile) => ({
+        label: codicon`${getIcon('vscode-key')} ${profile.label}`,
+        description: 'IAM Credential, sourced from IAM Identity Center',
+        data: profile,
+    }))
+
+    return await showQuickPick(items, {
+        title,
+        placeholder: 'Select an IAM role',
+        buttons: [createRefreshButton(), createExitButton()],
+    })
+}
+
 export function createConnectionPrompter(auth: Auth, type?: 'iam' | 'iam-only' | 'sso') {
     const addNewConnection = {
         label: codicon`${getIcon('vscode-plus')} Add New Connection`,
@@ -433,22 +476,22 @@ export function createConnectionPrompter(auth: Auth, type?: 'iam' | 'iam-only' |
         for await (const conn of connections) {
             if (conn.label.includes('profile:') && !hasShownEdit) {
                 hasShownEdit = true
-                yield [toPickerItem(conn), editCredentials]
+                yield [await toPickerItem(conn), editCredentials]
             } else {
-                yield [toPickerItem(conn)]
+                yield [await toPickerItem(conn)]
             }
         }
     }
 
-    function toPickerItem(conn: Connection): DataQuickPickItem<Connection> {
+    async function toPickerItem(conn: Connection): Promise<DataQuickPickItem<Connection>> {
         const state = auth.getConnectionState(conn)
         // Only allow SSO connections to be deleted
         const deleteButton: vscode.QuickInputButton[] = conn.type === 'sso' ? [createDeleteConnectionButton()] : []
         if (state === 'valid') {
             return {
                 data: conn,
                 label: codicon`${getConnectionIcon(conn)} ${conn.label}`,
-                description: getConnectionDescription(conn),
+                description: await getConnectionDescription(conn),
                 buttons: [...deleteButton],
             }
         }
@@ -502,7 +545,7 @@ export function createConnectionPrompter(auth: Auth, type?: 'iam' | 'iam-only' |
         }
     }
 
-    function getConnectionDescription(conn: Connection) {
+    async function getConnectionDescription(conn: Connection) {
         if (conn.type === 'iam') {
             // TODO: implement a proper `getConnectionSource` method to discover where a connection came from
             const descSuffix = conn.id.startsWith('profile:')
@@ -514,6 +557,14 @@ export function createConnectionPrompter(auth: Auth, type?: 'iam' | 'iam-only' |
             return `IAM Credential, ${descSuffix}`
         }
 
+        // If this is an SSO connection, check if it has linked IAM profiles
+        if (isSsoConnection(conn)) {
+            const linkedProfiles = await getLinkedIamProfiles(auth, conn)
+            if (linkedProfiles.length > 0) {
+                return `Has ${linkedProfiles.length} IAM role${linkedProfiles.length > 1 ? 's' : ''} (click to select)`
+            }
+        }
+
         const toolAuths = getDependentAuths(conn)
         if (toolAuths.length === 0) {
             return undefined

packages/core/src/test/credentials/auth.test.ts

@@ -570,6 +570,47 @@ describe('Auth', function () {
             assert.strictEqual((await promptForConnection(auth))?.id, conn.id)
         })
 
+        it('shows a second quickPick for linked IAM profiles when selecting an SSO connection', async function () {
+            let quickPickCount = 0
+            getTestWindow().onDidShowQuickPick(async (picker) => {
+                await picker.untilReady()
+                quickPickCount++
+
+                if (quickPickCount === 1) {
+                    // First picker: select the SSO connection
+                    const connItem = picker.findItemOrThrow(/IAM Identity Center/)
+                    picker.acceptItem(connItem)
+                } else if (quickPickCount === 2) {
+                    // Second picker: select the linked IAM profile
+                    const linkedItem = picker.findItemOrThrow(/TestRole/)
+                    picker.acceptItem(linkedItem)
+                }
+            })
+
+            const linkedSsoProfile = createSsoProfile({ scopes: scopesSsoAccountAccess })
+            const conn = await auth.createConnection(linkedSsoProfile)
+
+            // Mock the SSOClient to return account roles
+            auth.ssoClient.listAccounts.returns(
+                toCollection(async function* () {
+                    yield [{ accountId: '123456789012' }]
+                })
+            )
+            auth.ssoClient.listAccountRoles.callsFake(() =>
+                toCollection(async function* () {
+                    yield [{ accountId: '123456789012', roleName: 'TestRole' }]
+                })
+            )
+
+            // Should get a linked IAM profile back, not the SSO connection
+            const result = await promptForConnection(auth)
+            assert.ok(isIamConnection(result || undefined), 'Expected an IAM connection to be returned')
+            assert.ok(
+                result?.id.startsWith(`sso:${conn.id}#`),
+                'Expected the IAM connection to be linked to the SSO connection'
+            )
+        })
+
         it('refreshes when clicking the refresh button', async function () {
             getTestWindow().onDidShowQuickPick(async (picker) => {
                 await picker.untilReady()

packages/toolkit/.changes/next-release/Feature-852d6637-ac3c-4b7f-b846-649d23da87e3.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Improved connection actions for SSO"
+}