refactor: encapsulate model ID and token limits in session

Author: chungjac

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

@@ -3059,7 +3059,7 @@ ${' '.repeat(8)}}
             // Create a session and set initial model
             chatController.onTabAdd({ tabId: mockTabId })
             const session = chatSessionManagementService.getSession(mockTabId).data!
-            session.modelId = initialModelId
+            session.setModel(initialModelId, cachedModels)
 
             // Get initial token limits (default 200K)
             const initialLimits = session.tokenLimits
@@ -3137,7 +3137,7 @@ ${' '.repeat(8)}}
                 getCachedModelsStub.returns(cachedData)
 
                 const session = chatSessionManagementService.getSession(mockTabId).data!
-                session.modelId = 'model1'
+                session.setModel('model1', cachedData.models)
 
                 const result = await chatController.onListAvailableModels({ tabId: mockTabId })
 
@@ -3324,7 +3324,7 @@ ${' '.repeat(8)}}
 
             it('should use defaultModelId from cache when session has no modelId', async () => {
                 const session = chatSessionManagementService.getSession(mockTabId).data!
-                session.modelId = undefined
+                session.setModel(undefined, undefined)
 
                 const result = await chatController.onListAvailableModels({ tabId: mockTabId })
 
@@ -3341,7 +3341,7 @@ ${' '.repeat(8)}}
                 })
 
                 const session = chatSessionManagementService.getSession(mockTabId).data!
-                session.modelId = undefined
+                session.setModel(undefined, undefined)
 
                 const result = await chatController.onListAvailableModels({ tabId: mockTabId })
 

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

@@ -780,13 +780,12 @@ export class AgenticChatController implements ChatHandlers {
 
         // Handle error cases by returning default model
         if (!success || errorFromAPI) {
-            // Even in error cases, calculate token limits from the default/fallback model
+            // Even in error cases, set the model with token limits
             if (success && session) {
-                const fallbackModel = models.find(model => model.id === DEFAULT_MODEL_ID)
-                const maxInputTokens = TokenLimitsCalculator.extractMaxInputTokens(fallbackModel)
-                const tokenLimits = TokenLimitsCalculator.calculate(maxInputTokens)
-                session.setTokenLimits(tokenLimits)
-                this.#log(`Token limits calculated for fallback model (error case): ${JSON.stringify(tokenLimits)}`)
+                session.setModel(DEFAULT_MODEL_ID, models)
+                this.#log(
+                    `Model set for fallback (error case): ${DEFAULT_MODEL_ID}, tokenLimits: ${JSON.stringify(session.tokenLimits)}`
+                )
             }
             return {
                 tabId: params.tabId,
@@ -828,16 +827,10 @@ export class AgenticChatController implements ChatHandlers {
             selectedModelId = defaultModelId || getMappedModelId(DEFAULT_MODEL_ID)
         }
 
-        // Store the selected model in the session
-        session.modelId = selectedModelId
-
-        // Extract maxInputTokens from the selected model and calculate token limits
-        const selectedModel = models.find(model => model.id === selectedModelId)
-        const maxInputTokens = TokenLimitsCalculator.extractMaxInputTokens(selectedModel)
-        const tokenLimits = TokenLimitsCalculator.calculate(maxInputTokens)
-        session.setTokenLimits(tokenLimits)
+        // Store the selected model in the session (automatically calculates token limits)
+        session.setModel(selectedModelId, models)
         this.#log(
-            `Token limits calculated for initial model selection (${selectedModelId}): ${JSON.stringify(tokenLimits)}`
+            `Model set for initial selection: ${selectedModelId}, tokenLimits: ${JSON.stringify(session.tokenLimits)}`
         )
 
         return {
@@ -4678,17 +4671,13 @@ export class AgenticChatController implements ChatHandlers {
         session.pairProgrammingMode = params.optionsValues['pair-programmer-mode'] === 'true'
         const newModelId = params.optionsValues['model-selection']
 
-        // Recalculate token limits when model changes
-        if (newModelId && newModelId !== session.modelId) {
+        // Set model (automatically recalculates token limits)
+        if (newModelId !== session.modelId) {
             const cachedData = this.#chatHistoryDb.getCachedModels()
-            const selectedModel = cachedData?.models?.find(model => model.id === newModelId)
-            const maxInputTokens = TokenLimitsCalculator.extractMaxInputTokens(selectedModel)
-            const tokenLimits = TokenLimitsCalculator.calculate(maxInputTokens)
-            session.setTokenLimits(tokenLimits)
-            this.#log(`Token limits calculated for model switch (${newModelId}): ${JSON.stringify(tokenLimits)}`)
+            session.setModel(newModelId, cachedData?.models)
+            this.#log(`Model set for model switch: ${newModelId}, tokenLimits: ${JSON.stringify(session.tokenLimits)}`)
         }
 
-        session.modelId = newModelId
         this.#chatHistoryDb.setModelId(session.modelId)
         this.#chatHistoryDb.setPairProgrammingMode(session.pairProgrammingMode)
     }

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

@@ -329,6 +329,45 @@ describe('Chat Session Service', () => {
         })
     })
 
+    describe('setModel encapsulation', () => {
+        let chatSessionService: ChatSessionService
+
+        beforeEach(() => {
+            chatSessionService = new ChatSessionService()
+        })
+
+        it('should initialize with undefined modelId and default token limits', () => {
+            assert.strictEqual(chatSessionService.modelId, undefined)
+            assert.strictEqual(chatSessionService.tokenLimits.maxInputTokens, 200_000)
+        })
+
+        it('should set modelId and calculate token limits together', () => {
+            const models = [
+                { id: 'model-1', name: 'Model 1', description: 'Test', tokenLimits: { maxInputTokens: 300_000 } },
+            ]
+
+            chatSessionService.setModel('model-1', models)
+
+            assert.strictEqual(chatSessionService.modelId, 'model-1')
+            assert.strictEqual(chatSessionService.tokenLimits.maxInputTokens, 300_000)
+            assert.strictEqual(chatSessionService.tokenLimits.maxOverallCharacters, Math.floor(300_000 * 3.5))
+        })
+
+        it('should use default token limits when model not found in list', () => {
+            chatSessionService.setModel('unknown-model', [])
+
+            assert.strictEqual(chatSessionService.modelId, 'unknown-model')
+            assert.strictEqual(chatSessionService.tokenLimits.maxInputTokens, 200_000)
+        })
+
+        it('should use default token limits when models list is undefined', () => {
+            chatSessionService.setModel('some-model', undefined)
+
+            assert.strictEqual(chatSessionService.modelId, 'some-model')
+            assert.strictEqual(chatSessionService.tokenLimits.maxInputTokens, 200_000)
+        })
+    })
+
     describe('IAM client source property', () => {
         it('sets source to Origin.IDE when using StreamingClientServiceIAM', async () => {
             const codeWhispererStreamingClientIAM = stubInterface<StreamingClientServiceIAM>()

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

@@ -18,6 +18,7 @@ import { QErrorTransformer } from '../agenticChat/retry/errorTransformer'
 import { DelayNotification } from '../agenticChat/retry/delayInterceptor'
 import { MAX_REQUEST_ATTEMPTS } from '../agenticChat/constants/constants'
 import { TokenLimits, TokenLimitsCalculator } from '../agenticChat/utils/tokenLimitsCalculator'
+import { Model } from '@aws/language-server-runtimes/protocol'
 
 export type ChatSessionServiceConfig = CodeWhispererStreamingClientConfig
 type FileChange = { before?: string; after?: string }
@@ -29,8 +30,8 @@ type DeferredHandler = {
 export class ChatSessionService {
     public pairProgrammingMode: boolean = true
     public contextListSent: boolean = false
-    public modelId: string | undefined
     public isMemoryBankGeneration: boolean = false
+    #modelId: string | undefined
     #lsp?: Features['lsp']
     #abortController?: AbortController
     #currentPromptId?: string
@@ -145,6 +146,13 @@ export class ChatSessionService {
         this.#tokenLimits = TokenLimitsCalculator.calculate()
     }
 
+    /**
+     * Gets the model ID for this session
+     */
+    public get modelId(): string | undefined {
+        return this.#modelId
+    }
+
     /**
      * Gets the token limits for this session
      */
@@ -153,11 +161,15 @@ export class ChatSessionService {
     }
 
     /**
-     * Sets the token limits for this session
-     * @param limits The token limits to set
+     * Sets the model for this session, automatically calculating token limits.
+     * This encapsulates model ID and token limits as a single entity.
+     * @param modelId The model ID to set
+     * @param models Optional list of available models to look up token limits from
      */
-    public setTokenLimits(limits: TokenLimits): void {
-        this.#tokenLimits = limits
+    public setModel(modelId: string | undefined, models?: Model[]): void {
+        this.#modelId = modelId
+        const maxInputTokens = TokenLimitsCalculator.extractMaxInputTokens(models?.find(m => m.id === modelId))
+        this.#tokenLimits = TokenLimitsCalculator.calculate(maxInputTokens)
     }
 
     public async sendMessage(request: SendMessageCommandInput): Promise<SendMessageCommandOutput> {