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

## Problem
- Connect button only appeared for spaces with remote access enabled.
When users landed on the toolkit and saw running spaces without remote
access, there was no connect button visible. This created confusion
because:
-- Users expected to be able to connect to any running space
-- No visual indication that connection was possible
-- Users didn't understand why some spaces were "connectable" and others
weren't

## Solution
- Always show connect button for all spaces (running and stopped,
regardless of remote access status). When users click connect on a space
without remote access, show clear dialog explaining what needs to happen
and get their consent upfront.
- There is no breaking change.
- Tested for all use-cases SM-AI and SMUS

## Current User Exp
---

https://github.com/user-attachments/assets/9c5c8634-42df-4cd1-988b-12cc8008d69e


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

Author: bhavya2109sharma

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'
     }