fix(smus): Improve error handling and handle No projects and spaces scenarios (aws#2200)

## Problem
SageMaker Unified Studio explorer failed to handle scenarios when no
projects or spaces are available, causing poor user experience and
potential errors.

## Solution
- Enhanced DataZone client error handling for missing resources
- Added proper handling for empty project and space lists in explorer
nodes
- Improved error messages and fallback behavior in SageMaker space
integration
- Updated utility functions to gracefully handle null/undefined
scenarios
- Should fix the Space tooltip not to include the userId and SM
domainId, and fixes the the scenario of after sign out with project
already selected - If we sign in again with different domain, project
from other domain is still shows up. Thanks to Bhargava for the fix.
- Added/Updated the test cases accordingly

---
- 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.

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

Author: 3 people

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

@@ -131,8 +131,10 @@ export class SagemakerSpace {
         const domainId = this.spaceApp?.DomainId ?? '-'
         const owner = this.spaceApp?.OwnershipSettingsSummary?.OwnerUserProfileName || '-'
         const instanceType = this.spaceApp?.App?.ResourceSpec?.InstanceType ?? '-'
-
-        return `**Space:** ${spaceName} \n\n**Application:** ${appType} \n\n**Domain ID:** ${domainId} \n\n**User Profile:** ${owner} \n\n**Instance:** ${instanceType}`
+        if (this.isSMUSSpace) {
+            return `**Space:** ${spaceName} \n\n**Application:** ${appType} \n\n**Instance Type:** ${instanceType}`
+        }
+        return `**Space:** ${spaceName} \n\n**Application:** ${appType} \n\n**Domain ID:** ${domainId} \n\n**User Profile:** ${owner}`
     }
 
     public getAppIcon() {

packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioConnectionParentNode.ts

@@ -38,8 +38,16 @@ export class SageMakerUnifiedStudioConnectionParentNode implements TreeNode {
             this.connectionType
         )
         const childrenNodes = []
-        if (!this.connections?.items) {
-            return []
+        if (!this.connections?.items || this.connections.items.length === 0) {
+            return [
+                {
+                    id: 'smusNoConnections',
+                    resource: {},
+                    getTreeItem: () =>
+                        new vscode.TreeItem('[No connections found]', vscode.TreeItemCollapsibleState.None),
+                    getParent: () => this,
+                },
+            ]
         }
         for (const connection of this.connections.items) {
             childrenNodes.push(new SageMakerUnifiedStudioConnectionNode(this, connection))

packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioProjectNode.ts

@@ -9,11 +9,13 @@ import { getLogger } from '../../../shared/logger/logger'
 import { telemetry } from '../../../shared/telemetry/telemetry'
 import { AwsCredentialIdentity } from '@aws-sdk/types'
 import { SageMakerUnifiedStudioDataNode } from './sageMakerUnifiedStudioDataNode'
-import { DataZoneProject } from '../../shared/client/datazoneClient'
+import { DataZoneClient, DataZoneProject } from '../../shared/client/datazoneClient'
 import { SageMakerUnifiedStudioRootNode } from './sageMakerUnifiedStudioRootNode'
 import { SagemakerClient } from '../../../shared/clients/sagemaker'
 import { SmusAuthenticationProvider } from '../../auth/providers/smusAuthenticationProvider'
 import { SageMakerUnifiedStudioComputeNode } from './sageMakerUnifiedStudioComputeNode'
+import { getIcon } from '../../../shared/icons'
+import { SmusUtils } from '../../shared/smusUtils'
 
 /**
  * Tree node representing a SageMaker Unified Studio project
@@ -36,18 +38,25 @@ export class SageMakerUnifiedStudioProjectNode implements TreeNode {
 
     public async getTreeItem(): Promise<vscode.TreeItem> {
         if (this.project) {
-            const item = new vscode.TreeItem('Project: ' + this.project.name, vscode.TreeItemCollapsibleState.Collapsed)
+            const item = new vscode.TreeItem('Project: ' + this.project.name, vscode.TreeItemCollapsibleState.Expanded)
             item.contextValue = 'smusSelectedProject'
             item.tooltip = `Project: ${this.project.name}\nID: ${this.project.id}`
+            item.iconPath = getIcon('vscode-folder-opened')
             return item
         }
+
         const item = new vscode.TreeItem('Select a project', vscode.TreeItemCollapsibleState.None)
         item.contextValue = 'smusProjectSelectPicker'
         item.command = {
             command: 'aws.smus.projectView',
             title: 'Select Project',
             arguments: [this],
         }
+        item.iconPath = getIcon('vscode-folder-opened')
+
+        // Auto-invoke project selection after sign-in
+        void vscode.commands.executeCommand('aws.smus.projectView', this)
+
         return item
     }
 
@@ -61,6 +70,24 @@ export class SageMakerUnifiedStudioProjectNode implements TreeNode {
                 result: 'Succeeded',
                 passive: false,
             })
+            const hasAccess = await this.checkProjectAccess(this.project.id)
+            if (!hasAccess) {
+                return [
+                    {
+                        id: 'smusProjectAccessDenied',
+                        resource: {},
+                        getTreeItem: () => {
+                            const item = new vscode.TreeItem(
+                                'You are not a member of this project. Contact any of its owners to add you as a member.',
+                                vscode.TreeItemCollapsibleState.None
+                            )
+                            return item
+                        },
+                        getParent: () => this,
+                    },
+                ]
+            }
+
             const dataNode = new SageMakerUnifiedStudioDataNode(this)
             this.sagemakerClient = await this.initializeSagemakerClient(
                 this.authProvider.activeConnection?.ssoRegion || 'us-east-1'
@@ -87,17 +114,44 @@ export class SageMakerUnifiedStudioProjectNode implements TreeNode {
     }
 
     public async setProject(project: any): Promise<void> {
+        await this.cleanupProjectResources()
+        this.project = project
+    }
+
+    public getProject(): DataZoneProject | undefined {
+        return this.project
+    }
+
+    public async clearProject(): Promise<void> {
+        await this.cleanupProjectResources()
+        this.project = undefined
+        await this.refreshNode()
+    }
+
+    private async cleanupProjectResources(): Promise<void> {
         await this.authProvider.invalidateAllProjectCredentialsInCache()
         if (this.sagemakerClient) {
             this.sagemakerClient.dispose()
             this.sagemakerClient = undefined
         }
-        this.project = project
-        await this.refreshNode()
     }
 
-    public getProject(): DataZoneProject | undefined {
-        return this.project
+    private async checkProjectAccess(projectId: string): Promise<boolean> {
+        try {
+            const dzClient = await DataZoneClient.getInstance(this.authProvider)
+            const userId = await dzClient.getUserId()
+            if (!userId) {
+                return false
+            }
+            const ssoUserProfileId = SmusUtils.extractSSOIdFromUserId(userId)
+            const memberships = await dzClient.fetchAllProjectMemberships(projectId)
+            const hasAccess = memberships.some((member) => member.memberDetails?.user?.userId === ssoUserProfileId)
+            this.logger.debug(`Project access check for user ${ssoUserProfileId}: ${hasAccess}`)
+            return hasAccess
+        } catch (err) {
+            this.logger.error('Failed to check project access: %s', (err as Error).message)
+            return false
+        }
     }
 
     private async initializeSagemakerClient(regionCode: string): Promise<SagemakerClient> {

packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.ts

@@ -19,6 +19,8 @@ import { SmusAuthenticationProvider } from '../../auth/providers/smusAuthenticat
 const contextValueSmusRoot = 'sageMakerUnifiedStudioRoot'
 const contextValueSmusLogin = 'sageMakerUnifiedStudioLogin'
 const contextValueSmusLearnMore = 'sageMakerUnifiedStudioLearnMore'
+const projectPickerTitle = 'Select a SageMaker Unified Studio project you want to open'
+const projectPickerPlaceholder = 'Select project'
 
 export class SageMakerUnifiedStudioRootNode implements TreeNode {
     public readonly id = 'smusRootNode'
@@ -38,7 +40,9 @@ export class SageMakerUnifiedStudioRootNode implements TreeNode {
         this.projectNode = new SageMakerUnifiedStudioProjectNode(this, this.authProvider, this.extensionContext)
 
         // Subscribe to auth provider connection changes to refresh the node
-        this.authProvider.onDidChange(() => {
+        this.authProvider.onDidChange(async () => {
+            // Clear the project when connection changes
+            await this.projectNode.clearProject()
             this.onDidChangeEmitter.fire()
         })
     }
@@ -352,11 +356,7 @@ export async function selectSMUSProject(projectNode?: SageMakerUnifiedStudioProj
 
         const smusProjects = allProjects
 
-        if (smusProjects.length === 0) {
-            void vscode.window.showInformationMessage('No projects found in the domain')
-            return
-        }
-        // Process projects: sort by updatedAt, filter out current project, and map to quick pick items
+        // Process projects: sort by updatedAt, and map to quick pick items
         const items = [...smusProjects]
             .sort(
                 (a, b) =>
@@ -366,28 +366,45 @@ export async function selectSMUSProject(projectNode?: SageMakerUnifiedStudioProj
             .filter(
                 (project) =>
                     // Filter out the Generative AI Model Governance project that is part of QiuckStart resources
-                    project.name !== 'GenerativeAIModelGovernanceProject' &&
-                    (!projectNode?.getProject() || project.id !== projectNode.getProject()?.id)
+                    project.name !== 'GenerativeAIModelGovernanceProject'
             )
             .map((project) => ({
                 label: project.name,
                 detail: 'ID: ' + project.id,
                 description: project.description,
                 data: project,
             }))
-
-        const quickPick = createQuickPick(items, {
-            title: 'Select a SageMaker Unified Studio project you want to open',
-            placeholder: 'Select project',
-        })
-
-        const selectedProject = await quickPick.prompt()
-        if (selectedProject && projectNode) {
-            await projectNode.setProject(selectedProject)
-            // Refresh the entire tree view
-            await vscode.commands.executeCommand('aws.smus.rootView.refresh')
+        if (items.length === 0) {
+            logger.info('No projects found in the domain')
+            void vscode.window.showInformationMessage('No projects found in the domain')
+            // If no projects are found, show "No projects found" in the quick pick
+            const quickPickItem = [
+                {
+                    label: 'No projects found',
+                    detail: '',
+                    description: '',
+                    data: {},
+                },
+            ]
+            const quickPick = createQuickPick(quickPickItem, {
+                title: projectPickerTitle,
+                placeholder: projectPickerPlaceholder,
+            })
+            await quickPick.prompt()
+        } else {
+            const quickPick = createQuickPick(items, {
+                title: projectPickerTitle,
+                placeholder: projectPickerPlaceholder,
+            })
+
+            const selectedProject = await quickPick.prompt()
+            if (selectedProject && !('type' in selectedProject) && projectNode) {
+                await projectNode.setProject(selectedProject)
+                // Refresh the entire tree view
+                await vscode.commands.executeCommand('aws.smus.rootView.refresh')
+            }
+            return selectedProject
         }
-        return selectedProject
     } catch (err) {
         logger.error('Failed to select project: %s', (err as Error).message)
         void vscode.window.showErrorMessage(`Failed to select project: ${(err as Error).message}`)

packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpacesParentNode.ts

@@ -17,6 +17,7 @@ import { UserProfileMetadata } from '../../../awsService/sagemaker/explorer/sage
 import { SagemakerUnifiedStudioSpaceNode } from './sageMakerUnifiedStudioSpaceNode'
 import { PollingSet } from '../../../shared/utilities/pollingSet'
 import { SmusAuthenticationProvider } from '../../auth/providers/smusAuthenticationProvider'
+import { SmusUtils } from '../../shared/smusUtils'
 
 export class SageMakerUnifiedStudioSpacesParentNode implements TreeNode {
     public readonly id = 'smusSpacesParentNode'
@@ -55,11 +56,29 @@ export class SageMakerUnifiedStudioSpacesParentNode implements TreeNode {
     }
 
     public async getChildren(): Promise<TreeNode[]> {
-        await this.updateChildren()
+        try {
+            await this.updateChildren()
+        } catch (err) {
+            return this.getNoSpacesFoundChildren()
+        }
         const nodes = [...this.sagemakerSpaceNodes.values()]
+        if (nodes.length === 0) {
+            return this.getNoSpacesFoundChildren()
+        }
         return nodes
     }
 
+    private getNoSpacesFoundChildren(): TreeNode[] {
+        return [
+            {
+                id: 'smusNoSpaces',
+                resource: {},
+                getTreeItem: () => new vscode.TreeItem('[No Spaces found]', vscode.TreeItemCollapsibleState.None),
+                getParent: () => this,
+            },
+        ]
+    }
+
     public getParent(): TreeNode | undefined {
         return this.parent
     }
@@ -104,7 +123,7 @@ export class SageMakerUnifiedStudioSpacesParentNode implements TreeNode {
         const toolingEnvId = await datazoneClient
             .getToolingEnvironmentId(datazoneClient.getDomainId(), this.projectId)
             .catch((err) => {
-                this.logger.error('Failed to get tooling environment ID: %s', err as Error)
+                this.logger.error('Failed to get tooling environment ID for project %s', this.projectId)
                 throw new Error(`Failed to get tooling environment ID: ${err.message}`)
             })
         if (!toolingEnvId) {
@@ -144,7 +163,7 @@ export class SageMakerUnifiedStudioSpacesParentNode implements TreeNode {
         const datazoneClient = await DataZoneClient.getInstance(this.authProvider)
         // Will be of format: 'ABCA4NU3S7PEOLDQPLXYZ:user-12345678-d061-70a4-0bf2-eeee67a6ab12'
         const userId = await datazoneClient.getUserId()
-        const ssoUserProfileId = this.extractSSOIdFromUserId(userId || '')
+        const ssoUserProfileId = SmusUtils.extractSSOIdFromUserId(userId || '')
         const sagemakerDomainId = await this.getSageMakerDomainId()
         const [spaceApps, domains] = await this.sagemakerClient.fetchSpaceAppsAndDomains(
             sagemakerDomainId,
@@ -188,13 +207,4 @@ export class SageMakerUnifiedStudioSpacesParentNode implements TreeNode {
                 )
         )
     }
-
-    private extractSSOIdFromUserId(userId: string): string {
-        const match = userId.match(/user-(.+)$/)
-        if (!match) {
-            this.logger.error(`Invalid UserId format: ${userId}`)
-            throw new Error(`Invalid UserId format: ${userId}`)
-        }
-        return match[1]
-    }
 }