fix(sagemaker): Add proactive credential refresh for SSO-based SageMaker Space connections

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

@@ -27,13 +27,21 @@ import {
 } from './hyperpodCommands'
 import { SagemakerHyperpodNode } from './explorer/sagemakerHyperpodNode'
 import { getLogger } from '../../shared/logger/logger'
+import { stopAllSsoCredentialRefreshers } from './credentialMapping'
 
 let terminalActivityInterval: NodeJS.Timeout | undefined
 let backgroundStateInterval: NodeJS.Timeout | undefined
 
 export async function activate(ctx: ExtContext): Promise<void> {
+    getLogger().info('SageMaker: activation started')
     ctx.extensionContext.subscriptions.push(
         uriHandlers.register(ctx),
+        {
+            dispose: () => {
+                getLogger().debug('SageMaker: registered SSO credential refresher cleanup on deactivation')
+                stopAllSsoCredentialRefreshers()
+            },
+        },
         Commands.register('aws.sagemaker.openRemoteConnection', async (node: SagemakerSpaceNode) => {
             if (!validateNode(node)) {
                 return

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

@@ -11,6 +11,8 @@
 import { DevSettings } from '../../shared/settings'
 import { Auth } from '../../auth/auth'
 import { SpaceMappings, SsmConnectionInfo } from './types'
+import { getCredentialsFromStore } from '../../auth/credentials/store'
+import { CredentialsId, fromString } from '../../auth/providers/credentials'
 import { getLogger } from '../../shared/logger/logger'
 import { parseArn } from './utils'
 import { SagemakerUnifiedStudioSpaceNode } from '../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode'
@@ -20,6 +22,131 @@
 const mappingFileName = '.sagemaker-space-profiles'
 const mappingFilePath = path.join(os.homedir(), '.aws', mappingFileName)
 
+export interface SsoCachedCredentials {
+    credentials: {
+        accessKeyId: string
+        secretAccessKey: string
+        sessionToken?: string
+        expiration?: Date
+    }
+}
+
+/**
+ * Proactive credential refresh for SSO-based SageMaker Space connections.
+ *
+ * Follows the same pattern as SMUS `ProjectRoleCredentialsProvider.startProactiveCredentialRefresh()`:
+ * - Checks every 10 seconds using setTimeout (handles sleep/resume correctly)
+ * - Refreshes when credentials expire within 5 minutes (safety buffer)
+ * - Writes fresh credentials to the mapping file so the detached server always reads valid creds
+ *
+ * Without this, SSO connections disconnect after ~1 hour when the initial STS credentials expire
+ * because `persistLocalCredentials()` only writes once at connection time.
+ */
+export class SsoCredentialRefresher {
+    private refreshTimer?: ReturnType<typeof setTimeout>
+    private active = false
+    readonly checkIntervalMs: number
+    readonly safetyBufferMs: number
+
+    constructor(
+        private readonly spaceArn: string,
+        private readonly getCachedCredentials: () => SsoCachedCredentials | undefined,
+        private readonly credentialsId: CredentialsId,
+        options?: { checkIntervalMs?: number; safetyBufferMs?: number }
+    ) {
+        this.checkIntervalMs = options?.checkIntervalMs ?? 60_000
+        this.safetyBufferMs = options?.safetyBufferMs ?? 5 * 60_000
+    }
+
+    public start(): void {
+        if (this.active) {
+            getLogger().debug(`SSO refresh [${this.spaceArn}]: already active, skipping start`)
+            return
+        }
+        this.active = true
+        getLogger().info(`SSO refresh [${this.spaceArn}]: started (check every ${this.checkIntervalMs / 1000}s, buffer ${this.safetyBufferMs / 60000} min)`)
+        this.scheduleNextCheck()
+    }
+
+    public stop(): void {
+        this.active = false
+        if (this.refreshTimer !== undefined) {
+            clearTimeout(this.refreshTimer)
+            this.refreshTimer = undefined
+        }
+        getLogger().info(`SSO refresh [${this.spaceArn}]: stopped`)
+    }
+
+    public isActive(): boolean {
+        return this.active
+    }
+
+    private scheduleNextCheck(): void {
+        if (!this.active) {
+            return
+        }
+        this.refreshTimer = setTimeout(async () => {
+            try {
+                await this.refreshIfNeeded()
+            } catch (error) {
+                getLogger().error(`SSO credential refresh failed for ${this.spaceArn}: %O`, error)
+            }
+            if (this.active) {
+                this.scheduleNextCheck()
+            }
+        }, this.checkIntervalMs)
+    }
+
+    private async refreshIfNeeded(): Promise<void> {
+        const cached = this.getCachedCredentials()
+        const expiration = cached?.credentials.expiration?.getTime()
+        const now = Date.now()
+        const minutesLeft = expiration ? Math.round((expiration - now) / 60000) : 'unknown'
+
+        getLogger().debug(`SSO refresh check [${this.spaceArn}]: expiry in ${minutesLeft} min, buffer=${this.safetyBufferMs / 60000} min`)
+
+        if (expiration && expiration - now > this.safetyBufferMs) {
+            return // still fresh
+        }
+
+        getLogger().info(`SSO refresh [${this.spaceArn}]: credentials expiring soon (${minutesLeft} min left), fetching fresh via GetRoleCredentials`)
+
+        try {
+            const freshCreds = await getCredentialsFromStore(this.credentialsId, globals.loginManager.store)
+            if (!freshCreds) {
+                getLogger().warn(`SSO refresh [${this.spaceArn}]: getCredentialsFromStore returned undefined - bearer token may be expired`)
+                return
+            }
+
+            getLogger().debug(`SSO refresh [${this.spaceArn}]: got fresh creds, expiry=${freshCreds.expiration?.toISOString()}`)
+
+            await setSpaceSsoProfile(
+                this.spaceArn,
+                freshCreds.accessKeyId,
+                freshCreds.secretAccessKey,
+                freshCreds.sessionToken ?? ''
+            )
+
+            getLogger().info(`SSO refresh [${this.spaceArn}]: mapping file updated with fresh credentials`)
+        } catch (error) {
+            getLogger().error(`SSO refresh [${this.spaceArn}]: failed to refresh: %O`, error)
+        }
+    }
+}
+
+/** Active SSO credential refreshers, keyed by spaceArn */
+const activeSsoRefreshers = new Map<string, SsoCredentialRefresher>()
+
+/**
+ * Stops all active SSO credential refreshers. Call on extension deactivation or user logout.
+ */
+export function stopAllSsoCredentialRefreshers(): void {
+    for (const refresher of activeSsoRefreshers.values()) {
+        refresher.stop()
+    }
+    activeSsoRefreshers.clear()
+}
+
 export async function loadMappings(): Promise<SpaceMappings> {
     try {
         if (!(await fs.existsFile(mappingFilePath))) {
@@ -51,19 +178,36 @@
  */
 export async function persistLocalCredentials(spaceArn: string): Promise<void> {
     const currentProfileId = Auth.instance.getCurrentProfileId()
+    getLogger().info(`SageMaker persistLocalCredentials: called for space ${spaceArn}, profileId=${currentProfileId}`)
     if (!currentProfileId) {
         throw new ToolkitError('No current profile ID available for saving space credentials.')
     }
 
     if (currentProfileId.startsWith('sso:')) {
         const credentials = globals.loginManager.store.credentialsCache[currentProfileId]
+        getLogger().debug('SageMaker persistLocalCredentials: writing SSO credentials to mapping file')
         await setSpaceSsoProfile(
             spaceArn,
             credentials.credentials.accessKeyId,
             credentials.credentials.secretAccessKey,
             credentials.credentials.sessionToken ?? ''
         )
+
+        // Start proactive credential refresh for SSO connections.
+        // Without this, the mapping file goes stale after ~1h (default STS TTL)
+        // and the detached server reads expired credentials on reconnect.
+        const existing = activeSsoRefreshers.get(spaceArn)
+        if (existing) {
+            existing.stop()
+        }
+        const refresher = new SsoCredentialRefresher(spaceArn, () => {
+            return globals.loginManager.store.credentialsCache[currentProfileId] as SsoCachedCredentials | undefined
+        }, fromString(currentProfileId))
+        activeSsoRefreshers.set(spaceArn, refresher)
+        refresher.start()
+        getLogger().info(`SageMaker persistLocalCredentials: SSO credential refresher started for ${spaceArn}`)
     } else {
+        getLogger().debug(`SageMaker persistLocalCredentials: IAM profile ${currentProfileId}, no refresher needed`)
         await setSpaceIamProfile(spaceArn, currentProfileId)
     }
 }

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

@@ -20,6 +20,7 @@
  * @throws If the profile is missing, malformed, or unsupported.
  */
 export async function resolveCredentialsFor(connectionIdentifier: string): Promise<any> {
+    console.log('Detached server resolveCredentialsFor: resolving for', connectionIdentifier)
     const mapping = await readMapping()
     const profile = mapping.localCredential?.[connectionIdentifier] as LocalCredentialProfile
 
@@ -34,6 +35,7 @@
                 if (!name) {
                     throw new Error(`Invalid IAM profile name for "${connectionIdentifier}"`)
                 }
+                console.log('Detached server: using IAM profile', name, '(dynamic resolution via fromIni)')
                 return fromIni({ profile: name })
             } else if ('smusProjectId' in profile) {
                 const { accessKey, secret, token } = mapping.smusProjects?.[profile.smusProjectId] || {}
@@ -55,6 +57,7 @@
                 if (!accessKey || !secret || !token) {
                     throw new Error(`Missing SSO credentials for "${connectionIdentifier}"`)
                 }
+                console.log('Detached server: using SSO static credentials from mapping file')
                 return {
                     accessKeyId: accessKey,
                     secretAccessKey: secret,

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

@@ -44,9 +44,11 @@ export async function handleGetSession(req: IncomingMessage, res: ServerResponse
 
     let credentials
     try {
+        console.log('getSession: resolving credentials for', connectionIdentifier, 'attempt', count)
         credentials = await resolveCredentialsFor(connectionIdentifier)
     } catch (err) {
         console.error('Failed to resolve credentials:', err)
+        console.log('getSession: credential resolution failed - are mapping file creds stale?')
         res.writeHead(500, { 'Content-Type': 'text/plain' })
         res.end((err as Error).message)
         return
@@ -58,6 +60,7 @@ export async function handleGetSession(req: IncomingMessage, res: ServerResponse
 
     try {
         const session = await startSagemakerSession({ region, connectionIdentifier, credentials })
+        console.log('getSession: StartSession succeeded, sessionId=', session.SessionId)
         attemptCount.delete(connectionIdentifier)
         lastAttemptTime.delete(connectionIdentifier)
         res.writeHead(200, { 'Content-Type': 'application/json' })

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

@@ -84,10 +84,13 @@ export async function handleGetSessionAsync(req: IncomingMessage, res: ServerRes
                 `http://localhost:${serverInfo.port}/refresh_token`
             )}`
 
+            // Mark pending BEFORE opening browser to prevent race condition:
+            // browser callback can arrive and write "fresh" before markPending completes,
+            // then markPending overwrites "fresh" back to "pending".
+            await store.markPending(connectionIdentifier, requestId)
             await open(url)
             res.writeHead(202, { 'Content-Type': 'text/plain' })
             res.end('Session is not ready yet. Please retry in a few seconds.')
-            await store.markPending(connectionIdentifier, requestId)
             return
         }
     } catch (err) {

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

@@ -29,6 +29,28 @@ export const ideProcessPatterns = {
 
 const server = http.createServer((req: IncomingMessage, res: ServerResponse) => {
     const parsedUrl = url.parse(req.url || '', true)
+    const ts = new Date().toISOString()
+    const serverAddr = server.address()
+    const serverPort = serverAddr && typeof serverAddr === 'object' ? serverAddr.port : 'unknown'
+
+    // Log query params (redact sensitive values)
+    const safeParams: Record<string, string> = {}
+    for (const [k, v] of Object.entries(parsedUrl.query)) {
+        if (k === 'token' || k === 'ws_url' || k === 'session') {
+            safeParams[k] = `[REDACTED ${String(v).length} chars]`
+        } else {
+            safeParams[k] = String(v)
+        }
+    }
+    console.log(`[${ts}] REQUEST: ${req.method} ${parsedUrl.pathname} on port ${serverPort}`)
+    console.log(`[${ts}] PARAMS: ${JSON.stringify(safeParams)}`)
+
+    // Wrap res.writeHead to log response status
+    const origWriteHead = res.writeHead.bind(res)
+    res.writeHead = (statusCode: number, ...args: any[]) => {
+        console.log(`[${ts}] RESPONSE: ${statusCode} ${parsedUrl.pathname} (port ${serverPort}, pid ${process.pid})`)
+        return origWriteHead(statusCode, ...args)
+    }
 
     switch (parsedUrl.pathname) {
         case '/get_session':

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

@@ -99,6 +99,12 @@ export class SessionStore {
             throw new Error(`No mapping found for connectionId: "${connectionId}"`)
         }
 
+        // Never downgrade from "fresh" to "pending" — prevents race condition where
+        // the browser callback (setSession) writes "fresh" before this completes.
+        if (entry.requests[requestId]?.status === 'fresh') {
+            return
+        }
+
         entry.requests[requestId] = {
             sessionId: '',
             token: '',

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

@@ -193,8 +193,10 @@ export async function prepareDevEnvConnection(opts: DevEnvConnectionOptions) {
     // save space credential mapping
     if (connectionType === 'sm_lc') {
         if (!isSMUS) {
+            getLogger().info('SageMaker prepareDevEnvConnection: persisting credentials for space connection (SSO/IAM via persistLocalCredentials)')
             await persistLocalCredentials(spaceArn)
         } else {
+            getLogger().info('SageMaker prepareDevEnvConnection: persisting credentials for space connection (SMUS via persistSmusProjectCreds)')
             await persistSmusProjectCreds(spaceArn, node as SagemakerUnifiedStudioSpaceNode)
         }
     } else if (connectionType === 'sm_dl') {
@@ -447,6 +449,8 @@ export async function stopLocalServer(ctx: vscode.ExtensionContext): Promise<voi
         } catch (err: any) {
             if (err.code === 'ESRCH') {
                 logger.warn(`no process found with PID ${pid}. It may have already exited.`)
+            } else if (err.code === 'EPERM') {
+                logger.warn(`permission denied killing PID ${pid}. Process may be elevated or in a different security context. Continuing.`)
             } else {
                 throw ToolkitError.chain(err, 'failed to stop local server')
             }

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

@@ -83,6 +83,25 @@ describe('handleGetSessionAsync', () => {
         assert(ctx.storeStub.markPending.calledWith('abc', 'req123'))
     })
 
+    it('calls markPending before opening browser to prevent race condition', async () => {
+        ctx.req = { url: '/session_async?connection_identifier=abc&request_id=req123' }
+
+        ctx.storeStub.getFreshEntry.returns(Promise.resolve(undefined))
+        ctx.storeStub.getStatus.returns(Promise.resolve('not-started'))
+        ctx.storeStub.getRefreshUrl.returns(Promise.resolve('https://example.com/refresh'))
+
+        const callOrder: string[] = []
+        ctx.storeStub.markPending.callsFake(async () => { callOrder.push('markPending') })
+        sinon.stub(utils, 'open').callsFake(async () => { callOrder.push('open'); return undefined as any })
+        sinon.stub(utils, 'readServerInfo').resolves({ pid: 1234, port: 4567 })
+        sinon.stub(utils, 'parseArn').returns({ region: 'us-east-1', accountId: '123456789012', resourceName: 'test-space' })
+
+        await handleGetSessionAsync(ctx.req as http.IncomingMessage, ctx.res as http.ServerResponse)
+
+        assert.strictEqual(callOrder[0], 'markPending', 'markPending must be called before open')
+        assert.strictEqual(callOrder[1], 'open', 'open must be called after markPending')
+    })
+
     it('responds with 500 if unexpected error occurs', async () => {
         ctx.req = { url: '/session_async?connection_identifier=abc&request_id=req123' }
         ctx.storeStub.getFreshEntry.throws(new Error('fail'))

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

@@ -155,6 +155,25 @@ describe('SessionStore', () => {
         assert.strictEqual(updated.deepLink[connectionId].requests['newReq'].status, 'pending')
     })
 
+    it('does not overwrite fresh status with pending (race condition guard)', async () => {
+        const store = new SessionStore()
+        readMappingStub.returns({
+            deepLink: {
+                [connectionId]: {
+                    refreshUrl: 'https://refresh.url',
+                    requests: {
+                        [requestId]: { sessionId: 's1', token: 't1', url: 'u1', status: 'fresh' },
+                    },
+                },
+            },
+        })
+
+        await store.markPending(connectionId, requestId)
+
+        // Should NOT have written — fresh must not be downgraded to pending
+        assert(writeMappingStub.notCalled)
+    })
+
     it('sets session entry with default fresh status', async () => {
         const store = new SessionStore()
         const info: SsmConnectionInfo = {