fix(smsus): Update DER cred expiry time and project pick (aws#2206)

**Description**

Reduced the DER cred expiry time to 10 min default. The API is being
updated as well.

Also updated the auth logic to invoke project picker on sign in,
re-auth.

**Testing Done**

Unit tests, tested manually on VSCode as well - The signin, re-auth and
sign-out cases.

---

- 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: Bhargava Varadharajan <[email protected]>

Author: vpbhargav

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

@@ -16,7 +16,7 @@ import { SmusCredentialExpiry, SmusTimeouts, SmusErrorCodes, validateCredentialF
  * Credentials provider for SageMaker Unified Studio Domain Execution Role (DER)
  * Uses SSO tokens to get DER credentials via the /sso/redeem-token endpoint
  *
- * This provider implements internal caching with 55-minute expiry and handles
+ * This provider implements internal caching with 10-minute expiry and handles
  * its own credential lifecycle independently
  */
 export class DomainExecRoleCredentialsProvider implements CredentialsProvider {
@@ -122,7 +122,7 @@ export class DomainExecRoleCredentialsProvider implements CredentialsProvider {
     public async getCredentials(): Promise<AWS.Credentials> {
         this.logger.debug(`SMUS DER: Getting DER credentials for domain ${this.domainId}`)
 
-        // Check cache first (55-minute expiry with 5-minute buffer for proactive refresh)
+        // Check cache first (10-minute expiry with 5-minute buffer for proactive refresh)
         if (this.credentialCache && this.credentialCache.expiresAt > new Date()) {
             this.logger.debug(`SMUS DER: Using cached DER credentials for domain ${this.domainId}`)
             return this.credentialCache.credentials
@@ -207,6 +207,7 @@ export class DomainExecRoleCredentialsProvider implements CredentialsProvider {
                     accessKeyId: string
                     secretAccessKey: string
                     sessionToken: string
+                    expiration: string
                 }
             }
 
@@ -225,9 +226,23 @@ export class DomainExecRoleCredentialsProvider implements CredentialsProvider {
             validateCredentialFields(credentials, 'InvalidCredentialResponse', 'API response')
 
             // Create credentials with expiration
-            // Note: The response doesn't include expiration, so we set it to 55 minutes from now
-            // TODO: Update when the API provides actual expiration time
-            const credentialExpiresAt = new globals.clock.Date(Date.now() + SmusCredentialExpiry.derExpiryMs)
+            // Note: The response doesn't include expiration yet, so we set it to 10 minutes for now if it does't exist
+            let credentialExpiresAt: Date
+            if (credentials.expiration) {
+                // The API returns expiration as a string, convert to Date
+                const parsedExpiration = new Date(credentials.expiration)
+                // Check if the parsed date is valid
+                if (isNaN(parsedExpiration.getTime())) {
+                    this.logger.warn(
+                        `SMUS DER: Invalid expiration date string: ${credentials.expiration}, using default expiration`
+                    )
+                    credentialExpiresAt = new Date(Date.now() + SmusCredentialExpiry.derExpiryMs)
+                } else {
+                    credentialExpiresAt = parsedExpiration
+                }
+            } else {
+                credentialExpiresAt = new Date(Date.now() + SmusCredentialExpiry.derExpiryMs)
+            }
 
             const awsCredentials: AWS.Credentials = {
                 accessKeyId: credentials.accessKeyId as string,
@@ -236,7 +251,7 @@ export class DomainExecRoleCredentialsProvider implements CredentialsProvider {
                 expiration: credentialExpiresAt,
             }
 
-            // Cache DER credentials with 55-minute expiry (5-minute buffer for proactive refresh)
+            // Cache DER credentials with 10-minute expiry (5-minute buffer for proactive refresh)
             const cacheExpiresAt = new globals.clock.Date(Date.now() + SmusCredentialExpiry.derExpiryMs)
             this.credentialCache = {
                 credentials: awsCredentials,

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

@@ -175,6 +175,12 @@ export class SmusAuthenticationProvider {
 
                     // Use the existing connection
                     const result = await this.secondaryAuth.useNewConnection(existingConn)
+
+                    // Auto-invoke project selection after successful sign-in (but not in SMUS space environment)
+                    if (!SmusUtils.isInSmusSpaceEnvironment()) {
+                        void vscode.commands.executeCommand('aws.smus.switchProject')
+                    }
+
                     return result
                 }
 
@@ -192,6 +198,12 @@ export class SmusAuthenticationProvider {
 
                     const result = await this.secondaryAuth.useNewConnection(smusConn)
                     logger.debug(`SMUS: Reauthenticated connection successfully, id=${result.id}`)
+
+                    // Auto-invoke project selection after successful reauthentication (but not in SMUS space environment)
+                    if (!SmusUtils.isInSmusSpaceEnvironment()) {
+                        void vscode.commands.executeCommand('aws.smus.switchProject')
+                    }
+
                     return result
                 }
             }
@@ -214,6 +226,12 @@ export class SmusAuthenticationProvider {
             }
 
             const result = await this.secondaryAuth.useNewConnection(smusConn)
+
+            // Auto-invoke project selection after successful sign-in (but not in SMUS space environment)
+            if (!SmusUtils.isInSmusSpaceEnvironment()) {
+                void vscode.commands.executeCommand('aws.smus.switchProject')
+            }
+
             return result
         } catch (e) {
             throw ToolkitError.chain(e, 'Failed to connect to SageMaker Unified Studio', {

packages/core/src/sagemakerunifiedstudio/shared/smusUtils.ts

@@ -30,8 +30,8 @@ interface DataZoneSsoLoginResponse {
  * Credential expiry time constants for SMUS providers (in milliseconds)
  */
 export const SmusCredentialExpiry = {
-    /** Domain Execution Role (DER) credentials expiry time: 55 minutes */
-    derExpiryMs: 55 * 60 * 1000,
+    /** Domain Execution Role (DER) credentials expiry time: 10 minutes */
+    derExpiryMs: 10 * 60 * 1000,
     /** Project Role credentials expiry time: 10 minutes */
     projectExpiryMs: 10 * 60 * 1000,
     /** Connection credentials expiry time: 10 minutes */

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

@@ -303,12 +303,40 @@ describe('DomainExecRoleCredentialsProvider', function () {
         it('should set default expiration when not provided in response', async function () {
             const credentials = await derProvider.getCredentials()
 
-            // Should have expiration set to 55 mins from now
+            // Should have expiration set to 10 mins from now
             assert.ok(credentials.expiration)
             const expirationTime = credentials.expiration!.getTime()
-            const expectedTime = Date.now() + 55 * 60 * 1000 // 1 hour
+            const expectedTime = Date.now() + 10 * 60 * 1000 // 10 minutes
             const timeDiff = Math.abs(expirationTime - expectedTime)
-            assert.ok(timeDiff < 5000, 'Expiration should be 55 mins from now')
+            assert.ok(timeDiff < 5000, 'Expiration should be 10 mins from now')
+        })
+
+        it('should use expiration from API response when provided', async function () {
+            const futureExpiration = new Date(Date.now() + 2 * 60 * 60 * 1000) // 2 hours from now
+            const responseWithExpiration = {
+                credentials: {
+                    accessKeyId: 'AKIA-DER-KEY',
+                    secretAccessKey: 'der-secret-key',
+                    sessionToken: 'der-session-token',
+                    expiration: futureExpiration.toISOString(), // API returns expiration as ISO string
+                },
+            }
+
+            fetchStub.resolves({
+                ok: true,
+                status: 200,
+                statusText: 'OK',
+                json: sinon.stub().resolves(responseWithExpiration),
+            } as any)
+
+            const credentials = await derProvider.getCredentials()
+
+            // Should use the expiration from the API response
+            assert.ok(credentials.expiration)
+            const expirationTime = credentials.expiration!.getTime()
+            const expectedTime = futureExpiration.getTime()
+            const timeDiff = Math.abs(expirationTime - expectedTime)
+            assert.ok(timeDiff < 1000, 'Should use expiration from API response')
         })
 
         it('should handle JSON parsing errors', async function () {
@@ -326,6 +354,38 @@ describe('DomainExecRoleCredentialsProvider', function () {
                 }
             )
         })
+
+        it('should handle invalid expiration string in response', async function () {
+            const responseWithInvalidExpiration = {
+                credentials: {
+                    accessKeyId: 'AKIA-DER-KEY',
+                    secretAccessKey: 'der-secret-key',
+                    sessionToken: 'der-session-token',
+                    expiration: 'invalid-date-string', // Invalid date string
+                },
+            }
+
+            fetchStub.resolves({
+                ok: true,
+                status: 200,
+                statusText: 'OK',
+                json: sinon.stub().resolves(responseWithInvalidExpiration),
+            } as any)
+
+            const credentials = await derProvider.getCredentials()
+
+            // Should fall back to default expiration when date parsing fails
+            assert.ok(credentials.expiration)
+            const expirationTime = credentials.expiration!.getTime()
+
+            // Should be a valid timestamp (not NaN) using the default expiration
+            assert.ok(!isNaN(expirationTime), 'Should have valid expiration timestamp')
+
+            // Should be close to now + 10 minutes (default expiration)
+            const expectedTime = Date.now() + 10 * 60 * 1000
+            const timeDiff = Math.abs(expirationTime - expectedTime)
+            assert.ok(timeDiff < 5000, 'Should fall back to default expiration for invalid date string')
+        })
     })
 
     describe('invalidate', function () {

packages/core/src/test/sagemakerunifiedstudio/auth/smusAuthenticationProvider.test.ts

@@ -5,6 +5,7 @@
 
 import assert from 'assert'
 import sinon from 'sinon'
+import * as vscode from 'vscode'
 
 // Mock the setContext function BEFORE importing modules that use it
 const setContextModule = require('../../../shared/vscode/setContext')
@@ -24,6 +25,8 @@ describe('SmusAuthenticationProvider', function () {
     let smusAuthProvider: SmusAuthenticationProvider
     let extractDomainInfoStub: sinon.SinonStub
     let getSsoInstanceInfoStub: sinon.SinonStub
+    let isInSmusSpaceEnvironmentStub: sinon.SinonStub
+    let executeCommandStub: sinon.SinonStub
     let mockSecondaryAuthState: {
         activeConnection: SmusConnection | undefined
         hasSavedConnection: boolean
@@ -94,9 +97,14 @@ describe('SmusAuthenticationProvider', function () {
             .stub(SmusUtils, 'extractDomainInfoFromUrl')
             .returns({ domainId: testDomainId, region: testRegion })
         getSsoInstanceInfoStub = sinon.stub(SmusUtils, 'getSsoInstanceInfo').resolves(testSsoInstanceInfo)
+        isInSmusSpaceEnvironmentStub = sinon.stub(SmusUtils, 'isInSmusSpaceEnvironment').returns(false)
+        executeCommandStub = sinon.stub(vscode.commands, 'executeCommand').resolves()
         sinon.stub(require('../../../auth/secondaryAuth'), 'getSecondaryAuth').returns(mockSecondaryAuth)
 
         smusAuthProvider = new SmusAuthenticationProvider(mockAuth, mockSecondaryAuth)
+
+        // Reset the executeCommand stub for clean state
+        executeCommandStub.resetHistory()
     })
 
     afterEach(function () {
@@ -189,6 +197,7 @@ describe('SmusAuthenticationProvider', function () {
             assert.ok(getSsoInstanceInfoStub.calledWith(testDomainUrl))
             assert.ok(mockAuth.createConnection.called)
             assert.ok(mockSecondaryAuth.useNewConnection.called)
+            assert.ok(executeCommandStub.calledWith('aws.smus.switchProject'))
         })
 
         it('should reuse existing valid connection', async function () {
@@ -201,6 +210,7 @@ describe('SmusAuthenticationProvider', function () {
             assert.strictEqual(result, mockSmusConnection)
             assert.ok(mockAuth.createConnection.notCalled)
             assert.ok(mockSecondaryAuth.useNewConnection.calledWith(existingConnection))
+            assert.ok(executeCommandStub.calledWith('aws.smus.switchProject'))
         })
 
         it('should reauthenticate existing invalid connection', async function () {
@@ -213,6 +223,7 @@ describe('SmusAuthenticationProvider', function () {
             assert.strictEqual(result, mockSmusConnection)
             assert.ok(mockAuth.reauthenticate.calledWith(existingConnection))
             assert.ok(mockSecondaryAuth.useNewConnection.called)
+            assert.ok(executeCommandStub.calledWith('aws.smus.switchProject'))
         })
 
         it('should throw error for invalid domain URL', async function () {
@@ -225,6 +236,8 @@ describe('SmusAuthenticationProvider', function () {
                     return err.code === 'FailedToConnect' && (err.cause as any)?.code === 'InvalidDomainUrl'
                 }
             )
+            // Should not trigger project selection on error
+            assert.ok(executeCommandStub.notCalled)
         })
 
         it('should handle SmusUtils errors', async function () {
@@ -235,6 +248,8 @@ describe('SmusAuthenticationProvider', function () {
                 () => smusAuthProvider.connectToSmus(testDomainUrl),
                 (err: ToolkitError) => err.code === 'FailedToConnect'
             )
+            // Should not trigger project selection on error
+            assert.ok(executeCommandStub.notCalled)
         })
 
         it('should handle auth creation errors', async function () {
@@ -245,6 +260,47 @@ describe('SmusAuthenticationProvider', function () {
                 () => smusAuthProvider.connectToSmus(testDomainUrl),
                 (err: ToolkitError) => err.code === 'FailedToConnect'
             )
+            // Should not trigger project selection on error
+            assert.ok(executeCommandStub.notCalled)
+        })
+
+        it('should not trigger project selection in SMUS space environment', async function () {
+            isInSmusSpaceEnvironmentStub.returns(true)
+            mockAuth.listConnections.resolves([])
+
+            const result = await smusAuthProvider.connectToSmus(testDomainUrl)
+
+            assert.strictEqual(result, mockSmusConnection)
+            assert.ok(mockAuth.createConnection.called)
+            assert.ok(mockSecondaryAuth.useNewConnection.called)
+            assert.ok(executeCommandStub.notCalled)
+        })
+
+        it('should not trigger project selection when reusing connection in SMUS space environment', async function () {
+            isInSmusSpaceEnvironmentStub.returns(true)
+            const existingConnection = { ...mockSmusConnection, domainUrl: testDomainUrl.toLowerCase() }
+            mockAuth.listConnections.resolves([existingConnection])
+            mockAuth.getConnectionState.returns('valid')
+
+            const result = await smusAuthProvider.connectToSmus(testDomainUrl)
+
+            assert.strictEqual(result, mockSmusConnection)
+            assert.ok(mockSecondaryAuth.useNewConnection.calledWith(existingConnection))
+            assert.ok(executeCommandStub.notCalled)
+        })
+
+        it('should not trigger project selection when reauthenticating in SMUS space environment', async function () {
+            isInSmusSpaceEnvironmentStub.returns(true)
+            const existingConnection = { ...mockSmusConnection, domainUrl: testDomainUrl.toLowerCase() }
+            mockAuth.listConnections.resolves([existingConnection])
+            mockAuth.getConnectionState.returns('invalid')
+
+            const result = await smusAuthProvider.connectToSmus(testDomainUrl)
+
+            assert.strictEqual(result, mockSmusConnection)
+            assert.ok(mockAuth.reauthenticate.calledWith(existingConnection))
+            assert.ok(mockSecondaryAuth.useNewConnection.called)
+            assert.ok(executeCommandStub.notCalled)
         })
     })
 

packages/core/src/test/sagemakerunifiedstudio/shared/smusUtils.test.ts

@@ -164,7 +164,7 @@ describe('SmusUtils', () => {
         })
 
         it('should export SmusCredentialExpiry with correct values', () => {
-            assert.strictEqual(SmusCredentialExpiry.derExpiryMs, 55 * 60 * 1000)
+            assert.strictEqual(SmusCredentialExpiry.derExpiryMs, 10 * 60 * 1000)
             assert.strictEqual(SmusCredentialExpiry.projectExpiryMs, 10 * 60 * 1000)
             assert.strictEqual(SmusCredentialExpiry.connectionExpiryMs, 10 * 60 * 1000)
         })