fix(chat): Prune history if it's too large, populate statusCode in error message correctly

Author: jguoamz

packages/core/src/codewhispererChat/controllers/chat/controller.ts

@@ -987,6 +987,7 @@ export class ChatController {
     private processException(e: any, tabID: string) {
         let errorMessage = ''
         let requestID = undefined
+        let statusCode = undefined
         const defaultMessage = 'Failed to get response'
         if (typeof e === 'string') {
             errorMessage = e.toUpperCase()
@@ -996,6 +997,7 @@ export class ChatController {
         } else if (e instanceof CodeWhispererStreamingServiceException) {
             errorMessage = e.message
             requestID = e.$metadata.requestId
+            statusCode = e.$metadata.httpStatusCode
         } else if (e instanceof Error) {
             errorMessage = e.message
         }
@@ -1005,7 +1007,7 @@ export class ChatController {
             this.sessionStorage.setAgentLoopInProgress(tabID, false)
         }
 
-        this.messenger.sendErrorMessage(errorMessage, tabID, requestID)
+        this.messenger.sendErrorMessage(errorMessage, tabID, requestID, statusCode)
         getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)
 
         this.sessionStorage.deleteSession(tabID)

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

@@ -566,12 +566,17 @@ export class Messenger {
         )
     }
 
-    public sendErrorMessage(errorMessage: string | undefined, tabID: string, requestID: string | undefined) {
+    public sendErrorMessage(
+        errorMessage: string | undefined,
+        tabID: string,
+        requestID: string | undefined,
+        statusCode?: number
+    ) {
         this.showChatExceptionMessage(
             {
                 errorMessage: errorMessage,
                 sessionID: undefined,
-                statusCode: undefined,
+                statusCode: statusCode?.toString(),
             },
             tabID,
             requestID

packages/core/src/codewhispererChat/storages/chatHistory.ts

@@ -6,8 +6,8 @@ import { ChatMessage, ToolResult, ToolResultStatus, ToolUse } from '@amzn/codewh
 import { randomUUID } from '../../shared/crypto'
 import { getLogger } from '../../shared/logger/logger'
 
-// Maximum number of messages to keep in history
-const MaxConversationHistoryLength = 100
+// Maximum number of characters to keep in history
+const MaxConversationHistoryCharacters = 600_000
 
 /**
  * ChatHistoryManager handles the storage and manipulation of chat history
@@ -70,9 +70,6 @@ export class ChatHistoryManager {
      */
     public appendUserMessage(newMessage: ChatMessage): void {
         this.lastUserMessage = newMessage
-        if (!newMessage.userInputMessage?.content || newMessage.userInputMessage?.content.trim() === '') {
-            this.logger.warn('input must not be empty when adding new messages')
-        }
         this.history.push(this.formatChatHistoryMessage(this.lastUserMessage))
     }
 
@@ -118,22 +115,81 @@ export class ChatHistoryManager {
             this.clearRecentHistory()
         }
 
-        if (this.history.length <= MaxConversationHistoryLength) {
-            return
+        // Check if we need to trim based on character count
+        const totalCharacters = this.calculateHistoryCharacterCount()
+        if (totalCharacters > MaxConversationHistoryCharacters) {
+            this.logger.debug(
+                `History size (${totalCharacters} chars) exceeds limit of ${MaxConversationHistoryCharacters} chars`
+            )
+            // Keep removing messages from the beginning until we're under the limit
+            do {
+                // Find the next valid user message to start from
+                const indexToTrim = this.findIndexToTrim()
+                if (indexToTrim !== undefined && indexToTrim > 0) {
+                    this.logger.debug(
+                        `Removing the first ${indexToTrim} elements in the history due to character count limit`
+                    )
+                    this.history.splice(0, indexToTrim)
+                } else {
+                    // If we can't find a valid starting point, reset it
+                    this.logger.debug('Could not find a valid point to trim, reset history to reduce character count')
+                    this.history = []
+                }
+            } while (
+                this.calculateHistoryCharacterCount() > MaxConversationHistoryCharacters &&
+                this.history.length > 2
+            )
         }
+    }
 
-        const indexToTrim = this.findIndexToTrim()
-        if (indexToTrim !== undefined) {
-            this.logger.debug(`Removing the first ${indexToTrim} elements in the history`)
-            this.history.splice(0, indexToTrim)
-        } else {
-            this.logger.debug('No valid starting user message found in the history, clearing')
-            this.history = []
+    private calculateHistoryCharacterCount(): number {
+        let count = 0
+        for (const message of this.history) {
+            // Count characters in user messages
+            if (message.userInputMessage?.content) {
+                count += message.userInputMessage.content.length
+                count += message.userInputMessage.userInputMessageContext?.editorState?.document?.text?.length ?? 0
+                if (message.userInputMessage.userInputMessageContext?.editorState?.relevantDocuments) {
+                    for (const document of message.userInputMessage.userInputMessageContext.editorState
+                        .relevantDocuments) {
+                        count += document.text?.length ?? 0
+                    }
+                }
+                if (message.userInputMessage.userInputMessageContext?.additionalContext) {
+                    for (const document of message.userInputMessage.userInputMessageContext.additionalContext) {
+                        count += document.innerContext?.length ?? 0
+                    }
+                }
+            }
+
+            // Count characters in assistant messages
+            if (message.assistantResponseMessage?.content) {
+                count += message.assistantResponseMessage.content.length
+            }
+
+            try {
+                // Count characters in tool uses and results
+                if (message.assistantResponseMessage?.toolUses) {
+                    for (const toolUse of message.assistantResponseMessage.toolUses) {
+                        count += JSON.stringify(toolUse).length
+                    }
+                }
+
+                if (message.userInputMessage?.userInputMessageContext?.toolResults) {
+                    for (const toolResult of message.userInputMessage.userInputMessageContext.toolResults) {
+                        count += JSON.stringify(toolResult).length
+                    }
+                }
+            } catch (error: any) {
+                this.logger.error(`Error calculating character count for tool uses/results: ${error.message}`)
+            }
         }
+        this.logger.debug(`Current history characters: ${count}`)
+        return count
     }
 
     private findIndexToTrim(): number | undefined {
-        for (let i = 1; i < this.history.length; i++) {
+        for (let i = 2; i < this.history.length; i++) {
             const message = this.history[i]
             if (this.isValidUserMessageWithoutToolResults(message)) {
                 return i
@@ -162,6 +218,10 @@ export class ChatHistoryManager {
     private ensureCurrentMessageIsValid(newUserMessage: ChatMessage): void {
         const lastHistoryMessage = this.history[this.history.length - 1]
         if (!lastHistoryMessage) {
+            if (newUserMessage.userInputMessage?.userInputMessageContext?.toolResults) {
+                this.logger.debug('No history message found, but new user message has tool results.')
+                newUserMessage.userInputMessage.userInputMessageContext.toolResults = undefined
+            }
             return
         }