From 34e7f1bea5f84e6647077843cf6d885fac966597 Mon Sep 17 00:00:00 2001 From: Reed Hamilton Date: Fri, 18 Jul 2025 12:01:09 -0700 Subject: [PATCH 1/3] Make sso quickpick option display profiles --- packages/core/src/auth/utils.ts | 61 ++++++++++++++++++++++++++++++--- 1 file changed, 56 insertions(+), 5 deletions(-) diff --git a/packages/core/src/auth/utils.ts b/packages/core/src/auth/utils.ts index 910c5cee949..b455780f45d 100644 --- a/packages/core/src/auth/utils.ts +++ b/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 { + 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 { + const title = `Select an IAM Role for ${ssoConnection.label}` + + const items: DataQuickPickItem[] = 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,14 +476,14 @@ 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 { + async function toPickerItem(conn: Connection): Promise> { const state = auth.getConnectionState(conn) // Only allow SSO connections to be deleted const deleteButton: vscode.QuickInputButton[] = conn.type === 'sso' ? [createDeleteConnectionButton()] : [] @@ -448,7 +491,7 @@ export function createConnectionPrompter(auth: Auth, type?: 'iam' | 'iam-only' | 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 From efb2d498af6d414d2f375499de4358ad2750620c Mon Sep 17 00:00:00 2001 From: Reed Hamilton Date: Fri, 18 Jul 2025 12:01:22 -0700 Subject: [PATCH 2/3] test: make sso quickpick option display profiles --- .../core/src/test/credentials/auth.test.ts | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/packages/core/src/test/credentials/auth.test.ts b/packages/core/src/test/credentials/auth.test.ts index 022e2c5c6e7..a409f2ba2e2 100644 --- a/packages/core/src/test/credentials/auth.test.ts +++ b/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() From 9f082d9009b1aeb00a4d5e9181f43f964862ea63 Mon Sep 17 00:00:00 2001 From: Reed Hamilton Date: Fri, 18 Jul 2025 12:57:01 -0700 Subject: [PATCH 3/3] changelog update --- .../Feature-852d6637-ac3c-4b7f-b846-649d23da87e3.json | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 packages/toolkit/.changes/next-release/Feature-852d6637-ac3c-4b7f-b846-649d23da87e3.json diff --git a/packages/toolkit/.changes/next-release/Feature-852d6637-ac3c-4b7f-b846-649d23da87e3.json b/packages/toolkit/.changes/next-release/Feature-852d6637-ac3c-4b7f-b846-649d23da87e3.json new file mode 100644 index 00000000000..2d7652be636 --- /dev/null +++ b/packages/toolkit/.changes/next-release/Feature-852d6637-ac3c-4b7f-b846-649d23da87e3.json @@ -0,0 +1,4 @@ +{ + "type": "Feature", + "description": "Improved connection actions for SSO" +}