fix: show CW connection prompter when appropriate

Before this change we showed a connection prompter
that was not CW specific, with this commit it
now will be the correct prompter for CW specific scopes.

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

Author: nkomonen-amazon

src/auth/ui/vue/show.ts

@@ -46,7 +46,7 @@ import { CancellationError } from '../../../shared/utilities/timeoutUtils'
 import { validateSsoUrl, validateSsoUrlFormat } from '../../sso/validation'
 import { debounce } from '../../../shared/utilities/functionUtils'
 import { AuthError, ServiceItemId, isServiceItemId, userCancelled } from './types'
-import { awsIdSignIn } from '../../../codewhisperer/util/showSsoPrompt'
+import { awsIdSignIn, showCodeWhispererConnectionPrompt } from '../../../codewhisperer/util/showSsoPrompt'
 import { connectToEnterpriseSso } from '../../../codewhisperer/util/getStartUrl'
 import { trustedDomainCancellation } from '../../sso/model'
 import { FeatureId, CredentialSourceId, Result, telemetry } from '../../../shared/telemetry/telemetry'
@@ -741,16 +741,20 @@ export type AuthSource = (typeof AuthSources)[keyof typeof AuthSources]
 export const showManageConnections = Commands.declare(
     { id: showConnectionsPageCommand, compositeKey: { 1: 'source' } },
     (context: vscode.ExtensionContext) => (_: VsCodeCommandArg, source: AuthSource, serviceToShow?: ServiceItemId) => {
+        if (_ !== placeholder) {
+            source = 'vscodeComponent'
+        }
+
         // The auth webview page does not make sense to use in C9,
         // so show the auth quick pick instead.
         if (isCloud9('any')) {
+            if (source.toLowerCase().includes('codewhisperer')) {
+                // Show CW specific quick pick for CW connections
+                return showCodeWhispererConnectionPrompt()
+            }
             return addConnection.execute()
         }
 
-        if (_ !== placeholder) {
-            source = 'vscodeComponent'
-        }
-
         if (!isServiceItemId(serviceToShow)) {
             serviceToShow = undefined
         }

src/codewhisperer/commands/basicCommands.ts

@@ -13,7 +13,7 @@ import { startSecurityScanWithProgress, confirmStopSecurityScan } from './startS
 import { SecurityPanelViewProvider } from '../views/securityPanelViewProvider'
 import { CodeSuggestionsState, codeScanState } from '../models/model'
 import { connectToEnterpriseSso, getStartUrl } from '../util/getStartUrl'
-import { showConnectionPrompt } from '../util/showSsoPrompt'
+import { showCodeWhispererConnectionPrompt } from '../util/showSsoPrompt'
 import { ReferenceLogViewProvider } from '../service/referenceLogViewProvider'
 import { AuthUtil } from '../util/authUtil'
 import { isCloud9 } from '../../shared/extensionUtilities'
@@ -116,7 +116,7 @@ export const showManageCwConnections = Commands.declare(
 /** @deprecated in favor of the `Add Connection` page */
 export const showSsoSignIn = Commands.declare('aws.codeWhisperer.sso', () => async () => {
     telemetry.ui_click.emit({ elementId: 'cw_signUp_Cta' })
-    await showConnectionPrompt()
+    await showCodeWhispererConnectionPrompt()
 })
 
 // Shortcut command to directly connect to Identity Center or prompt start URL entry

src/codewhisperer/util/showSsoPrompt.ts

@@ -15,27 +15,14 @@ import { ToolkitError } from '../../shared/errors'
 import { createCommonButtons } from '../../shared/ui/buttons'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { isCloud9 } from '../../shared/extensionUtilities'
-import { isIamConnection } from '../../auth/connection'
 import { createBuilderIdItem, createSsoItem, createIamItem } from '../../auth/utils'
 
-export const showConnectionPrompt = async () => {
-    // Skip this prompt on C9 because:
-    // * The UI looks bad with C9's style of pickers
-    // * C9 will always start with _some_ form of auth so this prompt is less common
-    if (isCloud9()) {
-        if (isCloud9('classic')) {
-            const iamConn = (await AuthUtil.instance.auth.listConnections()).find(isIamConnection)
-            if (iamConn) {
-                await AuthUtil.instance.auth.useConnection(iamConn)
-            }
-        } else {
-            await AuthUtil.instance.connectToAwsBuilderId()
-        }
-
-        return
-    }
+export const showCodeWhispererConnectionPrompt = async () => {
+    const items = isCloud9('classic')
+        ? [createSsoItem(), createCodeWhispererIamItem()]
+        : [createBuilderIdItem(), createSsoItem(), createCodeWhispererIamItem()]
 
-    const resp = await showQuickPick([createBuilderIdItem(), createSsoItem(), createCodeWhispererIamItem()], {
+    const resp = await showQuickPick(items, {
         title: 'CodeWhisperer: Add Connection to AWS',
         placeholder: 'Select a connection option to start using CodeWhisperer',
         buttons: createCommonButtons() as vscode.QuickInputButton[],

src/test/codewhisperer/util/showSsoPrompt.test.ts

@@ -7,7 +7,7 @@ import * as vscode from 'vscode'
 import assert from 'assert'
 import * as sinon from 'sinon'
 import { resetCodeWhispererGlobalVariables } from '../testUtil'
-import { awsIdSignIn, showConnectionPrompt } from '../../../codewhisperer/util/showSsoPrompt'
+import { awsIdSignIn, showCodeWhispererConnectionPrompt } from '../../../codewhisperer/util/showSsoPrompt'
 import { getTestLogger } from '../../globalSetup.test'
 import { AuthUtil } from '../../../codewhisperer/util/authUtil'
 import { getTestWindow } from '../../shared/vscode/window'
@@ -30,7 +30,7 @@ describe('showConnectionPrompt', function () {
             picker.acceptItem(picker.items[0])
         })
 
-        await showConnectionPrompt()
+        await showCodeWhispererConnectionPrompt()
 
         assert.ok(authUtilSpy.called)
         const assertTelemetry = assertTelemetryCurried('ui_click')