fix(smus): Add proactive cred refresh for active SSH connections

**Description**

Added a proactive cred check and refresh when SSH connections
are established.

Also updated the error messages to be actionable for users.

**Motivation**

Bug : Previously once cred expired, we were throwing blanket
error which did not tell user what the issue was and there
was no path to recovery as well. Now with proactive cred refresh,
user should be able to retry in ~10-15 seconds.

**Testing Done**

Tested all flows manually. Unit tests partial, needs to be updated.

Author: Bhargava Varadharajan

packages/core/src/auth/auth.ts

@@ -233,7 +233,8 @@ export class Auth implements AuthService, ConnectionManager {
         }
 
         const provider = this.getSsoTokenProvider(connection.id, profile)
-        const token = await provider.getToken()
+        // Calling existing getToken private method - It will handle setting the connection state etc.
+        const token = await this._getToken(connection.id, provider)
 
         if (!token?.accessToken) {
             throw new Error(`No access token available for connection ${connection.id}`)
@@ -946,10 +947,22 @@ export class Auth implements AuthService, ConnectionManager {
         if (previousState === 'valid') {
             // Non-token expiration errors can happen. We must log it here, otherwise they are lost.
             getLogger().warn(`auth: valid connection became invalid. Last error: %s`, this.#validationErrors.get(id))
-
             const timeout = new Timeout(60000)
             this.#invalidCredentialsTimeouts.set(id, timeout)
 
+            // Check if this is a SMUS profile - if so, skip the generic prompt
+            // as SMUS has its own reauthentication flow
+            const isSmusConnection = profile.type === 'sso' && 'domainUrl' in profile && 'domainId' in profile
+            if (isSmusConnection) {
+                getLogger().debug(`auth: Skipping generic reauthentication prompt for SMUS connection ${id}`)
+                // For SMUS connections, just throw the InvalidConnection error
+                // The SMUS auth provider will handle showing the appropriate prompt
+                throw new ToolkitError('Connection is invalid or expired. Try logging in again.', {
+                    code: errorCode.invalidConnection,
+                    cause: this.#validationErrors.get(id),
+                })
+            }
+
             const connLabel = profile.metadata.label ?? (profile.type === 'sso' ? this.getSsoProfileLabel(profile) : id)
             const message = localize(
                 'aws.auth.invalidConnection',

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

@@ -74,9 +74,12 @@ export async function persistLocalCredentials(spaceArn: string): Promise<void> {
 export async function persistSmusProjectCreds(spaceArn: string, node: SagemakerUnifiedStudioSpaceNode): Promise<void> {
     const nodeParent = node.getParent() as SageMakerUnifiedStudioSpacesParentNode
     const authProvider = nodeParent.getAuthProvider()
-    const projectAuthProvider = await authProvider.getProjectCredentialProvider(nodeParent.getProjectId())
+    const projectId = nodeParent.getProjectId()
+    const projectAuthProvider = await authProvider.getProjectCredentialProvider(projectId)
     await projectAuthProvider.getCredentials()
-    await setSmusSpaceSsoProfile(spaceArn, nodeParent.getProjectId())
+    await setSmusSpaceSsoProfile(spaceArn, projectId)
+    // Trigger SSH credential refresh for the project
+    projectAuthProvider.startProactiveCredentialRefresh()
 }
 
 /**

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

@@ -15,6 +15,7 @@ import { open } from './utils'
 export enum ExceptionType {
     ACCESS_DENIED = 'AccessDeniedException',
     DEFAULT = 'Default',
+    EXPIRED_TOKEN = 'ExpiredTokenException',
     INTERNAL_FAILURE = 'InternalFailure',
     RESOURCE_LIMIT_EXCEEDED = 'ResourceLimitExceeded',
     THROTTLING = 'ThrottlingException',
@@ -31,13 +32,18 @@ export const getVSCodeErrorTitle = (error: SageMakerServiceException): string =>
     return ErrorText.StartSession[ExceptionType.DEFAULT].Title
 }
 
-export const getVSCodeErrorText = (error: SageMakerServiceException): string => {
+export const getVSCodeErrorText = (error: SageMakerServiceException, isSmus?: boolean): string => {
     const exceptionType = error.name as ExceptionType
 
     switch (exceptionType) {
         case ExceptionType.ACCESS_DENIED:
         case ExceptionType.VALIDATION:
             return ErrorText.StartSession[exceptionType].Text.replace('{message}', error.message)
+        case ExceptionType.EXPIRED_TOKEN:
+            // Use SMUS-specific message if in SMUS context
+            return isSmus
+                ? ErrorText.StartSession[ExceptionType.EXPIRED_TOKEN].SmusText
+                : ErrorText.StartSession[exceptionType].Text
         case ExceptionType.INTERNAL_FAILURE:
         case ExceptionType.RESOURCE_LIMIT_EXCEEDED:
         case ExceptionType.THROTTLING:
@@ -57,6 +63,12 @@ export const ErrorText = {
             Title: 'Unexpected system error',
             Text: 'We encountered an unexpected error: [{exceptionType}]. Please contact your administrator and provide them with this error so they can investigate the issue.',
         },
+        [ExceptionType.EXPIRED_TOKEN]: {
+            Title: 'Authentication expired',
+            Text: 'Your session has expired. Please refresh your credentials and try again.',
+            SmusText:
+                'Your session has expired. This is likely due to network connectivity issues after machine sleep/resume. Please wait 10-30 seconds for automatic credential refresh, then try again. If the issue persists, try reconnecting through AWS Toolkit.',
+        },
         [ExceptionType.INTERNAL_FAILURE]: {
             Title: 'Failed to connect remotely to VSCode',
             Text: 'Unable to establish remote connection to VSCode. This could be due to several factors. Please try again by clicking the VSCode button. If the problem persists, please contact your admin.',

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

@@ -6,7 +6,7 @@
 // Disabled: detached server files cannot import vscode.
 /* eslint-disable aws-toolkits/no-console-log */
 import { IncomingMessage, ServerResponse } from 'http'
-import { startSagemakerSession, parseArn } from '../utils'
+import { startSagemakerSession, parseArn, isSmusConnection } from '../utils'
 import { resolveCredentialsFor } from '../credentials'
 import url from 'url'
 import { SageMakerServiceException } from '@amzn/sagemaker-client'
@@ -33,6 +33,8 @@ export async function handleGetSession(req: IncomingMessage, res: ServerResponse
     }
 
     const { region } = parseArn(connectionIdentifier)
+    // Detect if this is a SMUS connection for specialized error handling
+    const isSmus = await isSmusConnection(connectionIdentifier)
 
     try {
         const session = await startSagemakerSession({ region, connectionIdentifier, credentials })
@@ -48,7 +50,7 @@ export async function handleGetSession(req: IncomingMessage, res: ServerResponse
         const error = err as SageMakerServiceException
         console.error(`Failed to start SageMaker session for ${connectionIdentifier}:`, err)
         const errorTitle = getVSCodeErrorTitle(error)
-        const errorText = getVSCodeErrorText(error)
+        const errorText = getVSCodeErrorText(error, isSmus)
         await openErrorPage(errorTitle, errorText)
         res.writeHead(500, { 'Content-Type': 'text/plain' })
         res.end('Failed to start SageMaker session')

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

@@ -121,6 +121,24 @@ async function processWriteQueue() {
     }
 }
 
+/**
+ * Detects if the connection identifier is using SMUS credentials
+ * @param connectionIdentifier - The connection identifier to check
+ * @returns Promise<boolean> - true if SMUS, false otherwise
+ */
+export async function isSmusConnection(connectionIdentifier: string): Promise<boolean> {
+    try {
+        const mapping = await readMapping()
+        const profile = mapping.localCredential?.[connectionIdentifier]
+
+        // Check if profile exists and has smusProjectId
+        return profile && 'smusProjectId' in profile
+    } catch (err) {
+        // If we can't read the mapping, assume not SMUS to avoid breaking existing functionality
+        return false
+    }
+}
+
 /**
  * Writes the mapping to a temp file and atomically renames it to the target path.
  * Uses a queue to prevent race conditions when multiple requests try to write simultaneously.

packages/core/src/sagemakerunifiedstudio/auth/providers/connectionCredentialsProvider.ts

@@ -218,4 +218,18 @@ export class ConnectionCredentialsProvider implements CredentialsProvider {
             `SMUS Connection: Successfully invalidated connection credentials cache for connection ${this.connectionId}`
         )
     }
+
+    /**
+     * Disposes of the provider and cleans up resources
+     */
+    public dispose(): void {
+        this.logger.debug(
+            `SMUS Connection: Disposing connection credentials provider for connection ${this.connectionId}`
+        )
+        // Clear cache to clean up resources
+        this.invalidate()
+        this.logger.debug(
+            `SMUS Connection: Successfully disposed connection credentials provider for connection ${this.connectionId}`
+        )
+    }
 }

packages/core/src/sagemakerunifiedstudio/auth/providers/domainExecRoleCredentialsProvider.ts

@@ -314,4 +314,12 @@ export class DomainExecRoleCredentialsProvider implements CredentialsProvider {
         this.credentialCache = undefined
         this.logger.debug(`SMUS DER: Successfully invalidated DER credentials cache for domain ${this.domainId}`)
     }
+    /**
+     * Disposes of the provider and cleans up resources
+     */
+    public dispose(): void {
+        this.logger.debug(`SMUS DER: Disposing DER credentials provider for domain ${this.domainId}`)
+        this.invalidate()
+        this.logger.debug(`SMUS DER: Successfully disposed DER credentials provider for domain ${this.domainId}`)
+    }
 }

packages/core/src/sagemakerunifiedstudio/auth/providers/projectRoleCredentialsProvider.ts

@@ -28,6 +28,11 @@ export class ProjectRoleCredentialsProvider implements CredentialsProvider {
         credentials: AWS.Credentials
         expiresAt: Date
     }
+    private refreshTimer?: NodeJS.Timeout
+    private readonly refreshInterval = 10 * 60 * 1000 // 10 minutes
+    private readonly checkInterval = 10 * 1000 // 10 seconds - check frequently, refresh based on actual time
+    private sshRefreshActive = false
+    private lastRefreshTime?: Date
 
     constructor(
         private readonly smusAuthProvider: SmusAuthenticationProvider,
@@ -162,11 +167,21 @@ export class ProjectRoleCredentialsProvider implements CredentialsProvider {
 
             return awsCredentials
         } catch (err) {
-            this.logger.error(
-                'SMUS Project: Failed to get project credentials for project %s: %s',
-                this.projectId,
-                (err as Error).message
-            )
+            this.logger.error('SMUS Project: Failed to get project credentials for project %s: %s', this.projectId, err)
+
+            // Handle InvalidGrantException specially - indicates need for reauthentication
+            if (err instanceof Error && err.name === 'InvalidGrantException') {
+                // Invalidate cache when authentication fails
+                this.invalidate()
+                throw new ToolkitError(
+                    `Failed to get project credentials for project ${this.projectId}: ${err.message}. Reauthentication required.`,
+                    {
+                        code: 'InvalidRefreshToken',
+                        cause: err,
+                    }
+                )
+            }
+
             throw new ToolkitError(`Failed to get project credentials for project ${this.projectId}: ${err}`, {
                 code: 'ProjectCredentialsFetchFailed',
                 cause: err instanceof Error ? err : undefined,
@@ -192,6 +207,139 @@ export class ProjectRoleCredentialsProvider implements CredentialsProvider {
         }
     }
 
+    /**
+     * Starts proactive credential refresh for SSH connections
+     *
+     * Uses an expiry-based approach with safety buffer:
+     * - Checks every 10 seconds using setTimeout
+     * - Refreshes when credentials expire within 5 minutes (safety buffer)
+     * - Falls back to 10-minute time-based refresh if no expiry information available
+     * - Handles sleep/resume because it uses wall-clock time for expiry checks
+     *
+     * This means credentials are refreshed just before they expire, reducing
+     * unnecessary API calls while ensuring credentials remain valid.
+     */
+    public startProactiveCredentialRefresh(): void {
+        if (this.sshRefreshActive) {
+            this.logger.debug(`SMUS Project: SSH refresh already active for project ${this.projectId}`)
+            return
+        }
+
+        this.logger.info(`SMUS Project: Starting SSH credential refresh for project ${this.projectId}`)
+        this.sshRefreshActive = true
+        this.lastRefreshTime = new Date() // Initialize refresh time
+
+        // Start the check timer (checks every 10 seconds, refreshes every 10 minutes based on actual time)
+        this.scheduleNextCheck()
+    }
+
+    /**
+     * Stops proactive credential refresh
+     * Called when SSH connection ends or SMUS disconnects
+     */
+    public stopProactiveCredentialRefresh(): void {
+        if (!this.sshRefreshActive) {
+            return
+        }
+
+        this.logger.info(`SMUS Project: Stopping SSH credential refresh for project ${this.projectId}`)
+        this.sshRefreshActive = false
+        this.lastRefreshTime = undefined
+
+        // Clean up timer
+        if (this.refreshTimer) {
+            clearTimeout(this.refreshTimer)
+            this.refreshTimer = undefined
+        }
+    }
+
+    /**
+     * Schedules the next credential check (every 10 seconds)
+     * Refreshes credentials when they expire within 5 minutes (safety buffer)
+     * Falls back to 10-minute time-based refresh if no expiry information available
+     * This handles sleep/resume scenarios correctly
+     */
+    private scheduleNextCheck(): void {
+        if (!this.sshRefreshActive) {
+            return
+        }
+        // Check every 10 seconds, but only refresh every 10 minutes based on actual time elapsed
+        this.refreshTimer = setTimeout(async () => {
+            try {
+                const now = new Date()
+                // Check if we need to refresh based on actual time elapsed
+                if (this.shouldPerformRefresh(now)) {
+                    await this.refresh()
+                }
+                // Schedule next check if still active
+                if (this.sshRefreshActive) {
+                    this.scheduleNextCheck()
+                }
+            } catch (error) {
+                this.logger.error(
+                    `SMUS Project: Failed to refresh credentials for project ${this.projectId}: %O`,
+                    error
+                )
+                // Continue trying even if refresh fails. Dispose will handle stopping the refresh.
+                if (this.sshRefreshActive) {
+                    this.scheduleNextCheck()
+                }
+            }
+        }, this.checkInterval)
+    }
+
+    /**
+     * Determines if a credential refresh should be performed based on credential expiration
+     * This handles sleep/resume scenarios properly and is more efficient than time-based refresh
+     */
+    private shouldPerformRefresh(now: Date): boolean {
+        if (!this.lastRefreshTime || !this.credentialCache) {
+            // First refresh or no cached credentials
+            this.logger.debug(`SMUS Project: First refresh - no previous credentials for ${this.projectId}`)
+            return true
+        }
+
+        // Check if credentials expire soon (with 5-minute safety buffer)
+        const safetyBufferMs = 5 * 60 * 1000 // 5 minutes before expiry
+        const expiryTime = this.credentialCache.credentials.expiration?.getTime()
+
+        if (!expiryTime) {
+            // No expiry info - fall back to time-based refresh as safety net
+            const timeSinceLastRefresh = now.getTime() - this.lastRefreshTime.getTime()
+            const shouldRefresh = timeSinceLastRefresh >= this.refreshInterval
+            return shouldRefresh
+        }
+
+        const timeUntilExpiry = expiryTime - now.getTime()
+        const shouldRefresh = timeUntilExpiry < safetyBufferMs
+        return shouldRefresh
+    }
+
+    /**
+     * Performs credential refresh by invalidating cache and fetching fresh credentials
+     */
+    private async refresh(): Promise<void> {
+        const now = new Date()
+        const expiryTime = this.credentialCache?.credentials.expiration?.getTime()
+
+        if (expiryTime) {
+            const minutesUntilExpiry = Math.round((expiryTime - now.getTime()) / 60000)
+            this.logger.debug(
+                `SMUS Project: Refreshing credentials for project ${this.projectId} - expires in ${minutesUntilExpiry} minutes`
+            )
+        } else {
+            const minutesSinceLastRefresh = this.lastRefreshTime
+                ? Math.round((now.getTime() - this.lastRefreshTime.getTime()) / 60000)
+                : 0
+            this.logger.debug(
+                `SMUS Project: Refreshing credentials for project ${this.projectId} - time-based refresh after ${minutesSinceLastRefresh} minutes`
+            )
+        }
+
+        await this.getCredentials()
+        this.lastRefreshTime = new Date()
+    }
+
     /**
      * Invalidates cached project credentials
      * Clears the internal cache without fetching new credentials
@@ -204,4 +352,12 @@ export class ProjectRoleCredentialsProvider implements CredentialsProvider {
             `SMUS Project: Successfully invalidated project credentials cache for project ${this.projectId}`
         )
     }
+
+    /**
+     * Disposes of the provider and cleans up resources
+     */
+    public dispose(): void {
+        this.stopProactiveCredentialRefresh()
+        this.invalidate()
+    }
 }