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

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

@@ -16,10 +16,19 @@ import _ from 'lodash'
 import { prepareDevEnvConnection, tryRemoteConnection } from './model'
 import { ExtContext } from '../../shared/extensions'
 import { SagemakerClient } from '../../shared/clients/sagemaker'
+import { AccessDeniedException } from '@amzn/sagemaker-client'
 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,
+    RemoteAccess,
+    RemoteAccessRequiredMessage,
+    SpaceStatus,
+} from './constants'
 import { SagemakerUnifiedStudioSpaceNode } from '../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode'
 
 const localize = nls.loadMessageBundle()
@@ -136,10 +145,10 @@ export async function stopSpace(
     sageMakerClient?: SagemakerClient
 ) {
     await tryRefreshNode(node)
-    if (node.getStatus() === 'Stopped' || node.getStatus() === 'Stopping') {
+    if (node.getStatus() === SpaceStatus.STOPPED || node.getStatus() === SpaceStatus.STOPPING) {
         void vscode.window.showWarningMessage(`Space ${node.spaceApp.SpaceName} is already in Stopped/Stopping state.`)
         return
-    } else if (node.getStatus() === 'Starting') {
+    } else if (node.getStatus() === SpaceStatus.STARTING) {
         void vscode.window.showWarningMessage(
             `Space ${node.spaceApp.SpaceName} is in Starting state. Wait until it is Running to attempt stop again.`
         )
@@ -167,7 +176,7 @@ export async function stopSpace(
         })
     } catch (err) {
         const error = err as Error
-        if (error.name === 'AccessDeniedException') {
+        if (error instanceof AccessDeniedException) {
             throw new ToolkitError('You do not have permission to stop spaces. Please contact your administrator', {
                 cause: error,
                 code: error.name,
@@ -195,56 +204,205 @@ 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 === SpaceStatus.RUNNING && remoteAccess !== RemoteAccess.ENABLED) {
+        return handleRunningSpaceWithDisabledAccess(node, ctx, spaceName, sageMakerClient)
+    } else if (nodeStatus === SpaceStatus.STOPPED) {
+        return handleStoppedSpace(node, ctx, spaceName, sageMakerClient)
+    } else if (nodeStatus === SpaceStatus.RUNNING) {
+        return handleRunningSpaceWithEnabledAccess(node, ctx, spaceName)
+    }
+}
+
+/**
+ * Checks if an instance type upgrade will be needed for remote access
+ */
+export 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!
+
+        // Get current instance type
+        const currentResourceSpec =
+            appType === 'JupyterLab'
+                ? spaceDetails.SpaceSettings!.JupyterLabAppSettings?.DefaultResourceSpec
+                : spaceDetails.SpaceSettings!.CodeEditorAppSettings?.DefaultResourceSpec
+
+        const currentInstanceType = currentResourceSpec?.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 instanceof AccessDeniedException) {
+            throw new ToolkitError('You do not have permission to describe spaces. Please contact your administrator', {
+                cause: error,
+                code: error.name,
+            })
+        }
+        throw err
+    }
+}
+
+/**
+ * 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 = RemoteAccessRequiredMessage
+    }
+
+    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,
+                })
 
-            // 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)
+                // 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)
+                await client.waitForAppInService(node.spaceApp.DomainId!, spaceName, node.spaceApp.App!.AppType!)
+                await tryRemoteConnection(node, ctx, progress)
+            } catch (err: any) {
+                // Handle user declining instance type upgrade
+                if (err.code === InstanceTypeError) {
+                    return
                 }
-            )
-        } 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
+                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,
-            })
         }
-    } else if (node.getStatus() === 'Running') {
-        // For running spaces, show progress
+    )
+}
+
+/**
+ * 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)
+
         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, node.spaceApp.App!.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)
+        }
+    )
+}

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

@@ -18,6 +18,19 @@ export const InstanceTypeInsufficientMemory: Record<string, string> = {
     'ml.c5.large': 'ml.c5.xlarge',
 }
 
+// Remote access constants
+export const RemoteAccess = {
+    ENABLED: 'ENABLED',
+    DISABLED: 'DISABLED',
+} as const
+
+export const SpaceStatus = {
+    RUNNING: 'Running',
+    STOPPED: 'Stopped',
+    STARTING: 'Starting',
+    STOPPING: 'Stopping',
+} as const
+
 export const InstanceTypeInsufficientMemoryMessage = (
     spaceName: string,
     chosenInstanceType: string,
@@ -29,3 +42,6 @@ export const InstanceTypeInsufficientMemoryMessage = (
 export const InstanceTypeNotSelectedMessage = (spaceName: string) => {
     return `No instanceType specified for [${spaceName}]. ${InstanceTypeMinimum} is the default instance type, which meets minimum 8 GiB memory requirements for remote access. Continuing will start your space with instanceType [${InstanceTypeMinimum}] and remotely connect.`
 }
+
+export const RemoteAccessRequiredMessage =
+    'This space requires remote access to be enabled.\nWould you like to restart the space and connect?\nAny unsaved work will be lost.'

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

@@ -11,6 +11,7 @@ import { getIcon, IconPath } from '../../shared/icons'
 import { generateSpaceStatus, updateIdleFile, startMonitoringTerminalActivity, ActivityCheckInterval } from './utils'
 import { UserActivity } from '../../shared/extensionUtilities'
 import { getLogger } from '../../shared/logger/logger'
+import { SpaceStatus, RemoteAccess } from './constants'
 
 export class SagemakerSpace {
     public label: string = ''
@@ -53,7 +54,7 @@ export class SagemakerSpace {
     }
 
     public isPending(): boolean {
-        return this.getStatus() !== 'Running' && this.getStatus() !== 'Stopped'
+        return this.getStatus() !== SpaceStatus.RUNNING && this.getStatus() !== SpaceStatus.STOPPED
     }
 
     public getStatus(): string {
@@ -148,8 +149,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 === RemoteAccess.ENABLED) {
+                baseTooltip += `\n\n**Remote Access:** Enabled`
+            } else if (remoteAccess === 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}`
     }
@@ -166,20 +177,12 @@ export class SagemakerSpace {
 
     public getContext(): string {
         const status = this.getStatus()
-        if (status === 'Running' && this.spaceApp.SpaceSettingsSummary?.RemoteAccess === 'ENABLED') {
-            return 'awsSagemakerSpaceRunningRemoteEnabledNode'
-        } else if (status === 'Running' && this.spaceApp.SpaceSettingsSummary?.RemoteAccess === 'DISABLED') {
-            return 'awsSagemakerSpaceRunningRemoteDisabledNode'
-        } else if (status === 'Running' && this.isSMUSSpace) {
+
+        // only distinguish between running and non-running states
+        if (status === SpaceStatus.RUNNING) {
             return 'awsSagemakerSpaceRunningNode'
-        } else if (status === 'Stopped' && this.spaceApp.SpaceSettingsSummary?.RemoteAccess === 'ENABLED') {
-            return 'awsSagemakerSpaceStoppedRemoteEnabledNode'
-        } else if (
-            (status === 'Stopped' && !this.spaceApp.SpaceSettingsSummary?.RemoteAccess) ||
-            this.spaceApp.SpaceSettingsSummary?.RemoteAccess === 'DISABLED'
-        ) {
-            return 'awsSagemakerSpaceStoppedRemoteDisabledNode'
         }
+
         return this.isSMUSSpace ? 'smusSpaceNode' : 'awsSagemakerSpaceNode'
     }