Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 20 additions & 8 deletions packages/core/src/awsService/sagemaker/sagemakerSpace.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 = ''
Expand All @@ -34,6 +37,10 @@ export class SagemakerSpace {
}

public updateSpace(spaceApp: SagemakerSpaceApp) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When all is this invoked and why does it help in this case?

// 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) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add a comment to clarify when this will happen, finding it hard to understand.

I know not from this change but spaceApp.App just feels very confusing, if there is scope to rename to better variables, we should.

Copy link
Contributor Author

@PotatoWKY PotatoWKY Nov 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

True, do you have any suggestions to replace "spaceApp", basically spaceApp is a space that have a App as one of the variables attached to it.

we would also need to change the name for the SagemakerSpaceApp type:

export interface SagemakerSpaceApp extends SpaceDetails {
    App?: AppDetails
    DomainSpaceKey: string
}

Copy link
Contributor Author

@PotatoWKY PotatoWKY Nov 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what about type: "SagemakerSpaceDetails", variable just: "space" or "spaceDetails"?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just noticed that SagemakerSpace is the class name, so the current naming can be more confusing, like:

SagemakerSpace.spaceApp.App.AppName

a lot of spaces and apps

what about:
SagemakerSpace.spaceDetails.appDetails.AppName

this.spaceApp.App = spaceApp.App
}
Comment on lines +41 to +43
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Under what circumstances will this apply, and what problem is this addressing?

Copy link
Contributor Author

@PotatoWKY PotatoWKY Nov 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this resolve the edge case when this.spaceApp.App is null.
where it is unable to update the status of the app if app is null because current logic to update the app is only updating it's status as a string, if the App object itself is null, then we can not update its status

public setSpaceStatus(spaceStatus: string, appStatus: string) {
    this.spaceApp.Status = spaceStatus
    if (this.spaceApp.App) {
        this.spaceApp.App.Status = appStatus
    }
}

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) {
Expand Down Expand Up @@ -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}]`
)
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

else? Are there other possible cases?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are no other cases, gonna add error handling for else case


// AWS DescribeSpace API returns full details with property names like 'SpaceSettings'
// but our internal SagemakerSpaceApp type expects 'SpaceSettingsSummary' (from ListSpaces API)
Expand Down Expand Up @@ -195,7 +208,6 @@ export class SagemakerSpace {
* Sets up user activity monitoring for SageMaker spaces
*/
export async function setupUserActivityMonitoring(extensionContext: vscode.ExtensionContext): Promise<void> {
const logger = getLogger()
logger.info('setupUserActivityMonitoring: Starting user activity monitoring setup')

const tmpDirectory = '/tmp/'
Expand Down
7 changes: 7 additions & 0 deletions packages/core/src/shared/clients/sagemaker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,13 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
return this.makeRequest(DeleteAppCommand, request)
}

public async listAppForSpace(domainId: string, spaceName: string): Promise<AppDetails | undefined> {
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

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
})
})
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down Expand Up @@ -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')
Expand All @@ -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)
})
})
})
35 changes: 35 additions & 0 deletions packages/core/src/test/shared/clients/sagemakerClient.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading