wip

justinmk3 · justinmk3 · commit abf4e15e7f41 · 2025-05-22T13:05:25.000-04:00
diff --git a/chat-client/src/client/mynahUi.ts b/chat-client/src/client/mynahUi.ts
@@ -57,7 +57,7 @@ import {
 } from './utils'
 import { ChatHistory, ChatHistoryList } from './features/history'
 import { pairProgrammingModeOff, pairProgrammingModeOn, programmerModeCard } from './texts/pairProgramming'
-import { paidTierCard, upgradeQButton } from './texts/paidTier'
+import { paidTierSuccessCard, freeTierLimitReachedCard, upgradeQButton } from './texts/paidTier'
 
 export interface InboundChatApi {
     addChatResponse(params: ChatResult, tabId: string, isPartialResult: boolean): void
@@ -822,20 +822,39 @@ export const createMynahUi = (
         })
     }
 
-    const updateChat = (params: ChatUpdateParams) => {
-        if (params.data?.placeholderText === 'upgrade-q') {
-            const tabId = params.tabId !== 'xxx' ? params.tabId : getOrCreateTabId()!
-            const upgradeQMode: 'paidtier' | 'freetier' | 'freetier-limit' = (params as any).upgradeQMode
-            // TODO: only do this for 'freetier-limit'
-            const needUpgrade = upgradeQMode === 'freetier' || upgradeQMode === 'freetier-limit'
+    /**
+     * Adjusts the UI when the user changes to/from free-tier/paid-tier.
+     * Shows a message if the user reaches free-tier limit.
+     * Shows a message if the user just upgraded to paid-tier.
+     */
+    const onPaidTierModeChange = (
+        tabId: string,
+        mode: 'paidtier' | 'paidtier-success' | 'freetier' | 'freetier-limit'
+    ) => {
+        if (!['paidtier', 'paidtier-success', 'freetier', 'freetier-limit'].includes(mode)) {
+            return // invalid mode
+        }
 
-            mynahUi.updateStore(tabId, {
-                promptInputButtons: needUpgrade ? [upgradeQButton] : [],
-                promptInputDisabledState: needUpgrade,
-            })
-            if (needUpgrade) {
-                mynahUi.addChatItem(tabId, paidTierCard)
-            }
+        tabId = tabId !== '' ? tabId : getOrCreateTabId()!
+
+        // Detect if the tab is already showing the "Upgrade Q" calls-to-action.
+        const didShowLimitReached = mynahUi.getTabData(tabId)?.getStore()?.promptInputButtons?.[0] === upgradeQButton
+        if (mode === 'freetier-limit' && !didShowLimitReached) {
+            mynahUi.addChatItem(tabId, freeTierLimitReachedCard)
+        } else if (mode === 'paidtier-success') {
+            mynahUi.addChatItem(tabId, paidTierSuccessCard)
+        }
+
+        mynahUi.updateStore(tabId, {
+            promptInputButtons: mode === 'freetier-limit' ? [upgradeQButton] : [],
+            promptInputDisabledState: mode === 'freetier-limit',
+        })
+    }
+
+    const updateChat = (params: ChatUpdateParams) => {
+        // HACK: Special field sent by `agenticChatController.ts:setPaidTierMode()`.
+        if ((params as any).paidTierMode) {
+            onPaidTierModeChange(params.tabId, (params as any).paidTierMode as any)
             return
         }
 
diff --git a/chat-client/src/client/tabs/tabFactory.ts b/chat-client/src/client/tabs/tabFactory.ts
@@ -10,7 +10,6 @@ import { disclaimerCard } from '../texts/disclaimer'
 import { ChatMessage } from '@aws/language-server-runtimes-types'
 import { ChatHistory } from '../features/history'
 import { pairProgrammingPromptInput, programmerModeCard } from '../texts/pairProgramming'
-import { paidTierCard, upgradeQButton } from '../texts/paidTier'
 
 export type DefaultTabData = MynahUIDataModel
 
@@ -50,15 +49,13 @@ export class TabFactory {
     public getChatItems(
         needWelcomeMessages: boolean,
         pairProgrammingCardActive: boolean,
-        // paidTierCardActive: boolean,
         chatMessages?: ChatMessage[]
     ): ChatItem[] {
         return [
             ...(this.bannerMessage ? [this.getBannerMessage() as ChatItem] : []),
             ...(needWelcomeMessages
                 ? [
                       ...(this.agenticMode && pairProgrammingCardActive ? [programmerModeCard] : []),
-                      // ...(paidTierCardActive ? [paidTierCard] : []),
                       {
                           type: ChatItemType.ANSWER,
                           body: `Hi, I'm Amazon Q. I can answer your software development questions. 
diff --git a/chat-client/src/client/texts/paidTier.ts b/chat-client/src/client/texts/paidTier.ts
@@ -1,21 +1,41 @@
-import { ChatItem, ChatItemButton, ChatItemFormItem, ChatItemType } from '@aws/mynah-ui'
+import { ChatItem, ChatItemButton, ChatItemFormItem, ChatItemType, TextBasedFormItem } from '@aws/mynah-ui'
 
-export const paidTierCard: ChatItem = {
+export const freeTierLimitReachedCard: ChatItem = {
     type: ChatItemType.ANSWER,
     title: 'FREE TIER LIMIT REACHED',
     header: {
         icon: 'q',
         iconStatus: 'primary',
         body: 'Upgrade to Amazon Q Pro',
     },
-    messageId: 'paidTierCardId',
+    messageId: 'freetier-limit',
     fullWidth: true,
     canBeDismissed: false,
     body: 'You have reached the free tier limit. Upgrade to Amazon Q Pro.\n\n[Learn More...](https://aws.amazon.com/q/pricing/)',
 }
 
-export const paidTierPromptInput: ChatItemFormItem = {
-    type: 'switch',
+export const paidTierSuccessCard: ChatItem = {
+    type: ChatItemType.ANSWER,
+    title: 'UPGRADED TO AMAZON Q PRO',
+    header: {
+        icon: 'q',
+        iconStatus: 'primary',
+        body: 'Welcome to Amazon Q Pro',
+        status: {
+            status: 'success',
+            icon: 'q',
+            text: 'Success',
+        },
+    },
+    messageId: 'paidtier-success',
+    fullWidth: true,
+    canBeDismissed: true,
+    body: 'Upgraded to Amazon Q Pro\n\n[Learn More...](https://aws.amazon.com/q/)',
+}
+
+export const paidTierPromptInput: TextBasedFormItem = {
+    placeholder: '111111111111',
+    type: 'textinput',
     id: 'paid-tier',
     tooltip: 'Upgrade to Amazon Q Pro',
     value: 'true',
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts
@@ -30,6 +30,7 @@ import {
     PromptInputOptionChangeParams,
     TextDocument,
     ChatUpdateParams,
+    MessageType,
 } from '@aws/language-server-runtimes/protocol'
 import {
     ApplyWorkspaceEditParams,
@@ -132,6 +133,7 @@ import { URI } from 'vscode-uri'
 import { AgenticChatError, customerFacingErrorCodes, isRequestAbortedError, unactionableErrorCodes } from './errors'
 import { CommandCategory } from './tools/executeBash'
 import { UserWrittenCodeTracker } from '../../shared/userWrittenCodeTracker'
+import { PaidTierMode } from '../paidTier/paidTier'
 
 type ChatHandlers = Omit<
     LspHandlers<Chat>,
@@ -163,7 +165,7 @@ export class AgenticChatController implements ChatHandlers {
     #userWrittenCodeTracker: UserWrittenCodeTracker | undefined
     #toolUseStartTimes: Record<string, number> = {}
     #toolUseLatencies: Array<{ toolName: string; toolUseId: string; latency: number }> = []
-    #freeTierLimit = false
+    #paidTierMode: PaidTierMode | undefined
 
     /**
      * Determines the appropriate message ID for a tool use based on tool type and name
@@ -277,29 +279,47 @@ export class AgenticChatController implements ChatHandlers {
                     .createSubscriptionToken({
                         accountId: awsAccountId,
                     })
+
                 if (!r.encodedVerificationUrl) {
                     this.#log('missing encodedVerificationUrl in server response')
+                    this.#features.lsp.window
+                        .showMessage({
+                            message: 'Subscription request failed. Check the account id.',
+                            type: MessageType.Error,
+                        })
+                        .catch(e => {
+                            this.#log(`showMessage failed: ${(e as Error).message}`)
+                        })
                     return {
                         success: false,
                         failureReason: 'missing encodedVerificationUrl in server response',
                     }
                 }
 
-                this.#log(`encodedVerificationUrl ${r.encodedVerificationUrl}`)
+                const uri = r.encodedVerificationUrl
 
                 try {
-                    URI.parse(r.encodedVerificationUrl)
+                    URI.parse(uri)
                 } catch (e) {
-                    this.#log(`invalid encodedVerificationUrl: '${r.encodedVerificationUrl}': ${(e as Error).message}`)
+                    this.#log(`invalid encodedVerificationUrl: '${uri}': ${(e as Error).message}`)
                     return {
                         success: false,
                         failureReason: 'invalid encodedVerificationUrl',
                     }
                 }
 
+                this.#features.lsp.window
+                    .showMessage({
+                        message: 'Upgraded to [Amazon Q Pro](https://aws.amazon.com/q/)',
+                        type: MessageType.Info,
+                    })
+                    .catch(e => {
+                        this.#log(`showMessage failed: ${(e as Error).message}`)
+                    })
+
                 this.#features.lsp.window.showDocument({
                     external: true, // Client is expected to open the URL in a web browser.
-                    uri: r.encodedVerificationUrl,
+                    uri: uri,
                 })
             } catch (e) {
                 return {
@@ -308,7 +328,7 @@ export class AgenticChatController implements ChatHandlers {
                 }
             }
 
-            this.setUpgradeQMode(params.tabId, 'paidtier')
+            this.setPaidTierMode(params.tabId, 'paidtier-success')
             return { success: true }
         } else {
             return {
@@ -1971,7 +1991,7 @@ export class AgenticChatController implements ChatHandlers {
 
         // TODO handle free tier limit exceeded
         if (isFreeTierLimitError(err)) {
-            this.setUpgradeQMode(tabId, 'freetier-limit')
+            this.setPaidTierMode(tabId, 'freetier-limit')
             // throw new AmazonQFreeTierLimitError()
         }
 
@@ -2239,12 +2259,6 @@ export class AgenticChatController implements ChatHandlers {
         } catch (error) {
             this.#log('Error initializing context commands: ' + error)
         }
-
-        try {
-            this.setUpgradeQMode()
-        } catch (err) {
-            this.#log('Error initializing Free Tier state: ' + (err as Error).message)
-        }
     }
 
     onSendFeedback({ tabId, feedbackPayload }: FeedbackParams) {
@@ -2271,7 +2285,7 @@ export class AgenticChatController implements ChatHandlers {
 
         this.#chatSessionManagementService.createSession(params.tabId)
 
-        this.setUpgradeQMode(params.tabId)
+        this.setPaidTierMode(params.tabId)
     }
 
     onTabChange(params: TabChangeParams) {
@@ -2434,57 +2448,44 @@ export class AgenticChatController implements ChatHandlers {
      * Updates the "Upgrade Q" (subscription tier) state of the UI in the chat component. If `mode` is not given, the user's subscription status is checked by calling the Q service.
      *
      * `mode` behavior:
-     * - 'freetier': always show "Upgrade Q" button.
+     * - 'freetier': treated as 'freetier-limit' if `this.#paidTierMode='freetier-limit'`.
      * - 'freetier-limit': also show "Free Tier limit reached" card in chat.
      *     - This mode is "sticky" until 'paidtier' is passed to override it.
-     * - 'paidtier': don't show "Upgrade Q" button.
+     * - 'paidtier': disable any "free-tier limit" UI.
      */
-    async setUpgradeQMode(
-        tabId?: string,
-        mode?: 'paidtier' | 'freetier' | 'freetier-limit' /*, session: ChatSessionService*/
-    ) {
-        this.#log(
-            `xxx setUpgradeQMode: mode=${mode} getCodewhispererService=${!!this.#serviceManager?.getCodewhispererService()}`
-        )
+    setPaidTierMode(tabId?: string, mode?: PaidTierMode) {
+        this.#log(`xxx setPaidTierMode: mode=${mode}`)
+
         if (mode === 'freetier-limit') {
-            this.#freeTierLimit = true // Sticky until 'paidtier' is sent.
+            this.#paidTierMode = mode // Sticky until 'paidtier' is sent.
         } else if (mode === 'paidtier') {
-            this.#freeTierLimit = false
-        } else if (this.#freeTierLimit && (!mode || mode === 'freetier')) {
+            this.#paidTierMode = mode
+        } else if (this.#paidTierMode === 'freetier-limit' && mode === 'freetier') {
             mode = 'freetier-limit'
         } else if (!mode) {
-            try {
-                // Note: intentionally async.
-                AmazonQTokenServiceManager.getInstance()
-                    ?.getCodewhispererService()
-                    .getSubscriptionStatus()
-                    .then(o => {
-                        this.#log(`xxx getSubscriptionStatus: ${o.status} ${o.encodedVerificationUrl}`)
-                        this.setUpgradeQMode(tabId, o.status === 'ACTIVE' ? 'paidtier' : 'freetier')
-                    })
-                    .catch(err => {
-                        this.#log(`xxx getSubscriptionStatus failed: ${JSON.stringify(err)}`)
-                    })
-
-                // const isFreeTierUser = getSsoConnectionType(this.#features.credentialsProvider) === 'builderId'
-                // mode = isFreeTierUser ? 'freetier' : 'paidtier'
-            } catch {
-                this.#log(`xxx yucky`)
-            }
+            // Note: intentionally async.
+            AmazonQTokenServiceManager.getInstance()
+                ?.getCodewhispererService()
+                .getSubscriptionStatus()
+                .then(o => {
+                    this.#log(`xxx getSubscriptionStatus: ${o.status} ${o.encodedVerificationUrl}`)
+                    this.setPaidTierMode(tabId, o.status === 'ACTIVE' ? 'paidtier' : 'freetier')
+                })
+                .catch(err => {
+                    this.#log(`xxx getSubscriptionStatus failed: ${JSON.stringify(err)}`)
+                })
+            // const isFreeTierUser = getSsoConnectionType(this.#features.credentialsProvider) === 'builderId'
+            // mode = isFreeTierUser ? 'freetier' : 'paidtier'
 
             return
         }
 
         const o: ChatUpdateParams = {
-            tabId: tabId ?? 'xxx',
-            state: { inProgress: false },
-            data: {
-                // Special flag recognized by `chat-client/src/client/mynahUi.ts`.
-                placeholderText: 'upgrade-q',
-                messages: [],
-            },
+            tabId: tabId ?? '',
+            // data: { messages: [] },
         }
-        ;(o as any).upgradeQMode = mode
+        // Special flag recognized by `chat-client/src/client/mynahUi.ts`.
+        ;(o as any).paidTierMode = mode
         this.#features.chat.sendChatUpdate(o)
     }
 
@@ -2735,6 +2736,9 @@ export class AgenticChatController implements ChatHandlers {
         const updatedOptOutPreference = newConfig.optOutTelemetryPreference
         this.#telemetryService.updateOptOutPreference(updatedOptOutPreference)
         this.#log(`Chat configuration telemetry preference to ${updatedOptOutPreference}`)
+
+        // Force a service request to get current Q user subscription status.
+        this.#paidTierMode = undefined
     }
 
     #getTools(session: ChatSessionService) {
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/qAgenticChatServer.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/qAgenticChatServer.ts
@@ -152,7 +152,7 @@ export const QAgenticChatServer =
         })
 
         // ;(chat as any).onPromptInputButtonClick((params: any) => {
-        //     chatController.setUpgradeQMode(params.tabId, 'paidtier')
+        //     chatController.setPaidTierMode(params.tabId, 'paidtier')
         // })
 
         chat.onButtonClick(params => {
diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts
@@ -24,8 +24,6 @@ type DeferredHandler = {
 export class ChatSessionService {
     public shareCodeWhispererContentWithAWS = false
     public pairProgrammingMode: boolean = true
-    /** Decides whether the Chat UI should show "Upgrade Q" button/card. */
-    public upgradeQMode: 'paidtier' | 'freetier' | 'freetier-limit' = 'paidtier'
     public contextListSent: boolean = false
     #abortController?: AbortController
     #currentPromptId?: string
diff --git a/server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.ts b/server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.ts
@@ -112,7 +112,6 @@ export const QConfigurationServerToken =
                 token: CancellationToken
             ): Promise<QConfigurationResponse | void> => {
                 const section = params.section
-                logging.info('zzzzzzzz1')
 
                 let customizations: Customizations | CustomizationWithMetadata[] = []
                 let developerProfiles: AmazonQDeveloperProfile[] = []
diff --git a/server/aws-lsp-codewhisperer/src/language-server/paidTier/paidTier.ts b/server/aws-lsp-codewhisperer/src/language-server/paidTier/paidTier.ts
@@ -0,0 +1 @@
+export type PaidTierMode = 'paidtier' | 'paidtier-success' | 'freetier' | 'freetier-limit'
diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/BaseAmazonQServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/BaseAmazonQServiceManager.ts
@@ -108,7 +108,6 @@ export abstract class BaseAmazonQServiceManager<
     abstract handleOnUpdateConfiguration(params: UpdateConfigurationParams, token: CancellationToken): Promise<void>
 
     public async handleDidChangeConfiguration(): Promise<void> {
-        this.logging.debug('xxxx 7 handleDidChangeConfiguration')
         if (this.isConfigChangeInProgress) {
             this.logging.debug(CONFIGURATION_CHANGE_IN_PROGRESS_MSG)
             return
@@ -130,7 +129,6 @@ export abstract class BaseAmazonQServiceManager<
     }
 
     protected updateCachedServiceConfig(): void {
-        this.logging.debug('xxxx 8 updateCachedServiceConfig')
         if (this.cachedCodewhispererService) {
             const customizationArn = this.configurationCache.getProperty('customizationArn')
             this.logging.debug(`Using customization=${customizationArn}`)
diff --git a/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.ts b/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.ts

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+export type PaidTierMode = 'paidtier' \| 'paidtier-success' \| 'freetier' \| 'freetier-limit'`