fix: improve history management (aws#2312) (aws#2357)

Co-authored-by: aws-toolkit-automation <[email protected]>

Author: jguoamz

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

@@ -137,7 +137,7 @@ import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/Ba
 import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager'
 import { AmazonQWorkspaceConfig } from '../../shared/amazonQServiceManager/configurationUtils'
 import { TabBarController } from './tabBarController'
-import { ChatDatabase, MaxOverallCharacters, ToolResultValidationError } from './tools/chatDb/chatDb'
+import { ChatDatabase, ToolResultValidationError } from './tools/chatDb/chatDb'
 import {
     AgenticChatEventParser,
     ChatResultWithMetadata as AgenticChatResultWithMetadata,
@@ -187,6 +187,8 @@ import {
     DEFAULT_WINDOW_REJECT_SHORTCUT,
     DEFAULT_MACOS_STOP_SHORTCUT,
     DEFAULT_WINDOW_STOP_SHORTCUT,
+    COMPACTION_CHARACTER_THRESHOLD,
+    MAX_OVERALL_CHARACTERS,
 } from './constants/constants'
 import {
     AgenticChatError,
@@ -1179,8 +1181,7 @@ export class AgenticChatController implements ChatHandlers {
      * Runs the compaction, making requests and processing tool uses until completion
      */
     #shouldCompact(currentRequestCount: number): boolean {
-        // 80% of 570K limit
-        if (currentRequestCount > 456_000) {
+        if (currentRequestCount > COMPACTION_CHARACTER_THRESHOLD) {
             this.#debug(`Current request total character count is: ${currentRequestCount}, prompting user to compact`)
             return true
         } else {
@@ -1225,7 +1226,7 @@ export class AgenticChatController implements ChatHandlers {
             //  Get and process the messages from history DB to maintain invariants for service requests
             try {
                 const { history: historyMessages, historyCount: historyCharCount } =
-                    this.#chatHistoryDb.fixAndGetHistory(tabId, conversationIdentifier ?? '', currentMessage, [])
+                    this.#chatHistoryDb.fixAndGetHistory(tabId, currentMessage, [])
                 messages = historyMessages
                 characterCount = historyCharCount
             } catch (err) {
@@ -1403,12 +1404,7 @@ export class AgenticChatController implements ChatHandlers {
                         history: historyMessages,
                         historyCount: historyCharacterCount,
                         currentCount: currentInputCount,
-                    } = this.#chatHistoryDb.fixAndGetHistory(
-                        tabId,
-                        conversationId,
-                        currentMessage,
-                        pinnedContextMessages
-                    )
+                    } = this.#chatHistoryDb.fixAndGetHistory(tabId, currentMessage, pinnedContextMessages)
                     messages = historyMessages
                     currentRequestCount = currentInputCount + historyCharacterCount
                     this.#debug(`Request total character count: ${currentRequestCount}`)
@@ -2779,7 +2775,7 @@ export class AgenticChatController implements ChatHandlers {
             body: COMPACTION_HEADER_BODY,
             buttons,
         } as any
-        const body = COMPACTION_BODY(Math.round((characterCount / MaxOverallCharacters) * 100))
+        const body = COMPACTION_BODY(Math.round((characterCount / MAX_OVERALL_CHARACTERS) * 100))
         return {
             type: 'tool',
             messageId,

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

@@ -1,5 +1,3 @@
-import { BedrockModel } from './modelSelection'
-
 // Error message constants
 export const GENERIC_ERROR_MS = 'An unexpected error occurred, check the logs for more information.'
 export const OUTPUT_LIMIT_EXCEEDS_PARTIAL_MSG = 'output exceeds maximum character limit of'
@@ -16,6 +14,10 @@ export const SERVICE_MANAGER_POLL_INTERVAL_MS = 100
 export const GENERATE_ASSISTANT_RESPONSE_INPUT_LIMIT = 500_000
 
 // Compaction
+// Maximum number of characters per request used for compaction prompt
+// 200K tokens * 3.5 = 700K characters, intentionally overestimating with 3.5:1 ratio
+export const MAX_OVERALL_CHARACTERS = 700_000
+export const COMPACTION_CHARACTER_THRESHOLD = 0.7 * MAX_OVERALL_CHARACTERS
 export const COMPACTION_BODY = (threshold: number) =>
     `The context window is almost full (${threshold}%) and exceeding it will clear your history. Amazon Q can compact your history instead.`
 export const COMPACTION_HEADER_BODY = 'Compact chat history?'

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.ts

@@ -45,11 +45,6 @@ export class ToolResultValidationError extends Error {
 }
 
 export const EMPTY_CONVERSATION_LIST_ID = 'empty'
-// Maximum number of characters to keep in request
-// (200K tokens - 8K output tokens - 2k system prompt) * 3 = 570K characters, intentionally overestimating with 3:1 ratio
-export const MaxOverallCharacters = 570_000
-// Maximum number of history messages to include in each request to the LLM
-const maxConversationHistoryMessages = 250
 
 /**
  * A singleton database class that manages chat history persistence using LokiJS.
@@ -708,15 +703,12 @@ export class ChatDatabase {
 
     /**
      * Prepare the history messages for service request and fix the persisted history in DB to maintain the following invariants:
-     * 1. The history contains at most MaxConversationHistoryMessages messages. Oldest messages are dropped.
-     * 2. The first message is from the user and without any tool usage results, and the last message is from the assistant.
+     * 1. The first message is from the user and without any tool usage results, and the last message is from the assistant.
      *    The history contains alternating sequene of userMessage followed by assistantMessages
-     * 3. The toolUse and toolResult relationship is valid
-     * 4. The history character length is <= MaxConversationHistoryCharacters - newUserMessageCharacterCount. Oldest messages are dropped.
+     * 2. The toolUse and toolResult relationship is valid
      */
     fixAndGetHistory(
         tabId: string,
-        conversationId: string,
         newUserMessage: ChatMessage,
         pinnedContextMessages: ChatMessage[]
     ): MessagesWithCharacterCount {
@@ -732,18 +724,19 @@ export class ChatDatabase {
 
         this.#features.logging.info(`Fixing history: tabId=${tabId}`)
 
-        // 1. Make sure the length of the history messages don't exceed MaxConversationHistoryMessages
-        let allMessages = this.getMessages(tabId, maxConversationHistoryMessages)
+        let allMessages = this.getMessages(tabId)
         if (allMessages.length > 0) {
-            // 2. Fix history: Ensure messages in history is valid for server side checks
+            // 1. Fix history: Ensure messages in history is valid for server side checks
             this.ensureValidMessageSequence(tabId, allMessages)
 
-            // 3. Fix new user prompt: Ensure lastMessage in history toolUse and newMessage toolResult relationship is valid
+            // 2. Fix new user prompt: Ensure lastMessage in history toolUse and newMessage toolResult relationship is valid
             this.validateAndFixNewMessageToolResults(allMessages, newUserMessage)
 
-            // 4. NOTE: Keep this trimming logic at the end of the preprocess.
-            // Make sure max characters ≤ remaining Character Budget, must be put at the end of preprocessing
-            messagesWithCount = this.trimMessagesToMaxLength(allMessages, newUserInputCount, tabId, conversationId)
+            messagesWithCount = {
+                history: allMessages,
+                historyCount: this.calculateMessagesCharacterCount(allMessages),
+                currentCount: newUserInputCount,
+            }
 
             // Edge case: If the history is empty and the next message contains tool results, then we have to just abandon them.
             if (
@@ -772,74 +765,11 @@ export class ChatDatabase {
         return messagesWithCount
     }
 
-    /**
-     * Finds a suitable "break point" index in the message sequence.
-     *
-     * It ensures that the "break point" is at a clean conversation boundary where:
-     * 1. The message is from a user (type === 'prompt')
-     * 2. The message doesn't contain tool results that would break tool use/result pairs
-     * 3. The message has a non-empty body
-     *
-     * @param allMessages The array of conversation messages to search through
-     * @returns The index to trim from, or undefined if no suitable trimming point is found
-     */
-    private findIndexToTrim(allMessages: Message[]): number | undefined {
-        for (let i = 2; i < allMessages.length; i++) {
-            const message = allMessages[i]
-            if (message.type === ('prompt' as ChatItemType) && this.isValidUserMessageWithoutToolResults(message)) {
-                return i
-            }
-        }
-        return undefined
-    }
-
     private isValidUserMessageWithoutToolResults(message: Message): boolean {
         const ctx = message.userInputMessageContext
         return !!ctx && (!ctx.toolResults || ctx.toolResults.length === 0) && message.body !== ''
     }
 
-    private trimMessagesToMaxLength(
-        messages: Message[],
-        newUserInputCount: number,
-        tabId: string,
-        conversationId: string
-    ): MessagesWithCharacterCount {
-        let historyCharacterCount = this.calculateMessagesCharacterCount(messages)
-        const maxHistoryCharacterSize = Math.max(0, MaxOverallCharacters - newUserInputCount)
-        let trimmedHistory = false
-        this.#features.logging.debug(
-            `Current history character count: ${historyCharacterCount}, remaining history character budget: ${maxHistoryCharacterSize}`
-        )
-        while (historyCharacterCount > maxHistoryCharacterSize && messages.length > 2) {
-            trimmedHistory = true
-            // Find the next valid user message to start from
-            const indexToTrim = this.findIndexToTrim(messages)
-            if (indexToTrim !== undefined && indexToTrim > 0) {
-                this.#features.logging.debug(
-                    `Removing the first ${indexToTrim} elements in the history due to character count limit`
-                )
-                messages.splice(0, indexToTrim)
-            } else {
-                this.#features.logging.debug(
-                    'Could not find a valid point to trim, reset history to reduce character count'
-                )
-                this.replaceHistory(tabId, 'cwc', conversationId, [])
-                return { history: [], historyCount: 0, currentCount: newUserInputCount }
-            }
-            historyCharacterCount = this.calculateMessagesCharacterCount(messages)
-            this.#features.logging.debug(`History character count post trimming: ${historyCharacterCount}`)
-        }
-
-        if (trimmedHistory) {
-            this.replaceHistory(tabId, 'cwc', conversationId, messages)
-        }
-        return {
-            history: messages,
-            historyCount: historyCharacterCount,
-            currentCount: newUserInputCount,
-        }
-    }
-
     private calculateToolSpecCharacterCount(currentMessage: ChatMessage): number {
         let count = 0
         if (currentMessage.userInputMessage?.userInputMessageContext?.tools) {