fix(sagemaker): GetStatus error when refreshing large number of spaces and fix deeplink reconnect (#2161)

## Problem

- When there are a large number of SageMaker space nodes, refreshing the
Explorer tree can take time. During this window, the underlying node
objects may be temporarily undefined, causing actions like "Connect" or
"Stop" to fail. This happens because all space nodes were previously
being refreshed whenever an action was taken on any one node, or when
the Explorer tree was refreshed.

- Deeplink reconnect can fail if the user is not logged in with the same
credentials originally used to describe the space. This is because the
Toolkit currently makes a `DescribeSpace` call using local credentials.
Ideally, Studio UI should pass the `appType` directly in the deeplink,
avoiding the need for Toolkit to make this call. However, since changes
to Studio UI can’t be deployed before the NY Summit, this update will be
handled in a future MFE release.

## Solution

- Updated the logic to **refresh only the specific space node** being
acted on, instead of refreshing all nodes. This avoids unnecessary
delays and reduces the likelihood of undefined node states during
actions.

- Added a **warning message** when the full space list is being
refreshed. If a user tries to interact with a space during this time,
they will see a message indicating that space information is still
loading and to try again shortly.

- Temporarily **hardcoded the `appType` to `JupyterLab`** in the
reconnect URI for all app types. Reconnection will still work for both
Code Editor and JupyterLab, although the URL path will always show
`/jupyterlab`. This is a temporary workaround until Studio UI can send
the correct `appType`.

- Added **telemetry for deeplink connect**


---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

Author: aws-asolidu

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

@@ -4,6 +4,7 @@
  */
 
 import * as path from 'path'
+import * as vscode from 'vscode'
 import { Commands } from '../../shared/vscode/commands2'
 import { SagemakerSpaceNode } from './explorer/sagemakerSpaceNode'
 import { SagemakerParentNode } from './explorer/sagemakerParentNode'
@@ -20,6 +21,9 @@ export async function activate(ctx: ExtContext): Promise<void> {
     ctx.extensionContext.subscriptions.push(
         uriHandlers.register(ctx),
         Commands.register('aws.sagemaker.openRemoteConnection', async (node: SagemakerSpaceNode) => {
+            if (!validateNode(node)) {
+                return
+            }
             await telemetry.sagemaker_openRemoteConnection.run(async () => {
                 await openRemoteConnect(node, ctx.extensionContext)
             })
@@ -32,6 +36,9 @@ export async function activate(ctx: ExtContext): Promise<void> {
         }),
 
         Commands.register('aws.sagemaker.stopSpace', async (node: SagemakerSpaceNode) => {
+            if (!validateNode(node)) {
+                return
+            }
             await telemetry.sagemaker_stopSpace.run(async () => {
                 await stopSpace(node, ctx.extensionContext)
             })
@@ -62,3 +69,14 @@ export async function activate(ctx: ExtContext): Promise<void> {
         })
     }
 }
+
+/**
+ * Checks if a node  is undefined and shows a warning message if so.
+ */
+function validateNode(node: unknown): boolean {
+    if (!node) {
+        void vscode.window.showWarningMessage('Space information is being refreshed. Please try again shortly.')
+        return false
+    }
+    return true
+}

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

@@ -12,8 +12,6 @@ import { DevSettings } from '../../shared/settings'
 import { Auth } from '../../auth/auth'
 import { SpaceMappings, SsmConnectionInfo } from './types'
 import { getLogger } from '../../shared/logger/logger'
-import { SagemakerClient } from '../../shared/clients/sagemaker'
-import { AppType } from '@amzn/sagemaker-client'
 import { parseArn } from './detached-server/utils'
 
 const mappingFileName = '.sagemaker-space-profiles'
@@ -83,25 +81,13 @@ export async function persistSSMConnection(
     wsUrl?: string,
     token?: string
 ): Promise<void> {
-    const { region, spaceName } = parseArn(appArn)
+    const { region } = parseArn(appArn)
     const endpoint = DevSettings.instance.get('endpoints', {})['sagemaker'] ?? ''
-    const client = new SagemakerClient(region)
 
-    const spaceDetails = await client.describeSpace({
-        DomainId: domain,
-        SpaceName: spaceName,
-    })
-
-    let appSubDomain: string
-    if (spaceDetails.SpaceSettings?.AppType === AppType.JupyterLab) {
-        appSubDomain = 'jupyterlab'
-    } else if (spaceDetails.SpaceSettings?.AppType === AppType.CodeEditor) {
-        appSubDomain = 'code-editor'
-    } else {
-        throw new ToolkitError(
-            `Unsupported or missing app type for space. Expected JupyterLab or CodeEditor, got: ${spaceDetails.SpaceSettings?.AppType ?? 'undefined'}`
-        )
-    }
+    // TODO: Hardcoded to 'jupyterlab' due to a bug in Studio that only supports refreshing
+    // the token for both CodeEditor and JupyterLab Apps in the jupyterlab subdomain.
+    // This will be fixed shortly after NYSummit launch to support refresh URL in CodeEditor subdomain.
+    const appSubDomain = 'jupyterlab'
 
     let envSubdomain: string
 

packages/core/src/awsService/sagemaker/explorer/sagemakerSpaceNode.ts

@@ -162,15 +162,17 @@ export class SagemakerSpaceNode extends AWSTreeNodeBase implements AWSResourceNo
 
     public async refreshNode(): Promise<void> {
         await this.updateSpaceAppStatus()
-        await tryRefreshNode(this)
+        await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', this)
     }
 }
 
 export async function tryRefreshNode(node?: SagemakerSpaceNode) {
     if (node) {
-        const n = node instanceof SagemakerSpaceNode ? node.parent : node
         try {
-            await n.refreshNode()
+            // For SageMaker spaces, refresh just the individual space node to avoid expensive
+            // operation of refreshing all spaces in the domain
+            await node.updateSpaceAppStatus()
+            await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', node)
         } catch (e) {
             getLogger().error('refreshNode failed: %s', (e as Error).message)
         }

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

@@ -121,7 +121,7 @@ export async function startLocalServer(ctx: vscode.ExtensionContext) {
     const errLog = path.join(storagePath, 'sagemaker-local-server.err.log')
     const infoFilePath = path.join(storagePath, 'sagemaker-local-server-info.json')
 
-    logger.info(`local server logs at ${storagePath}/sagemaker-local-server.*.log`)
+    logger.info(`sagemaker-local-server.*.log at ${storagePath}`)
 
     const customEndpoint = DevSettings.instance.get('endpoints', {})['sagemaker']
 

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

@@ -7,17 +7,20 @@ import * as vscode from 'vscode'
 import { SearchParams } from '../../shared/vscode/uriHandler'
 import { deeplinkConnect } from './commands'
 import { ExtContext } from '../../shared/extensions'
+import { telemetry } from '../../shared/telemetry/telemetry'
 
 export function register(ctx: ExtContext) {
     async function connectHandler(params: ReturnType<typeof parseConnectParams>) {
-        await deeplinkConnect(
-            ctx,
-            params.connection_identifier,
-            params.session,
-            `${params.ws_url}&cell-number=${params['cell-number']}`,
-            params.token,
-            params.domain
-        )
+        await telemetry.sagemaker_deeplinkConnect.run(async () => {
+            await deeplinkConnect(
+                ctx,
+                params.connection_identifier,
+                params.session,
+                `${params.ws_url}&cell-number=${params['cell-number']}`,
+                params.token,
+                params.domain
+            )
+        })
     }
 
     return vscode.Disposable.from(ctx.uriHandler.onPath('/connect/sagemaker', connectHandler, parseConnectParams))

packages/core/src/shared/telemetry/vscodeTelemetry.json

@@ -268,6 +268,15 @@
                 }
             ]
         },
+        {
+            "name": "sagemaker_deeplinkConnect",
+            "description": "Connect to SageMake Space via a deeplink",
+            "metadata": [
+                {
+                    "type": "result"
+                }
+            ]
+        },
         {
             "name": "amazonq_didSelectProfile",
             "description": "Emitted after the user's Q Profile has been set, whether the user was prompted with a dialog, or a profile was automatically assigned after signing in.",

packages/core/src/test/awsService/sagemaker/credentialMapping.test.ts

@@ -181,7 +181,12 @@ describe('credentialMapping', () => {
             )
         })
 
-        it('throws error when app type is unsupported', async () => {
+        // TODO: Skipped due to hardcoded appSubDomain. Currently hardcoded to 'jupyterlab' due to
+        // a bug in Studio that only supports refreshing the token for both CodeEditor and JupyterLab
+        // Apps in the jupyterlab subdomain. This will be fixed shortly after NYSummit launch to
+        // support refresh URL in CodeEditor subdomain. Additionally, appType will be determined by
+        // the deeplink URI rather than the describeSpace call from the toolkit.
+        it.skip('throws error when app type is unsupported', async () => {
             sandbox.stub(DevSettings.instance, 'get').returns({})
             sandbox.stub(fs, 'existsFile').resolves(false)