feat: sm ide #3837

Problem:
New SM VSCode IDE with IAM role, we want to use the IAM role for Code Whisperer
auth. Currently Code Whisperer auth used Builder ID/SSO/Cloud9(IAM). Current IAM
auth is for Cloud9 IDE and does not allow for paginated results.

Solution:
Add new IDE type for SM VSCode that can be used for this feature and future
modifications. Use the new IDE type for conditional on connection type to use,
defaulting to IAM. Use IAM for Code Whisperer and make sure we use the correct,
`generateRecommendations`, endpoint to allow for auth with IAM and paginated
results.

Author: JedML

src/auth/auth.ts

@@ -53,6 +53,7 @@ import {
     loadLinkedProfilesIntoStore,
     ssoAccountAccessScopes,
 } from './connection'
+import { isSageMaker, isCloud9 } from '../shared/extensionUtilities'
 
 interface AuthService {
     /**
@@ -777,3 +778,15 @@ export class Auth implements AuthService, ConnectionManager {
             : `${localizedText.iamIdentityCenter} (${truncatedUrl})`
     }
 }
+/**
+ * Returns true if credentials are provided by the environment (ex. via ~/.aws/)
+ *
+ * @param isC9 boolean for if Cloud9 is host
+ * @param isSM boolean for if SageMaker is host
+ * @returns boolean for if C9 "OR" SM
+ */
+export function hasVendedIamCredentials(isC9?: boolean, isSM?: boolean) {
+    isC9 ??= isCloud9()
+    isSM ??= isSageMaker()
+    return isSM || isC9
+}

src/codewhisperer/client/codewhisperer.ts

@@ -9,7 +9,7 @@ import * as CodeWhispererUserClient from './codewhispereruserclient'
 import { ListAvailableCustomizationsResponse, SendTelemetryEventRequest } from './codewhispereruserclient'
 import * as CodeWhispererConstants from '../models/constants'
 import { ServiceOptions } from '../../shared/awsClientBuilder'
-import { isCloud9 } from '../../shared/extensionUtilities'
+import { hasVendedIamCredentials } from '../../auth/auth'
 import { CodeWhispererSettings } from '../util/codewhispererSettings'
 import { PromiseResult } from 'aws-sdk/lib/request'
 import { AuthUtil } from '../util/authUtil'
@@ -85,7 +85,8 @@ export class DefaultCodeWhispererClient {
                         }
                         // This logic is for backward compatability with legacy SDK v2 behavior for refreshing
                         // credentials. Once the Toolkit adds a file watcher for credentials it won't be needed.
-                        if (isCloud9()) {
+
+                        if (hasVendedIamCredentials()) {
                             req.on('retry', resp => {
                                 if (
                                     resp.error?.code === 'AccessDeniedException' &&

src/codewhisperer/commands/basicCommands.ts

@@ -60,14 +60,15 @@ export const createGettingStartedNode = () =>
 
 export const enableCodeSuggestions = Commands.declare(
     'aws.codeWhisperer.enableCodeSuggestions',
-    (context: ExtContext) => async () => {
-        await set(CodeWhispererConstants.autoTriggerEnabledKey, true, context.extensionContext.globalState)
-        await vscode.commands.executeCommand('setContext', 'CODEWHISPERER_ENABLED', true)
-        await vscode.commands.executeCommand('aws.codeWhisperer.refresh')
-        if (!isCloud9()) {
-            await vscode.commands.executeCommand('aws.codeWhisperer.refreshStatusBar')
+    (context: ExtContext) =>
+        async (isAuto: boolean = true) => {
+            await set(CodeWhispererConstants.autoTriggerEnabledKey, isAuto, context.extensionContext.globalState)
+            await vscode.commands.executeCommand('setContext', 'CODEWHISPERER_ENABLED', true)
+            await vscode.commands.executeCommand('aws.codeWhisperer.refresh')
+            if (!isCloud9()) {
+                await vscode.commands.executeCommand('aws.codeWhisperer.refreshStatusBar')
+            }
         }
-    }
 )
 
 export const showReferenceLog = Commands.declare(

src/codewhisperer/explorer/codewhispererNode.ts

@@ -19,7 +19,7 @@ import {
 import { createGettingStartedNode } from '../commands/basicCommands'
 import { Commands } from '../../shared/vscode/commands2'
 import { RootNode } from '../../awsexplorer/localExplorer'
-import { isCloud9 } from '../../shared/extensionUtilities'
+import { hasVendedIamCredentials } from '../../auth/auth'
 import { AuthUtil } from '../util/authUtil'
 import { TreeNode } from '../../shared/treeview/resourceTreeDataProvider'
 
@@ -57,9 +57,13 @@ export class CodeWhispererNode implements RootNode {
 
     private getDescription(): string {
         if (AuthUtil.instance.isConnectionValid()) {
-            return AuthUtil.instance.isEnterpriseSsoInUse()
-                ? 'IAM Identity Center Connected'
-                : 'AWS Builder ID Connected'
+            if (AuthUtil.instance.isEnterpriseSsoInUse()) {
+                return 'IAM Identity Center Connected'
+            } else if (AuthUtil.instance.isBuilderIdInUse()) {
+                return 'AWS Builder ID Connected'
+            } else {
+                return 'IAM Connected'
+            }
         } else if (AuthUtil.instance.isConnectionExpired()) {
             return 'Expired Connection'
         }
@@ -76,13 +80,13 @@ export class CodeWhispererNode implements RootNode {
             return [createSsoSignIn(), createLearnMore()]
         }
         if (this._showFreeTierLimitReachedNode) {
-            if (isCloud9()) {
+            if (hasVendedIamCredentials()) {
                 return [createFreeTierLimitMetNode(), createOpenReferenceLogNode()]
             } else {
                 return [createFreeTierLimitMetNode(), createSecurityScanNode(), createOpenReferenceLogNode()]
             }
         } else {
-            if (isCloud9()) {
+            if (hasVendedIamCredentials()) {
                 return [createAutoSuggestionsNode(autoTriggerEnabled), createOpenReferenceLogNode()]
             } else {
                 if (AuthUtil.instance.isValidEnterpriseSsoInUse() && AuthUtil.instance.isCustomizationFeatureEnabled) {

src/codewhisperer/service/recommendationHandler.ts

@@ -14,7 +14,8 @@ import { AWSError } from 'aws-sdk'
 import { isAwsError } from '../../shared/errors'
 import { TelemetryHelper } from '../util/telemetryHelper'
 import { getLogger } from '../../shared/logger'
-import { isCloud9 } from '../../shared/extensionUtilities'
+import { isCloud9, isSageMaker } from '../../shared/extensionUtilities'
+import { hasVendedIamCredentials } from '../../auth/auth'
 import {
     asyncCallWithTimeout,
     isInlineCompletionEnabled,
@@ -40,6 +41,7 @@ import { AuthUtil } from '../util/authUtil'
 import { CodeWhispererUserGroupSettings } from '../util/userGroupUtil'
 import { CWInlineCompletionItemProvider } from './inlineCompletionItemProvider'
 import { application } from '../util/codeWhispererApplication'
+import { openUrl } from '../../shared/utilities/vsCodeUtils'
 import { indent } from '../../shared/utilities/textUtilities'
 
 /**
@@ -110,8 +112,10 @@ export class RecommendationHandler {
         isFirstPaginationCall: boolean,
         promise: Promise<any>
     ): Promise<any> {
-        const timeoutMessage = isCloud9() ? `Generate recommendation timeout.` : `List recommendation timeout`
-        if (isManualTriggerOn && triggerType === 'OnDemand' && (isCloud9() || isFirstPaginationCall)) {
+        const timeoutMessage = hasVendedIamCredentials()
+            ? 'Generate recommendation timeout.'
+            : 'List recommendation timeout'
+        if (isManualTriggerOn && triggerType === 'OnDemand' && (hasVendedIamCredentials() || isFirstPaginationCall)) {
             return vscode.window.withProgress(
                 {
                     location: vscode.ProgressLocation.Notification,
@@ -154,6 +158,7 @@ export class RecommendationHandler {
         autoTriggerType?: CodewhispererAutomatedTriggerType,
         pagination: boolean = true,
         page: number = 0,
+        isSM: boolean = isSageMaker(),
         retry: boolean = false
     ): Promise<GetRecommendationsResponse> {
         let invocationResult: 'Succeeded' | 'Failed' = 'Failed'
@@ -227,9 +232,8 @@ export class RecommendationHandler {
             startTime = performance.now()
             this.lastInvocationTime = startTime
             const mappedReq = runtimeLanguageContext.mapToRuntimeLanguage(request)
-            const codewhispererPromise = pagination
-                ? client.listRecommendations(mappedReq)
-                : client.generateRecommendations(mappedReq)
+            const codewhispererPromise =
+                pagination && !isSM ? client.listRecommendations(mappedReq) : client.generateRecommendations(mappedReq)
             const resp = await this.getServerResponse(
                 triggerType,
                 config.isManualTriggerEnabled,
@@ -270,6 +274,18 @@ export class RecommendationHandler {
                 errorCode = error.code
                 reason = `CodeWhisperer Invocation Exception: ${error?.code ?? error?.name ?? 'unknown'}`
                 await this.onThrottlingException(error, triggerType)
+
+                if (error?.code === 'AccessDeniedException' && errorMessage?.includes('no identity-based policy')) {
+                    getLogger().error('CodeWhisperer AccessDeniedException : %s', (error as Error).message)
+                    vscode.window
+                        .showErrorMessage(`CodeWhisperer: ${error?.message}`, CodeWhispererConstants.settingsLearnMore)
+                        .then(async resp => {
+                            if (resp === CodeWhispererConstants.settingsLearnMore) {
+                                openUrl(vscode.Uri.parse(CodeWhispererConstants.learnMoreUri))
+                            }
+                        })
+                    await vscode.commands.executeCommand('aws.codeWhisperer.enableCodeSuggestions', false)
+                }
             } else {
                 errorMessage = error as string
                 reason = error ? String(error) : 'unknown'

src/codewhisperer/util/authUtil.ts

@@ -8,7 +8,7 @@ import * as CodeWhispererConstants from '../models/constants'
 import { Auth } from '../../auth/auth'
 import { ToolkitError } from '../../shared/errors'
 import { getSecondaryAuth } from '../../auth/secondaryAuth'
-import { isCloud9 } from '../../shared/extensionUtilities'
+import { isCloud9, isSageMaker } from '../../shared/extensionUtilities'
 import { PromptSettings } from '../../shared/settings'
 import {
     ssoAccountAccessScopes,
@@ -32,6 +32,10 @@ export const isValidCodeWhispererConnection = (conn?: Connection): conn is Conne
         return isIamConnection(conn)
     }
 
+    if (isSageMaker()) {
+        return isIamConnection(conn)
+    }
+
     return (
         (isCloud9('codecatalyst') && isIamConnection(conn)) ||
         (isSsoConnection(conn) && hasScopes(conn, codewhispererScopes))

src/extension.ts

@@ -27,6 +27,7 @@ import {
     getToolkitEnvironmentDetails,
     initializeComputeRegion,
     isCloud9,
+    isSageMaker,
     showQuickStartWebview,
     showWelcomeMessage,
 } from './shared/extensionUtilities'
@@ -209,7 +210,10 @@ export async function activate(context: vscode.ExtensionContext) {
             })
         )
 
-        await codecatalyst.activate(extContext)
+        // do not enable codecatalyst for sagemaker
+        if (!isSageMaker()) {
+            await codecatalyst.activate(extContext)
+        }
 
         await activateCloudFormationTemplateRegistry(context)
 

src/shared/extensionUtilities.ts

@@ -22,13 +22,15 @@ const localize = nls.loadMessageBundle()
 const vscodeAppname = 'Visual Studio Code'
 const cloud9Appname = 'AWS Cloud9'
 const cloud9CnAppname = 'Amazon Cloud9'
+const sageMakerAppname = 'SageMaker Code Editor'
 const notInitialized = 'notInitialized'
 
 export const mostRecentVersionKey: string = 'globalsMostRecentVersion'
 
 export enum IDE {
     vscode,
     cloud9,
+    sagemaker,
     unknown,
 }
 
@@ -44,7 +46,11 @@ export function getIdeType(): IDE {
         return IDE.cloud9
     }
 
-    // Theia doesn't necessarily have all env propertie
+    if (vscode.env.appName === sageMakerAppname) {
+        return IDE.sagemaker
+    }
+
+    // Theia doesn't necessarily have all env properties
     // so we should be defensive and assume appName is nullable.
     if (vscode.env.appName?.startsWith(vscodeAppname)) {
         return IDE.vscode
@@ -80,12 +86,29 @@ export function getIdeProperties(): IdeProperties {
                 return createCloud9Properties(localize('AWS.title.cn', 'Amazon'))
             }
             return createCloud9Properties(company)
+        case IDE.sagemaker:
+            if (isCn()) {
+                // check for cn region
+                return createSageMakerProperties(localize('AWS.title.cn', 'Amazon'))
+            }
+            return createSageMakerProperties(company)
         // default is IDE.vscode
         default:
             return vscodeVals
     }
 }
 
+function createSageMakerProperties(company: string): IdeProperties {
+    return {
+        shortName: localize('AWS.vscode.shortName', '{0} Code Editor', company),
+        longName: localize('AWS.vscode.longName', '{0} SageMaker Code Editor', company),
+        commandPalette: localize('AWS.vscode.commandPalette', 'Command Palette'),
+        codelens: localize('AWS.vscode.codelens', 'CodeLens'),
+        codelenses: localize('AWS.vscode.codelenses', 'CodeLenses'),
+        company,
+    }
+}
+
 function createCloud9Properties(company: string): IdeProperties {
     return {
         shortName: localize('AWS.cloud9.shortName', 'Cloud9'),
@@ -109,6 +132,10 @@ export function isCloud9(flavor: 'classic' | 'codecatalyst' | 'any' = 'any'): bo
     return (flavor === 'classic' && !codecat) || (flavor === 'codecatalyst' && codecat)
 }
 
+export function isSageMaker(): boolean {
+    return vscode.env.appName === sageMakerAppname
+}
+
 export function isCn(): boolean {
     return getComputeRegion()?.startsWith('cn') ?? false
 }
@@ -338,10 +365,17 @@ export function getToolkitEnvironmentDetails(): string {
 }
 
 /**
- * Returns the Cloud9 compute region or 'unknown' if we can't pull a region, or `undefined` if this is not Cloud9.
+ * Returns the Cloud9/SageMaker compute region or 'unknown' if we can't pull a region, or `undefined` if this is not Cloud9 or SageMaker.
  */
-export async function initializeComputeRegion(metadata?: Ec2MetadataClient, isC9: boolean = isCloud9()): Promise<void> {
-    if (isC9) {
+
+export async function initializeComputeRegion(
+    metadata?: Ec2MetadataClient,
+    isC9?: boolean,
+    isSM?: boolean
+): Promise<void> {
+    isC9 ??= isCloud9()
+    isSM ??= isSageMaker()
+    if (isC9 || isSM) {
         metadata ??= new DefaultEc2MetadataClient()
         try {
             const identity = await metadata.getInstanceIdentity()

src/test/codewhisperer/explorer/codewhispererNode.test.ts

@@ -39,6 +39,7 @@ describe('codewhispererNode', function () {
 
         it('should create a node showing AWS Builder ID connection', function () {
             sinon.stub(AuthUtil.instance, 'isUsingSavedConnection').get(() => true)
+            sinon.stub(AuthUtil.instance, 'isBuilderIdInUse').resolves(true)
             isConnectionValid.returns(true)
 
             const node = codewhispererNode
@@ -50,6 +51,20 @@ describe('codewhispererNode', function () {
             assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.Collapsed)
         })
 
+        it('should create a node showing IAM connection', function () {
+            sinon.stub(AuthUtil.instance, 'isUsingSavedConnection').get(() => true)
+            //sinon.stub(AuthUtil.instance, 'isBuilderIdInUse').resolves(false)
+            isConnectionValid.returns(true)
+
+            const node = codewhispererNode
+            const treeItem = node.getTreeItem()
+
+            assert.strictEqual(treeItem.label, 'CodeWhisperer')
+            assert.strictEqual(treeItem.contextValue, 'awsCodeWhispererNodeSaved')
+            assert.strictEqual(treeItem.description, 'IAM Connected')
+            assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.Collapsed)
+        })
+
         it('should create a node showing enterprise SSO connection', function () {
             sinon.stub(AuthUtil.instance, 'isUsingSavedConnection').get(() => true)
             sinon.stub(AuthUtil.instance, 'isEnterpriseSsoInUse').resolves(true)

src/test/shared/extensionUtilities.test.ts

@@ -166,24 +166,45 @@ describe('initializeComputeRegion, getComputeRegion', async function () {
         assert.strictEqual(getComputeRegion(), 'us-weast-1')
     })
 
+    it('returns a compute region for sagemaker', async function () {
+        sandbox.stub(metadataService, 'getInstanceIdentity').resolves({ region: 'us-weast-1' })
+
+        await initializeComputeRegion(metadataService, false, true)
+        assert.strictEqual(getComputeRegion(), 'us-weast-1')
+    })
+
     it('returns "unknown" if cloud9 and the MetadataService request fails', async function () {
         sandbox.stub(metadataService, 'getInstanceIdentity').rejects({} as AWSError)
 
         await initializeComputeRegion(metadataService, true)
         assert.strictEqual(getComputeRegion(), 'unknown')
     })
 
+    it('returns "unknown" if sagemaker and the MetadataService request fails', async function () {
+        sandbox.stub(metadataService, 'getInstanceIdentity').rejects({} as AWSError)
+
+        await initializeComputeRegion(metadataService, false, true)
+        assert.strictEqual(getComputeRegion(), 'unknown')
+    })
+
     it('returns "unknown" if cloud9 and can not find a region', async function () {
         sandbox.stub(metadataService, 'getInstanceIdentity').resolves({} as InstanceIdentity)
 
         await initializeComputeRegion(metadataService, true)
         assert.strictEqual(getComputeRegion(), 'unknown')
     })
 
-    it('returns undefined if not cloud9', async function () {
+    it('returns "unknown" if sagemaker and can not find a region', async function () {
+        sandbox.stub(metadataService, 'getInstanceIdentity').resolves({} as InstanceIdentity)
+
+        await initializeComputeRegion(metadataService, false, true)
+        assert.strictEqual(getComputeRegion(), 'unknown')
+    })
+
+    it('returns undefined if not cloud9 or sagemaker', async function () {
         sandbox.stub(metadataService, 'getInstanceIdentity').callsArgWith(1, undefined, 'lol')
 
-        await initializeComputeRegion(metadataService, false)
+        await initializeComputeRegion(metadataService, false, false)
         assert.strictEqual(getComputeRegion(), undefined)
     })