fix(smus): Improve error handling when the Space takes too long to start

Author: PotatoWKY

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

@@ -332,7 +332,8 @@ async function handleRunningSpaceWithDisabledAccess(
                 await client.waitForAppInService(
                     node.spaceApp.DomainId!,
                     spaceName,
-                    node.spaceApp.SpaceSettingsSummary!.AppType!
+                    node.spaceApp.SpaceSettingsSummary!.AppType!,
+                    progress
                 )
                 await tryRemoteConnection(node, ctx, progress)
             } catch (err: any) {
@@ -376,7 +377,8 @@ async function handleStoppedSpace(
                 await client.waitForAppInService(
                     node.spaceApp.DomainId!,
                     spaceName,
-                    node.spaceApp.SpaceSettingsSummary!.AppType!
+                    node.spaceApp.SpaceSettingsSummary!.AppType!,
+                    progress
                 )
                 await tryRemoteConnection(node, ctx, progress)
             }

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

@@ -364,10 +364,13 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
         domainId: string,
         spaceName: string,
         appType: string,
-        maxRetries = 30,
-        intervalMs = 5000
+        progress?: vscode.Progress<{ message?: string; increment?: number }>
     ): Promise<void> {
-        for (let attempt = 0; attempt < maxRetries; attempt++) {
+        const softTimeoutRetries = 12 // 1 minute
+        const hardTimeoutRetries = 120 // 10 minutes
+        const intervalMs = 5000
+
+        for (let attempt = 0; attempt < hardTimeoutRetries; attempt++) {
             const { Status } = await this.describeApp({
                 DomainId: domainId,
                 SpaceName: spaceName,
@@ -383,6 +386,12 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
                 throw new ToolkitError(`App failed to start. Status: ${Status}`)
             }
 
+            if (attempt === softTimeoutRetries) {
+                progress?.report({
+                    message: `Starting the space is taking longer than usual. The space will connect when ready`,
+                })
+            }
+
             await sleep(intervalMs)
         }
 

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

@@ -249,12 +249,31 @@ describe('SagemakerClient.waitForAppInService', function () {
     })
 
     it('times out after max retries', async function () {
-        describeAppStub.resolves({ Status: 'Pending' })
+        let callCount = 0
+        describeAppStub.callsFake(async () => {
+            callCount++
+            return { Status: 'Pending' }
+        })
 
-        await assert.rejects(
-            client.waitForAppInService('domain1', 'space1', 'CodeEditor', 2, 10),
-            /Timed out waiting for app/
-        )
+        // Stub globals.clock.setTimeout to make sleep instant
+        const globalsModule = await import('../../../shared/extensionGlobals.js')
+        const globals = globalsModule.default as any
+        const originalSetTimeout = globals.clock.setTimeout
+        globals.clock.setTimeout = ((cb: any) => {
+            // eslint-disable-next-line @typescript-eslint/no-implied-eval
+            setImmediate(cb)
+            return 0 as any
+        }) as any
+
+        try {
+            await assert.rejects(
+                client.waitForAppInService('domain1', 'space1', 'CodeEditor'),
+                /Timed out waiting for app/
+            )
+            assert.strictEqual(callCount, 120)
+        } finally {
+            globals.clock.setTimeout = originalSetTimeout
+        }
     })
 })