feat(smus): Add Project Space management functionality (aws#2190)

## Problem
SMUS users need Project Space management functionality

## Solution

- Add SageMakerUnifiedStudioSpaceNode for individual space
representation
- Add SageMakerUnifiedStudioSpacesParentNode for space container
management
- Add SagemakerSpace class for unified space operations and status
management
- Enhance SageMaker credential mapping with SMUS project support
- Add space-specific icons and detached server credential resolution
- Update SageMaker commands and model types for space functionality
- Added test cases.
---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.


###Note: Due to lot of constriants - I apologies that I had to raise
such a big PR.

Co-authored-by: guntamb <[email protected]>

Author: bharathGuntamadugu

packages/core/package.json

@@ -421,26 +421,40 @@
                     "fontCharacter": "\\f1e0"
                 }
             },
-            "aws-schemas-registry": {
+            "aws-sagemakerunifiedstudio-spaces": {
                 "description": "AWS Contributed Icon",
                 "default": {
                     "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
                     "fontCharacter": "\\f1e1"
                 }
             },
-            "aws-schemas-schema": {
+            "aws-sagemakerunifiedstudio-spaces-dark": {
                 "description": "AWS Contributed Icon",
                 "default": {
                     "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
                     "fontCharacter": "\\f1e2"
                 }
             },
-            "aws-stepfunctions-preview": {
+            "aws-schemas-registry": {
                 "description": "AWS Contributed Icon",
                 "default": {
                     "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
                     "fontCharacter": "\\f1e3"
                 }
+            },
+            "aws-schemas-schema": {
+                "description": "AWS Contributed Icon",
+                "default": {
+                    "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+                    "fontCharacter": "\\f1e4"
+                }
+            },
+            "aws-stepfunctions-preview": {
+                "description": "AWS Contributed Icon",
+                "default": {
+                    "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+                    "fontCharacter": "\\f1e5"
+                }
             }
         }
     },

packages/core/package.nls.json

@@ -231,6 +231,7 @@
     "AWS.command.s3.uploadFile": "Upload Files...",
     "AWS.command.s3.uploadFileToParent": "Upload to Parent...",
     "AWS.command.smus.switchProject": "Switch Project",
+    "AWS.command.smus.refreshProject": "Refresh Project",
     "AWS.command.smus.signOut": "Sign Out",
     "AWS.command.sagemaker.filterSpaces": "Filter Sagemaker Spaces",
     "AWS.command.stepFunctions.createStateMachineFromTemplate": "Create a new Step Functions state machine",
@@ -483,6 +484,5 @@
     "AWS.toolkit.lambda.walkthrough.step1.description": "Locally test and debug your code.",
     "AWS.toolkit.lambda.walkthrough.step2.title": "Deploy to the cloud",
     "AWS.toolkit.lambda.walkthrough.step2.description": "Test your application in the cloud from within VS Code. \n\nNote: The AWS CLI and the SAM CLI require AWS Credentials to interact with the cloud. For information on setting up your credentials, see [Authentication and access credentials](https://docs.aws.amazon.com/cli/latest/userguide/cli-chap-configure.html). \n\n[Configure credentials](command:aws.toolkit.lambda.walkthrough.credential)",
-    "AWS.toolkit.lambda.serverlessLand.quickpickTitle": "Create application with Serverless template",
-    "AWS.command.smus.signout": "Sign Out"
+    "AWS.toolkit.lambda.serverlessLand.quickpickTitle": "Create application with Serverless template"
 }

packages/core/resources/icons/aws/sagemakerunifiedstudio/spaces-dark.svg

@@ -20,6 +20,7 @@ import { ToolkitError } from '../../shared/errors'
 import { showConfirmationMessage } from '../../shared/utilities/messages'
 import { RemoteSessionError } from '../../shared/remoteSession'
 import { ConnectFromRemoteWorkspaceMessage, InstanceTypeError } from './constants'
+import { SagemakerUnifiedStudioSpaceNode } from '../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode'
 
 const localize = nls.loadMessageBundle()
 
@@ -101,6 +102,8 @@ export async function deeplinkConnect(
             connectionIdentifier,
             ctx.extensionContext,
             'sm_dl',
+            false /* isSMUS */,
+            undefined /* node */,
             session,
             wsUrl,
             token,
@@ -125,7 +128,11 @@ export async function deeplinkConnect(
     }
 }
 
-export async function stopSpace(node: SagemakerSpaceNode, ctx: vscode.ExtensionContext) {
+export async function stopSpace(
+    node: SagemakerSpaceNode | SagemakerUnifiedStudioSpaceNode,
+    ctx: vscode.ExtensionContext,
+    sageMakerClient?: SagemakerClient
+) {
     const spaceName = node.spaceApp.SpaceName!
     const confirmed = await showConfirmationMessage({
         prompt: `You are about to stop this space. Any active resource will also be stopped. Are you sure you want to stop the space?`,
@@ -137,8 +144,8 @@ export async function stopSpace(node: SagemakerSpaceNode, ctx: vscode.ExtensionC
     if (!confirmed) {
         return
     }
-
-    const client = new SagemakerClient(node.regionCode)
+    //  In case of SMUS, we pass in a SM Client and for SM AI, it creates a new SM Client.
+    const client = sageMakerClient ? sageMakerClient : new SagemakerClient(node.regionCode)
     try {
         await client.deleteApp({
             DomainId: node.spaceApp.DomainId!,
@@ -159,14 +166,19 @@ export async function stopSpace(node: SagemakerSpaceNode, ctx: vscode.ExtensionC
     await tryRefreshNode(node)
 }
 
-export async function openRemoteConnect(node: SagemakerSpaceNode, ctx: vscode.ExtensionContext) {
+export async function openRemoteConnect(
+    node: SagemakerSpaceNode | SagemakerUnifiedStudioSpaceNode,
+    ctx: vscode.ExtensionContext,
+    sageMakerClient?: SagemakerClient
+) {
     if (isRemoteWorkspace()) {
         void vscode.window.showErrorMessage(ConnectFromRemoteWorkspaceMessage)
         return
     }
 
     if (node.getStatus() === 'Stopped') {
-        const client = new SagemakerClient(node.regionCode)
+        //  In case of SMUS, we pass in a SM Client and for SM AI, it creates a new SM Client.
+        const client = sageMakerClient ? sageMakerClient : new SagemakerClient(node.regionCode)
 
         try {
             await client.startSpace(node.spaceApp.SpaceName!, node.spaceApp.DomainId!)

packages/core/resources/icons/aws/sagemakerunifiedstudio/spaces.svg

@@ -13,6 +13,8 @@ import { Auth } from '../../auth/auth'
 import { SpaceMappings, SsmConnectionInfo } from './types'
 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'
 
 const mappingFileName = '.sagemaker-space-profiles'
 const mappingFilePath = path.join(os.homedir(), '.aws', mappingFileName)
@@ -44,9 +46,9 @@ export async function saveMappings(data: SpaceMappings): Promise<void> {
 
 /**
  * Persists the current profile to the appropriate space mapping based on connection type and profile format.
- * @param appArn - The identifier for the SageMaker space.
+ * @param spaceArn - The arn for the SageMaker space.
  */
-export async function persistLocalCredentials(appArn: string): Promise<void> {
+export async function persistLocalCredentials(spaceArn: string): Promise<void> {
     const currentProfileId = Auth.instance.getCurrentProfileId()
     if (!currentProfileId) {
         throw new ToolkitError('No current profile ID available for saving space credentials.')
@@ -55,33 +57,45 @@ export async function persistLocalCredentials(appArn: string): Promise<void> {
     if (currentProfileId.startsWith('sso:')) {
         const credentials = globals.loginManager.store.credentialsCache[currentProfileId]
         await setSpaceSsoProfile(
-            appArn,
+            spaceArn,
             credentials.credentials.accessKeyId,
             credentials.credentials.secretAccessKey,
             credentials.credentials.sessionToken ?? ''
         )
     } else {
-        await setSpaceIamProfile(appArn, currentProfileId)
+        await setSpaceIamProfile(spaceArn, currentProfileId)
     }
 }
 
+/**
+ * Persists the current selected SMUS Project Role creds to the appropriate space mapping.
+ * @param spaceArn - The identifier for the SageMaker Space.
+ */
+export async function persistSmusProjectCreds(spaceArn: string, node: SagemakerUnifiedStudioSpaceNode): Promise<void> {
+    const nodeParent = node.getParent() as SageMakerUnifiedStudioSpacesParentNode
+    const authProvider = nodeParent.getAuthProvider()
+    const projectAuthProvider = await authProvider.getProjectCredentialProvider(nodeParent.getProjectId())
+    await projectAuthProvider.getCredentials()
+    await setSmusSpaceSsoProfile(spaceArn, nodeParent.getProjectId())
+}
+
 /**
  * Persists deep link credentials for a SageMaker space using a derived refresh URL based on environment.
  *
- * @param appArn - ARN of the SageMaker space.
+ * @param spaceArn - ARN of the SageMaker space.
  * @param domain - The domain ID associated with the space.
  * @param session - SSM session ID.
  * @param wsUrl - SSM WebSocket URL.
  * @param token - Bearer token for the session.
  */
 export async function persistSSMConnection(
-    appArn: string,
+    spaceArn: string,
     domain: string,
     session?: string,
     wsUrl?: string,
     token?: string
 ): Promise<void> {
-    const { region } = parseArn(appArn)
+    const { region } = parseArn(spaceArn)
     const endpoint = DevSettings.instance.get('endpoints', {})['sagemaker'] ?? ''
 
     // TODO: Hardcoded to 'jupyterlab' due to a bug in Studio that only supports refreshing
@@ -107,7 +121,7 @@ export async function persistSSMConnection(
             : `${envSubdomain}.studio.${region}.asfiovnxocqpcry.com`
 
     const refreshUrl = `https://studio-${domain}.${baseDomain}/${appSubDomain}`
-    await setSpaceCredentials(appArn, refreshUrl, {
+    await setSpaceCredentials(spaceArn, refreshUrl, {
         sessionId: session ?? '-',
         url: wsUrl ?? '-',
         token: token ?? '-',
@@ -116,51 +130,63 @@ export async function persistSSMConnection(
 
 /**
  * Sets or updates an IAM credential profile for a given space.
- * @param spaceName - The name of the SageMaker space.
+ * @param spaceArn - The name of the SageMaker space.
  * @param profileName - The local AWS profile name to associate.
  */
-export async function setSpaceIamProfile(spaceName: string, profileName: string): Promise<void> {
+export async function setSpaceIamProfile(spaceArn: string, profileName: string): Promise<void> {
     const data = await loadMappings()
     data.localCredential ??= {}
-    data.localCredential[spaceName] = { type: 'iam', profileName }
+    data.localCredential[spaceArn] = { type: 'iam', profileName }
     await saveMappings(data)
 }
 
 /**
  * Sets or updates an SSO credential profile for a given space.
- * @param spaceName - The name of the SageMaker space.
+ * @param spaceArn - The arn of the SageMaker space.
  * @param accessKey - Temporary access key from SSO.
  * @param secret - Temporary secret key from SSO.
  * @param token - Session token from SSO.
  */
 export async function setSpaceSsoProfile(
-    spaceName: string,
+    spaceArn: string,
     accessKey: string,
     secret: string,
     token: string
 ): Promise<void> {
     const data = await loadMappings()
     data.localCredential ??= {}
-    data.localCredential[spaceName] = { type: 'sso', accessKey, secret, token }
+    data.localCredential[spaceArn] = { type: 'sso', accessKey, secret, token }
+    await saveMappings(data)
+}
+
+/**
+ * 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.
+ */
+export async function setSmusSpaceSsoProfile(spaceArn: string, projectId: string): Promise<void> {
+    const data = await loadMappings()
+    data.localCredential ??= {}
+    data.localCredential[spaceArn] = { type: 'sso', smusProjectId: projectId }
     await saveMappings(data)
 }
 
 /**
  * Stores SSM connection information for a given space, typically from a deep link session.
  * This initializes the request as 'fresh' and includes a refresh URL if provided.
- * @param spaceName - The name of the SageMaker space.
+ * @param spaceArn - The arn of the SageMaker space.
  * @param refreshUrl - URL to use for refreshing session tokens.
  * @param credentials - The session information used to initiate the connection.
  */
 export async function setSpaceCredentials(
-    spaceName: string,
+    spaceArn: string,
     refreshUrl: string,
     credentials: SsmConnectionInfo
 ): Promise<void> {
     const data = await loadMappings()
     data.deepLink ??= {}
 
-    data.deepLink[spaceName] = {
+    data.deepLink[spaceArn] = {
         refreshUrl,
         requests: {
             'initial-connection': {

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

@@ -36,15 +36,30 @@ export async function resolveCredentialsFor(connectionIdentifier: string): Promi
             return fromIni({ profile: name })
         }
         case 'sso': {
-            const { accessKey, secret, token } = profile
-            if (!accessKey || !secret || !token) {
+            if ('accessKey' in profile && 'secret' in profile && 'token' in profile) {
+                const { accessKey, secret, token } = profile
+                if (!accessKey || !secret || !token) {
+                    throw new Error(`Missing SSO credentials for "${connectionIdentifier}"`)
+                }
+                return {
+                    accessKeyId: accessKey,
+                    secretAccessKey: secret,
+                    sessionToken: token,
+                }
+            } else if ('smusProjectId' in profile) {
+                // Handle SMUS project ID case
+                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 SSO credentials for "${connectionIdentifier}"`)
             }
-            return {
-                accessKeyId: accessKey,
-                secretAccessKey: secret,
-                sessionToken: token,
-            }
         }
         default:
             throw new Error(`Unsupported profile type "${profile}"`)

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

@@ -96,7 +96,6 @@ export async function readMapping() {
     try {
         const content = await fs.readFile(mappingFilePath, 'utf-8')
         console.log(`Mapping file path: ${mappingFilePath}`)
-        console.log(`Conents: ${content}`)
         return JSON.parse(content)
     } catch (err) {
         throw new Error(`Failed to read mapping file: ${err instanceof Error ? err.message : String(err)}`)