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

docs/build.md

@@ -1,11 +1,23 @@
 # Build
 
-The AmazonQ features rely on the `codewhisperer-streaming` service, who's client
-is generated from the service's smithy models and placed in
-`src.gen/@amzn/codewhisperer-streaming` (For more
-information about this client and how it is generated, please see this
+The AmazonQ features rely on the `codewhisperer-streaming` service, to support both Sigv4 and Bearer token modes of this service,
+two clients are generated from the service's smithy models and placed in
+`src.gen/@amzn/amazon-q-developer-streaming-client` and `src.gen/@amzn/codewhisperer-streaming` respectively (For more
+information about these clients and how they are generated, please see this
 [quip document](https://quip-amazon.com/2dAWAvTIYXXr/Build-instructions-for-AWS-CodeWhisperer-Streaming-Typescript-client)).
 
+## @amzn/amazon-q-developer-streaming client
+
+This client is a standalone npm project in typescript, and it is added to
+the project as a workspace in the project's root `package.json` with the line `"workspaces": [ ..., "src.gen/@amzn/amazon-q-developer-streaming" ]`.
+The client may be manually built using `npm run build -w @amzn/amazon-q-developer-streaming"`.
+The `generateClients` run script ensures that this dependency is
+built before the toolkit project itself. Workspaces are automatically ready to
+be imported in the root toolkit project by their declared package.json name,
+(`@amzn/amazon-q-developer-streaming` in this case).
+
+## @amzn/codewhisperer-streaming client
+
 This client is a standalone npm project in typescript, and it is added to
 the project as a workspace in the project's root `package.json` with the line `"workspaces": [ ..., "src.gen/@amzn/codewhisperer-streaming" ]`.
 The client may be manually built using `npm run build -w @amzn/codewhisperer-streaming"`.

package.json

@@ -20,7 +20,7 @@
     },
     "scripts": {
         "prepare": "ts-node ./scripts/prepare.ts",
-        "preinstall": "cd src.gen/@amzn/codewhisperer-streaming && npm i",
+        "preinstall": "cd src.gen/@amzn/codewhisperer-streaming && npm i && cd ../amazon-q-developer-streaming-client && npm i",
         "postinstall": "npm run buildCustomLintPlugin && npm run postinstall -ws --if-present",
         "buildCustomLintPlugin": "npm run build -w plugins/eslint-plugin-aws-toolkits",
         "compile": "npm run compile -w packages/",

packages/amazonq/.changes/next-release/Bug Fix-0ee7be24-b0f8-460b-a852-a2d08b6d289e.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Use Sagemaker environment IAM Credentials for Code Completion when they're available"
+}

packages/amazonq/.changes/next-release/Bug Fix-4db8f989-da38-467f-bf2e-598c9ed9af81.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Inline: Code completion not working for Sagemaker Pro Tier users."
+}

packages/amazonq/.changes/next-release/Bug Fix-cbd8faed-6375-40e7-9f82-0e402e8fc0ad.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Disable /transform and /dev commands for sagemaker users as they're not supported"
+}

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

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Enable Free Tier Chat for IAM users"
+}

packages/amazonq/package.json

@@ -196,13 +196,13 @@
                     "type": "webview",
                     "id": "aws.amazonq.AmazonCommonAuth",
                     "name": "%AWS.amazonq.login%",
-                    "when": "!aws.isSageMaker && !aws.isWebExtHost && aws.amazonq.showLoginView"
+                    "when": "!aws.isWebExtHost && aws.amazonq.showLoginView"
                 },
                 {
                     "type": "webview",
                     "id": "aws.AmazonQChatView",
                     "name": "%AWS.amazonq.chat%",
-                    "when": "!aws.isSageMaker && !aws.isWebExtHost && !aws.amazonq.showLoginView"
+                    "when": "!aws.isWebExtHost && !aws.amazonq.showLoginView"
                 },
                 {
                     "id": "aws.AmazonQNeverShowBadge",

packages/amazonq/src/api.ts

@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { SendMessageCommandOutput, SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
 import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseRequest } from '@amzn/codewhisperer-streaming'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { ChatSession } from 'aws-core-vscode/codewhispererChat'
@@ -12,7 +13,11 @@ export default {
     chatApi: {
         async chat(request: GenerateAssistantResponseRequest): Promise<GenerateAssistantResponseCommandOutput> {
             const chatSession = new ChatSession()
-            return chatSession.chat(request)
+            return chatSession.chatSso(request)
+        },
+        async chatIam(request: SendMessageRequest): Promise<SendMessageCommandOutput> {
+            const chatSession = new ChatSession()
+            return chatSession.chatIam(request)
         },
     },
     authApi: {

packages/amazonq/src/extension.ts

@@ -44,6 +44,7 @@ import {
     setContext,
     setupUninstallHandler,
     maybeShowMinVscodeWarning,
+    isSageMaker,
 } from 'aws-core-vscode/shared'
 import { ExtStartUpSources, telemetry } from 'aws-core-vscode/telemetry'
 import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
@@ -193,7 +194,7 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
                 }
             }
             const currConn = AuthUtil.instance.conn
-            if (currConn !== undefined && !isAnySsoConnection(currConn)) {
+            if (currConn !== undefined && !(isAnySsoConnection(currConn) || isSageMaker())) {
                 getLogger().error(`Current Amazon Q connection is not SSO, type is: %s`, currConn?.type)
             }
 

packages/amazonq/src/inlineChat/controller/inlineChatController.ts

@@ -13,7 +13,16 @@ import { computeDecorations } from '../decorations/computeDecorations'
 import { CodelensProvider } from '../codeLenses/codeLenseProvider'
 import { PromptMessage, ReferenceLogController } from 'aws-core-vscode/codewhispererChat'
 import { CodeWhispererSettings } from 'aws-core-vscode/codewhisperer'
-import { codicon, getIcon, getLogger, messages, setContext, Timeout, textDocumentUtil } from 'aws-core-vscode/shared'
+import {
+    codicon,
+    getIcon,
+    getLogger,
+    messages,
+    setContext,
+    Timeout,
+    textDocumentUtil,
+    isSageMaker,
+} from 'aws-core-vscode/shared'
 import { InlineLineAnnotationController } from '../decorations/inlineLineAnnotationController'
 
 export class InlineChatController {
@@ -162,6 +171,11 @@ export class InlineChatController {
             return
         }
 
+        if (isSageMaker()) {
+            void vscode.window.showWarningMessage('Amazon Q: Inline chat is not supported in Sagemaker')
+            return
+        }
+
         if (this.task && this.task.isActiveState()) {
             void vscode.window.showWarningMessage(
                 'Amazon Q: Reject or Accept the current suggestion before creating a new one'

packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts

@@ -120,7 +120,7 @@ export class InlineChatProvider {
         let response: GenerateAssistantResponseCommandOutput | undefined = undefined
         session.createNewTokenSource()
         try {
-            response = await session.chat(request)
+            response = await session.chatSso(request)
             getLogger().info(
                 `response to tab: ${tabID} conversationID: ${session.sessionIdentifier} requestID: ${response.$metadata.requestId} metadata: ${JSON.stringify(
                     response.$metadata

packages/core/package.json

@@ -457,6 +457,7 @@
         "webfont": "^11.2.26"
     },
     "dependencies": {
+        "@amzn/amazon-q-developer-streaming-client": "file:../../src.gen/@amzn/amazon-q-developer-streaming-client",
         "@amzn/codewhisperer-streaming": "file:../../src.gen/@amzn/codewhisperer-streaming",
         "@aws-sdk/client-cognito-identity": "^3.637.0",
         "@aws-sdk/client-lambda": "^3.637.0",

packages/core/src/amazonq/extApi.ts

@@ -5,6 +5,7 @@
 
 import vscode from 'vscode'
 import { VSCODE_EXTENSION_ID } from '../shared/utilities'
+import { SendMessageCommandOutput, SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
 import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseRequest } from '@amzn/codewhisperer-streaming'
 import { FeatureAuthState } from '../codewhisperer/util/authUtil'
 import { ToolkitError } from '../shared'
@@ -16,6 +17,7 @@ import { ToolkitError } from '../shared'
 export interface api {
     chatApi: {
         chat(request: GenerateAssistantResponseRequest): Promise<GenerateAssistantResponseCommandOutput>
+        chatIam(request: SendMessageRequest): Promise<SendMessageCommandOutput>
     }
     authApi: {
         reauthIfNeeded(): Promise<void>

packages/core/src/amazonq/index.ts

@@ -42,11 +42,12 @@ import { FeatureContext } from '../shared'
 export function createMynahUI(
     ideApi: any,
     amazonQEnabled: boolean,
-    featureConfigsSerialized: [string, FeatureContext][]
+    featureConfigsSerialized: [string, FeatureContext][],
+    disabledCommands?: string[]
 ) {
     if (typeof window !== 'undefined') {
         const mynahUI = require('./webview/ui/main')
-        return mynahUI.createMynahUI(ideApi, amazonQEnabled, featureConfigsSerialized)
+        return mynahUI.createMynahUI(ideApi, amazonQEnabled, featureConfigsSerialized, disabledCommands)
     }
     throw new Error('Not implemented for node')
 }

packages/core/src/amazonq/webview/generators/webViewContent.ts

@@ -6,7 +6,7 @@
 import path from 'path'
 import { Uri, Webview } from 'vscode'
 import { AuthUtil } from '../../../codewhisperer/util/authUtil'
-import { FeatureConfigProvider, FeatureContext, globals } from '../../../shared'
+import { FeatureConfigProvider, FeatureContext, globals, isSageMaker } from '../../../shared'
 
 export class WebViewContentGenerator {
     private async generateFeatureConfigsData(): Promise<string> {
@@ -77,14 +77,16 @@ export class WebViewContentGenerator {
         // Fetch featureConfigs and use it within the script
         const featureConfigsString = await this.generateFeatureConfigsData()
 
+        const disabledCommandsString = isSageMaker() ? `['/dev', '/transform']` : '[]'
+
         return `
         <script type="text/javascript" src="${javascriptEntrypoint.toString()}" defer onload="init()"></script>
         ${cssLinks}
         <script type="text/javascript">
             const init = () => {
                 createMynahUI(acquireVsCodeApi(), ${
                     (await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected'
-                },${featureConfigsString});
+                },${featureConfigsString},${disabledCommandsString});
             }
         </script>
         `

packages/core/src/amazonq/webview/ui/main.ts

@@ -31,7 +31,8 @@ import { tryNewMap } from '../../util/functionUtils'
 export const createMynahUI = (
     ideApi: any,
     amazonQEnabled: boolean,
-    featureConfigsSerialized: [string, FeatureContext][]
+    featureConfigsSerialized: [string, FeatureContext][],
+    disabledCommands?: string[]
 ) => {
     // eslint-disable-next-line prefer-const
     let mynahUI: MynahUI
@@ -77,6 +78,7 @@ export const createMynahUI = (
     let tabDataGenerator = new TabDataGenerator({
         isFeatureDevEnabled,
         isGumbyEnabled,
+        disabledCommands,
     })
 
     // eslint-disable-next-line prefer-const
@@ -118,11 +120,13 @@ export const createMynahUI = (
                 tabsStorage,
                 isFeatureDevEnabled,
                 isGumbyEnabled,
+                disabledCommands,
             })
 
             tabDataGenerator = new TabDataGenerator({
                 isFeatureDevEnabled,
                 isGumbyEnabled,
+                disabledCommands,
             })
 
             featureConfigs = tryNewMap(featureConfigsSerialized)

packages/core/src/amazonq/webview/ui/messages/controller.ts

@@ -15,6 +15,7 @@ export interface MessageControllerProps {
     tabsStorage: TabsStorage
     isFeatureDevEnabled: boolean
     isGumbyEnabled: boolean
+    disabledCommands?: string[]
 }
 
 export class MessageController {
@@ -30,6 +31,7 @@ export class MessageController {
         this.tabDataGenerator = new TabDataGenerator({
             isFeatureDevEnabled: props.isFeatureDevEnabled,
             isGumbyEnabled: props.isGumbyEnabled,
+            disabledCommands: props.disabledCommands,
         })
     }
 

packages/core/src/amazonq/webview/ui/quickActions/generator.ts

@@ -9,22 +9,25 @@ import { TabType } from '../storages/tabsStorage'
 export interface QuickActionGeneratorProps {
     isFeatureDevEnabled: boolean
     isGumbyEnabled: boolean
+    disableCommands?: string[]
 }
 
 export class QuickActionGenerator {
     public isFeatureDevEnabled: boolean
     private isGumbyEnabled: boolean
+    private disabledCommands: string[]
 
     constructor(props: QuickActionGeneratorProps) {
         this.isFeatureDevEnabled = props.isFeatureDevEnabled
         this.isGumbyEnabled = props.isGumbyEnabled
+        this.disabledCommands = props.disableCommands ?? []
     }
 
     public generateForTab(tabType: TabType): QuickActionCommandGroup[] {
         const quickActionCommands = [
             {
                 commands: [
-                    ...(this.isFeatureDevEnabled
+                    ...(this.isFeatureDevEnabled && !this.disabledCommands.includes('/dev')
                         ? [
                               {
                                   command: '/dev',
@@ -33,7 +36,7 @@ export class QuickActionGenerator {
                               },
                           ]
                         : []),
-                    ...(this.isGumbyEnabled
+                    ...(this.isGumbyEnabled && !this.disabledCommands.includes('/transform')
                         ? [
                               {
                                   command: '/transform',
@@ -55,7 +58,7 @@ export class QuickActionGenerator {
                     },
                 ],
             },
-        ]
+        ].filter((section) => section.commands.length > 0)
 
         const commandUnavailability: Record<
             TabType,

packages/core/src/amazonq/webview/ui/quickActions/handler.ts

@@ -15,6 +15,7 @@ export interface QuickActionsHandlerProps {
     tabsStorage: TabsStorage
     isFeatureDevEnabled: boolean
     isGumbyEnabled: boolean
+    disabledCommands?: string[]
 }
 
 export class QuickActionHandler {
@@ -32,6 +33,7 @@ export class QuickActionHandler {
         this.tabDataGenerator = new TabDataGenerator({
             isFeatureDevEnabled: props.isFeatureDevEnabled,
             isGumbyEnabled: props.isGumbyEnabled,
+            disabledCommands: props.disabledCommands,
         })
         this.isFeatureDevEnabled = props.isFeatureDevEnabled
         this.isGumbyEnabled = props.isGumbyEnabled

packages/core/src/amazonq/webview/ui/tabs/generator.ts

@@ -12,6 +12,7 @@ import { TabTypeDataMap } from './constants'
 export interface TabDataGeneratorProps {
     isFeatureDevEnabled: boolean
     isGumbyEnabled: boolean
+    disabledCommands?: string[]
 }
 
 export class TabDataGenerator {
@@ -23,6 +24,7 @@ export class TabDataGenerator {
         this.quickActionsGenerator = new QuickActionGenerator({
             isFeatureDevEnabled: props.isFeatureDevEnabled,
             isGumbyEnabled: props.isGumbyEnabled,
+            disableCommands: props.disabledCommands,
         })
     }
 

packages/core/src/auth/activation.ts

@@ -3,16 +3,28 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import * as vscode from 'vscode'
 import { Auth } from './auth'
 import { LoginManager } from './deprecated/loginManager'
 import { fromString } from './providers/credentials'
 import { getLogger } from '../shared/logger'
-import { ExtensionUse } from './utils'
-import { isCloud9 } from '../shared/extensionUtilities'
+import { ExtensionUse, initializeCredentialsProviderManager } from './utils'
+import { isAmazonQ, isCloud9, isSageMaker } from '../shared/extensionUtilities'
 import { isInDevEnv } from '../shared/vscode/env'
 import { isWeb } from '../shared/extensionGlobals'
 
+interface SagemakerCookie {
+    authMode?: 'Sso' | 'Iam'
+}
+
 export async function initialize(loginManager: LoginManager): Promise<void> {
+    if (isAmazonQ() && isSageMaker()) {
+        // 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()
+        }
+    }
     Auth.instance.onDidChangeActiveConnection(async (conn) => {
         // This logic needs to be moved to `Auth.useConnection` to correctly record `passive`
         if (conn?.type === 'iam' && conn.state === 'valid') {

packages/core/src/auth/index.ts

@@ -17,6 +17,7 @@ export {
     isAnySsoConnection,
     isBuilderIdConnection,
     getTelemetryMetadataForConn,
+    isIamConnection,
 } from './connection'
 export { Auth } from './auth'
 export { CredentialsStore } from './credentials/store'