feat(smus): improve project selection error handling and remove auto-selection (aws#2205)

## Problem:
Project selection had poor error handling and automatically invoked
after sign-in causing some unintended UX issues.

## Solution:
- Add access denied error handling with user-friendly messages
- Remove automatic project selection after sign-in
- Refactor selection logic with helper functions for better error
management
- Updated the 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.

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

Author: bharathGuntamadugu

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

@@ -69,9 +69,6 @@ export class SageMakerUnifiedStudioProjectNode implements TreeNode {
         }
         item.iconPath = getIcon('vscode-folder-opened')
 
-        // Auto-invoke project selection after sign-in
-        void vscode.commands.executeCommand('aws.smus.projectView', this)
-
         return item
     }
 

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

@@ -340,74 +340,84 @@ export const smusSignOutCommand = Commands.declare('aws.smus.signOut', () => asy
     }
 })
 
+function isAccessDenied(error: Error): boolean {
+    return error.name.includes('AccessDenied')
+}
+
+function createProjectQuickPickItems(projects: any[]) {
+    return projects
+        .sort(
+            (a, b) =>
+                (b.updatedAt ? new Date(b.updatedAt).getTime() : 0) -
+                (a.updatedAt ? new Date(a.updatedAt).getTime() : 0)
+        )
+        .filter((project) => project.name !== 'GenerativeAIModelGovernanceProject')
+        .map((project) => ({
+            label: project.name,
+            detail: 'ID: ' + project.id,
+            description: project.description,
+            data: project,
+        }))
+}
+
+async function showQuickPick(items: any[]) {
+    const quickPick = createQuickPick(items, {
+        title: projectPickerTitle,
+        placeholder: projectPickerPlaceholder,
+    })
+    return await quickPick.prompt()
+}
+
 export async function selectSMUSProject(projectNode?: SageMakerUnifiedStudioProjectNode) {
     const logger = getLogger()
+
     try {
         const authProvider = SmusAuthenticationProvider.fromContext()
-        const activeConnection = authProvider.activeConnection
-        if (!activeConnection) {
-            logger.error('There is no active connection to display project view')
+        if (!authProvider.activeConnection) {
+            logger.error('No active connection to display project view')
             return
         }
-        const authenticatedDataZoneClient = await DataZoneClient.getInstance(authProvider)
-        logger.debug('SMUS: DataZone client instance obtained successfully')
 
-        // Fetching all projects in the specified domain using the client's fetchAllProjects method
-        const allProjects = await authenticatedDataZoneClient.fetchAllProjects()
+        const client = await DataZoneClient.getInstance(authProvider)
+        logger.debug('DataZone client instance obtained successfully')
 
-        const smusProjects = allProjects
+        const allProjects = await client.fetchAllProjects()
+        const items = createProjectQuickPickItems(allProjects)
 
-        // Process projects: sort by updatedAt, and map to quick pick items
-        const items = [...smusProjects]
-            .sort(
-                (a, b) =>
-                    (b.updatedAt ? new Date(b.updatedAt).getTime() : 0) -
-                    (a.updatedAt ? new Date(a.updatedAt).getTime() : 0)
-            )
-            .filter(
-                (project) =>
-                    // Filter out the Generative AI Model Governance project that is part of QiuckStart resources
-                    project.name !== 'GenerativeAIModelGovernanceProject'
-            )
-            .map((project) => ({
-                label: project.name,
-                detail: 'ID: ' + project.id,
-                description: project.description,
-                data: project,
-            }))
         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 = [
+            await showQuickPick([{ label: 'No projects found', detail: '', description: '', data: {} }])
+            return
+        }
+
+        const selectedProject = await showQuickPick(items)
+        if (
+            selectedProject &&
+            typeof selectedProject === 'object' &&
+            selectedProject !== null &&
+            !('type' in selectedProject) &&
+            projectNode
+        ) {
+            await projectNode.setProject(selectedProject)
+            await vscode.commands.executeCommand('aws.smus.rootView.refresh')
+        }
+
+        return selectedProject
+    } catch (err) {
+        const error = err as Error
+
+        if (isAccessDenied(error)) {
+            await showQuickPick([
                 {
-                    label: 'No projects found',
-                    detail: '',
-                    description: '',
-                    data: {},
+                    label: '$(error)',
+                    description: "You don't have permissions to view projects. Please contact your administrator",
                 },
-            ]
-            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
         }
-    } 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}`)
+
+        logger.error('Failed to select project: %s', error.message)
+        void vscode.window.showErrorMessage(`Failed to select project: ${error.message}`)
     }
 }

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

@@ -18,6 +18,7 @@ import { SagemakerUnifiedStudioSpaceNode } from './sageMakerUnifiedStudioSpaceNo
 import { PollingSet } from '../../../shared/utilities/pollingSet'
 import { SmusAuthenticationProvider } from '../../auth/providers/smusAuthenticationProvider'
 import { SmusUtils } from '../../shared/smusUtils'
+import { getIcon } from '../../../shared/icons'
 
 export class SageMakerUnifiedStudioSpacesParentNode implements TreeNode {
     public readonly id = 'smusSpacesParentNode'
@@ -59,6 +60,10 @@ export class SageMakerUnifiedStudioSpacesParentNode implements TreeNode {
         try {
             await this.updateChildren()
         } catch (err) {
+            const error = err as Error
+            if (error.name === 'AccessDeniedException') {
+                return this.getAccessDeniedChildren()
+            }
             return this.getNoSpacesFoundChildren()
         }
         const nodes = [...this.sagemakerSpaceNodes.values()]
@@ -79,6 +84,24 @@ export class SageMakerUnifiedStudioSpacesParentNode implements TreeNode {
         ]
     }
 
+    private getAccessDeniedChildren(): TreeNode[] {
+        return [
+            {
+                id: 'smusAccessDenied',
+                resource: {},
+                getTreeItem: () => {
+                    const item = new vscode.TreeItem(
+                        "You don't have permission to view spaces. Please contact your administrator.",
+                        vscode.TreeItemCollapsibleState.None
+                    )
+                    item.iconPath = getIcon('vscode-error')
+                    return item
+                },
+                getParent: () => this,
+            },
+        ]
+    }
+
     public getParent(): TreeNode | undefined {
         return this.parent
     }

packages/core/src/sagemakerunifiedstudio/shared/client/datazoneClient.ts

@@ -342,7 +342,7 @@ export class DataZoneClient {
             )
             return { memberships: response.members, nextToken: response.nextToken }
         } catch (err) {
-            this.logger.error('DataZoneClient: Failed to list project memberships: %s', err as Error)
+            this.logger.error('DataZoneClient: Failed to list project memberships: %s', (err as Error).message)
             throw err
         }
     }
@@ -370,7 +370,7 @@ export class DataZoneClient {
             this.logger.debug(`DataZoneClient: Fetched a total of ${allMemberships.length} project memberships`)
             return allMemberships
         } catch (err) {
-            this.logger.error('DataZoneClient: Failed to fetch all project memberships: %s', err as Error)
+            this.logger.error('DataZoneClient: Failed to fetch all project memberships: %s', (err as Error).message)
             throw err
         }
     }
@@ -420,7 +420,7 @@ export class DataZoneClient {
             this.logger.debug(`DataZoneClient: Found ${projects.length} projects for domain ${this.domainId}`)
             return { projects, nextToken: response.nextToken }
         } catch (err) {
-            this.logger.error('DataZoneClient: Failed to list projects: %s', err as Error)
+            this.logger.error('DataZoneClient: Failed to list projects: %s', (err as Error).message)
             throw err
         }
     }

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

@@ -117,46 +117,26 @@ describe('SageMakerUnifiedStudioProjectNode', function () {
     })
 
     describe('setProject', function () {
-        it('updates the project without firing change event', async function () {
+        it('updates the project and calls cleanupProjectResources', async function () {
+            const cleanupSpy = sinon.spy(projectNode as any, 'cleanupProjectResources')
             await projectNode.setProject(mockProject)
             assert.strictEqual(projectNode['project'], mockProject)
-        })
-
-        it('invalidates credentials and disposes existing sagemaker client', async function () {
-            // Set up existing sagemaker client with mock
-            const mockClient = { dispose: sinon.stub() } as any
-            projectNode['sagemakerClient'] = mockClient
-
-            await projectNode.setProject(mockProject)
-
-            assert((projectNode['authProvider'].invalidateAllProjectCredentialsInCache as sinon.SinonStub).calledOnce)
-            assert(mockClient.dispose.calledOnce)
-            assert.strictEqual(projectNode['sagemakerClient'], undefined)
+            assert(cleanupSpy.calledOnce)
         })
     })
 
     describe('clearProject', function () {
-        it('clears the project and fires change event', async function () {
+        it('clears the project, calls cleanupProjectResources and fires change event', async function () {
             await projectNode.setProject(mockProject)
+            const cleanupSpy = sinon.spy(projectNode as any, 'cleanupProjectResources')
             const emitterSpy = sinon.spy(projectNode['onDidChangeEmitter'], 'fire')
 
             await projectNode.clearProject()
 
             assert.strictEqual(projectNode['project'], undefined)
+            assert(cleanupSpy.calledOnce)
             assert(emitterSpy.calledOnce)
         })
-
-        it('invalidates credentials and disposes existing sagemaker client', async function () {
-            // Set up existing sagemaker client with mock
-            const mockClient = { dispose: sinon.stub() } as any
-            projectNode['sagemakerClient'] = mockClient
-
-            await projectNode.clearProject()
-
-            assert((projectNode['authProvider'].invalidateAllProjectCredentialsInCache as sinon.SinonStub).calledOnce)
-            assert(mockClient.dispose.calledOnce)
-            assert.strictEqual(projectNode['sagemakerClient'], undefined)
-        })
     })
 
     describe('getProject', function () {
@@ -320,4 +300,26 @@ describe('SageMakerUnifiedStudioProjectNode', function () {
             assert.strictEqual(hasAccess, false)
         })
     })
+
+    describe('cleanupProjectResources', function () {
+        it('invalidates credentials and disposes existing sagemaker client', async function () {
+            // Set up existing sagemaker client with mock
+            const mockClient = { dispose: sinon.stub() } as any
+            projectNode['sagemakerClient'] = mockClient
+
+            await projectNode['cleanupProjectResources']()
+
+            assert((projectNode['authProvider'].invalidateAllProjectCredentialsInCache as sinon.SinonStub).calledOnce)
+            assert(mockClient.dispose.calledOnce)
+            assert.strictEqual(projectNode['sagemakerClient'], undefined)
+        })
+
+        it('handles case when no sagemaker client exists', async function () {
+            projectNode['sagemakerClient'] = undefined
+
+            await projectNode['cleanupProjectResources']()
+
+            assert((projectNode['authProvider'].invalidateAllProjectCredentialsInCache as sinon.SinonStub).calledOnce)
+        })
+    })
 })