wip: control MynahUI from chat server

Author: justinmk3

chat-client/src/client/mynahUi.ts

@@ -57,6 +57,7 @@ import {
 } from './utils'
 import { ChatHistory, ChatHistoryList } from './features/history'
 import { pairProgrammingModeOff, pairProgrammingModeOn, programmerModeCard } from './texts/pairProgramming'
+import { paidTierCard, upgradeQButton } from './texts/paidTier'
 
 export interface InboundChatApi {
     addChatResponse(params: ChatResult, tabId: string, isPartialResult: boolean): void
@@ -171,8 +172,6 @@ export const createMynahUi = (
 ): [MynahUI, InboundChatApi] => {
     let disclaimerCardActive = !disclaimerAcknowledged
     let programmingModeCardActive = !pairProgrammingCardAcknowledged
-    const isFreeTierLimitReached = true
-    let paidTierCardActive = isFreeTierLimitReached
     let contextCommandGroups: ContextCommandGroups | undefined
 
     let chatEventHandlers: ChatEventHandler = {
@@ -265,12 +264,7 @@ export const createMynahUi = (
             // We check if tabMetadata.openTabKey exists - if it does and is set to true, we skip showing welcome messages
             // since this indicates we're loading a previous chat session rather than starting a new one.
             if (!tabStore?.tabMetadata || !tabStore.tabMetadata.openTabKey) {
-                defaultTabConfig.chatItems = tabFactory.getChatItems(
-                    true,
-                    programmingModeCardActive,
-                    paidTierCardActive,
-                    []
-                )
+                defaultTabConfig.chatItems = tabFactory.getChatItems(true, programmingModeCardActive, false, [])
             }
             mynahUi.updateStore(tabId, defaultTabConfig)
             messager.onTabAdd(tabId)
@@ -499,7 +493,7 @@ export const createMynahUi = (
     const mynahUiProps: MynahUIProps = {
         tabs: {},
         defaults: {
-            store: tabFactory.createTab(false, paidTierCardActive),
+            store: tabFactory.createTab(false, false),
         },
         config: {
             maxTabs: 10,
@@ -538,7 +532,7 @@ export const createMynahUi = (
     // This distinction helps maintain consistent tab behavior between fresh conversations and restored sessions.
     const createTabId = (openTab?: boolean) => {
         const tabId = mynahUi.updateStore('', {
-            ...tabFactory.createTab(disclaimerCardActive, paidTierCardActive),
+            ...tabFactory.createTab(disclaimerCardActive, false),
             tabMetadata: { openTabKey: openTab ? true : false },
         })
         if (tabId === undefined) {
@@ -818,6 +812,25 @@ 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
+
+            // const chatItem: ChatItem = {
+            //     type: ChatItemType.DIRECTIVE,
+            //     contentHorizontalAlignment: 'center',
+            //     fullWidth: true,
+            //     body: `Upgrade Q: ${upgradeQMode}`,
+            // }
+            // mynahUi.addChatItem(tabId, chatItem)
+            upgradeQButton.description = `Upgrade Q: ${upgradeQMode}`
+            mynahUi.updateStore(tabId, {
+                promptInputButtons: upgradeQMode === 'paidtier' ? [] : [upgradeQButton],
+                chatItems: upgradeQMode === 'freetier-limit' ? [paidTierCard] : [],
+            })
+            return
+        }
+
         const isChatLoading = params.state?.inProgress
         mynahUi.updateStore(params.tabId, {
             loadingChat: isChatLoading,
@@ -1005,7 +1018,7 @@ ${params.message}`,
                     chatItems: tabFactory.getChatItems(
                         messages ? false : true,
                         programmingModeCardActive,
-                        paidTierCardActive,
+                        false,
                         messages
                     ),
                 })

chat-client/src/client/tabs/tabFactory.ts

@@ -5,31 +5,17 @@ import {
     MynahUIDataModel,
     QuickActionCommandGroup,
     TabBarMainAction,
-    ChatItemButton,
 } from '@aws/mynah-ui'
 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 } from '../texts/paidTier'
+import { paidTierCard, upgradeQButton } from '../texts/paidTier'
 
 export type DefaultTabData = MynahUIDataModel
 
 export const ExportTabBarButtonId = 'export'
 
-const upgradeToProButton: ChatItemButton = {
-    flash: 'once',
-    fillState: 'hover',
-    position: 'outside',
-    id: 'upgrade-q',
-    // https://github.com/aws/mynah-ui/blob/main/src/components/icon/icons/q.svg
-    // https://github.com/aws/mynah-ui/blob/main/src/components/icon/icons/rocket.svg
-    // icon: MynahIcons.Q,
-    description: 'Upgrade to Amazon Q Pro',
-    text: 'Upgrade Q',
-    status: 'info',
-}
-
 export class TabFactory {
     private history: boolean = false
     private export: boolean = false
@@ -52,7 +38,7 @@ export class TabFactory {
         const tabData: MynahUIDataModel = {
             ...this.getDefaultTabData(),
             ...(disclaimerCardActive ? { promptInputStickyCard: disclaimerCard } : {}),
-            promptInputButtons: showUpgradeButton ? [upgradeToProButton] : [],
+            promptInputButtons: showUpgradeButton ? [upgradeQButton] : [],
             promptInputOptions: this.agenticMode ? [pairProgrammingPromptInput] : [],
             cancelButtonWhenLoading: this.agenticMode, // supported for agentic chat only
         }

chat-client/src/client/texts/paidTier.ts

@@ -1,4 +1,4 @@
-import { ChatItem, ChatItemFormItem, ChatItemType } from '@aws/mynah-ui'
+import { ChatItem, ChatItemButton, ChatItemFormItem, ChatItemType } from '@aws/mynah-ui'
 
 export const paidTierCard: ChatItem = {
     type: ChatItemType.ANSWER,
@@ -31,3 +31,16 @@ export const paidTierStep1: ChatItem = {
     type: ChatItemType.DIRECTIVE,
     body: 'You have upgraded to Amazon Q Pro',
 }
+
+export const upgradeQButton: ChatItemButton = {
+    flash: 'once',
+    fillState: 'hover',
+    position: 'outside',
+    id: 'upgrade-q',
+    // https://github.com/aws/mynah-ui/blob/main/src/components/icon/icons/q.svg
+    // https://github.com/aws/mynah-ui/blob/main/src/components/icon/icons/rocket.svg
+    // icon: MynahIcons.Q,
+    description: 'Upgrade to Amazon Q Pro',
+    text: 'Upgrade Q',
+    status: 'info',
+}

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

@@ -29,6 +29,7 @@ import {
     InlineChatResultParams,
     PromptInputOptionChangeParams,
     TextDocument,
+    ChatUpdateParams,
 } from '@aws/language-server-runtimes/protocol'
 import {
     ApplyWorkspaceEditParams,
@@ -73,7 +74,14 @@ import { ChatSessionManagementService } from '../chat/chatSessionManagementServi
 import { ChatTelemetryController } from '../chat/telemetry/chatTelemetryController'
 import { QuickAction } from '../chat/quickActions'
 import { Metric } from '../../shared/telemetry/metric'
-import { getErrorMessage, getHttpStatusCode, getRequestID, isAwsError, isNullish, isObject } from '../../shared/utils'
+import {
+    getErrorMessage,
+    getHttpStatusCode,
+    getRequestID,
+    getSsoConnectionType,
+    isFreeTierLimitError,
+    isNullish,
+} from '../../shared/utils'
 import { HELP_MESSAGE, loadingMessage } from '../chat/constants'
 import { TelemetryService } from '../../shared/telemetry/telemetryService'
 import {
@@ -147,6 +155,7 @@ export class AgenticChatController implements ChatHandlers {
     #additionalContextProvider: AdditionalContextProvider
     #contextCommandsProvider: ContextCommandsProvider
     #stoppedToolUses = new Set<string>()
+    #freeTierLimit = false
 
     /**
      * Determines the appropriate message ID for a tool use based on tool type and name
@@ -239,7 +248,15 @@ export class AgenticChatController implements ChatHandlers {
             await this.#renderStoppedShellCommand(params.tabId, params.messageId)
             return { success: true }
         } else if (params.buttonId === 'upgrade-q') {
-            this.#features.logging.warn('clicked upgrade-to-pro-button')
+            const awsAccountId = (params as any).awsAccountId
+            if (typeof awsAccountId !== 'string') {
+                this.#log(`invalid awsAccountId: ${awsAccountId}`)
+                return {
+                    success: false,
+                    failureReason: 'invalid awsAccountId',
+                }
+            }
+            this.setUpgradeQMode(params.tabId, 'paidtier')
             return { success: true }
         } else {
             return {
@@ -1574,6 +1591,12 @@ export class AgenticChatController implements ChatHandlers {
         metric.setDimension('cwsprChatResponseCode', getHttpStatusCode(err) ?? 0)
         metric.setDimension('languageServerVersion', this.#features.runtime.serverInfo.version)
 
+        // TODO handle free tier limit exceeded
+        if (isFreeTierLimitError(err)) {
+            this.setUpgradeQMode(tabId, 'freetier-limit')
+            // throw new AmazonQFreeTierLimitError()
+        }
+
         // use custom error message for unactionable errors (user-dependent errors like PromptCharacterLimit)
         if (err.code && err.code in unactionableErrorCodes) {
             const customErrMessage = unactionableErrorCodes[err.code as keyof typeof unactionableErrorCodes]
@@ -1604,7 +1627,7 @@ export class AgenticChatController implements ChatHandlers {
             return createAuthFollowUpResult(authFollowType)
         }
 
-        if (customerFacingErrorCodes.includes(err.code)) {
+        if (isFreeTierLimitError(err) || customerFacingErrorCodes.includes(err.code)) {
             this.#features.logging.error(`${loggingUtils.formatErr(err)}`)
             if (err.code === 'InputTooLong') {
                 // Clear the chat history in the database for this tab
@@ -1813,6 +1836,12 @@ export class AgenticChatController implements ChatHandlers {
 
     async onReady() {
         await this.#tabBarController.loadChats()
+        try {
+            this.setUpgradeQMode()
+        } catch (err) {
+            this.#log('Error initializing Free Tier state: ' + (err as Error).message)
+        }
+
         try {
             const localProjectContextController = await LocalProjectContextController.getInstance()
             const contextItems = await localProjectContextController.getContextCommandItems()
@@ -1846,6 +1875,8 @@ export class AgenticChatController implements ChatHandlers {
         this.#telemetryController.activeTabId = params.tabId
 
         this.#chatSessionManagementService.createSession(params.tabId)
+
+        this.setUpgradeQMode(params.tabId)
     }
 
     onTabChange(params: TabChangeParams) {
@@ -1860,6 +1891,8 @@ export class AgenticChatController implements ChatHandlers {
             name: ChatTelemetryEventName.EnterFocusConversation,
             data: {},
         })
+
+        this.setUpgradeQMode(params.tabId)
     }
 
     onTabRemove(params: TabRemoveParams) {
@@ -2004,6 +2037,43 @@ export class AgenticChatController implements ChatHandlers {
         }
     }
 
+    /**
+     * Enables or disables the "Upgrade Q" UI in the chat component.
+     *
+     * `mode` behavior:
+     * - 'freetier': always show "Upgrade Q" button.
+     * - '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.
+     */
+    async setUpgradeQMode(
+        tabId?: string,
+        mode?: 'paidtier' | 'freetier' | 'freetier-limit' /*, session: ChatSessionService*/
+    ) {
+        if (mode === 'freetier-limit') {
+            this.#freeTierLimit = true // Sticky until 'paidtier' is sent.
+        } else if (mode === 'paidtier') {
+            this.#freeTierLimit = false
+        } else if (this.#freeTierLimit && (!mode || mode === 'freetier')) {
+            mode = 'freetier-limit'
+        } else if (!mode) {
+            const isFreeTierUser = getSsoConnectionType(this.#features.credentialsProvider) === 'builderId'
+            mode = isFreeTierUser ? 'freetier' : 'paidtier'
+        }
+
+        const o: ChatUpdateParams = {
+            tabId: tabId ?? 'xxx',
+            state: { inProgress: false },
+            data: {
+                // Special flag recognized by `chat-client/src/client/mynahUi.ts`.
+                placeholderText: 'upgrade-q',
+                messages: [],
+            },
+        }
+        ;(o as any).upgradeQMode = mode
+        this.#features.chat.sendChatUpdate(o)
+    }
+
     async #processGenerateAssistantResponseResponseWithTimeout(
         response: GenerateAssistantResponseCommandOutput,
         metric: Metric<AddMessageEvent>,

server/aws-lsp-codewhisperer/src/language-server/agenticChat/qAgenticChatServer.ts

@@ -51,8 +51,6 @@ export const QAgenticChatServer =
             chatController.updateConfiguration(updatedConfig)
         }
 
-        const isFreeTierUser = getSsoConnectionType(credentialsProvider)
-
         lsp.onInitialized(async () => {
             // Get initialized service manager and inject it to chatSessionManagementService to pass it down
             amazonQServiceManager = AmazonQTokenServiceManager.getInstance()
@@ -153,6 +151,10 @@ export const QAgenticChatServer =
             return chatController.onPromptInputOptionChange(params)
         })
 
+        // ;(chat as any).onPromptInputButtonClick((params: any) => {
+        //     chatController.setUpgradeQMode(params.tabId, 'paidtier')
+        // })
+
         chat.onButtonClick(params => {
             return chatController.onButtonClick(params)
         })

server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts

@@ -24,6 +24,8 @@ 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
     #conversationId?: string