diff --git a/packages/core/src/awsService/sagemaker/sagemakerSpace.ts b/packages/core/src/awsService/sagemaker/sagemakerSpace.ts index 24618966e46..5eeed4e8551 100644 --- a/packages/core/src/awsService/sagemaker/sagemakerSpace.ts +++ b/packages/core/src/awsService/sagemaker/sagemakerSpace.ts @@ -11,8 +11,11 @@ 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 { ToolkitError } from '../../shared/errors' import { SpaceStatus, RemoteAccess } from './constants' +const logger = getLogger('sagemaker') + export class SagemakerSpace { public label: string = '' public contextValue: string = '' @@ -34,6 +37,10 @@ export class SagemakerSpace { } public updateSpace(spaceApp: SagemakerSpaceApp) { + // Edge case when this.spaceApp.App is null, returned by ListApp API for a Space that is not connected to for over 24 hours + if (!this.spaceApp.App) { + this.spaceApp.App = spaceApp.App + } this.setSpaceStatus(spaceApp.Status ?? '', spaceApp.App?.Status ?? '') // Only update RemoteAccess property to minimize impact due to minor structural differences between variables if (this.spaceApp.SpaceSettingsSummary && spaceApp.SpaceSettingsSummary?.RemoteAccess) { @@ -107,13 +114,19 @@ export class SagemakerSpace { DomainId: this.spaceApp.DomainId, SpaceName: this.spaceApp.SpaceName, }) - - const app = await this.client.describeApp({ - DomainId: this.spaceApp.DomainId, - AppName: this.spaceApp.App?.AppName, - AppType: this.spaceApp?.SpaceSettingsSummary?.AppType, - SpaceName: this.spaceApp.SpaceName, - }) + // get app using ListApps API, with given DomainId and SpaceName + const app = + this.spaceApp.DomainId && this.spaceApp.SpaceName + ? await this.client.listAppForSpace(this.spaceApp.DomainId, this.spaceApp.SpaceName) + : undefined + if (!app) { + logger.error( + `updateSpaceAppStatus: unable to get app, [DomainId: ${this.spaceApp.DomainId}], [SpaceName: ${this.spaceApp.SpaceName}]` + ) + throw new ToolkitError( + `Cannot update app status without [DomainId: ${this.spaceApp.DomainId} and SpaceName: ${this.spaceApp.SpaceName}]` + ) + } // AWS DescribeSpace API returns full details with property names like 'SpaceSettings' // but our internal SagemakerSpaceApp type expects 'SpaceSettingsSummary' (from ListSpaces API) @@ -195,7 +208,6 @@ export class SagemakerSpace { * Sets up user activity monitoring for SageMaker spaces */ export async function setupUserActivityMonitoring(extensionContext: vscode.ExtensionContext): Promise { - const logger = getLogger() logger.info('setupUserActivityMonitoring: Starting user activity monitoring setup') const tmpDirectory = '/tmp/' diff --git a/packages/core/src/shared/clients/sagemaker.ts b/packages/core/src/shared/clients/sagemaker.ts index 14d8947c896..165a51c5a08 100644 --- a/packages/core/src/shared/clients/sagemaker.ts +++ b/packages/core/src/shared/clients/sagemaker.ts @@ -128,6 +128,13 @@ export class SagemakerClient extends ClientWrapper { return this.makeRequest(DeleteAppCommand, request) } + public async listAppForSpace(domainId: string, spaceName: string): Promise { + const appsList = await this.listApps({ DomainIdEquals: domainId, SpaceNameEquals: spaceName }) + .flatten() + .promise() + return appsList[0] // At most one App for one SagemakerSpace + } + public async startSpace(spaceName: string, domainId: string, skipInstanceTypePrompts: boolean = false) { let spaceDetails: DescribeSpaceCommandOutput diff --git a/packages/core/src/test/awsService/sagemaker/explorer/sagemakerSpaceNode.test.ts b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerSpaceNode.test.ts index b0fc6d78c0f..81d8d844bf9 100644 --- a/packages/core/src/test/awsService/sagemaker/explorer/sagemakerSpaceNode.test.ts +++ b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerSpaceNode.test.ts @@ -112,11 +112,13 @@ describe('SagemakerSpaceNode', function () { it('updates space app status', async function () { const describeSpaceStub = sinon.stub(SagemakerClient.prototype, 'describeSpace') describeSpaceStub.resolves({ SpaceName: 'TestSpace', Status: 'InService', $metadata: {} }) - describeAppStub.resolves({ AppName: 'TestApp', Status: 'InService', $metadata: {} }) + + const listAppForSpaceStub = sinon.stub(SagemakerClient.prototype, 'listAppForSpace') + listAppForSpaceStub.resolves({ AppName: 'TestApp', Status: 'InService' }) await testSpaceAppNode.updateSpaceAppStatus() sinon.assert.calledOnce(describeSpaceStub) - sinon.assert.calledOnce(describeAppStub) + sinon.assert.calledOnce(listAppForSpaceStub) }) }) diff --git a/packages/core/src/test/awsService/sagemaker/sagemakerSpace.test.ts b/packages/core/src/test/awsService/sagemaker/sagemakerSpace.test.ts index 2a52b08a3a6..12db8b9c6f6 100644 --- a/packages/core/src/test/awsService/sagemaker/sagemakerSpace.test.ts +++ b/packages/core/src/test/awsService/sagemaker/sagemakerSpace.test.ts @@ -63,6 +63,7 @@ describe('SagemakerSpace', function () { mockClient.describeSpace.resolves(mockDescribeSpaceResponse) mockClient.describeApp.resolves(mockDescribeAppResponse) + mockClient.listAppForSpace.resolves(mockDescribeAppResponse) const space = new SagemakerSpace(mockClient as any, 'us-east-1', mockSpaceApp) const updateSpaceSpy = sinon.spy(space, 'updateSpace') @@ -107,9 +108,8 @@ describe('SagemakerSpace', function () { Status: 'InService', $metadata: { requestId: 'test-request-id' }, } - + mockClient.listAppForSpace.resolves(mockDescribeAppResponse) mockClient.describeSpace.resolves(mockDescribeSpaceResponse) - mockClient.describeApp.resolves(mockDescribeAppResponse) const space = new SagemakerSpace(mockClient as any, 'us-east-1', mockSpaceApp) const updateSpaceSpy = sinon.spy(space, 'updateSpace') @@ -125,5 +125,37 @@ describe('SagemakerSpace', function () { assert.strictEqual(updateSpaceArgs.OwnershipSettingsSummary, undefined) assert.strictEqual(updateSpaceArgs.SpaceSharingSettingsSummary, undefined) }) + + it('should update app status using listAppForSpace', async function () { + const mockDescribeSpaceResponse = { + SpaceName: 'test-space', + Status: 'InService', + DomainId: 'test-domain', + $metadata: { requestId: 'test-request-id' }, + } + + const mockAppFromList = { + AppName: 'listed-app', + Status: 'InService', + $metadata: { requestId: 'test-request-id' }, + } + + mockClient.describeSpace.resolves(mockDescribeSpaceResponse) + mockClient.listAppForSpace.resolves(mockAppFromList) + + // Create space without App.AppName + const spaceWithoutAppName: SagemakerSpaceApp = { + ...mockSpaceApp, + App: undefined, + } + + const space = new SagemakerSpace(mockClient as any, 'us-east-1', spaceWithoutAppName) + await space.updateSpaceAppStatus() + + // Verify listAppForSpace was called instead of describeApp + assert.ok(mockClient.listAppForSpace.calledOnce) + assert.ok(mockClient.listAppForSpace.calledWith('test-domain', 'test-space')) + assert.ok(mockClient.describeApp.notCalled) + }) }) }) diff --git a/packages/core/src/test/shared/clients/sagemakerClient.test.ts b/packages/core/src/test/shared/clients/sagemakerClient.test.ts index f748c066b23..0ebda0eeb83 100644 --- a/packages/core/src/test/shared/clients/sagemakerClient.test.ts +++ b/packages/core/src/test/shared/clients/sagemakerClient.test.ts @@ -173,6 +173,41 @@ describe('SagemakerClient.listSpaceApps', function () { }) }) +describe('SagemakerClient.listAppForSpace', function () { + const region = 'test-region' + let client: SagemakerClient + let listAppsStub: sinon.SinonStub + + beforeEach(function () { + client = new SagemakerClient(region) + listAppsStub = sinon.stub(client, 'listApps') + }) + + afterEach(function () { + sinon.restore() + }) + + it('returns first app for given domain and space', async function () { + const appDetails: AppDetails[] = [ + { AppName: 'app1', DomainId: 'domain1', SpaceName: 'space1', AppType: AppType.CodeEditor }, + ] + listAppsStub.returns(intoCollection([appDetails])) + + const result = await client.listAppForSpace('domain1', 'space1') + + assert.strictEqual(result?.AppName, 'app1') + sinon.assert.calledWith(listAppsStub, { DomainIdEquals: 'domain1', SpaceNameEquals: 'space1' }) + }) + + it('returns undefined when no apps found', async function () { + listAppsStub.returns(intoCollection([[]])) + + const result = await client.listAppForSpace('domain1', 'space1') + + assert.strictEqual(result, undefined) + }) +}) + describe('SagemakerClient.waitForAppInService', function () { const region = 'test-region' let client: SagemakerClient