fix(chat): prune history if it's too large

Author: jguoamz

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
@@ -118,22 +118,68 @@ 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
+            while (
+                this.calculateHistoryCharacterCount() > MaxConversationHistoryCharacters &&
+                this.history.length > 2
+            ) {
+                // 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, remove the oldest message
+                    this.logger.debug('Removing oldest message to reduce character count')
+                    this.history.shift()
+                }
+            }
         }
+    }
 
-        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 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}`)
+            }
         }
+        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