From 1553470c83e85c485fe0cd6743da810d7ca8933e Mon Sep 17 00:00:00 2001 From: Alvin Solidum Date: Wed, 6 Aug 2025 07:57:42 -0700 Subject: [PATCH 1/4] add codeeditor subdomain --- packages/core/src/awsService/sagemaker/commands.ts | 6 ++++-- .../src/awsService/sagemaker/credentialMapping.ts | 11 ++++++----- packages/core/src/awsService/sagemaker/model.ts | 5 +++-- packages/core/src/awsService/sagemaker/uriHandlers.ts | 6 ++++-- 4 files changed, 17 insertions(+), 11 deletions(-) diff --git a/packages/core/src/awsService/sagemaker/commands.ts b/packages/core/src/awsService/sagemaker/commands.ts index 0075d7e5dff..9b08fb0c66f 100644 --- a/packages/core/src/awsService/sagemaker/commands.ts +++ b/packages/core/src/awsService/sagemaker/commands.ts @@ -85,7 +85,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}` @@ -104,7 +105,8 @@ export async function deeplinkConnect( session, wsUrl, token, - domain + domain, + appType ) await startVscodeRemote( diff --git a/packages/core/src/awsService/sagemaker/credentialMapping.ts b/packages/core/src/awsService/sagemaker/credentialMapping.ts index 60d4e94260e..6567a3fa223 100644 --- a/packages/core/src/awsService/sagemaker/credentialMapping.ts +++ b/packages/core/src/awsService/sagemaker/credentialMapping.ts @@ -79,15 +79,16 @@ export async function persistSSMConnection( domain: string, session?: string, wsUrl?: string, - token?: string + token?: string, + appType?: string ): Promise { const { region } = parseArn(appArn) 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 diff --git a/packages/core/src/awsService/sagemaker/model.ts b/packages/core/src/awsService/sagemaker/model.ts index 20a667a0bfa..041cb155692 100644 --- a/packages/core/src/awsService/sagemaker/model.ts +++ b/packages/core/src/awsService/sagemaker/model.ts @@ -50,7 +50,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() @@ -72,7 +73,7 @@ export async function prepareDevEnvConnection( if (connectionType === 'sm_lc') { await persistLocalCredentials(appArn) } else if (connectionType === 'sm_dl') { - await persistSSMConnection(appArn, domain ?? '', session, wsUrl, token) + await persistSSMConnection(appArn, domain ?? '', session, wsUrl, token, appType) } await startLocalServer(ctx) diff --git a/packages/core/src/awsService/sagemaker/uriHandlers.ts b/packages/core/src/awsService/sagemaker/uriHandlers.ts index 17c3c512272..f40ea301510 100644 --- a/packages/core/src/awsService/sagemaker/uriHandlers.ts +++ b/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 ) }) } @@ -34,7 +35,8 @@ export function parseConnectParams(query: SearchParams) { 'session', 'ws_url', 'cell-number', - 'token' + 'token', + 'app_type' ) return params } From 5fff98a3f3b70bb7c7909cc92b21a2215c123fe5 Mon Sep 17 00:00:00 2001 From: Alvin Solidum Date: Wed, 6 Aug 2025 08:15:14 -0700 Subject: [PATCH 2/4] fix test --- packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts b/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts index 9ff24b2a3f9..caae053443e 100644 --- a/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts +++ b/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) From 9e98f85d8a741478d051b1d6988f7ae1e7bd5be9 Mon Sep 17 00:00:00 2001 From: Alvin Solidum Date: Wed, 6 Aug 2025 08:26:24 -0700 Subject: [PATCH 3/4] make app_type optional --- .../core/src/awsService/sagemaker/commands.ts | 2 +- .../src/awsService/sagemaker/uriHandlers.ts | 9 +++---- .../awsService/sagemaker/uriHandlers.test.ts | 24 +++++++++++++++++++ 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/packages/core/src/awsService/sagemaker/commands.ts b/packages/core/src/awsService/sagemaker/commands.ts index 9b08fb0c66f..8ed485dc9dc 100644 --- a/packages/core/src/awsService/sagemaker/commands.ts +++ b/packages/core/src/awsService/sagemaker/commands.ts @@ -86,7 +86,7 @@ export async function deeplinkConnect( wsUrl: string, token: string, domain: string, - appType: string + appType?: string ) { getLogger().debug( `sm:deeplinkConnect: connectionIdentifier: ${connectionIdentifier} session: ${session} wsUrl: ${wsUrl} token: ${token}` diff --git a/packages/core/src/awsService/sagemaker/uriHandlers.ts b/packages/core/src/awsService/sagemaker/uriHandlers.ts index f40ea301510..6f1143d9054 100644 --- a/packages/core/src/awsService/sagemaker/uriHandlers.ts +++ b/packages/core/src/awsService/sagemaker/uriHandlers.ts @@ -28,15 +28,16 @@ export function register(ctx: ExtContext) { } export function parseConnectParams(query: SearchParams) { - const params = query.getFromKeysOrThrow( + const requiredParams = query.getFromKeysOrThrow( 'connection_identifier', 'domain', 'user_profile', 'session', 'ws_url', 'cell-number', - 'token', - 'app_type' + 'token' ) - return params + const optionalParams = query.getFromKeys('app_type') + + return { ...requiredParams, ...optionalParams } } diff --git a/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts b/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts index caae053443e..07e20e424b6 100644 --- a/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts +++ b/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts @@ -56,5 +56,29 @@ 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[1], 'abc123') + assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[2], 'sess-xyz') + 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], undefined) }) }) From 9f2101f1eaa3167a76245537861d9b47bee5696a Mon Sep 17 00:00:00 2001 From: Alvin Solidum Date: Wed, 27 Aug 2025 17:18:04 -0700 Subject: [PATCH 4/4] remove dupe --- .../core/src/test/awsService/sagemaker/uriHandlers.test.ts | 5 ----- 1 file changed, 5 deletions(-) diff --git a/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts b/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts index 07e20e424b6..f27df1fcb11 100644 --- a/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts +++ b/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts @@ -74,11 +74,6 @@ describe('SageMaker URI handler', function () { await handler.handleUri(uri) assert.ok(deeplinkConnectStub.calledOnce) - assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[1], 'abc123') - assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[2], 'sess-xyz') - 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], undefined) }) })