Merge master into feature/emr

Author: aws-toolkit-automation

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

@@ -191,32 +191,60 @@ export async function openRemoteConnect(
         void vscode.window.showErrorMessage(ConnectFromRemoteWorkspaceMessage)
         return
     }
+
+    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(node.spaceApp.SpaceName!, node.spaceApp.DomainId!)
+            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',
                 })
             }
-            await client.waitForAppInService(node.spaceApp.DomainId!, node.spaceApp.SpaceName!, appType)
-            await tryRemoteConnection(node, ctx)
+
+            // 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)
+                }
+            )
         } 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: ${(err as Error).message}`, {
-                    cause: err as Error,
-                    code: err.code,
-                })
+            // just return without showing progress
+            if (err.code === InstanceTypeError) {
+                return
             }
+            throw new ToolkitError(`Remote connection failed: ${(err as Error).message}`, {
+                cause: err as Error,
+                code: err.code,
+            })
         }
     } else if (node.getStatus() === 'Running') {
-        await tryRemoteConnection(node, ctx)
+        // For running spaces, show progress
+        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/model.ts

@@ -27,12 +27,14 @@ const logger = getLogger('sagemaker')
 
 export async function tryRemoteConnection(
     node: SagemakerSpaceNode | SagemakerUnifiedStudioSpaceNode,
-    ctx: vscode.ExtensionContext
+    ctx: vscode.ExtensionContext,
+    progress: vscode.Progress<{ message?: string; increment?: number }>
 ) {
     const spaceArn = (await node.getSpaceArn()) as string
     const isSMUS = node instanceof SagemakerUnifiedStudioSpaceNode
     const remoteEnv = await prepareDevEnvConnection(spaceArn, ctx, 'sm_lc', isSMUS, node)
     try {
+        progress.report({ message: 'Opening remote session' })
         await startVscodeRemote(
             remoteEnv.SessionProcess,
             remoteEnv.hostname,