feat(sagemakerunifiedstudio): Fix list spaces in cross region cross account cases (#2209)

## Problem
1. In a cross region cross account set up. seeing error fetching space
2. The error message when cannot connect is not meaningful. I had to
check on telemetry log to see the error message

## Solution
Fixed the above issue


---

- 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: Zulin Liu <[email protected]>

Author: liuzulin

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

@@ -161,7 +161,7 @@ export async function stopSpace(
                 code: error.name,
             })
         } else {
-            throw new ToolkitError(`Failed to stop space: ${spaceName}`, {
+            throw new ToolkitError(`Failed to stop space ${spaceName}: ${(error as Error).message}`, {
                 cause: error,
                 code: error.name,
             })
@@ -198,7 +198,7 @@ export async function openRemoteConnect(
         } catch (err: any) {
             // Ignore InstanceTypeError since it means the user decided not to use an instanceType with more memory
             if (err.code !== InstanceTypeError) {
-                throw new ToolkitError('Remote connection failed.', {
+                throw new ToolkitError(`Remote connection failed: ${(err as Error).message}`, {
                     cause: err as Error,
                     code: err.code,
                 })

packages/core/src/sagemakerunifiedstudio/auth/providers/domainExecRoleCredentialsProvider.ts

@@ -202,15 +202,16 @@ export class DomainExecRoleCredentialsProvider implements CredentialsProvider {
                 )
             }
 
-            const data = (await response.json()) as {
+            const responseText = await response.text()
+
+            const data = JSON.parse(responseText) as {
                 credentials: {
                     accessKeyId: string
                     secretAccessKey: string
                     sessionToken: string
                     expiration: string
                 }
             }
-
             this.logger.debug(`SMUS DER: Successfully received credentials from API for domain ${this.domainId}`)
 
             // Validate the response data structure
@@ -226,21 +227,50 @@ export class DomainExecRoleCredentialsProvider implements CredentialsProvider {
             validateCredentialFields(credentials, 'InvalidCredentialResponse', 'API response')
 
             // Create credentials with expiration
-            // Note: The response doesn't include expiration yet, so we set it to 10 minutes for now if it does't exist
             let credentialExpiresAt: Date
             if (credentials.expiration) {
-                // The API returns expiration as a string, convert to Date
-                const parsedExpiration = new Date(credentials.expiration)
+                // Handle both epoch timestamps and ISO date strings
+                let parsedExpiration: Date
+
+                // Check if expiration is a numeric string (epoch timestamp)
+                const expirationNum = Number(credentials.expiration)
+                if (!isNaN(expirationNum) && expirationNum > 0) {
+                    // Treat as epoch timestamp in seconds and convert to milliseconds
+                    const timestampMs = expirationNum * 1000
+                    parsedExpiration = new Date(timestampMs)
+                    this.logger.debug(
+                        `SMUS DER: Parsed epoch timestamp ${credentials.expiration} (seconds) as ${parsedExpiration.toISOString()}`
+                    )
+                } else {
+                    // Treat as ISO date string
+                    parsedExpiration = new Date(credentials.expiration)
+                    if (!isNaN(parsedExpiration.getTime())) {
+                        this.logger.debug(
+                            `SMUS DER: Parsed ISO date string ${credentials.expiration} as ${parsedExpiration.toISOString()}`
+                        )
+                    } else {
+                        this.logger.debug(
+                            `SMUS DER: Failed to parse ISO date string ${credentials.expiration} - invalid date format`
+                        )
+                    }
+                }
+
                 // Check if the parsed date is valid
                 if (isNaN(parsedExpiration.getTime())) {
                     this.logger.warn(
-                        `SMUS DER: Invalid expiration date string: ${credentials.expiration}, using default expiration`
+                        `SMUS DER: Invalid expiration value: ${credentials.expiration}, using default expiration`
                     )
                     credentialExpiresAt = new Date(Date.now() + SmusCredentialExpiry.derExpiryMs)
                 } else {
                     credentialExpiresAt = parsedExpiration
                 }
+                if (!isNaN(credentialExpiresAt.getTime())) {
+                    this.logger.debug(`SMUS DER: Credential expires at ${credentialExpiresAt.toISOString()}`)
+                } else {
+                    this.logger.debug(`SMUS DER: Invalid credential expiration date, using default`)
+                }
             } else {
+                this.logger.debug(`SMUS DER: No expiration provided, using default`)
                 credentialExpiresAt = new Date(Date.now() + SmusCredentialExpiry.derExpiryMs)
             }
 

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

@@ -98,12 +98,13 @@ export class SageMakerUnifiedStudioDataNode implements TreeNode {
             (conn) => (conn.type as ConnectionType) === ConnectionType.LAKEHOUSE
         )
 
-        // Create Bucket parent node if there are S3 connections
-        if (s3Connections.length > 0) {
-            const bucketNode = this.createBucketParentNode(project, s3Connections, region)
-            dataNodes.push(bucketNode)
+        // Add Lakehouse nodes first
+        for (const connection of lakehouseConnections) {
+            const node = await this.createLakehouseNode(project, connection, region)
+            dataNodes.push(node)
         }
 
+        // Add Redshift nodes second
         for (const connection of redshiftConnections) {
             if (connection.name.startsWith('project.lakehouse')) {
                 continue
@@ -115,9 +116,10 @@ export class SageMakerUnifiedStudioDataNode implements TreeNode {
             dataNodes.push(node)
         }
 
-        for (const connection of lakehouseConnections) {
-            const node = await this.createLakehouseNode(project, connection, region)
-            dataNodes.push(node)
+        // Add S3 Bucket parent node last
+        if (s3Connections.length > 0) {
+            const bucketNode = this.createBucketParentNode(project, s3Connections, region)
+            dataNodes.push(bucketNode)
         }
 
         this.logger.info(`Created ${dataNodes.length} total connection nodes`)

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

@@ -117,9 +117,17 @@ export class SageMakerUnifiedStudioProjectNode implements TreeNode {
                     return [dataNode]
                 }
 
-                this.sagemakerClient = await this.initializeSagemakerClient(
-                    this.authProvider.activeConnection?.ssoRegion || 'us-east-1'
-                )
+                const dzClient = await DataZoneClient.getInstance(this.authProvider)
+                if (!this.project?.id) {
+                    throw new Error('Project ID is required')
+                }
+                const toolingEnv = await dzClient.getToolingEnvironment(this.project.id)
+                const spaceAwsAccountRegion = toolingEnv.awsAccountRegion
+
+                if (!spaceAwsAccountRegion) {
+                    throw new Error('No AWS account region found in tooling environment')
+                }
+                this.sagemakerClient = await this.initializeSagemakerClient(spaceAwsAccountRegion)
                 const computeNode = new SageMakerUnifiedStudioComputeNode(
                     this,
                     this.extensionContext,

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

@@ -8,7 +8,6 @@ import { SageMakerUnifiedStudioComputeNode } from './sageMakerUnifiedStudioCompu
 import { updateInPlace } from '../../../shared/utilities/collectionUtils'
 import { DataZoneClient } from '../../shared/client/datazoneClient'
 import { DescribeDomainResponse } from '@amzn/sagemaker-client'
-import { GetEnvironmentCommandOutput } from '@aws-sdk/client-datazone'
 import { getDomainUserProfileKey } from '../../../awsService/sagemaker/utils'
 import { getLogger } from '../../../shared/logger/logger'
 import { TreeNode } from '../../../shared/treeview/resourceTreeDataProvider'
@@ -31,6 +30,7 @@ export class SageMakerUnifiedStudioSpacesParentNode implements TreeNode {
     public readonly onDidChangeTreeItem = this.onDidChangeEmitter.event
     public readonly onDidChangeChildren = this.onDidChangeEmitter.event
     public readonly pollingSet: PollingSet<string> = new PollingSet(5, this.updatePendingNodes.bind(this))
+    private spaceAwsAccountRegion: string | undefined
 
     public constructor(
         private readonly parent: SageMakerUnifiedStudioComputeNode,
@@ -143,16 +143,9 @@ export class SageMakerUnifiedStudioSpacesParentNode implements TreeNode {
         if (!datazoneClient) {
             throw new Error('DataZone client is not initialized')
         }
-        const toolingEnvId = await datazoneClient
-            .getToolingEnvironmentId(datazoneClient.getDomainId(), this.projectId)
-            .catch((err) => {
-                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) {
-            throw new Error('No default environment found for project')
-        }
-        const toolingEnv: GetEnvironmentCommandOutput = await datazoneClient.getEnvironmentDetails(toolingEnvId)
+
+        const toolingEnv = await datazoneClient.getToolingEnvironment(this.projectId)
+        this.spaceAwsAccountRegion = toolingEnv.awsAccountRegion
         if (toolingEnv.provisionedResources) {
             for (const resource of toolingEnv.provisionedResources) {
                 if (resource.name === 'sageMakerDomainId') {
@@ -224,7 +217,10 @@ export class SageMakerUnifiedStudioSpacesParentNode implements TreeNode {
                 new SagemakerUnifiedStudioSpaceNode(
                     this as any,
                     this.sagemakerClient,
-                    datazoneClient.getRegion(),
+                    this.spaceAwsAccountRegion ||
+                        (() => {
+                            throw new Error('No AWS account region found in tooling environment')
+                        })(),
                     this.spaceApps.get(key)!,
                     true /* isSMUSSpace */
                 )

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

@@ -15,6 +15,7 @@ import {
     S3PropertiesOutput,
     ConnectionType,
     GluePropertiesOutput,
+    GetEnvironmentCommandOutput,
 } from '@aws-sdk/client-datazone'
 import { getLogger } from '../../../shared/logger/logger'
 import type { SmusAuthenticationProvider } from '../../auth/providers/smusAuthenticationProvider'
@@ -753,6 +754,33 @@ export class DataZoneClient {
         }
     }
 
+    /**
+     * Gets the tooling environment details for a project
+     * @param projectId The project ID
+     * @returns The tooling environment details
+     */
+    public async getToolingEnvironment(projectId: string): Promise<GetEnvironmentCommandOutput> {
+        const logger = getLogger()
+
+        const datazoneClient = await DataZoneClient.getInstance(this.authProvider)
+        if (!datazoneClient) {
+            throw new Error('DataZone client is not initialized')
+        }
+
+        const toolingEnvId = await datazoneClient
+            .getToolingEnvironmentId(datazoneClient.getDomainId(), projectId)
+            .catch((err) => {
+                logger.error('Failed to get tooling environment ID for project %s', projectId)
+                throw new Error(`Failed to get tooling environment ID: ${err.message}`)
+            })
+
+        if (!toolingEnvId) {
+            throw new Error('No default environment found for project')
+        }
+
+        return await datazoneClient.getEnvironmentDetails(toolingEnvId)
+    }
+
     public async getUserId(): Promise<string | undefined> {
         const derCredProvider = await this.authProvider.getDerCredentialsProvider()
         this.logger.debug(`Calling STS GetCallerIdentity using DER credentials of ${this.getDomainId()}`)