From c9bd715324e4a94e31fdb22171f1da133f73a96a Mon Sep 17 00:00:00 2001 From: Bhavya Sharma Date: Tue, 11 Nov 2025 11:53:06 -0800 Subject: [PATCH 1/3] fix(sagemaker): Disable start/stop button for intermediate state of the space --- packages/core/src/awsService/sagemaker/sagemakerSpace.ts | 6 +++++- packages/toolkit/package.json | 4 ++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/packages/core/src/awsService/sagemaker/sagemakerSpace.ts b/packages/core/src/awsService/sagemaker/sagemakerSpace.ts index 5eeed4e8551..f4bcfdd952f 100644 --- a/packages/core/src/awsService/sagemaker/sagemakerSpace.ts +++ b/packages/core/src/awsService/sagemaker/sagemakerSpace.ts @@ -191,11 +191,15 @@ export class SagemakerSpace { public getContext(): string { const status = this.getStatus() - // only distinguish between running and non-running states if (status === SpaceStatus.RUNNING) { return 'awsSagemakerSpaceRunningNode' } + if (status === SpaceStatus.STOPPED) { + return 'awsSagemakerSpaceStoppedNode' + } + + // For all other states (STARTING, STOPPING, etc.), return base context return this.isSMUSSpace ? 'smusSpaceNode' : 'awsSagemakerSpaceNode' } diff --git a/packages/toolkit/package.json b/packages/toolkit/package.json index a7de5f18113..d55e8596e6d 100644 --- a/packages/toolkit/package.json +++ b/packages/toolkit/package.json @@ -1514,12 +1514,12 @@ { "command": "aws.sagemaker.openRemoteConnection", "group": "inline@1", - "when": "view != aws.smus.rootView && viewItem =~ /^(awsSagemakerSpaceRunningNode|awsSagemakerSpaceNode)$/" + "when": "view != aws.smus.rootView && viewItem =~ /^(awsSagemakerSpaceRunningNode|awsSagemakerSpaceStoppedNode)$/" }, { "command": "aws.smus.openRemoteConnection", "group": "inline@1", - "when": "view == aws.smus.rootView && viewItem =~ /^(awsSagemakerSpaceRunningNode|smusSpaceNode)$/" + "when": "view == aws.smus.rootView && viewItem =~ /^(awsSagemakerSpaceRunningNode|awsSagemakerSpaceStoppedNode)$/" }, { "command": "_aws.toolkit.notifications.dismiss", From 7b80f7db0b1ef5306bc5259ed779d6cad0129076 Mon Sep 17 00:00:00 2001 From: Bhavya Sharma Date: Wed, 12 Nov 2025 09:35:21 -0800 Subject: [PATCH 2/3] chore: trigger CI From 7abfff7af8f15243ea0a911f82a1e98bad23f576 Mon Sep 17 00:00:00 2001 From: Bhavya Sharma Date: Fri, 14 Nov 2025 08:44:19 -0800 Subject: [PATCH 3/3] fix(sagemaker): Fix various SMUS bugs --- packages/core/src/awsService/sagemaker/commands.ts | 2 +- packages/core/src/awsService/sagemaker/constants.ts | 2 +- .../nodes/sageMakerUnifiedStudioProjectNode.ts | 2 +- packages/core/src/shared/clients/sagemaker.ts | 2 +- .../src/test/awsService/sagemaker/commands.test.ts | 12 ++++++------ .../src/test/shared/clients/sagemakerClient.test.ts | 11 +++++++---- 6 files changed, 17 insertions(+), 14 deletions(-) diff --git a/packages/core/src/awsService/sagemaker/commands.ts b/packages/core/src/awsService/sagemaker/commands.ts index 66ffe35fbee..7573c51020b 100644 --- a/packages/core/src/awsService/sagemaker/commands.ts +++ b/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', }) diff --git a/packages/core/src/awsService/sagemaker/constants.ts b/packages/core/src/awsService/sagemaker/constants.ts index 1e0875cd385..35cc8b5c938 100644 --- a/packages/core/src/awsService/sagemaker/constants.ts +++ b/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) => { diff --git a/packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioProjectNode.ts b/packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioProjectNode.ts index 8097ceed9e7..114ffe77212 100644 --- a/packages/core/src/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioProjectNode.ts +++ b/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) diff --git a/packages/core/src/shared/clients/sagemaker.ts b/packages/core/src/shared/clients/sagemaker.ts index 165a51c5a08..1d92c629b61 100644 --- a/packages/core/src/shared/clients/sagemaker.ts +++ b/packages/core/src/shared/clients/sagemaker.ts @@ -175,7 +175,7 @@ export class SagemakerClient extends ClientWrapper { instanceType, InstanceTypeInsufficientMemory[instanceType] ), - confirm: 'Restart and Connect', + confirm: 'Restart Space and Connect', cancel: 'Cancel', type: 'warning', }) diff --git a/packages/core/src/test/awsService/sagemaker/commands.test.ts b/packages/core/src/test/awsService/sagemaker/commands.test.ts index fd835cfe79e..8b7a445b1b4 100644 --- a/packages/core/src/test/awsService/sagemaker/commands.test.ts +++ b/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 diff --git a/packages/core/src/test/shared/clients/sagemakerClient.test.ts b/packages/core/src/test/shared/clients/sagemakerClient.test.ts index 0ebda0eeb83..3d92c176d42 100644 --- a/packages/core/src/test/shared/clients/sagemakerClient.test.ts +++ b/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.')