feat(sagemaker): free tier Q Chat with auto-login for iam users and login option for pro tier users #5858

ahusseinali · 2024-10-24T18:37:17Z

Problem

AmazonQ

IAM users didn't have access to Q Chat (Free Tier)
Sagemaker IAM users weren't automatically logged in using their Sagemaker studio credentials.
For Sagemaker users, login screen was disabled & users didn't have access to chat or code completion.
Sagemaker Pro Tier users code completion was broken & didn't generate any recommendations even after successful login.
Q Icon hides from Activity Bar when Toolkit extension is installed

Solution

Add new Q Developer Client to enable Free tier chat.
Terminate all the conditions that used to force chat to get disabled for IAM users
Terminate the conditions that used to hide login screen for Sagemaker users
Login Sagemaker IAM users automatically using their environment credentials (Free tier)
Allow Sagemaker SSO pro tier users to login using the login screen.

License: I confirm that my contribution is made under the terms of the Apache 2.0 license.

github-actions · 2024-10-24T18:37:32Z

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.

nkomonen-amazon · 2024-10-24T19:09:50Z

package-lock.json

There is a large change here and it is not clear what triggered it. Are you able to extract this related change to an isolated commit so we can better understand

I'm not sure what triggerred it either, but I believe it's the dependencies of the new Q developer client I added in src.gen

That new generated client won't result in new package-lock.json.

At this point I'm clueless as to why the package.json file was changed. But I see many of the dependencies added there are also part of the newly added client package-lock.json. this is why I suspected they're related.

I inspected the dependencies added / changed further and they're only relevant to the new client code

packages/amazonq/.changes/next-release/Feature-a007f02d-d737-4429-9b7d-245aa30d879c.json

hayemaxi · 2024-10-24T20:45:14Z

src.gen/@amzn/amazon-q-developer-streaming-client/package.json

@@ -0,0 +1,91 @@
+{
+    "name": "@amzn/amazon-q-developer-streaming-client",


What is the purpose of this package? Why can we not use src.gen/@amzn/codewhisperer-streaming?

To answer this, please update https://github.com/aws/aws-toolkit-vscode/blob/master/docs/build.md

As discussed offline, this client supports Sigv4 credentials instead of bearer token. this is mandatory to support both IAM and SSO users.

The instructions to generate these clients will be updated in the same quip doc referenced in this .md file
I updated the wording to reflect the new changes

packages/toolkit/package.json

packages/core/src/shared/clients/qDeveloperChatClient.ts

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

hayemaxi · 2024-10-24T21:15:33Z

packages/core/src/auth/activation.ts


 export async function initialize(loginManager: LoginManager): Promise<void> {
+    if (isAmazonQ() && isSageMaker()) {
+        const result = (await vscode.commands.executeCommand('sagemaker.parseCookies')) as SagemakerCookie


Where does this command come from? Is it already registered in VSC by sagemaker? If so add a comment to mention this.

Correct, this command is registered in Sagemaker code editor. I'll add a comment

Lets make a new sagemaker folder in src/shared/sagemaker for all custom sagemaker code so that it is contained within a single place.

That's a big refactor that I'd rather defer to a later followup to reduce risk

This isn't a long-term approach. The Toolkit/Q extensions already have an api to accept connections. If needed, they can be enhanced to accept tokens/credentials, so the caller (SM in this case) can inject credentials into the extension.

packages/core/src/auth/activation.ts

hayemaxi · 2024-10-24T21:27:23Z

packages/core/src/codewhisperer/util/authUtil.ts


    if (isSageMaker()) {
-        return isIamConnection(conn)
+        return isIamConnection(conn) || (isSsoConnection(conn) && hasScopes(conn, codeWhispererCoreScopes))


Can we use amazonQScopes? Does sagemaker SSO auth allow for the complete Q scope set?

Suggested change

return isIamConnection(conn) || (isSsoConnection(conn) && hasScopes(conn, codeWhispererCoreScopes))

return isIamConnection(conn) || (isSsoConnection(conn) && hasScopes(conn, amazonQScopes))

It is code smell that we have these scope subsets, so it is easy to confuse which one needs to be used/checked. We have plans to fix this.

No it doesn't, but tackling the scopes is a separate issue that I'll address later on. For now, if user login with their SSO credentials, they will go through the normal login flow which appends the full set of credentials for them. I can omit the scope check from here though as it doesn't make a difference for Sagemaker users.

For now, I'll simplify the condition and put in a separate function for readability. it ignores any scope checking as we're handling deactivating features not supported by sagemaker in a separate change

packages/core/src/codewhisperer/util/authUtil.ts

hayemaxi · 2024-10-24T21:33:02Z

packages/core/src/codewhisperer/util/authUtil.ts

-        if (!isSsoConnection(conn)) {
-            throw new ToolkitError(`Connection "${conn.id}" is not a valid type: ${conn.type}`)
-        }


Instead of deleting this you can add the check

if (!isSsoConnection(conn) && !isSageMaker()) { ... }

but this change is not just for Sagemaker, the new client enables chat for both Sso and Iam connections, so this restriction no longer applies with this change

but this change is not just for Sagemaker

But this isn't a supported feature yet in all of the code yet right? That is a big product change to the Q extension. AFAIK there is no way to get IAM credentials outside of sagemaker. Until it's supported throughout we should still have these sanity checks. Also, for a change that big I would expect unit or integ test updates/additions if possible.

Ack, will adjust it now

packages/core/src/codewhisperer/util/authUtil.ts

packages/core/src/codewhispererChat/controllers/chat/controller.ts

packages/core/src/codewhisperer/util/authUtil.ts

hayemaxi

lgtm pending team approval

leigaol · 2024-10-28T19:19:53Z

packages/core/src/codewhisperer/service/recommendationHandler.ts

        page: number = 0,
-        isSM: boolean = isSageMaker(),
-        retry: boolean = false
+        generate: boolean = isIamConnection(AuthUtil.instance.conn)


This should no longer be a single function for both IAM and bearer token code path. The number of if-else is too much to comprehend. Not a blocker for this PR

Please do lots of regression testing post merge

Test /dev and /transform

bhadrip · 2024-10-28T19:42:22Z

packages/core/src/codewhispererChat/controllers/chat/controller.ts

+                    message: generateAssistantResponseResponse,
+                }
+            } else {
+                const { $metadata, sendMessageResponse } = await session.chatIam(request as SendMessageRequest)


source="IDE"
we need to set the request for the right routing in Q Dev

bhadrip · 2024-10-28T19:45:32Z

packages/core/src/codewhispererChat/controllers/chat/controller.ts

+                    message: sendMessageResponse,
+                }
+            }
            this.telemetryHelper.recordEnterFocusConversation(triggerEvent.tabID)


what kind of telemetry or error mechanism we can use to understand customers facing issues with the chat window ?

bhadrip · 2024-10-28T19:51:41Z

packages/core/src/auth/activation.ts

+        // The command `sagemaker.parseCookies` is registered in VS Code Sagemaker environment.
+        const result = (await vscode.commands.executeCommand('sagemaker.parseCookies')) as SagemakerCookie
+        if (result.authMode !== 'Sso') {
+            initializeCredentialsProviderManager()


how does this refactor affect toolkit ? In sagemaker studio sso mode, toolkit will still work in IAM mode.

The higher level condition is scoped to amazonq, so it shouldn't impact aws toolkit

bhadrip

how do we handle graceful failure of features that are not supported in free tier IAM mode

bhadrip

are there no unit tests for the changes ?

justinmk3 · 2024-10-28T20:01:03Z

please add a src.gen/@amzn/.gitignore file that will avoid mistakes like 972e2ad in the future.

ahusseinali · 2024-10-28T23:37:11Z

are there no unit tests for the changes ?

All the files I touched didn't have unit tests related to the changes. I plan to improve the coverage of this package as the code paths (especially for auth and API calls) are very critical

ahusseinali · 2024-10-28T23:45:25Z

please add a src.gen/@amzn/.gitignore file that will avoid mistakes like 972e2ad in the future.

Added rule

hayemaxi · 2024-10-29T03:42:53Z

CI is failing:

Error: src/inlineChat/provider/inlineChatProvider.ts(123,38): error TS2339: Property 'chat' does not exist on type 'ChatSession'.
Error: src/inlineChat/provider/inlineChatProvider.ts(125,102): error TS18048: 'response' is possibly 'undefined'.
Error: src/inlineChat/provider/inlineChatProvider.ts(126,21): error TS18048: 'response' is possibly 'undefined'.

another PR merged to master is using old function chat which is replaced in this PR by chatSso and chatIam. The logic will need to be updated to determine which one to call.

SM IAM permissions have been working as of v1.99 and now it's broken. This change brings it back to support code completion for SM code editor

gogakoreli

Verified in Sagemaker

…rs and login option for pro tier users" (#5884) Reverts #5858

ahusseinali requested review from a team as code owners October 24, 2024 18:37

nkomonen-amazon reviewed Oct 24, 2024

View reviewed changes