feat(amazonq): paid tier

Author: justinmk3

packages/amazonq/package.json

@@ -392,6 +392,11 @@
                     "when": "view == aws.amazonq.AmazonQChatView && aws.amazonq.connectedSsoIdc == true",
                     "group": "1_amazonQ@1"
                 },
+                {
+                    "command": "aws.amazonq.manageSubscription",
+                    "when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected",
+                    "group": "2_amazonQ@4"
+                },
                 {
                     "command": "aws.amazonq.signout",
                     "when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected && !aws.isSageMakerUnifiedStudio",
@@ -673,6 +678,13 @@
                 "category": "%AWS.amazonq.title%",
                 "icon": "$(question)"
             },
+            {
+                "command": "aws.amazonq.manageSubscription",
+                "title": "%AWS.command.manageSubscription%",
+                "category": "%AWS.amazonq.title%",
+                "icon": "$(gear)",
+                "enablement": "aws.codewhisperer.connected"
+            },
             {
                 "command": "aws.amazonq.signout",
                 "title": "%AWS.command.codewhisperer.signout%",

packages/amazonq/src/lsp/auth.ts

@@ -96,6 +96,8 @@ export class AmazonQLspAuth {
             token,
         })
 
+        // "aws/credentials/token/update"
+        // https://github.com/aws/language-servers/blob/44d81f0b5754747d77bda60b40cc70950413a737/core/aws-lsp-core/src/credentials/credentialsProvider.ts#L27
         await this.client.sendRequest(bearerCredentialsUpdateRequestType.method, request)
 
         this.client.info(`UpdateBearerToken: ${JSON.stringify(request)}`)

packages/amazonq/src/lsp/chat/activation.ts

@@ -6,7 +6,7 @@
 import { window } from 'vscode'
 import { LanguageClient } from 'vscode-languageclient'
 import { AmazonQChatViewProvider } from './webviewProvider'
-import { registerCommands } from './commands'
+import { focusAmazonQPanel, registerCommands } from './commands'
 import { registerLanguageServerEventListener, registerMessageListeners } from './messages'
 import { Commands, getLogger, globals, undefinedIfEmpty } from 'aws-core-vscode/shared'
 import { activate as registerLegacyChatListeners } from '../../app/chat/activation'
@@ -73,6 +73,14 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu
                 customization: undefinedIfEmpty(getSelectedCustomization().arn),
             })
         }),
+        Commands.register('aws.amazonq.manageSubscription', () => {
+            focusAmazonQPanel().catch((e) => languageClient.error(`[VSCode Client] focusAmazonQPanel() failed`))
+
+            languageClient.sendRequest('workspace/executeCommand', {
+                command: 'aws/chat/manageSubscription',
+                // arguments: [],
+            })
+        }),
         globals.logOutputChannel.onDidChangeLogLevel((logLevel) => {
             getLogger('amazonqLsp').info(`Local log level changed to ${logLevel}, notifying LSP`)
             void pushConfigUpdate(languageClient, {

packages/amazonq/src/lsp/chat/commands.ts

@@ -125,7 +125,7 @@ function registerGenericCommand(commandName: string, genericCommand: string, pro
  *
  * Instead, we just create our own as a temporary solution
  */
-async function focusAmazonQPanel() {
+export async function focusAmazonQPanel() {
     await Commands.tryExecute('aws.amazonq.AmazonQChatView.focus')
     await Commands.tryExecute('aws.amazonq.AmazonCommonAuth.focus')
 }

packages/amazonq/src/lsp/chat/messages.ts

@@ -61,6 +61,7 @@ import * as jose from 'jose'
 import { AmazonQChatViewProvider } from './webviewProvider'
 import { AuthUtil, ReferenceLogViewProvider } from 'aws-core-vscode/codewhisperer'
 import { amazonQDiffScheme, AmazonQPromptSettings, messages, openUrl } from 'aws-core-vscode/shared'
+import { credentialsValidation } from 'aws-core-vscode/auth'
 import {
     DefaultAmazonQAppInitContext,
     messageDispatcher,
@@ -70,6 +71,7 @@ import {
 } from 'aws-core-vscode/amazonq'
 import { telemetry, TelemetryBase } from 'aws-core-vscode/telemetry'
 import { isValidResponseError } from './error'
+import { focusAmazonQPanel } from './commands'
 
 export function registerLanguageServerEventListener(languageClient: LanguageClient, provider: AmazonQChatViewProvider) {
     languageClient.info(
@@ -326,13 +328,33 @@ export function registerMessageListeners(
                 }
                 break
             case buttonClickRequestType.method: {
+                if (message.params.buttonId === 'paidtier-upgrade-q') {
+                    focusAmazonQPanel().catch((e) => languageClient.error(`[VSCode Client] focusAmazonQPanel() failed`))
+
+                    const accountId = await vscode.window.showInputBox({
+                        title: 'Upgrade Amazon Q',
+                        prompt: 'Enter your 12-digit AWS account ID',
+                        placeHolder: '111111111111',
+                        validateInput: credentialsValidation.validateAwsAccount,
+                    })
+
+                    if (accountId) {
+                        languageClient.sendRequest('workspace/executeCommand', {
+                            command: 'aws/chat/manageSubscription',
+                            arguments: [accountId],
+                        })
+                    } else {
+                        languageClient.error('[VSCode Client] user canceled or did not input AWS account id')
+                    }
+                }
+
                 const buttonResult = await languageClient.sendRequest<ButtonClickResult>(
                     buttonClickRequestType.method,
                     message.params
                 )
                 if (!buttonResult.success) {
                     languageClient.error(
-                        `[VSCode Client] Failed to execute action associated with button with reason: ${buttonResult.failureReason}`
+                        `[VSCode Client] Failed to execute button action: ${buttonResult.failureReason}`
                     )
                 }
                 break
@@ -431,6 +453,8 @@ export function registerMessageListeners(
     languageClient.onRequest<ShowDocumentParams, ShowDocumentResult>(
         ShowDocumentRequest.method,
         async (params: ShowDocumentParams): Promise<ShowDocumentParams | ResponseError<ShowDocumentResult>> => {
+            focusAmazonQPanel().catch((e) => languageClient.error(`[VSCode Client] focusAmazonQPanel() failed`))
+
             try {
                 const uri = vscode.Uri.parse(params.uri)
 

packages/core/package.nls.json

@@ -137,6 +137,7 @@
     "AWS.command.codecatalyst.login": "Connect to CodeCatalyst",
     "AWS.command.codecatalyst.logout": "Sign out of CodeCatalyst",
     "AWS.command.codecatalyst.signout": "Sign Out",
+    "AWS.command.manageSubscription": "Manage Subscription",
     "AWS.command.amazonq.explainCode": "Explain",
     "AWS.command.amazonq.refactorCode": "Refactor",
     "AWS.command.amazonq.fixCode": "Fix",

packages/core/src/auth/credentials/validation.ts

@@ -127,6 +127,14 @@ async function validateProfileName(profileName: SectionName) {
     }
 }
 
+export function validateAwsAccount(s: string): string | undefined {
+    // AWS account IDs are exactly 12 digits
+    if (!/^\d{12}$/.test(s)) {
+        return 'Enter a valid 12-digit AWS account ID'
+    }
+    return undefined
+}
+
 // All shared credentials keys
 const sharedCredentialsKeysSet = new Set(Object.values(SharedCredentialsKeys))
 

packages/core/src/auth/index.ts

@@ -24,3 +24,4 @@ export { Auth } from './auth'
 export { CredentialsStore } from './credentials/store'
 export { LoginManager } from './deprecated/loginManager'
 export * as AuthUtils from './utils'
+export * as credentialsValidation from './credentials/validation'

packages/core/src/codewhisperer/ui/codeWhispererNodes.ts

@@ -175,6 +175,19 @@ export function createGettingStarted(): DataQuickPickItem<'gettingStarted'> {
     } as DataQuickPickItem<'gettingStarted'>
 }
 
+export function createManageSubscription(): DataQuickPickItem<'manageSubscription'> {
+    const label = localize('AWS.command.manageSubscription', 'Manage Subscription')
+    // const kind = AuthUtil.instance.isBuilderIdInUse() ? 'AWS Builder ID' : 'IAM Identity Center'
+
+    return {
+        data: 'manageSubscription',
+        label: label,
+        iconPath: getIcon('vscode-link-external'),
+        // TODO
+        onClick: () => Commands.tryExecute('aws.amazonq.manageSubscription'),
+    } as DataQuickPickItem<'manageSubscription'>
+}
+
 export function createSignout(): DataQuickPickItem<'signout'> {
     const label = localize('AWS.codewhisperer.signoutNode.label', 'Sign Out')
     const icon = getIcon('vscode-export')

packages/core/src/codewhisperer/ui/statusBarMenu.ts

@@ -11,6 +11,7 @@ import {
     createSelectCustomization,
     createReconnect,
     createGettingStarted,
+    createManageSubscription,
     createSignout,
     createSeparator,
     createSettingsNode,
@@ -106,7 +107,7 @@ export function getQuickPickItems(): DataQuickPickItem<string>[] {
         createSettingsNode(),
         ...(isUsingEnterpriseSso && regionProfile ? [createSelectRegionProfileNode(regionProfile)] : []),
         ...(AuthUtil.instance.isConnected() && !hasVendedIamCredentials() && !hasVendedCredentialsFromMetadata()
-            ? [createSignout()]
+            ? [createManageSubscription(), createSignout()]
             : []),
     ]