fix(sagemaker): Fix various SMUS bugs

Author: bhavya2109sharma

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

@@ -292,7 +292,7 @@ async function handleRunningSpaceWithDisabledAccess(
 
     const confirmed = await showConfirmationMessage({
         prompt,
-        confirm: 'Restart and Connect',
+        confirm: 'Restart Space and Connect',
         cancel: 'Cancel',
         type: 'warning',
     })

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

@@ -36,7 +36,7 @@ export const InstanceTypeInsufficientMemoryMessage = (
     chosenInstanceType: string,
     recommendedInstanceType: string
 ) => {
-    return `Unable to create app for [${spaceName}] because instanceType [${chosenInstanceType}] is not supported for remote access enabled spaces. Use instanceType with at least 8 GiB memory. Would you like to start your space with instanceType [${recommendedInstanceType}]?`
+    return `[${chosenInstanceType}] does not support remote access. Use an instanceType with at least 8 GiB memory. Would you like to start your space with instanceType [${recommendedInstanceType}]?`
 }
 
 export const InstanceTypeNotSelectedMessage = (spaceName: string) => {

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

@@ -132,7 +132,7 @@ export class SageMakerUnifiedStudioProjectNode implements TreeNode {
                 if (this.isFirstTimeSelection && !this.hasShownFirstTimeMessage) {
                     this.hasShownFirstTimeMessage = true
                     void vscode.window.showInformationMessage(
-                        'Find your space in the Explorer panel under SageMaker Unified Studio. Hover over any space and click the connection icon to connect remotely.'
+                        'Find your space in the Explorer panel under SageMaker Unified Studio. Hover over a space and click the connection icon to connect remotely.'
                     )
                 }
                 this.sagemakerClient = await this.initializeSagemakerClient(spaceAwsAccountRegion)

packages/core/src/shared/clients/sagemaker.ts

@@ -175,7 +175,7 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
                         instanceType,
                         InstanceTypeInsufficientMemory[instanceType]
                     ),
-                    confirm: 'Restart and Connect',
+                    confirm: 'Restart Space and Connect',
                     cancel: 'Cancel',
                     type: 'warning',
                 })

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

@@ -124,7 +124,7 @@ describe('SageMaker Commands', () => {
                 // Setup test window to handle confirmation dialog
                 getTestWindow().onDidShowMessage((message) => {
                     if (message.message.includes(RemoteAccessRequiredMessage)) {
-                        message.selectItem('Restart and Connect')
+                        message.selectItem('Restart Space and Connect')
                     }
                 })
 
@@ -182,7 +182,7 @@ describe('SageMaker Commands', () => {
                             InstanceTypeInsufficientMemoryMessage('test-space', 'ml.t3.medium', 'ml.t3.large')
                         )
                     ) {
-                        message.selectItem('Restart and Connect')
+                        message.selectItem('Restart Space and Connect')
                     }
                 })
 
@@ -236,8 +236,8 @@ describe('SageMaker Commands', () => {
 
                 // Setup test window to confirm
                 getTestWindow().onDidShowMessage((message) => {
-                    if (message.items.some((item) => item.title === 'Restart and Connect')) {
-                        message.selectItem('Restart and Connect')
+                    if (message.items.some((item) => item.title === 'Restart Space and Connect')) {
+                        message.selectItem('Restart Space and Connect')
                     }
                 })
 
@@ -337,7 +337,7 @@ describe('SageMaker Commands', () => {
 
                 // Verify no confirmation dialog shown for stopped space
                 const confirmMessages = getTestWindow().shownMessages.filter((m) =>
-                    m.message.includes('Restart and Connect')
+                    m.message.includes('Restart Space and Connect')
                 )
                 assert.strictEqual(confirmMessages.length, 0, 'Should not show confirmation for stopped space')
 
@@ -388,7 +388,7 @@ describe('SageMaker Commands', () => {
                 assert(mockTryRefreshNode.calledOnce)
                 // Verify no confirmation needed
                 const confirmMessages = getTestWindow().shownMessages.filter((m) =>
-                    m.message.includes('Restart and Connect')
+                    m.message.includes('Restart Space and Connect')
                 )
                 assert.strictEqual(confirmMessages.length, 0)
                 // Verify no space operations performed

packages/core/src/test/shared/clients/sagemakerClient.test.ts

@@ -11,6 +11,7 @@ import { DescribeDomainResponse } from '@amzn/sagemaker-client'
 import { intoCollection } from '../../../shared/utilities/collectionUtils'
 import { ToolkitError } from '../../../shared/errors'
 import { getTestWindow } from '../vscode/window'
+import { InstanceTypeInsufficientMemoryMessage } from '../../../awsService/sagemaker/constants'
 
 describe('SagemakerClient.fetchSpaceAppsAndDomains', function () {
     const region = 'test-region'
@@ -390,9 +391,10 @@ describe('SagemakerClient.startSpace', function () {
 
         const promise = client.startSpace('my-space', 'my-domain')
 
-        // Wait for the error message to appear and select "Restart and Connect"
-        await getTestWindow().waitForMessage(/not supported for remote access/)
-        getTestWindow().getFirstMessage().selectItem('Restart and Connect')
+        // Wait for the error message to appear and select "Restart Space and Connect"
+        const expectedMessage = InstanceTypeInsufficientMemoryMessage('my-space', 'ml.t3.medium', 'ml.t3.large')
+        await getTestWindow().waitForMessage(new RegExp(expectedMessage.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')))
+        getTestWindow().getFirstMessage().selectItem('Restart Space and Connect')
 
         await promise
         sinon.assert.calledOnce(updateSpaceStub)
@@ -416,7 +418,8 @@ describe('SagemakerClient.startSpace', function () {
         const promise = client.startSpace('my-space', 'my-domain')
 
         // Wait for the error message to appear and select "Cancel"
-        await getTestWindow().waitForMessage(/not supported for remote access/)
+        const expectedMessage = InstanceTypeInsufficientMemoryMessage('my-space', 'ml.t3.medium', 'ml.t3.large')
+        await getTestWindow().waitForMessage(new RegExp(expectedMessage.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')))
         getTestWindow().getFirstMessage().selectItem('Cancel')
 
         await assert.rejects(promise, (err: ToolkitError) => err.message === 'InstanceType has insufficient memory.')