Merge master into feature/emr

Author: aws-toolkit-automation

.gitignore

@@ -34,6 +34,7 @@ src.gen/*
 **/src/auth/sso/oidcclientpkce.d.ts
 **/src/sagemakerunifiedstudio/shared/client/gluecatalogapi.d.ts
 **/src/sagemakerunifiedstudio/shared/client/sqlworkbench.d.ts
+**/src/sagemakerunifiedstudio/shared/client/datazonecustomclient.d.ts
 
 # Generated by tests
 **/src/testFixtures/**/bin

packages/core/package.nls.json

@@ -231,6 +231,7 @@
     "AWS.command.s3.uploadFileToParent": "Upload to Parent...",
     "AWS.command.smus.switchProject": "Switch Project",
     "AWS.command.smus.refreshProject": "Refresh Project",
+    "AWS.command.smus.refresh": "Refresh",
     "AWS.command.smus.signOut": "Sign Out",
     "AWS.command.sagemaker.filterSpaces": "Filter Sagemaker Spaces",
     "AWS.command.stepFunctions.createStateMachineFromTemplate": "Create a new Step Functions state machine",

packages/core/scripts/build/generateServiceClient.ts

@@ -249,6 +249,10 @@ void (async () => {
             serviceJsonPath: 'src/sagemakerunifiedstudio/shared/client/sqlworkbench.json',
             serviceName: 'SQLWorkbench',
         },
+        {
+            serviceJsonPath: 'src/sagemakerunifiedstudio/shared/client/datazonecustomclient.json',
+            serviceName: 'DataZoneCustomClient',
+        },
     ]
     await generateServiceClients(serviceClientDefinitions)
 })()

packages/core/src/auth/auth.ts

@@ -598,7 +598,7 @@ export class Auth implements AuthService, ConnectionManager {
     }
 
     @withTelemetryContext({ name: 'updateConnectionState', class: authClassName })
-    private async updateConnectionState(id: Connection['id'], connectionState: ProfileMetadata['connectionState']) {
+    public async updateConnectionState(id: Connection['id'], connectionState: ProfileMetadata['connectionState']) {
         getLogger().info(`auth: Updating connection state of ${id} to ${connectionState}`)
 
         if (connectionState === 'authenticating') {

packages/core/src/awsService/sagemaker/credentialMapping.ts

@@ -15,6 +15,7 @@ import { getLogger } from '../../shared/logger/logger'
 import { parseArn } from './detached-server/utils'
 import { SagemakerUnifiedStudioSpaceNode } from '../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode'
 import { SageMakerUnifiedStudioSpacesParentNode } from '../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpacesParentNode'
+import { isSmusSsoConnection } from '../../sagemakerunifiedstudio/auth/model'
 
 const mappingFileName = '.sagemaker-space-profiles'
 const mappingFilePath = path.join(os.homedir(), '.aws', mappingFileName)
@@ -74,10 +75,11 @@ export async function persistLocalCredentials(spaceArn: string): Promise<void> {
 export async function persistSmusProjectCreds(spaceArn: string, node: SagemakerUnifiedStudioSpaceNode): Promise<void> {
     const nodeParent = node.getParent() as SageMakerUnifiedStudioSpacesParentNode
     const authProvider = nodeParent.getAuthProvider()
+    const activeConnection = authProvider.activeConnection
     const projectId = nodeParent.getProjectId()
     const projectAuthProvider = await authProvider.getProjectCredentialProvider(projectId)
     await projectAuthProvider.getCredentials()
-    await setSmusSpaceSsoProfile(spaceArn, projectId)
+    await setSmusSpaceProfile(spaceArn, projectId, isSmusSsoConnection(activeConnection) ? 'sso' : 'iam')
     // Trigger SSH credential refresh for the project
     projectAuthProvider.startProactiveCredentialRefresh()
 }
@@ -177,11 +179,16 @@ export async function setSpaceSsoProfile(
  * Sets the SM Space to map to SageMaker Unified Studio Project.
  * @param spaceArn - The arn of the SageMaker Unified Studio space.
  * @param projectId - The project ID associated with the SageMaker Unified Studio space.
+ * @param credentialType - The type of credential ('sso' or 'iam').
  */
-export async function setSmusSpaceSsoProfile(spaceArn: string, projectId: string): Promise<void> {
+export async function setSmusSpaceProfile(
+    spaceArn: string,
+    projectId: string,
+    credentialType: 'iam' | 'sso'
+): Promise<void> {
     const data = await loadMappings()
     data.localCredential ??= {}
-    data.localCredential[spaceArn] = { type: 'sso', smusProjectId: projectId }
+    data.localCredential[spaceArn] = { type: credentialType, smusProjectId: projectId }
     await saveMappings(data)
 }
 

packages/core/src/awsService/sagemaker/detached-server/credentials.ts

@@ -29,11 +29,25 @@ export async function resolveCredentialsFor(connectionIdentifier: string): Promi
 
     switch (profile.type) {
         case 'iam': {
-            const name = profile.profileName?.split(':')[1]
-            if (!name) {
-                throw new Error(`Invalid IAM profile name for "${connectionIdentifier}"`)
+            if ('profileName' in profile) {
+                const name = profile.profileName?.split(':')[1]
+                if (!name) {
+                    throw new Error(`Invalid IAM profile name for "${connectionIdentifier}"`)
+                }
+                return fromIni({ profile: name })
+            } else if ('smusProjectId' in profile) {
+                const { accessKey, secret, token } = mapping.smusProjects?.[profile.smusProjectId] || {}
+                if (!accessKey || !secret || !token) {
+                    throw new Error(`Missing ProjectRole credentials for SMUS Space "${connectionIdentifier}"`)
+                }
+                return {
+                    accessKeyId: accessKey,
+                    secretAccessKey: secret,
+                    sessionToken: token,
+                }
+            } else {
+                throw new Error(`Missing IAM credentials for "${connectionIdentifier}"`)
             }
-            return fromIni({ profile: name })
         }
         case 'sso': {
             if ('accessKey' in profile && 'secret' in profile && 'token' in profile) {

packages/core/src/awsService/sagemaker/detached-server/errorPage.ts

@@ -32,7 +32,11 @@ export const getVSCodeErrorTitle = (error: SageMakerServiceException): string =>
     return ErrorText.StartSession[ExceptionType.DEFAULT].Title
 }
 
-export const getVSCodeErrorText = (error: SageMakerServiceException, isSmus?: boolean): string => {
+export const getVSCodeErrorText = (
+    error: SageMakerServiceException,
+    isSmus?: boolean,
+    isSmusIamConn?: boolean
+): string => {
     const exceptionType = error.name as ExceptionType
 
     switch (exceptionType) {
@@ -41,9 +45,12 @@ export const getVSCodeErrorText = (error: SageMakerServiceException, isSmus?: bo
             return ErrorText.StartSession[exceptionType].Text.replace('{message}', error.message)
         case ExceptionType.EXPIRED_TOKEN:
             // Use SMUS-specific message if in SMUS context
-            return isSmus
-                ? ErrorText.StartSession[ExceptionType.EXPIRED_TOKEN].SmusText
-                : ErrorText.StartSession[exceptionType].Text
+            if (isSmus) {
+                return isSmusIamConn
+                    ? ErrorText.StartSession[ExceptionType.EXPIRED_TOKEN].SmusIamText
+                    : ErrorText.StartSession[ExceptionType.EXPIRED_TOKEN].SmusSsoText
+            }
+            return ErrorText.StartSession[exceptionType].Text
         case ExceptionType.INTERNAL_FAILURE:
         case ExceptionType.RESOURCE_LIMIT_EXCEEDED:
         case ExceptionType.THROTTLING:
@@ -66,8 +73,10 @@ export const ErrorText = {
         [ExceptionType.EXPIRED_TOKEN]: {
             Title: 'Authentication expired',
             Text: 'Your session has expired. Please refresh your credentials and try again.',
-            SmusText:
-                'Your session has expired. This is likely due to network connectivity issues after machine sleep/resume. Please wait 10-30 seconds for automatic credential refresh, then try again. If the issue persists, try reconnecting through AWS Toolkit.',
+            SmusSsoText:
+                'Your session has expired. This is likely due to network connectivity issues after machine sleep/resume. Wait 10-30 seconds for automatic credential refresh, then try again. If the issue persists, try reconnecting through AWS Toolkit.',
+            SmusIamText:
+                'Your session has expired. Update the credentials associated with the IAM profile or use a valid IAM profile, then try again.',
         },
         [ExceptionType.INTERNAL_FAILURE]: {
             Title: 'Failed to connect remotely to VSCode',

packages/core/src/awsService/sagemaker/detached-server/routes/getSession.ts

@@ -6,7 +6,7 @@
 // Disabled: detached server files cannot import vscode.
 /* eslint-disable aws-toolkits/no-console-log */
 import { IncomingMessage, ServerResponse } from 'http'
-import { startSagemakerSession, parseArn, isSmusConnection } from '../utils'
+import { startSagemakerSession, parseArn, isSmusConnection, isSmusIamConnection } from '../utils'
 import { resolveCredentialsFor } from '../credentials'
 import url from 'url'
 import { SageMakerServiceException } from '@amzn/sagemaker-client'
@@ -35,6 +35,7 @@ export async function handleGetSession(req: IncomingMessage, res: ServerResponse
     const { region } = parseArn(connectionIdentifier)
     // Detect if this is a SMUS connection for specialized error handling
     const isSmus = await isSmusConnection(connectionIdentifier)
+    const isSmusIamConn = await isSmusIamConnection(connectionIdentifier)
 
     try {
         const session = await startSagemakerSession({ region, connectionIdentifier, credentials })
@@ -50,7 +51,7 @@ export async function handleGetSession(req: IncomingMessage, res: ServerResponse
         const error = err as SageMakerServiceException
         console.error(`Failed to start SageMaker session for ${connectionIdentifier}:`, err)
         const errorTitle = getVSCodeErrorTitle(error)
-        const errorText = getVSCodeErrorText(error, isSmus)
+        const errorText = getVSCodeErrorText(error, isSmus, isSmusIamConn)
         await openErrorPage(errorTitle, errorText)
         res.writeHead(500, { 'Content-Type': 'text/plain' })
         res.end('Failed to start SageMaker session')

packages/core/src/awsService/sagemaker/detached-server/utils.ts

@@ -147,6 +147,24 @@ export async function isSmusConnection(connectionIdentifier: string): Promise<bo
     }
 }
 
+/**
+ * Detects if the connection identifier is using SMUS IAM credentials
+ * @param connectionIdentifier - The connection identifier to check
+ * @returns Promise<boolean> - true if SMUS IAM connection, false otherwise
+ */
+export async function isSmusIamConnection(connectionIdentifier: string): Promise<boolean> {
+    try {
+        const mapping = await readMapping()
+        const profile = mapping.localCredential?.[connectionIdentifier]
+
+        // Check if profile exists, has smusProjectId, and type is 'iam'
+        return profile && 'smusProjectId' in profile && profile.type === 'iam'
+    } catch (err) {
+        // If we can't detect it is iam connection, assume not SMUS IAM to avoid breaking existing functionality
+        return false
+    }
+}
+
 /**
  * Writes the mapping to a temp file and atomically renames it to the target path.
  * Uses a queue to prevent race conditions when multiple requests try to write simultaneously.

packages/core/src/awsService/sagemaker/types.ts

@@ -12,7 +12,7 @@ export interface SpaceMappings {
 export type LocalCredentialProfile =
     | { type: 'iam'; profileName: string }
     | { type: 'sso'; accessKey: string; secret: string; token: string }
-    | { type: 'sso'; smusProjectId: string }
+    | { type: 'sso' | 'iam'; smusProjectId: string }
 
 export interface DeeplinkSession {
     requests: Record<string, SsmConnectionInfo>