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

@@ -1,5 +1,4 @@
-/*!
- * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+/*! * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -59,6 +58,12 @@ const appTypeSettingsMap: Record<string, string> = {
     [AppType.CodeEditor as string]: 'CodeEditorAppSettings',
 } as const
 
+export const waitForAppConfig = {
+    softTimeoutRetries: 12,
+    hardTimeoutRetries: 120,
+    intervalMs: 5000,
+}
+
 export interface SagemakerSpaceApp extends SpaceDetails {
     App?: AppDetails
     DomainSpaceKey: string
@@ -364,10 +369,9 @@ 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++) {
+        for (let attempt = 0; attempt < waitForAppConfig.hardTimeoutRetries; attempt++) {
             const { Status } = await this.describeApp({
                 DomainId: domainId,
                 SpaceName: spaceName,
@@ -383,7 +387,13 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
                 throw new ToolkitError(`App failed to start. Status: ${Status}`)
             }
 
-            await sleep(intervalMs)
+            if (attempt === waitForAppConfig.softTimeoutRetries) {
+                progress?.report({
+                    message: `Starting the space is taking longer than usual. The space will connect when ready`,
+                })
+            }
+
+            await sleep(waitForAppConfig.intervalMs)
         }
 
         throw new ToolkitError(`Timed out waiting for app "${spaceName}" to reach "InService" status.`)

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

@@ -251,10 +251,18 @@ describe('SagemakerClient.waitForAppInService', function () {
     it('times out after max retries', async function () {
         describeAppStub.resolves({ Status: 'Pending' })
 
-        await assert.rejects(
-            client.waitForAppInService('domain1', 'space1', 'CodeEditor', 2, 10),
-            /Timed out waiting for app/
-        )
+        const sagemakerModule = await import('../../../shared/clients/sagemaker.js')
+        const originalValue = sagemakerModule.waitForAppConfig.hardTimeoutRetries
+        sagemakerModule.waitForAppConfig.hardTimeoutRetries = 3
+
+        try {
+            await assert.rejects(
+                client.waitForAppInService('domain1', 'space1', 'CodeEditor'),
+                /Timed out waiting for app/
+            )
+        } finally {
+            sagemakerModule.waitForAppConfig.hardTimeoutRetries = originalValue
+        }
     })
 })