Merge branch 'master' into fix/SSH-host

Author: bhavya2109sharma

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

@@ -115,10 +115,16 @@ export async function deeplinkConnect(
     wsUrl: string,
     token: string,
     domain: string,
-    appType?: string
+    appType?: string,
+    isSMUS: boolean = false
 ) {
     getLogger().debug(
-        `sm:deeplinkConnect: connectionIdentifier: ${connectionIdentifier} session: ${session} wsUrl: ${wsUrl} token: ${token}`
+        'sm:deeplinkConnect: connectionIdentifier: %s session: %s wsUrl: %s token: %s isSMUS: %s',
+        connectionIdentifier,
+        session,
+        wsUrl,
+        token,
+        isSMUS
     )
 
     if (isRemoteWorkspace()) {
@@ -134,7 +140,7 @@ export async function deeplinkConnect(
             connectionIdentifier,
             ctx.extensionContext,
             'sm_dl',
-            false /* isSMUS */,
+            isSMUS,
             undefined /* node */,
             session,
             wsUrl,
@@ -152,7 +158,10 @@ export async function deeplinkConnect(
         )
     } catch (err: any) {
         getLogger().error(
-            `sm:OpenRemoteConnect: Unable to connect to target space with arn: ${connectionIdentifier} error: ${err}`
+            'sm:OpenRemoteConnect: Unable to connect to target space with arn: %s error: %s isSMUS: %s',
+            connectionIdentifier,
+            err,
+            isSMUS
         )
 
         if (![RemoteSessionError.MissingExtension, RemoteSessionError.ExtensionVersionTooLow].includes(err.code)) {
@@ -357,7 +366,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) {
@@ -401,7 +411,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/awsService/sagemaker/constants.ts

@@ -61,3 +61,10 @@ export const SshConfigRemovalFailedMessage = (configHostName: string) =>
 
 export const SshConfigUpdateFailedMessage = (configPath: string, configHostName: string) =>
     `Failed to update SSH config section. Fix your ${configPath} file manually or remove the outdated ${configHostName} section.`
+export const SmusDeeplinkSessionExpiredError = {
+    title: 'Session Disconnected',
+    message:
+        'Your SageMaker Unified Studio session has been disconnected. Select a local (non-remote) VS Code window and use the SageMaker Unified Studio portal to connect again.',
+    code: 'SMUS_SESSION_DISCONNECTED',
+    shortMessage: 'Session disconnected, re-connect from SageMaker Unified Studio portal.',
+} as const

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

@@ -90,41 +90,51 @@ export async function persistSmusProjectCreds(spaceArn: string, node: SagemakerU
  * @param session - SSM session ID.
  * @param wsUrl - SSM WebSocket URL.
  * @param token - Bearer token for the session.
+ * @param appType - Application type (e.g., 'jupyterlab', 'codeeditor').
+ * @param isSMUS - If true, skip refreshUrl construction (SMUS connections cannot refresh).
  */
 export async function persistSSMConnection(
     spaceArn: string,
     domain: string,
     session?: string,
     wsUrl?: string,
     token?: string,
-    appType?: string
+    appType?: string,
+    isSMUS?: boolean
 ): Promise<void> {
-    const { region } = parseArn(spaceArn)
-    const endpoint = DevSettings.instance.get('endpoints', {})['sagemaker'] ?? ''
+    let refreshUrl: string | undefined
 
-    let appSubDomain = 'jupyterlab'
-    if (appType && appType.toLowerCase() === 'codeeditor') {
-        appSubDomain = 'code-editor'
-    }
+    if (!isSMUS) {
+        // Construct refreshUrl for SageMaker AI connections
+        const { region } = parseArn(spaceArn)
+        const endpoint = DevSettings.instance.get('endpoints', {})['sagemaker'] ?? ''
 
-    let envSubdomain: string
+        let appSubDomain = 'jupyterlab'
+        if (appType && appType.toLowerCase() === 'codeeditor') {
+            appSubDomain = 'code-editor'
+        }
 
-    if (endpoint.includes('beta')) {
-        envSubdomain = 'devo'
-    } else if (endpoint.includes('gamma')) {
-        envSubdomain = 'loadtest'
-    } else {
-        envSubdomain = 'studio'
-    }
+        let envSubdomain: string
 
-    // Use the standard AWS domain for 'studio' (prod).
-    // For non-prod environments, use the obfuscated domain 'asfiovnxocqpcry.com'.
-    const baseDomain =
-        envSubdomain === 'studio'
-            ? `studio.${region}.sagemaker.aws`
-            : `${envSubdomain}.studio.${region}.asfiovnxocqpcry.com`
+        if (endpoint.includes('beta')) {
+            envSubdomain = 'devo'
+        } else if (endpoint.includes('gamma')) {
+            envSubdomain = 'loadtest'
+        } else {
+            envSubdomain = 'studio'
+        }
+
+        // Use the standard AWS domain for 'studio' (prod).
+        // For non-prod environments, use the obfuscated domain 'asfiovnxocqpcry.com'.
+        const baseDomain =
+            envSubdomain === 'studio'
+                ? `studio.${region}.sagemaker.aws`
+                : `${envSubdomain}.studio.${region}.asfiovnxocqpcry.com`
+
+        refreshUrl = `https://studio-${domain}.${baseDomain}/${appSubDomain}`
+    }
+    // For SMUS connections, refreshUrl remains undefined
 
-    const refreshUrl = `https://studio-${domain}.${baseDomain}/${appSubDomain}`
     await setSpaceCredentials(spaceArn, refreshUrl, {
         sessionId: session ?? '-',
         url: wsUrl ?? '-',
@@ -179,12 +189,12 @@ export async function setSmusSpaceSsoProfile(spaceArn: string, projectId: string
  * Stores SSM connection information for a given space, typically from a deep link session.
  * This initializes the request as 'fresh' and includes a refresh URL if provided.
  * @param spaceArn - The arn of the SageMaker space.
- * @param refreshUrl - URL to use for refreshing session tokens.
+ * @param refreshUrl - URL to use for refreshing session tokens (undefined for SMUS connections).
  * @param credentials - The session information used to initiate the connection.
  */
 export async function setSpaceCredentials(
     spaceArn: string,
-    refreshUrl: string,
+    refreshUrl: string | undefined,
     credentials: SsmConnectionInfo
 ): Promise<void> {
     const data = await loadMappings()

packages/core/src/awsService/sagemaker/detached-server/routes/getSessionAsync.ts

@@ -9,6 +9,8 @@ import { IncomingMessage, ServerResponse } from 'http'
 import url from 'url'
 import { SessionStore } from '../sessionStore'
 import { open, parseArn, readServerInfo } from '../utils'
+import { openErrorPage } from '../errorPage'
+import { SmusDeeplinkSessionExpiredError } from '../../constants'
 
 export async function handleGetSessionAsync(req: IncomingMessage, res: ServerResponse): Promise<void> {
     const parsedUrl = url.parse(req.url || '', true)
@@ -46,8 +48,34 @@ export async function handleGetSessionAsync(req: IncomingMessage, res: ServerRes
             res.end()
             return
         } else if (status === 'not-started') {
-            const serverInfo = await readServerInfo()
             const refreshUrl = await store.getRefreshUrl(connectionIdentifier)
+
+            // Check if this is a SMUS connection (no refreshUrl available)
+            if (refreshUrl === undefined) {
+                console.log(`SMUS session expired for connection: ${connectionIdentifier}`)
+
+                // Clean up the expired connection entry
+                try {
+                    await store.cleanupExpiredConnection(connectionIdentifier)
+                    console.log(`Cleaned up expired connection: ${connectionIdentifier}`)
+                } catch (cleanupErr) {
+                    console.error(`Failed to cleanup expired connection: ${cleanupErr}`)
+                    // Continue with error response even if cleanup fails
+                }
+
+                await openErrorPage(SmusDeeplinkSessionExpiredError.title, SmusDeeplinkSessionExpiredError.message)
+                res.writeHead(400, { 'Content-Type': 'application/json' })
+                res.end(
+                    JSON.stringify({
+                        error: SmusDeeplinkSessionExpiredError.code,
+                        message: SmusDeeplinkSessionExpiredError.shortMessage,
+                    })
+                )
+                return
+            }
+
+            // Continue with existing SageMaker AI refresh flow
+            const serverInfo = await readServerInfo()
             const { spaceName } = parseArn(connectionIdentifier)
 
             const url = `${refreshUrl}/${encodeURIComponent(spaceName)}?remote_access_token_refresh=true&reconnect_identifier=${encodeURIComponent(

packages/core/src/awsService/sagemaker/detached-server/sessionStore.ts

@@ -9,7 +9,7 @@ import { readMapping, writeMapping } from './utils'
 export type SessionStatus = 'pending' | 'fresh' | 'consumed' | 'not-started'
 
 export class SessionStore {
-    async getRefreshUrl(connectionId: string) {
+    async getRefreshUrl(connectionId: string): Promise<string | undefined> {
         const mapping = await readMapping()
 
         if (!mapping.deepLink) {
@@ -21,10 +21,6 @@ export class SessionStore {
             throw new Error(`No mapping found for connectionId: "${connectionId}"`)
         }
 
-        if (!entry.refreshUrl) {
-            throw new Error(`No refreshUrl found for connectionId: "${connectionId}"`)
-        }
-
         return entry.refreshUrl
     }
 
@@ -113,6 +109,20 @@ export class SessionStore {
         await writeMapping(mapping)
     }
 
+    async cleanupExpiredConnection(connectionId: string) {
+        const mapping = await readMapping()
+
+        if (!mapping.deepLink) {
+            throw new Error('No deepLink mapping found')
+        }
+
+        // Remove the entire connection entry for the expired space
+        if (mapping.deepLink[connectionId]) {
+            delete mapping.deepLink[connectionId]
+            await writeMapping(mapping)
+        }
+    }
+
     async setSession(connectionId: string, requestId: string, ssmConnectionInfo: SsmConnectionInfo) {
         const mapping = await readMapping()
 

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

@@ -84,7 +84,7 @@ export async function prepareDevEnvConnection(
             await persistSmusProjectCreds(spaceArn, node as SagemakerUnifiedStudioSpaceNode)
         }
     } else if (connectionType === 'sm_dl') {
-        await persistSSMConnection(spaceArn, domain ?? '', session, wsUrl, token, appType)
+        await persistSSMConnection(spaceArn, domain ?? '', session, wsUrl, token, appType, isSMUS)
     }
 
     await startLocalServer(ctx)

packages/core/src/extensionNode.ts

@@ -199,7 +199,7 @@ export async function activate(context: vscode.ExtensionContext) {
             await handleAmazonQInstall()
         }
 
-        await activateSageMakerUnifiedStudio(context)
+        await activateSageMakerUnifiedStudio(extContext)
 
         await activateApplicationComposer(context)
         await activateThreatComposerEditor(context)

packages/core/src/sagemakerunifiedstudio/activation.ts

@@ -3,21 +3,25 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import * as vscode from 'vscode'
 import { activate as activateConnectionMagicsSelector } from './connectionMagicsSelector/activation'
 import { activate as activateExplorer } from './explorer/activation'
 import { isSageMaker } from '../shared/extensionUtilities'
 import { initializeResourceMetadata } from './shared/utils/resourceMetadataUtils'
 import { setContext } from '../shared/vscode/setContext'
 import { SmusUtils } from './shared/smusUtils'
+import * as smusUriHandlers from './uriHandlers'
+import { ExtContext } from '../shared/extensions'
 
-export async function activate(extensionContext: vscode.ExtensionContext): Promise<void> {
+export async function activate(ctx: ExtContext): Promise<void> {
     // Only run when environment is a SageMaker Unified Studio space
     if (isSageMaker('SMUS') || isSageMaker('SMUS-SPACE-REMOTE-ACCESS')) {
         await initializeResourceMetadata()
         // Setting context before any getContext calls to avoid potential race conditions.
         await setContext('aws.smus.inSmusSpaceEnvironment', SmusUtils.isInSmusSpaceEnvironment())
-        await activateConnectionMagicsSelector(extensionContext)
+        await activateConnectionMagicsSelector(ctx.extensionContext)
     }
-    await activateExplorer(extensionContext)
+    await activateExplorer(ctx.extensionContext)
+
+    // Register SMUS URI handler for deeplink connections
+    ctx.extensionContext.subscriptions.push(smusUriHandlers.register(ctx))
 }