feat(smus): UX improvements for connecting running spaces from toolkit

bhavya2109sharma · bhavya2109sharma · commit 1a5a7cfcb7ac · 2025-11-05T15:51:43.000-08:00
diff --git a/packages/core/src/awsService/sagemaker/commands.ts b/packages/core/src/awsService/sagemaker/commands.ts
@@ -19,7 +19,12 @@ import { SagemakerClient } from '../../shared/clients/sagemaker'
 import { ToolkitError } from '../../shared/errors'
 import { showConfirmationMessage } from '../../shared/utilities/messages'
 import { RemoteSessionError } from '../../shared/remoteSession'
-import { ConnectFromRemoteWorkspaceMessage, InstanceTypeError } from './constants'
+import {
+    ConnectFromRemoteWorkspaceMessage,
+    InstanceTypeError,
+    InstanceTypeInsufficientMemory,
+    InstanceTypeInsufficientMemoryMessage,
+} from './constants'
 import { SagemakerUnifiedStudioSpaceNode } from '../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode'
 
 const localize = nls.loadMessageBundle()
@@ -195,56 +200,228 @@ export async function openRemoteConnect(
     const spaceName = node.spaceApp.SpaceName!
     await tryRefreshNode(node)
 
-    // for Stopped SM spaces - check instance type before showing progress
-    if (node.getStatus() === 'Stopped') {
-        //  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(spaceName, node.spaceApp.DomainId!)
-            await tryRefreshNode(node)
-            const appType = node.spaceApp.SpaceSettingsSummary?.AppType
-            if (!appType) {
-                throw new ToolkitError('AppType is undefined for the selected space. Cannot start remote connection.', {
-                    code: 'undefinedAppType',
-                })
+    const remoteAccess = node.spaceApp.SpaceSettingsSummary?.RemoteAccess
+    const nodeStatus = node.getStatus()
+
+    // Route to appropriate handler based on space state
+    if (nodeStatus === 'Running' && remoteAccess === 'DISABLED') {
+        return handleRunningSpaceWithDisabledAccess(node, ctx, spaceName, sageMakerClient)
+    } else if (nodeStatus === 'Stopped') {
+        return handleStoppedSpace(node, ctx, spaceName, sageMakerClient)
+    } else if (nodeStatus === 'Running') {
+        return handleRunningSpaceWithEnabledAccess(node, ctx, spaceName)
+    }
+}
+
+/**
+ * Checks if an instance type upgrade will be needed for remote access
+ */
+async function checkInstanceTypeUpgradeNeeded(
+    node: SagemakerSpaceNode | SagemakerUnifiedStudioSpaceNode,
+    sageMakerClient?: SagemakerClient
+): Promise<{ upgradeNeeded: boolean; currentType?: string; recommendedType?: string }> {
+    const client = sageMakerClient || new SagemakerClient(node.regionCode)
+
+    try {
+        const spaceDetails = await client.describeSpace({
+            DomainId: node.spaceApp.DomainId!,
+            SpaceName: node.spaceApp.SpaceName!,
+        })
+
+        const appType = spaceDetails.SpaceSettings?.AppType
+        if (!appType) {
+            return { upgradeNeeded: false }
+        }
+
+        // Get current instance type
+        const requestedResourceSpec =
+            appType === 'JupyterLab'
+                ? spaceDetails.SpaceSettings?.JupyterLabAppSettings?.DefaultResourceSpec
+                : spaceDetails.SpaceSettings?.CodeEditorAppSettings?.DefaultResourceSpec
+
+        const currentInstanceType = requestedResourceSpec?.InstanceType
+
+        // Check if upgrade is needed
+        if (currentInstanceType && currentInstanceType in InstanceTypeInsufficientMemory) {
+            // Current type has insufficient memory
+            return {
+                upgradeNeeded: true,
+                currentType: currentInstanceType,
+                recommendedType: InstanceTypeInsufficientMemory[currentInstanceType],
             }
+        }
+
+        return { upgradeNeeded: false, currentType: currentInstanceType }
+    } catch (err) {
+        const error = err as Error
+        if (error.name === 'AccessDeniedException') {
+            throw new ToolkitError('You do not have permission to describe spaces. Please contact your administrator', {
+                cause: error,
+                code: error.name,
+            })
+        }
+        // For other errors, assume no upgrade needed to avoid blocking the flow
+        return { upgradeNeeded: false }
+    }
+}
+
+/**
+ * Handles connecting to a running space with disabled remote access
+ * Requires stopping the space, enabling remote access, and restarting
+ */
+async function handleRunningSpaceWithDisabledAccess(
+    node: SagemakerSpaceNode | SagemakerUnifiedStudioSpaceNode,
+    ctx: vscode.ExtensionContext,
+    spaceName: string,
+    sageMakerClient?: SagemakerClient
+) {
+    // Check if instance type upgrade will be needed
+    const instanceTypeInfo = await checkInstanceTypeUpgradeNeeded(node, sageMakerClient)
+
+    let prompt: string
+    if (instanceTypeInfo.upgradeNeeded) {
+        prompt = InstanceTypeInsufficientMemoryMessage(
+            spaceName,
+            instanceTypeInfo.currentType!,
+            instanceTypeInfo.recommendedType!
+        )
+    } else {
+        // Only remote access needs to be enabled
+        prompt = `This space requires remote access to be enabled.\nWould you like to restart the space and connect?\nAny unsaved work will be lost.`
+    }
+
+    const confirmed = await showConfirmationMessage({
+        prompt,
+        confirm: 'Restart and Connect',
+        cancel: 'Cancel',
+        type: 'warning',
+    })
+
+    if (!confirmed) {
+        return
+    }
+
+    // Enable remote access and connect
+    const client = sageMakerClient || new SagemakerClient(node.regionCode)
+
+    return await vscode.window.withProgress(
+        {
+            location: vscode.ProgressLocation.Notification,
+            cancellable: false,
+            title: `Connecting to ${spaceName}`,
+        },
+        async (progress) => {
+            try {
+                // Show initial progress message
+                progress.report({ message: 'Stopping the space' })
+
+                // Stop the running space
+                await client.deleteApp({
+                    DomainId: node.spaceApp.DomainId!,
+                    SpaceName: spaceName,
+                    AppType: node.spaceApp.App!.AppType!,
+                    AppName: node.spaceApp.App?.AppName,
+                })
+
+                // Update progress message
+                progress.report({ message: 'Starting the space' })
+
+                // Start the space with remote access enabled (skip prompts since user already consented)
+                await client.startSpace(spaceName, node.spaceApp.DomainId!, true)
+                await tryRefreshNode(node)
 
-            // Only start showing progress after instance type validation
-            return await vscode.window.withProgress(
-                {
-                    location: vscode.ProgressLocation.Notification,
-                    cancellable: false,
-                    title: `Connecting to ${spaceName}`,
-                },
-                async (progress) => {
-                    progress.report({ message: 'Starting the space.' })
-                    await client.waitForAppInService(node.spaceApp.DomainId!, spaceName, appType)
-                    await tryRemoteConnection(node, ctx, progress)
+                const appType = node.spaceApp.SpaceSettingsSummary?.AppType
+                if (!appType) {
+                    throw new ToolkitError(
+                        'AppType is undefined for the selected space. Cannot start remote connection.',
+                        {
+                            code: 'undefinedAppType',
+                        }
+                    )
                 }
-            )
-        } catch (err: any) {
-            // Ignore InstanceTypeError since it means the user decided not to use an instanceType with more memory
-            // just return without showing progress
-            if (err.code === InstanceTypeError) {
-                return
+
+                progress.report({ message: 'Starting the space' })
+                await client.waitForAppInService(node.spaceApp.DomainId!, spaceName, appType)
+                await tryRemoteConnection(node, ctx, progress)
+            } catch (err: any) {
+                // Handle user declining instance type upgrade
+                if (err.code === InstanceTypeError) {
+                    return
+                }
+                throw new ToolkitError(`Remote connection failed: ${err.message}`, {
+                    cause: err,
+                    code: err.code,
+                })
             }
-            throw new ToolkitError(`Remote connection failed: ${(err as Error).message}`, {
-                cause: err as Error,
-                code: err.code,
+        }
+    )
+}
+
+/**
+ * Handles connecting to a stopped space
+ * Starts the space and connects (remote access enabled automatically if needed)
+ */
+async function handleStoppedSpace(
+    node: SagemakerSpaceNode | SagemakerUnifiedStudioSpaceNode,
+    ctx: vscode.ExtensionContext,
+    spaceName: string,
+    sageMakerClient?: SagemakerClient
+) {
+    const client = sageMakerClient || new SagemakerClient(node.regionCode)
+
+    try {
+        await client.startSpace(spaceName, node.spaceApp.DomainId!)
+        await tryRefreshNode(node)
+
+        const appType = node.spaceApp.SpaceSettingsSummary?.AppType
+        if (!appType) {
+            throw new ToolkitError('AppType is undefined for the selected space. Cannot start remote connection.', {
+                code: 'undefinedAppType',
             })
         }
-    } else if (node.getStatus() === 'Running') {
-        // For running spaces, show progress
+
         return await vscode.window.withProgress(
             {
                 location: vscode.ProgressLocation.Notification,
                 cancellable: false,
                 title: `Connecting to ${spaceName}`,
             },
             async (progress) => {
+                progress.report({ message: 'Starting the space' })
+                await client.waitForAppInService(node.spaceApp.DomainId!, spaceName, appType)
                 await tryRemoteConnection(node, ctx, progress)
             }
         )
+    } catch (err: any) {
+        // Handle user declining instance type upgrade
+        if (err.code === InstanceTypeError) {
+            return
+        }
+        throw new ToolkitError(`Remote connection failed: ${(err as Error).message}`, {
+            cause: err as Error,
+            code: err.code,
+        })
     }
 }
+
+/**
+ * Handles connecting to a running space with enabled remote access
+ * Direct connection without any space modifications
+ */
+async function handleRunningSpaceWithEnabledAccess(
+    node: SagemakerSpaceNode | SagemakerUnifiedStudioSpaceNode,
+    ctx: vscode.ExtensionContext,
+    spaceName: string,
+    sageMakerClient?: SagemakerClient
+) {
+    return await vscode.window.withProgress(
+        {
+            location: vscode.ProgressLocation.Notification,
+            cancellable: false,
+            title: `Connecting to ${spaceName}`,
+        },
+        async (progress) => {
+            await tryRemoteConnection(node, ctx, progress)
+        }
+    )
+}
diff --git a/packages/core/src/awsService/sagemaker/sagemakerSpace.ts b/packages/core/src/awsService/sagemaker/sagemakerSpace.ts
@@ -148,8 +148,18 @@ export class SagemakerSpace {
         const domainId = this.spaceApp?.DomainId ?? '-'
         const owner = this.spaceApp?.OwnershipSettingsSummary?.OwnerUserProfileName || '-'
         const instanceType = this.spaceApp?.App?.ResourceSpec?.InstanceType ?? '-'
+        const remoteAccess = this.spaceApp?.SpaceSettingsSummary?.RemoteAccess
+
+        let baseTooltip = ''
         if (this.isSMUSSpace) {
-            return `**Space:** ${spaceName} \n\n**Application:** ${appType} \n\n**Instance Type:** ${instanceType}`
+            baseTooltip = `**Space:** ${spaceName} \n\n**Application:** ${appType} \n\n**Instance Type:** ${instanceType}`
+            if (remoteAccess === 'ENABLED') {
+                baseTooltip += `\n\n**Remote Access:** Enabled`
+            } else if (remoteAccess === 'DISABLED') {
+                baseTooltip += `\n\n**Remote Access:** Disabled`
+            }
+
+            return baseTooltip
         }
         return `**Space:** ${spaceName} \n\n**Application:** ${appType} \n\n**Domain ID:** ${domainId} \n\n**User Profile:** ${owner}`
     }
diff --git a/packages/core/src/shared/clients/sagemaker.ts b/packages/core/src/shared/clients/sagemaker.ts
@@ -48,7 +48,8 @@ import {
 import { getDomainSpaceKey } from '../../awsService/sagemaker/utils'
 import { getLogger } from '../logger/logger'
 import { ToolkitError } from '../errors'
-import { yes, no, continueText, cancel } from '../localizedText'
+import { continueText, cancel } from '../localizedText'
+import { showConfirmationMessage } from '../utilities/messages'
 import { AwsCredentialIdentity } from '@aws-sdk/types'
 import globals from '../extensionGlobals'
 
@@ -126,7 +127,7 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
         return this.makeRequest(DeleteAppCommand, request)
     }
 
-    public async startSpace(spaceName: string, domainId: string) {
+    public async startSpace(spaceName: string, domainId: string, skipInstanceTypePrompts: boolean = false) {
         let spaceDetails: DescribeSpaceCommandOutput
 
         // Get existing space details
@@ -155,35 +156,47 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
 
         // Is InstanceType defined and has enough memory?
         if (instanceType && instanceType in InstanceTypeInsufficientMemory) {
-            // Prompt user to select one with sufficient memory (1 level up from their chosen one)
-            const response = await vscode.window.showErrorMessage(
-                InstanceTypeInsufficientMemoryMessage(
-                    spaceDetails.SpaceName || '',
-                    instanceType,
-                    InstanceTypeInsufficientMemory[instanceType]
-                ),
-                yes,
-                no
-            )
+            if (skipInstanceTypePrompts) {
+                // User already consented, upgrade automatically
+                instanceType = InstanceTypeInsufficientMemory[instanceType]
+            } else {
+                // Prompt user to select one with sufficient memory (1 level up from their chosen one)
+                const confirmed = await showConfirmationMessage({
+                    prompt: InstanceTypeInsufficientMemoryMessage(
+                        spaceDetails.SpaceName || '',
+                        instanceType,
+                        InstanceTypeInsufficientMemory[instanceType]
+                    ),
+                    confirm: 'Restart and Connect',
+                    cancel: 'Cancel',
+                    type: 'warning',
+                })
 
-            if (response === no) {
-                throw new ToolkitError('InstanceType has insufficient memory.', { code: InstanceTypeError })
-            }
+                if (!confirmed) {
+                    throw new ToolkitError('InstanceType has insufficient memory.', { code: InstanceTypeError })
+                }
 
-            instanceType = InstanceTypeInsufficientMemory[instanceType]
+                instanceType = InstanceTypeInsufficientMemory[instanceType]
+            }
         } else if (!instanceType) {
-            // Prompt user to select the minimum supported instance type
-            const response = await vscode.window.showErrorMessage(
-                InstanceTypeNotSelectedMessage(spaceDetails.SpaceName || ''),
-                continueText,
-                cancel
-            )
+            if (skipInstanceTypePrompts) {
+                // User already consented, use minimum
+                instanceType = InstanceTypeMinimum
+            } else {
+                // Prompt user to select the minimum supported instance type
+                const confirmed = await showConfirmationMessage({
+                    prompt: InstanceTypeNotSelectedMessage(spaceDetails.SpaceName || ''),
+                    confirm: continueText,
+                    cancel: cancel,
+                    type: 'warning',
+                })
 
-            if (response === cancel) {
-                throw new ToolkitError('InstanceType not defined.', { code: InstanceTypeError })
-            }
+                if (!confirmed) {
+                    throw new ToolkitError('InstanceType not defined.', { code: InstanceTypeError })
+                }
 
-            instanceType = InstanceTypeMinimum
+                instanceType = InstanceTypeMinimum
+            }
         }
 
         // First, update the space if needed
diff --git a/packages/toolkit/package.json b/packages/toolkit/package.json
