fix: outdated history when trimming happens, add missing metric for compaction (aws#2047)

jguoamz · web-flow · commit 8390f6686c80 · 2025-08-05T09:18:12.000-07:00
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts
@@ -1022,11 +1022,8 @@ export class AgenticChatController implements ChatHandlers {
         if (currentMessage) {
             //  Get and process the messages from history DB to maintain invariants for service requests
             try {
-                const { messages: historyMessages, count: historyCharCount } = this.#chatHistoryDb.fixAndGetHistory(
-                    tabId,
-                    currentMessage,
-                    []
-                )
+                const { history: historyMessages, historyCount: historyCharCount } =
+                    this.#chatHistoryDb.fixAndGetHistory(tabId, conversationIdentifier ?? '', currentMessage, [])
                 messages = historyMessages
                 characterCount = historyCharCount
             } catch (err) {
@@ -1200,19 +1197,18 @@ export class AgenticChatController implements ChatHandlers {
             if (currentMessage) {
                 //  Get and process the messages from history DB to maintain invariants for service requests
                 try {
-                    const newUserInputCount = this.#chatHistoryDb.calculateNewMessageCharacterCount(
+                    const {
+                        history: historyMessages,
+                        historyCount: historyCharacterCount,
+                        currentCount: currentInputCount,
+                    } = this.#chatHistoryDb.fixAndGetHistory(
+                        tabId,
+                        conversationId,
                         currentMessage,
                         pinnedContextMessages
                     )
-                    const { messages: historyMessages, count: historyCharacterCount } =
-                        this.#chatHistoryDb.fixAndGetHistory(
-                            tabId,
-                            currentMessage,
-                            pinnedContextMessages,
-                            newUserInputCount
-                        )
                     messages = historyMessages
-                    currentRequestCount = newUserInputCount + historyCharacterCount
+                    currentRequestCount = currentInputCount + historyCharacterCount
                     this.#debug(`Request total character count: ${currentRequestCount}`)
                 } catch (err) {
                     if (err instanceof ToolResultValidationError) {
@@ -1438,6 +1434,10 @@ export class AgenticChatController implements ChatHandlers {
         }
 
         if (this.#shouldCompact(currentRequestCount)) {
+            this.#telemetryController.emitCompactNudge(
+                currentRequestCount,
+                this.#features.runtime.serverInfo.version ?? ''
+            )
             const messageId = this.#getMessageIdForCompact(uuid())
             const confirmationResult = this.#processCompactConfirmation(messageId, currentRequestCount)
             const cachedButtonBlockId = await chatResultStream.writeResultBlock(confirmationResult)
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.test.ts
@@ -91,6 +91,30 @@ describe('ChatDatabase', () => {
         })
     })
 
+    describe('replaceHistory', () => {
+        it('should replace history with messages', async () => {
+            await chatDb.databaseInitialize(0)
+            const tabId = 'tab-1'
+            const tabType = 'cwc'
+            const conversationId = 'conv-1'
+            const messages = [
+                { body: 'Test', type: 'prompt' as any, timestamp: new Date() },
+                { body: 'Thinking...', type: 'answer', timestamp: new Date() },
+            ]
+
+            // Call the method
+            chatDb.replaceHistory(tabId, tabType, conversationId, messages)
+
+            // Verify the messages array contains the summary and a dummy response
+            const messagesFromDb = chatDb.getMessages(tabId, 250)
+            assert.strictEqual(messagesFromDb.length, 2)
+            assert.strictEqual(messagesFromDb[0].body, 'Test')
+            assert.strictEqual(messagesFromDb[0].type, 'prompt')
+            assert.strictEqual(messagesFromDb[1].body, 'Thinking...')
+            assert.strictEqual(messagesFromDb[1].type, 'answer')
+        })
+    })
+
     describe('ensureValidMessageSequence', () => {
         it('should preserve valid alternating sequence', () => {
             const messages: Message[] = [
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.ts
@@ -634,6 +634,50 @@ export class ChatDatabase {
         }
     }
 
+    /**
+     * Replace history with summary/dummyResponse pair within a specified tab.
+     *
+     * This method manages chat messages by creating a new history with compacted summary and dummy response pairs
+     */
+    replaceHistory(tabId: string, tabType: TabType, conversationId: string, messages: Message[]) {
+        if (this.isInitialized()) {
+            const clientType = this.#features.lsp.getClientInitializeParams()?.clientInfo?.name || 'unknown'
+            const tabCollection = this.#db.getCollection<Tab>(TabCollection)
+
+            this.#features.logging.log(
+                `Update history with new messages: tabId=${tabId}, tabType=${tabType}, conversationId=${conversationId}`
+            )
+
+            const oldHistoryId = this.getOrCreateHistoryId(tabId)
+            // create a new historyId to start fresh
+            const historyId = this.createHistoryId(tabId)
+
+            const tabData = historyId ? tabCollection.findOne({ historyId }) : undefined
+            const tabTitle = tabData?.title || 'Amazon Q Chat'
+            messages = messages.map(msg => this.formatChatHistoryMessage(msg))
+            this.#features.logging.log(`Overriding tab with new historyId=${historyId}`)
+            tabCollection.insert({
+                historyId,
+                updatedAt: new Date(),
+                isOpen: true,
+                tabType: tabType,
+                title: tabTitle,
+                conversations: [
+                    {
+                        conversationId,
+                        clientType,
+                        updatedAt: new Date(),
+                        messages: messages,
+                    },
+                ],
+            })
+
+            if (oldHistoryId) {
+                tabCollection.findAndRemove({ historyId: oldHistoryId })
+            }
+        }
+    }
+
     formatChatHistoryMessage(message: Message): Message {
         if (message.type === ('prompt' as ChatItemType)) {
             let hasToolResults = false
@@ -663,11 +707,16 @@ export class ChatDatabase {
      */
     fixAndGetHistory(
         tabId: string,
+        conversationId: string,
         newUserMessage: ChatMessage,
-        pinnedContextMessages: ChatMessage[],
-        newUserInputCount?: number
+        pinnedContextMessages: ChatMessage[]
     ): MessagesWithCharacterCount {
-        let messagesWithCount: MessagesWithCharacterCount = { messages: [], count: 0 }
+        let newUserInputCount = this.calculateNewMessageCharacterCount(newUserMessage, pinnedContextMessages)
+        let messagesWithCount: MessagesWithCharacterCount = {
+            history: [],
+            historyCount: 0,
+            currentCount: newUserInputCount,
+        }
         if (!this.isInitialized()) {
             return messagesWithCount
         }
@@ -683,32 +732,32 @@ export class ChatDatabase {
             // 3. Fix new user prompt: Ensure lastMessage in history toolUse and newMessage toolResult relationship is valid
             this.validateAndFixNewMessageToolResults(allMessages, newUserMessage)
 
-            if (!newUserInputCount) {
-                newUserInputCount = this.calculateNewMessageCharacterCount(newUserMessage, pinnedContextMessages)
-            }
-
             // 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)
-            allMessages = messagesWithCount.messages
+            messagesWithCount = this.trimMessagesToMaxLength(allMessages, newUserInputCount, tabId, conversationId)
+
             // Edge case: If the history is empty and the next message contains tool results, then we have to just abandon them.
             if (
-                allMessages.length === 0 &&
+                messagesWithCount.history.length === 0 &&
                 newUserMessage.userInputMessage?.userInputMessageContext?.toolResults?.length &&
                 newUserMessage.userInputMessage?.userInputMessageContext?.toolResults?.length > 0
             ) {
                 this.#features.logging.warn('History overflow: abandoning dangling toolResults.')
                 newUserMessage.userInputMessage.userInputMessageContext.toolResults = []
                 newUserMessage.userInputMessage.content = 'The conversation history has overflowed, clearing state'
+                // Update character count for current message
+                this.#features.logging.debug(`Updating input character with pinnedContext`)
+                messagesWithCount.currentCount = this.calculateNewMessageCharacterCount(
+                    newUserMessage,
+                    pinnedContextMessages
+                )
             }
         }
 
         // Prepend pinned context fake message pair to beginning of history
         if (pinnedContextMessages.length === 2) {
-            messagesWithCount.messages = [
-                ...pinnedContextMessages.map(msg => chatMessageToMessage(msg)),
-                ...allMessages,
-            ]
+            const pinnedMessages = pinnedContextMessages.map(msg => chatMessageToMessage(msg))
+            messagesWithCount.history = [...pinnedMessages, ...messagesWithCount.history]
         }
 
         return messagesWithCount
@@ -740,11 +789,20 @@ export class ChatDatabase {
         return !!ctx && (!ctx.toolResults || ctx.toolResults.length === 0) && message.body !== ''
     }
 
-    private trimMessagesToMaxLength(messages: Message[], newUserInputCount: number): MessagesWithCharacterCount {
+    private trimMessagesToMaxLength(
+        messages: Message[],
+        newUserInputCount: number,
+        tabId: string,
+        conversationId: string
+    ): MessagesWithCharacterCount {
         let historyCharacterCount = this.calculateMessagesCharacterCount(messages)
         const maxHistoryCharacterSize = Math.max(0, MaxOverallCharacters - newUserInputCount)
-        this.#features.logging.debug(`Current remaining character budget: ${maxHistoryCharacterSize}`)
+        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) {
@@ -756,13 +814,20 @@ export class ChatDatabase {
                 this.#features.logging.debug(
                     'Could not find a valid point to trim, reset history to reduce character count'
                 )
-                return { messages: [], count: 0 }
+                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 {
-            messages,
-            count: historyCharacterCount,
+            history: messages,
+            historyCount: historyCharacterCount,
+            currentCount: newUserInputCount,
         }
     }
 
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/util.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/util.ts
@@ -126,8 +126,9 @@ export type TabWithDbMetadata = {
 export type DbReference = { collection: Collection<Tab>; db: Loki }
 
 export type MessagesWithCharacterCount = {
-    messages: Message[]
-    count: number
+    history: Message[]
+    historyCount: number
+    currentCount: number
 }
 
 /**
diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/telemetry/chatTelemetryController.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/telemetry/chatTelemetryController.ts
@@ -226,6 +226,18 @@ export class ChatTelemetryController {
             data: {
                 type,
                 characters,
+                credentialStartUrl: this.#credentialsProvider.getConnectionMetadata()?.sso?.startUrl,
+                languageServerVersion: languageServerVersion,
+            },
+        })
+    }
+
+    public emitCompactNudge(characters: number, languageServerVersion: string) {
+        this.#telemetry.emitMetric({
+            name: ChatTelemetryEventName.CompactNudge,
+            data: {
+                characters,
+                credentialStartUrl: this.#credentialsProvider.getConnectionMetadata()?.sso?.startUrl,
                 languageServerVersion: languageServerVersion,
             },
         })
diff --git a/server/aws-lsp-codewhisperer/src/shared/telemetry/types.ts b/server/aws-lsp-codewhisperer/src/shared/telemetry/types.ts
@@ -206,6 +206,7 @@ export enum ChatTelemetryEventName {
     MCPServerInit = 'amazonq_mcpServerInit',
     LoadHistory = 'amazonq_loadHistory',
     CompactHistory = 'amazonq_compactHistory',
+    CompactNudge = 'amazonq_compactNudge',
     ChatHistoryAction = 'amazonq_performChatHistoryAction',
     ExportTab = 'amazonq_exportTab',
     UiClick = 'ui_click',
@@ -231,6 +232,7 @@ export interface ChatTelemetryEventMap {
     [ChatTelemetryEventName.MCPServerInit]: MCPServerInitializeEvent
     [ChatTelemetryEventName.LoadHistory]: LoadHistoryEvent
     [ChatTelemetryEventName.CompactHistory]: CompactHistoryEvent
+    [ChatTelemetryEventName.CompactNudge]: CompactNudgeEvent
     [ChatTelemetryEventName.ChatHistoryAction]: ChatHistoryActionEvent
     [ChatTelemetryEventName.ExportTab]: ExportTabEvent
     [ChatTelemetryEventName.UiClick]: UiClickEvent
@@ -392,6 +394,13 @@ export type LoadHistoryEvent = {
 export type CompactHistoryEvent = {
     type: CompactHistoryActionType
     characters: number
+    credentialStartUrl?: string
+    languageServerVersion?: string
+}
+
+export type CompactNudgeEvent = {
+    characters: number
+    credentialStartUrl?: string
     languageServerVersion?: string
 }
 
