feat(amazonq): Enable SageMaker SSO Users to use AmazonQ Chat & Code completion using their Credentials & Q Pro-tier profile #6338

ahusseinali · 2025-01-10T21:00:46Z

Problem

SageMaker Amazon Q Pro-tier users didn't have the ability to use their pro-tier profile ARN when using Amazon Q Chat & Code Completion. Previously, we had SageMaker pro-tier SSO users unable to automatically use their Code Editor login credentials to create an auth connection with Amazon Q APIs. This meant they have to re-login in order to use AmazonQ chat & Code Completion.

One big obstacle to enable auto-login for SSO users was that their credentials are stored in HttpOnly Cookie and a change had to be set in place in SageMaker Code Editor environment in order for AmazonQ extension to be able to read the cookie value and send it as bearer token.

Solution

Along with a change implemented for SageMaker Code Editor in https://github.com/aws/sagemaker-code-editor/pull/99/files
This PR implements the following:

SageMaker client to read SageMaker user domain status to identify whether pro-tier is enabled or not & read the cookie value through custom SageMaker code editor command.
1.1. In case Pro-tier is not enabled, the login flow is similar to IAM SageMaker users (using Container environment variables)
1.2. The custom command execution stores cookie values into a file and the client reads the value from that file.
SageMakerSSOTokenProvider: Custom token provider parallel to the original SSO Token provider. It doesn't go through any creation UI logic as the token already exists in file. This provider polls for token refresh every 5 minutes integrating with the client mentioned in Support listing multiple user-selected regions in explorer #1
Pass Q Pro-tier ProfileARN (if it is enabled) to API calls that need it for SSO users
Change the connection validation logic to account for valid pre-created SSO connections for SageMaker environment.
Add Auto connect logic for SSO user flow in SageMaker

Followup steps

Simplify the logic for communicating the SSO credentials and creating a connection based on such credentials, by letting SageMaker Code Editor create a standard AWS Credentials File & adding a watcher logic in AWS toolkit codebase to automatically sync with such file existence

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.
License: I confirm that my contribution is made under the terms of the Apache 2.0 license.

…redentials & Pro Tier profile ARN for AmazonQ features

github-actions · 2025-01-10T21:01:02Z

This pull request modifies code in src/* but no tests were added/updated.
- Confirm whether tests should be added or ensure the PR description explains why tests are not required.

hayemaxi · 2025-01-13T16:31:03Z

packages/core/src/auth/sso/sageMakerAccessTokenProvider.ts

+        let scopes = [...scopesSsoAccountAccess]
+        if (isAmazonQ()) {
+            scopes = [...scopes, ...SageMakerSpaceClient.getQConnectionScopes()]
+        }


This implies that this token provider would work with toolkit. Is that intended?

Are the SSO scopes required? If not, do not attach them.

Suggested change

let scopes = [...scopesSsoAccountAccess]

if (isAmazonQ()) {

scopes = [...scopes, ...SageMakerSpaceClient.getQConnectionScopes()]

}

const scopes = [isAmazonQ() ? ...SageMakerSpaceClient.getQConnectionScopes() : ...scopesSsoAccountAccess]

hayemaxi · 2025-01-13T16:32:06Z

packages/core/src/shared/sagemaker/client/sagemaker.ts

+    }
+
+    public static getQConnectionScopes() {
+        return ['codewhisperer:completions', 'codewhisperer:conversations']


Can we give users all amazon Q scopes? Why are we only giving this subset

hayemaxi · 2025-01-13T16:41:48Z

packages/amazonq/.changes/next-release/Feature-2ca899db-fff6-469c-8294-f93fc965cfa7.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Enable SageMaker SSO user access token & Profile ARN to be used for accessing AmazonQ & CodeWhisperer features"


"CodeWhisperer" is no more

Suggested change

"description": "Enable SageMaker SSO user access token & Profile ARN to be used for accessing AmazonQ & CodeWhisperer features"

"description": "Auth: Enable SSO access for SageMaker users."

hayemaxi · 2025-01-13T16:43:50Z

packages/core/src/auth/auth.ts

@@ -963,6 +1003,20 @@ export class Auth implements AuthService, ConnectionManager {
        return this.authenticate(id, refresh)
    }

+    public async tryAutoConnectSageMaker(): Promise<StatefulConnection | undefined> {
+        try {
+            const sagemakerConnection = await this.createSageMakerSsoConnection()


Could the connection already exist in a saved state from a prior session?

hayemaxi · 2025-01-13T17:25:52Z

packages/core/src/shared/sagemaker/client/sagemaker.ts

+    private constructor(sageMakerClient?: DefaultSageMakerClient) {
+        this.sageMakerClient = sageMakerClient ?? new DefaultSageMakerClient(this.getDefaultRegion())
+    }


I don't think this pattern makes sense. This class is a singleton accessible via getInstance only, so you can never pass a custom sageMakerClient. The constructor is private.

hayemaxi · 2025-01-13T17:28:36Z

packages/core/src/shared/sagemaker/client/sagemaker.ts

+    public static getInstance(): SageMakerSpaceClient {
+        if (!this._instance) {
+            this._instance = new SageMakerSpaceClient()
+        }
+        return this._instance
+    }


Follow the same format that we have in other singletons for consistency

aws-toolkit-vscode/packages/core/src/codewhisperer/util/authUtil.ts

Lines 234 to 241 in e6c3101

public static get instance() {

if (this.#instance !== undefined) {

return this.#instance

}

const self = (this.#instance = new this())

return self

}

Suggested change

public static getInstance(): SageMakerSpaceClient {

if (!this._instance) {

this._instance = new SageMakerSpaceClient()

}

return this._instance

}

public static get instance() {

if (!this.#instance) {

this.#instance = new SageMakerSpaceClient()

}

return this.#instance

}

hayemaxi · 2025-01-13T17:31:11Z

packages/core/src/shared/sagemaker/client/sagemaker.ts

+        return this.cachedCookies
+    }
+
+    public async getAmazonQProfileArn(): Promise<string | undefined> {


What is a profile ARN and why do we want it? Please add a docstring

hayemaxi · 2025-01-13T17:33:04Z

packages/core/src/shared/sagemaker/client/sagemaker.ts

+}
+
+// Manually expire the token every 5 minutes as token refresh is handled externally by Sagemaker
+const TOKEN_EXPIRY = 5 * 60 * 1000


Use constants in

aws-toolkit-vscode/packages/core/src/shared/datetime.ts

Line 10 in e6c3101

export const oneMinute = oneSecond * 60

e.g.

Suggested change

const TOKEN_EXPIRY = 5 * 60 * 1000

const TOKEN_EXPIRY = oneMinute * 5

hayemaxi · 2025-01-13T17:34:29Z

packages/core/src/shared/sagemaker/client/sagemaker.ts

+    authMode?: 'Sso' | 'Iam'
+    expiryTime?: number
+    ssoExpiryTimestamp?: number
+    studioUserProfileName?: string
+    redirectURL?: string
+    AccessToken?: string
+    StudioSessionToken?: string


Nothing is guaranteed to exist in the sagemaker cookie? Would be nice if we had some documentation about sagemaker expectations (if it doesn't already exist)

hayemaxi · 2025-01-13T17:37:10Z

packages/core/src/shared/sagemaker/client/sagemaker.ts

+
+// Manually expire the token every 5 minutes as token refresh is handled externally by Sagemaker
+const TOKEN_EXPIRY = 5 * 60 * 1000
+const DEFAULT_REGION = 'us-east-1'


We already have logic for this. Use

aws-toolkit-vscode/packages/core/src/shared/regions/regionProvider.ts

Line 46 in e6c3101

return this.awsContext.getCredentialDefaultRegion() ?? defaultRegion

in globals, e.g.:

globals.regionProvider

hayemaxi · 2025-01-13T17:40:19Z

packages/core/src/shared/index.ts

@@ -59,3 +59,4 @@ export { i18n } from './i18n-helper'
 export * from './icons'
 export * as textDocumentUtil from './utilities/textDocumentUtilities'
 export { TabTypeDataMap } from '../amazonq/webview/ui/tabs/constants'
+export * from './sagemaker/client/sagemaker'


Is packages/amazonq going to use this soon? Otherwise this isn't needed.

hayemaxi · 2025-01-13T17:43:13Z

packages/core/src/auth/utils.ts

+    let useDefaultCredentials = true
+
+    if (isSageMaker() && isAmazonQ()) {
+        useDefaultCredentials = !(await SageMakerSpaceClient.getInstance().getAmazonQProfileArn())
+    }
+
+    if (useDefaultCredentials) {


Suggested change

let useDefaultCredentials = true

if (isSageMaker() && isAmazonQ()) {

useDefaultCredentials = !(await SageMakerSpaceClient.getInstance().getAmazonQProfileArn())

}

if (useDefaultCredentials) {

const useDefaultCredentials = isSageMaker() && isAmazonQ() && !(await SageMakerSpaceClient.getInstance().getAmazonQProfileArn())

if (useDefaultCredentials) {

hayemaxi · 2025-01-13T17:47:39Z

packages/core/src/auth/sso/sageMakerAccessTokenProvider.ts

+
+        return {
+            scopes,
+            startUrl: 'sagemaker',


This isn't a URL... can we use sagemakerCookie.redirectURL? I'm worried about callers using this field with functions expecting a real URL.

ahusseinali added 2 commits January 10, 2025 12:46

feat(amazonq): Enable SageMaker SSO user to automatically their SSO c…

00378ee

…redentials & Pro Tier profile ARN for AmazonQ features

Enable file reading from command execution

0b6de5d

ahusseinali requested review from a team as code owners January 10, 2025 21:00

ahusseinali changed the title ~~Enable SageMaker SSO Users to use AmazonQ Chat & Code completion using their Credentials & Q Pro-tier profile~~ feat(amazonq): Enable SageMaker SSO Users to use AmazonQ Chat & Code completion using their Credentials & Q Pro-tier profile Jan 10, 2025

hayemaxi reviewed Jan 13, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(amazonq): Enable SageMaker SSO Users to use AmazonQ Chat & Code completion using their Credentials & Q Pro-tier profile #6338

feat(amazonq): Enable SageMaker SSO Users to use AmazonQ Chat & Code completion using their Credentials & Q Pro-tier profile #6338

ahusseinali commented Jan 10, 2025

github-actions bot commented Jan 10, 2025

hayemaxi Jan 13, 2025

hayemaxi Jan 13, 2025

hayemaxi Jan 13, 2025

hayemaxi Jan 13, 2025

hayemaxi Jan 13, 2025

hayemaxi Jan 13, 2025

hayemaxi Jan 13, 2025

hayemaxi Jan 13, 2025 •

edited

Loading

hayemaxi Jan 13, 2025

hayemaxi Jan 13, 2025 •

edited

Loading

hayemaxi Jan 13, 2025

hayemaxi Jan 13, 2025

hayemaxi Jan 13, 2025

	"description": "Enable SageMaker SSO user access token & Profile ARN to be used for accessing AmazonQ & CodeWhisperer features"
	"description": "Auth: Enable SSO access for SageMaker users."

	public static get instance() {
	if (this.#instance !== undefined) {
	return this.#instance
	}

	const self = (this.#instance = new this())
	return self
	}

	const TOKEN_EXPIRY = 5 * 60 * 1000
	const TOKEN_EXPIRY = oneMinute * 5

feat(amazonq): Enable SageMaker SSO Users to use AmazonQ Chat & Code completion using their Credentials & Q Pro-tier profile #6338

Are you sure you want to change the base?

feat(amazonq): Enable SageMaker SSO Users to use AmazonQ Chat & Code completion using their Credentials & Q Pro-tier profile #6338

Conversation

ahusseinali commented Jan 10, 2025

Problem

Solution

Followup steps

github-actions bot commented Jan 10, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hayemaxi Jan 13, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hayemaxi Jan 13, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hayemaxi Jan 13, 2025 •

edited

Loading

hayemaxi Jan 13, 2025 •

edited

Loading