fix(smus): unble to refresh status on newly created App bug

Author: PotatoWKY

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,18 @@ 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
+        let app
+        if (this.spaceApp.DomainId && this.spaceApp.SpaceName) {
+            app = await this.client.listAppForSpace(this.spaceApp.DomainId, this.spaceApp.SpaceName)
+        } else {
+            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 +207,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/'

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

@@ -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
 

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)
     })
 })

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')
@@ -125,5 +126,37 @@ describe('SagemakerSpace', function () {
             assert.strictEqual(updateSpaceArgs.OwnershipSettingsSummary, undefined)
             assert.strictEqual(updateSpaceArgs.SpaceSharingSettingsSummary, undefined)
         })
+
+        it('should use listAppForSpace when AppName is not available', 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)
+        })
     })
 })

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