feat: Implement dynamic model selection based on extension capabilities and improve error handling (#1737)

* fix: temporarily fix version showing error throttling without model selection

* feat: get modelSelection from extention qcapabilities and handle error

* chore: prettier fix

* fix: fix the test related to modelId update

* fix: simplify error handling and add util for model selection

* fix: delete the test code

* fix: leave fonstants and fix test

* fix: give general error message for every existing users

* fix: revert back to default model and check the error message based on enabled model selection

* fix: delete unrelated comments and pass features.lsp rather than features

Author: XiaoxuanLu

chat-client/src/client/mynahUi.ts

@@ -64,7 +64,7 @@ import {
 import { ChatHistory, ChatHistoryList } from './features/history'
 import { pairProgrammingModeOff, pairProgrammingModeOn, programmerModeCard } from './texts/pairProgramming'
 import { ContextRule, RulesList } from './features/rules'
-import { getModelSelectionChatItem, modelUnavailableBanner } from './texts/modelSelection'
+import { getModelSelectionChatItem, modelUnavailableBanner, modelThrottledBanner } from './texts/modelSelection'
 import {
     freeTierLimitSticky,
     upgradeSuccessSticky,
@@ -1016,6 +1016,13 @@ export const createMynahUi = (
                     return
                 }
 
+                if (updatedMessage.messageId === 'modelThrottled') {
+                    mynahUi.updateStore(tabId, {
+                        promptInputStickyCard: modelThrottledBanner,
+                    })
+                    return
+                }
+
                 const oldMessage = chatItems.find(ci => ci.messageId === updatedMessage.messageId)
                 if (!oldMessage) return
 

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

@@ -54,3 +54,14 @@ export const modelUnavailableBanner: Partial<ChatItem> = {
     body: `The model you've selected is experiencing high load. Please switch to another model and try again.`,
     canBeDismissed: true,
 }
+
+export const modelThrottledBanner: Partial<ChatItem> = {
+    messageId: 'model-throttled-banner',
+    header: {
+        icon: 'warning',
+        iconStatus: 'warning',
+        body: '### Model Unavailable',
+    },
+    body: `I am experiencing high traffic, please try again shortly.`,
+    canBeDismissed: true,
+}

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

@@ -354,7 +354,6 @@ describe('AgenticChatController', () => {
     it('onTabAdd updates model ID in chat options and session', () => {
         const modelId = 'test-model-id'
         sinon.stub(ChatDatabase.prototype, 'getModelId').returns(modelId)
-
         chatController.onTabAdd({ tabId: mockTabId })
 
         sinon.assert.calledWithExactly(testFeatures.chat.chatOptionsUpdate, { modelId, tabId: mockTabId })

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

@@ -93,6 +93,7 @@ import {
     getSsoConnectionType,
     isUsageLimitError,
     isNullish,
+    enabledModelSelection,
 } from '../../shared/utils'
 import { HELP_MESSAGE, loadingMessage } from '../chat/constants'
 import { TelemetryService } from '../../shared/telemetry/telemetryService'
@@ -2352,6 +2353,19 @@ export class AgenticChatController implements ChatHandlers {
                     }
                     return emptyChatResult
                 }
+                if (err.message === `I am experiencing high traffic, please try again shortly.`) {
+                    this.#features.chat.sendChatUpdate({
+                        tabId: tabId,
+                        data: { messages: [{ messageId: 'modelThrottled' }] },
+                    })
+                    const emptyChatResult: ChatResult = {
+                        type: 'answer',
+                        body: '',
+                        messageId: errorMessageId,
+                        buttons: [],
+                    }
+                    return emptyChatResult
+                }
                 return responseData
             }
             return new ResponseError<ChatResult>(LSPErrorCodes.RequestFailed, err.message, responseData)

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

@@ -11,11 +11,15 @@ import { ChatSessionManagementService } from '../chat/chatSessionManagementServi
 import { QAgenticChatServer } from './qAgenticChatServer'
 import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager'
 import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager'
+import { Features } from '../types'
 
 describe('QAgenticChatServer', () => {
     const mockTabId = 'mockTabId'
     let disposeStub: sinon.SinonStub
-    let withAmazonQServiceSpy: sinon.SinonSpy<[amazonQService: AmazonQBaseServiceManager], ChatSessionManagementService>
+    let withAmazonQServiceSpy: sinon.SinonSpy<
+        [serviceManager: AmazonQBaseServiceManager, lsp?: Features['lsp'] | undefined],
+        ChatSessionManagementService
+    >
     let testFeatures: TestFeatures
     let amazonQServiceManager: AmazonQTokenServiceManager
     let disposeServer: () => void
@@ -72,7 +76,7 @@ describe('QAgenticChatServer', () => {
     })
 
     it('should initialize ChatSessionManagementService with AmazonQTokenServiceManager instance', () => {
-        sinon.assert.calledOnceWithExactly(withAmazonQServiceSpy, amazonQServiceManager)
+        sinon.assert.calledOnceWithExactly(withAmazonQServiceSpy, amazonQServiceManager, testFeatures.lsp)
     })
 
     it('dispose should dispose all chat session services', () => {

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

@@ -13,7 +13,7 @@ import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/A
 import { AmazonQWorkspaceConfig } from '../../shared/amazonQServiceManager/configurationUtils'
 import { TabBarController } from './tabBarController'
 import { AmazonQServiceInitializationError } from '../../shared/amazonQServiceManager/errors'
-import { safeGet } from '../../shared/utils'
+import { enabledModelSelection, safeGet } from '../../shared/utils'
 import { enabledMCP } from './tools/mcp/mcpUtils'
 
 export const QAgenticChatServer =
@@ -47,6 +47,7 @@ export const QAgenticChatServer =
                             ],
                         },
                         mcpServers: enabledMCP(params),
+                        // we should set it as true for current VSC and VS clients
                         modelSelection: true,
                         history: true,
                         export: TabBarController.enableChatExport(params)
@@ -64,7 +65,7 @@ export const QAgenticChatServer =
             // Get initialized service manager and inject it to chatSessionManagementService to pass it down
             amazonQServiceManager = AmazonQTokenServiceManager.getInstance()
             chatSessionManagementService =
-                ChatSessionManagementService.getInstance().withAmazonQServiceManager(amazonQServiceManager)
+                ChatSessionManagementService.getInstance().withAmazonQServiceManager(amazonQServiceManager, features.lsp)
 
             telemetryService = new TelemetryService(amazonQServiceManager, credentialsProvider, telemetry, logging)
 
@@ -144,7 +145,7 @@ export const QAgenticChatServer =
             return chatController.onListConversations(params)
         })
 
-         chat.onListRules(params => {
+        chat.onListRules(params => {
             return chatController.onListRules(params)
         })
 

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

@@ -1,11 +1,12 @@
 import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager'
-import { Result } from '../types'
+import { Result, Features } from '../types'
 import { ChatSessionService } from './chatSessionService'
 
 export class ChatSessionManagementService {
     static #instance?: ChatSessionManagementService
     #sessionByTab: Map<string, ChatSessionService> = new Map<string, any>()
     #serviceManager?: AmazonQBaseServiceManager
+    #lsp?: Features['lsp']
 
     public static getInstance() {
         if (!ChatSessionManagementService.#instance) {
@@ -21,8 +22,9 @@ export class ChatSessionManagementService {
 
     private constructor() {}
 
-    public withAmazonQServiceManager(serviceManager: AmazonQBaseServiceManager) {
+    public withAmazonQServiceManager(serviceManager: AmazonQBaseServiceManager, lsp?: Features['lsp']) {
         this.#serviceManager = serviceManager
+        this.#lsp = lsp
 
         return this
     }
@@ -39,7 +41,7 @@ export class ChatSessionManagementService {
             }
         }
 
-        const newSession = new ChatSessionService(this.#serviceManager)
+        const newSession = new ChatSessionService(this.#serviceManager, this.#lsp)
 
         this.#sessionByTab.set(tabId, newSession)
 

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

@@ -21,7 +21,9 @@ import {
 import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager'
 import { loggingUtils } from '@aws/lsp-core'
 import { Logging } from '@aws/language-server-runtimes/server-interface'
+import { Features } from '../types'
 import { getRequestID, isUsageLimitError } from '../../shared/utils'
+import { enabledModelSelection } from '../../shared/utils'
 
 export type ChatSessionServiceConfig = CodeWhispererStreamingClientConfig
 type FileChange = { before?: string; after?: string }
@@ -35,6 +37,7 @@ export class ChatSessionService {
     public pairProgrammingMode: boolean = true
     public contextListSent: boolean = false
     public modelId: string | undefined
+    #lsp?: Features['lsp']
     #abortController?: AbortController
     #currentPromptId?: string
     #conversationId?: string
@@ -121,8 +124,9 @@ export class ChatSessionService {
         this.#approvedPaths.add(normalizedPath)
     }
 
-    constructor(serviceManager?: AmazonQBaseServiceManager, logging?: Logging) {
+    constructor(serviceManager?: AmazonQBaseServiceManager, lsp?: Features['lsp'], logging?: Logging) {
         this.#serviceManager = serviceManager
+        this.#lsp = lsp
         this.#logging = logging
     }
 
@@ -144,6 +148,10 @@ export class ChatSessionService {
         return response
     }
 
+    private isModelSelectionEnabled(): boolean {
+        return enabledModelSelection(this.#lsp?.getClientInitializeParams())
+    }
+
     public async generateAssistantResponse(
         request: GenerateAssistantResponseCommandInput
     ): Promise<GenerateAssistantResponseCommandOutput> {
@@ -208,8 +216,11 @@ export class ChatSessionService {
                     error.message ===
                         'Encountered unexpectedly high load when processing the request, please try again.'
                 ) {
-                    error.message = `The model you selected is temporarily unavailable. Please switch to a different model and try again.`
+                    error.message = this.isModelSelectionEnabled()
+                        ? `The model you selected is temporarily unavailable. Please switch to a different model and try again.`
+                        : `I am experiencing high traffic, please try again shortly.`
                 }
+
                 throw error
             }
         } else {

server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.ts

@@ -39,6 +39,7 @@ export interface QClientCapabilities {
     developerProfiles?: boolean
     customizationsWithMetadata?: boolean
     mcp?: boolean
+    modelSelection?: boolean
 }
 
 type QConfigurationResponse =

server/aws-lsp-codewhisperer/src/shared/utils.ts

@@ -26,6 +26,8 @@ import { ServiceException } from '@smithy/smithy-client'
 import * as ignoreWalk from 'ignore-walk'
 import { getAuthFollowUpType } from '../language-server/chat/utils'
 import ignore = require('ignore')
+import { InitializeParams } from '@aws/language-server-runtimes/server-interface'
+import { QClientCapabilities } from '../language-server/configuration/qConfigurationServer'
 export type SsoConnectionType = 'builderId' | 'identityCenter' | 'none'
 
 export function isAwsError(error: unknown): error is AWSError {
@@ -316,6 +318,13 @@ export function getCompletionType(suggestion: Suggestion): CodewhispererCompleti
     return nonBlankLines > 1 ? 'Block' : 'Line'
 }
 
+export function enabledModelSelection(params: InitializeParams | undefined): boolean {
+    const qCapabilities = params?.initializationOptions?.aws?.awsClientCapabilities?.q as
+        | QClientCapabilities
+        | undefined
+    return qCapabilities?.modelSelection || false
+}
+
 export function parseJson(jsonString: string) {
     try {
         return JSON.parse(jsonString)