fix(sagemaker): Add code-editor subdomain to refreshUrl (#7826)

## Problem
-All space reconnection redirect links were hardcoded to route users to
the JupyterLab details page, due to a bug in the Studio web application.


## Solution
- With the bug now fixed in SageMaker Studio, we can support redirecting
to the Code Editor details page via refreshUrl.
- For backward compatibility, the default remains set to JupyterLab.
Once the fix is verified and rolled out across all regions, we can
remove the default behavior and instead explicitly check for the appType
(e.g., JupyterLab) before setting the redirect.

---

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

---------

Co-authored-by: Newton Der <[email protected]>
Co-authored-by: Newton Der <[email protected]>
Co-authored-by: Laxman Reddy <[email protected]>

Author: 4 people

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

@@ -86,7 +86,8 @@ export async function deeplinkConnect(
     session: string,
     wsUrl: string,
     token: string,
-    domain: string
+    domain: string,
+    appType?: string
 ) {
     getLogger().debug(
         `sm:deeplinkConnect: connectionIdentifier: ${connectionIdentifier} session: ${session} wsUrl: ${wsUrl} token: ${token}`
@@ -107,7 +108,8 @@ export async function deeplinkConnect(
             session,
             wsUrl,
             token,
-            domain
+            domain,
+            appType
         )
 
         await startVscodeRemote(

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

@@ -96,15 +96,16 @@ export async function persistSSMConnection(
     domain: string,
     session?: string,
     wsUrl?: string,
-    token?: string
+    token?: string,
+    appType?: string
 ): Promise<void> {
     const { region } = parseArn(spaceArn)
     const endpoint = DevSettings.instance.get('endpoints', {})['sagemaker'] ?? ''
 
-    // 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 appSubDomain = 'jupyterlab'
+    if (appType && appType.toLowerCase() === 'codeeditor') {
+        appSubDomain = 'code-editor'
+    }
 
     let envSubdomain: string
 

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

@@ -56,7 +56,8 @@ export async function prepareDevEnvConnection(
     session?: string,
     wsUrl?: string,
     token?: string,
-    domain?: string
+    domain?: string,
+    appType?: string
 ) {
     const remoteLogger = configureRemoteConnectionLogger()
     const { ssm, vsc, ssh } = (await ensureDependencies()).unwrap()
@@ -82,7 +83,7 @@ export async function prepareDevEnvConnection(
             await persistSmusProjectCreds(spaceArn, node as SagemakerUnifiedStudioSpaceNode)
         }
     } else if (connectionType === 'sm_dl') {
-        await persistSSMConnection(spaceArn, domain ?? '', session, wsUrl, token)
+        await persistSSMConnection(spaceArn, domain ?? '', session, wsUrl, token, appType)
     }
 
     await startLocalServer(ctx)

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

@@ -18,7 +18,8 @@ export function register(ctx: ExtContext) {
                 params.session,
                 `${params.ws_url}&cell-number=${params['cell-number']}`,
                 params.token,
-                params.domain
+                params.domain,
+                params.app_type
             )
         })
     }
@@ -27,7 +28,7 @@ export function register(ctx: ExtContext) {
 }
 
 export function parseConnectParams(query: SearchParams) {
-    const params = query.getFromKeysOrThrow(
+    const requiredParams = query.getFromKeysOrThrow(
         'connection_identifier',
         'domain',
         'user_profile',
@@ -36,5 +37,7 @@ export function parseConnectParams(query: SearchParams) {
         'cell-number',
         'token'
     )
-    return params
+    const optionalParams = query.getFromKeys('app_type')
+
+    return { ...requiredParams, ...optionalParams }
 }

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

@@ -44,6 +44,7 @@ describe('SageMaker URI handler', function () {
             ws_url: 'wss://example.com',
             'cell-number': '4',
             token: 'my-token',
+            app_type: 'jupyterlab',
         }
 
         const uri = createConnectUri(params)
@@ -55,5 +56,24 @@ describe('SageMaker URI handler', function () {
         assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[3], 'wss://example.com&cell-number=4')
         assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[4], 'my-token')
         assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[5], 'my-domain')
+        assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[6], 'jupyterlab')
+    })
+
+    it('calls deeplinkConnect with undefined app_type when not provided', async function () {
+        const params = {
+            connection_identifier: 'abc123',
+            domain: 'my-domain',
+            user_profile: 'me',
+            session: 'sess-xyz',
+            ws_url: 'wss://example.com',
+            'cell-number': '4',
+            token: 'my-token',
+        }
+
+        const uri = createConnectUri(params)
+        await handler.handleUri(uri)
+
+        assert.ok(deeplinkConnectStub.calledOnce)
+        assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[6], undefined)
     })
 })