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,12 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
         domainId: string,
         spaceName: string,
         appType: string,
-        maxRetries = 30,
+        progress?: vscode.Progress<{ message?: string; increment?: number }>,
+        softTimeoutRetries = 12, // 1 minute
+        hardTimeoutRetries = 120, // 10 minutes
         intervalMs = 5000
     ): Promise<void> {
-        for (let attempt = 0; attempt < maxRetries; attempt++) {
+        for (let attempt = 0; attempt < hardTimeoutRetries; attempt++) {
             const { Status } = await this.describeApp({
                 DomainId: domainId,
                 SpaceName: spaceName,
@@ -383,6 +385,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. Will connect when ready.`,
+                })
+            }
+
             await sleep(intervalMs)
         }
 

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

@@ -252,7 +252,7 @@ describe('SagemakerClient.waitForAppInService', function () {
         describeAppStub.resolves({ Status: 'Pending' })
 
         await assert.rejects(
-            client.waitForAppInService('domain1', 'space1', 'CodeEditor', 2, 10),
+            client.waitForAppInService('domain1', 'space1', 'CodeEditor', undefined, 2, 2, 10),
             /Timed out waiting for app/
         )
     })