From e58911bd04e2dd69d601b0464777fefaacf75c42 Mon Sep 17 00:00:00 2001 From: Tyrone Smith Date: Wed, 9 Apr 2025 23:14:04 -0700 Subject: [PATCH] feat(amazonq): Use Chat History DB instead of in-memory store --- .../codewhispererChat/clients/chat/v0/chat.ts | 24 +- .../controllers/chat/controller.ts | 29 +- .../controllers/chat/messenger/messenger.ts | 13 - .../codewhispererChat/storages/chatHistory.ts | 273 ------------------ .../storages/chatHistoryStorage.ts | 43 --- packages/core/src/shared/db/chatDb/chatDb.ts | 239 ++++++++++++++- 6 files changed, 251 insertions(+), 370 deletions(-) delete mode 100644 packages/core/src/codewhispererChat/storages/chatHistory.ts delete mode 100644 packages/core/src/codewhispererChat/storages/chatHistoryStorage.ts diff --git a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts index 53551b254f5..13b972fe101 100644 --- a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts +++ b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts @@ -16,16 +16,15 @@ import { createQDeveloperStreamingClient } from '../../../../shared/clients/qDev import { UserWrittenCodeTracker } from '../../../../codewhisperer/tracker/userWrittenCodeTracker' import { PromptMessage } from '../../../controllers/chat/model' import { FsWriteBackup } from '../../../../codewhispererChat/tools/fsWrite' +import { randomUUID } from '../../../../shared/crypto' export type ToolUseWithError = { toolUse: ToolUse error: Error | undefined } -import { getLogger } from '../../../../shared/logger/logger' -import { randomUUID } from '../../../../shared/crypto' export class ChatSession { - private sessionId?: string + private sessionId: string /** * _readFiles = list of files read from the project to gather context before generating response. * _showDiffOnFileWrite = Controls whether to show diff view (true) or file context view (false) to the user @@ -46,7 +45,7 @@ export class ChatSession { // TODO: doesn't handle the edge case when two files share the same relativePath string but from different root // e.g. root_a/file1 vs root_b/file1 relativePathToWorkspaceRoot: Map = new Map() - public get sessionIdentifier(): string | undefined { + public get sessionIdentifier(): string { return this.sessionId } @@ -86,13 +85,14 @@ export class ChatSession { constructor() { this.createNewTokenSource() + this.sessionId = randomUUID() } createNewTokenSource() { this.tokenSource = new vscode.CancellationTokenSource() } - public setSessionID(id?: string) { + public setSessionID(id: string) { this.sessionId = id } public get readFiles(): string[] { @@ -120,14 +120,6 @@ export class ChatSession { ) } - const responseStream = response.sendMessageResponse - for await (const event of responseStream) { - if ('messageMetadataEvent' in event) { - this.sessionId = event.messageMetadataEvent?.conversationId - break - } - } - UserWrittenCodeTracker.instance.onQFeatureInvoked() return response } @@ -142,12 +134,6 @@ export class ChatSession { ) } - this.sessionId = response.conversationId - if (this.sessionId?.length === 0) { - getLogger().debug(`Session ID: ${this.sessionId} is empty. Generating random UUID`) - this.sessionId = randomUUID() - } - UserWrittenCodeTracker.instance.onQFeatureInvoked() return response diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts index 8b1f733d0e2..571580e33d5 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/controller.ts @@ -97,10 +97,10 @@ import { amazonQTabSuffix } from '../../../shared/constants' import { OutputKind } from '../../tools/toolShared' import { ToolUtils, Tool, ToolType } from '../../tools/toolUtils' import { ChatStream } from '../../tools/chatStream' -import { ChatHistoryStorage } from '../../storages/chatHistoryStorage' import { tempDirPath } from '../../../shared/filesystemUtilities' import { Database } from '../../../shared/db/chatDb/chatDb' import { TabBarController } from './tabBarController' +import { messageToChatMessage } from '../../../shared/db/chatDb/util' export interface ChatControllerMessagePublishers { readonly processPromptChatMessage: MessagePublisher @@ -175,7 +175,6 @@ export class ChatController { private readonly userIntentRecognizer: UserIntentRecognizer private readonly telemetryHelper: CWCTelemetryHelper private userPromptsWatcher: vscode.FileSystemWatcher | undefined - private readonly chatHistoryStorage: ChatHistoryStorage private chatHistoryDb = Database.getInstance() private cancelTokenSource: vscode.CancellationTokenSource = new vscode.CancellationTokenSource() @@ -195,7 +194,6 @@ export class ChatController { this.editorContentController = new EditorContentController() this.promptGenerator = new PromptsGenerator() this.userIntentRecognizer = new UserIntentRecognizer() - this.chatHistoryStorage = new ChatHistoryStorage() this.tabBarController = new TabBarController(this.messenger) onDidChangeAmazonQVisibility((visible) => { @@ -417,7 +415,7 @@ export class ChatController { const session = this.sessionStorage.getSession(message.tabID) session.tokenSource.cancel() this.messenger.sendEmptyMessage(message.tabID, '', undefined) - this.chatHistoryStorage.getTabHistory(message.tabID).clearRecentHistory() + this.chatHistoryDb.clearRecentHistory(message.tabID) this.telemetryHelper.recordInteractionWithAgenticChat(AgenticChatInteractionType.StopChat, message) } @@ -475,7 +473,6 @@ export class ChatController { private async processTabCloseMessage(message: TabClosedMessage) { this.sessionStorage.deleteSession(message.tabID) - this.chatHistoryStorage.deleteHistory(message.tabID) this.triggerEventsStorage.removeTabEvents(message.tabID) // this.telemetryHelper.recordCloseChat(message.tabID) this.chatHistoryDb.updateTabOpenState(message.tabID, false) @@ -1039,7 +1036,7 @@ export class ChatController { getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`) this.sessionStorage.deleteSession(tabID) - this.chatHistoryStorage.getTabHistory(tabID).clear() + this.chatHistoryDb.clearTab(tabID) } private async processContextMenuCommand(command: EditorContextCommand) { @@ -1161,7 +1158,6 @@ export class ChatController { switch (message.command) { case 'clear': this.sessionStorage.deleteSession(message.tabID) - this.chatHistoryStorage.getTabHistory(message.tabID).clear() this.triggerEventsStorage.removeTabEvents(message.tabID) recordTelemetryChatRunCommand('clear') this.chatHistoryDb.clearTab(message.tabID) @@ -1452,10 +1448,6 @@ export class ChatController { } const session = this.sessionStorage.getSession(tabID) - if (!session.localHistoryHydrated) { - triggerPayload.history = this.chatHistoryDb.getMessages(triggerEvent.tabID, 10) - session.localHistoryHydrated = true - } await this.resolveContextCommandPayload(triggerPayload, session) triggerPayload.useRelevantDocuments = triggerPayload.context.some( (context) => typeof context !== 'string' && context.command === '@workspace' @@ -1494,16 +1486,14 @@ export class ChatController { const request = triggerPayloadToChatRequest(triggerPayload) - const chatHistory = this.chatHistoryStorage.getTabHistory(tabID) const currentMessage = request.conversationState.currentMessage if (currentMessage) { - chatHistory.fixHistory(currentMessage) + this.chatHistoryDb.fixHistory(tabID, currentMessage) } - request.conversationState.history = chatHistory.getHistory() - - const conversationId = chatHistory.getConversationId() || randomUUID() - chatHistory.setConversationId(conversationId) - request.conversationState.conversationId = conversationId + request.conversationState.history = this.chatHistoryDb + .getMessages(tabID) + .map((chat) => messageToChatMessage(chat)) + request.conversationState.conversationId = session.sessionIdentifier triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments) @@ -1562,7 +1552,6 @@ export class ChatController { this.telemetryHelper.recordStartConversation(triggerEvent, triggerPayload) if (currentMessage && session.sessionIdentifier) { - chatHistory.appendUserMessage(currentMessage) this.chatHistoryDb.addMessage(tabID, 'cwc', session.sessionIdentifier, { body: triggerPayload.message, type: 'prompt' as any, @@ -1577,7 +1566,6 @@ export class ChatController { response.$metadata.requestId } metadata: ${inspect(response.$metadata, { depth: 12 })}` ) - this.cancelTokenSource = new vscode.CancellationTokenSource() await this.messenger.sendAIResponse( response, @@ -1585,7 +1573,6 @@ export class ChatController { tabID, triggerID, triggerPayload, - chatHistory, this.cancelTokenSource.token ) diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts index f8306b00874..d45ed522708 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts @@ -57,7 +57,6 @@ import { } from '@aws/mynah-ui' import { Database } from '../../../../shared/db/chatDb/chatDb' import { TabType } from '../../../../amazonq/webview/ui/storages/tabsStorage' -import { ChatHistoryManager } from '../../../storages/chatHistory' import { ToolType, ToolUtils } from '../../../tools/toolUtils' import { ChatStream } from '../../../tools/chatStream' import path from 'path' @@ -185,7 +184,6 @@ export class Messenger { tabID: string, triggerID: string, triggerPayload: TriggerPayload, - chatHistoryManager: ChatHistoryManager, cancelToken: vscode.CancellationToken ) { let message = '' @@ -508,17 +506,6 @@ export class Messenger { ) ) - chatHistoryManager.pushAssistantMessage({ - assistantResponseMessage: { - messageId: messageID, - content: message, - references: codeReference, - ...(toolUse && - toolUse.input !== undefined && - toolUse.input !== '' && { toolUses: [{ ...toolUse }] }), - }, - }) - getLogger().info( `All events received. requestId=%s counts=%s`, response.$metadata.requestId, diff --git a/packages/core/src/codewhispererChat/storages/chatHistory.ts b/packages/core/src/codewhispererChat/storages/chatHistory.ts deleted file mode 100644 index d140582b189..00000000000 --- a/packages/core/src/codewhispererChat/storages/chatHistory.ts +++ /dev/null @@ -1,273 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -import { ChatMessage, ToolResult, ToolResultStatus, ToolUse } from '@amzn/codewhisperer-streaming' -import { randomUUID } from '../../shared/crypto' -import { getLogger } from '../../shared/logger/logger' - -// Maximum number of characters to keep in history -const MaxConversationHistoryCharacters = 600_000 - -/** - * ChatHistoryManager handles the storage and manipulation of chat history - * for CodeWhisperer Chat sessions. - */ -export class ChatHistoryManager { - private conversationId: string - private tabId: string - private history: ChatMessage[] = [] - private logger = getLogger() - private lastUserMessage?: ChatMessage - - constructor(tabId?: string) { - this.conversationId = randomUUID() - this.tabId = tabId ?? randomUUID() - } - - /** - * Get the conversation ID - */ - public getConversationId(): string { - return this.conversationId - } - - public setConversationId(conversationId: string) { - this.conversationId = conversationId - } - - /** - * Get the tab ID - */ - public getTabId(): string { - return this.tabId - } - - /** - * Set the tab ID - */ - public setTabId(tabId: string) { - this.tabId = tabId - } - - /** - * Get the full chat history - */ - public getHistory(): ChatMessage[] { - return [...this.history] - } - - /** - * Clear the conversation history - */ - public clear(): void { - this.history = [] - this.conversationId = '' - } - - /** - * Append a new user message to be sent - */ - public appendUserMessage(newMessage: ChatMessage): void { - this.lastUserMessage = newMessage - this.history.push(this.formatChatHistoryMessage(this.lastUserMessage)) - } - - /** - * Push an assistant message to the history - */ - public pushAssistantMessage(newMessage: ChatMessage): void { - if (newMessage !== undefined && this.lastUserMessage === undefined) { - this.logger.warn('first assistant response should always come after user input message') - return - } - // check if last message in histroy is assistant message and now replace it in that case - if (this.history.length > 0 && this.history.at(-1)?.assistantResponseMessage) { - this.history.pop() - } - this.history.push(newMessage) - } - - /** - * Fixes the history to maintain the following invariants: - * 1. The history length is <= MAX_CONVERSATION_HISTORY_LENGTH. Oldest messages are dropped. - * 2. The first message is from the user. Oldest messages are dropped if needed. - * 3. The last message is from the assistant. The last message is dropped if it is from the user. - * 4. If the last message is from the assistant and it contains tool uses, and a next user - * message is set without tool results, then the user message will have cancelled tool results. - */ - public fixHistory(newUserMessage: ChatMessage): void { - this.trimConversationHistory() - this.ensureLastMessageFromAssistant() - this.ensureCurrentMessageIsValid(newUserMessage) - } - - private trimConversationHistory(): void { - // make sure the UseInputMessage is the first stored message - if (this.history.length === 1 && this.history[0].assistantResponseMessage) { - this.history = [] - } - - if ( - this.history.at(-1)?.assistantResponseMessage?.content === '' && - this.history.at(-1)?.assistantResponseMessage?.toolUses === undefined - ) { - this.clearRecentHistory() - } - - // 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 - ) - } - } - - 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}`) - } - } - this.logger.debug(`Current history characters: ${count}`) - return count - } - - private findIndexToTrim(): number | undefined { - for (let i = 2; i < this.history.length; i++) { - const message = this.history[i] - if (this.isValidUserMessageWithoutToolResults(message)) { - return i - } - } - return undefined - } - - private isValidUserMessageWithoutToolResults(message: ChatMessage): boolean { - if (!message.userInputMessage) { - return false - } - const ctx = message.userInputMessage.userInputMessageContext - return Boolean( - ctx && (!ctx.toolResults || ctx.toolResults.length === 0) && message.userInputMessage.content !== '' - ) - } - - private ensureLastMessageFromAssistant(): void { - if (this.history.length > 0 && this.history[this.history.length - 1].userInputMessage !== undefined) { - this.logger.debug('Last message in history is from the user, dropping') - this.history.pop() - } - } - - 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 - // tool results are empty, so content must not be empty - newUserMessage.userInputMessage.content = 'Conversation history was too large, so it was cleared.' - } - return - } - - if (lastHistoryMessage.assistantResponseMessage?.toolUses?.length) { - const toolResults = newUserMessage.userInputMessage?.userInputMessageContext?.toolResults - if (!toolResults || toolResults.length === 0) { - const abandonedToolResults = this.createAbandonedToolResults( - lastHistoryMessage.assistantResponseMessage.toolUses - ) - - if (newUserMessage.userInputMessage?.userInputMessageContext) { - newUserMessage.userInputMessage.userInputMessageContext.toolResults = abandonedToolResults - } - } - } - } - - private createAbandonedToolResults(toolUses: ToolUse[]): ToolResult[] { - return toolUses.map((toolUse) => ({ - toolUseId: toolUse.toolUseId, - content: [ - { - type: 'Text', - text: 'Tool use was cancelled by the user', - }, - ], - status: ToolResultStatus.ERROR, - })) - } - - private formatChatHistoryMessage(message: ChatMessage): ChatMessage { - if (message.userInputMessage !== undefined) { - return { - userInputMessage: { - ...message.userInputMessage, - userInputMessageContext: { - // Only keep toolResults in history - toolResults: message.userInputMessage.userInputMessageContext?.toolResults, - }, - }, - } - } - return message - } - - public clearRecentHistory(): void { - if (this.history.length === 0) { - return - } - - const lastHistoryMessage = this.history[this.history.length - 1] - - if (lastHistoryMessage.userInputMessage?.userInputMessageContext) { - this.history.pop() - } else if (lastHistoryMessage.assistantResponseMessage) { - this.history.splice(-2) - } - } -} diff --git a/packages/core/src/codewhispererChat/storages/chatHistoryStorage.ts b/packages/core/src/codewhispererChat/storages/chatHistoryStorage.ts deleted file mode 100644 index a01edb035a1..00000000000 --- a/packages/core/src/codewhispererChat/storages/chatHistoryStorage.ts +++ /dev/null @@ -1,43 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import { ChatHistoryManager } from './chatHistory' - -/** - * ChatHistoryStorage manages ChatHistoryManager instances for multiple tabs. - * Each tab has its own ChatHistoryManager to maintain separate chat histories. - */ -export class ChatHistoryStorage { - private histories: Map = new Map() - - /** - * Gets the ChatHistoryManager for a specific tab. - * If no history exists for the tab, creates a new one. - * - * @param tabId The ID of the tab - * @returns The ChatHistoryManager for the specified tab - */ - public getTabHistory(tabId: string): ChatHistoryManager { - const historyFromStorage = this.histories.get(tabId) - if (historyFromStorage !== undefined) { - return historyFromStorage - } - - // Create a new ChatHistoryManager with the tabId - const newHistory = new ChatHistoryManager(tabId) - this.histories.set(tabId, newHistory) - - return newHistory - } - - /** - * Deletes the ChatHistoryManager for a specific tab. - * - * @param tabId The ID of the tab - */ - public deleteHistory(tabId: string) { - this.histories.delete(tabId) - } -} diff --git a/packages/core/src/shared/db/chatDb/chatDb.ts b/packages/core/src/shared/db/chatDb/chatDb.ts index cbacba4bb11..2b949b23869 100644 --- a/packages/core/src/shared/db/chatDb/chatDb.ts +++ b/packages/core/src/shared/db/chatDb/chatDb.ts @@ -20,6 +20,10 @@ import crypto from 'crypto' import path from 'path' import { fs } from '../../fs/fs' import { getLogger } from '../../logger/logger' +import { ChatMessage, ToolResultStatus } from '@amzn/codewhisperer-streaming' + +// Maximum number of characters to keep in history +const MaxConversationHistoryCharacters = 600_000 /** * A singleton database class that manages chat history persistence using LokiJS. @@ -140,6 +144,18 @@ export class Database { return undefined } + getActiveConversation(historyId: string): Conversation | undefined { + const tabCollection = this.db.getCollection(TabCollection) + const tabData = tabCollection.findOne({ historyId }) + + if (!tabData?.conversations.length) { + this.logger.debug('No active conversations found') + return undefined + } + + return tabData.conversations[0] + } + clearTab(tabId: string) { if (this.initialized) { const tabCollection = this.db.getCollection(TabCollection) @@ -154,6 +170,34 @@ export class Database { } } + // Removes the most recent message(s) from the chat history for a given tab + clearRecentHistory(tabId: string): void { + if (this.initialized) { + const historyId = this.historyIdMapping.get(tabId) + this.logger.info(`Clearing recent history: tabId=${tabId}, historyId=${historyId || 'undefined'}`) + if (historyId) { + const tabCollection = this.db.getCollection(TabCollection) + const tabData = tabCollection.findOne({ historyId }) + if (tabData) { + const activeConversation = tabData.conversations[0] + const allMessages = this.getMessages(tabId) + const lastMessage = allMessages[allMessages.length - 1] + this.logger.debug(`Last message type: ${lastMessage.type}`) + + if (lastMessage.type === ('prompt' as ChatItemType)) { + allMessages.pop() + this.logger.info(`Removed last user message`) + } else { + allMessages.splice(-2) + this.logger.info(`Removed last assistant message and user message`) + } + activeConversation.messages = allMessages + tabCollection.update(tabData) + } + } + } + } + updateTabOpenState(tabId: string, isOpen: boolean) { if (this.initialized) { const tabCollection = this.db.getCollection(TabCollection) @@ -209,7 +253,7 @@ export class Database { * @param tabId The ID of the tab to get messages from * @param numMessages Optional number of most recent messages to return. If not provided, returns all messages. */ - getMessages(tabId: string, numMessages?: number) { + getMessages(tabId: string, numMessages?: number): Message[] { if (this.initialized) { const tabCollection = this.db.getCollection(TabCollection) const historyId = this.historyIdMapping.get(tabId) @@ -268,6 +312,7 @@ export class Database { message.type === ('prompt' as ChatItemType) && message.body.trim().length > 0 ? message.body : tabData?.title || 'Amazon Q Chat' + message = this.formatChatHistoryMessage(message) if (tabData) { this.logger.info(`Found existing tab data, updating conversations`) tabData.conversations = updateOrCreateConversation(tabData.conversations, conversationId, message) @@ -287,4 +332,196 @@ export class Database { } } } + + private formatChatHistoryMessage(message: Message): Message { + if (message.type === ('prompt' as ChatItemType)) { + return { + ...message, + userInputMessageContext: { + // Only keep toolResults in history + toolResults: message.userInputMessageContext?.toolResults, + }, + } + } + return message + } + + /** + * Fixes the history to maintain the following invariants: + * 1. The history character length is <= MAX_CONVERSATION_HISTORY_CHARACTERS. Oldest messages are dropped. + * 2. The first message is from the user. Oldest messages are dropped if needed. + * 3. The last message is from the assistant. The last message is dropped if it is from the user. + * 4. If the last message is from the assistant and it contains tool uses, and a next user + * message is set without tool results, then the user message will have cancelled tool results. + */ + fixHistory(tabId: string, newUserMessage: ChatMessage): void { + if (!this.initialized) { + return + } + const historyId = this.historyIdMapping.get(tabId) + this.logger.info(`Fixing history: tabId=${tabId}, historyId=${historyId || 'undefined'}`) + + if (!historyId) { + return + } + + const tabCollection = this.db.getCollection(TabCollection) + const tabData = tabCollection.findOne({ historyId }) + if (!tabData) { + return + } + + const activeConversation = tabData.conversations[0] + let allMessages = activeConversation.messages + this.logger.info(`Found ${allMessages.length} messages in conversation`) + + // Drop empty assistant partial if it’s the last message + this.handleEmptyAssistantMessage(allMessages) + + // Make sure max characters ≤ MaxConversationHistoryCharacters + allMessages = this.trimMessagesToMaxLength(allMessages) + + // Ensure messages in history a valid for server side checks + this.ensureValidMessageSequence(allMessages) + + // If the last message is from the assistant and it contains tool uses, and a next user message is set without tool results, then the user message will have cancelled tool results. + this.handleToolUses(allMessages, newUserMessage) + + activeConversation.messages = allMessages + tabCollection.update(tabData) + this.logger.info(`Updated tab data in collection`) + } + + private handleEmptyAssistantMessage(messages: Message[]): void { + if (messages.length === 0) { + return + } + + const lastMsg = messages[messages.length - 1] + if ( + lastMsg.type === ('answer' as ChatItemType) && + (!lastMsg.body || lastMsg.body.trim().length === 0) && + (!lastMsg.toolUses || lastMsg.toolUses.length === 0) + ) { + this.logger.debug( + 'Last message is empty partial assistant. Removed last assistant message and user message' + ) + messages.splice(-2) + } + } + + private trimMessagesToMaxLength(messages: Message[]): Message[] { + let totalCharacters = this.calculateCharacterCount(messages) + while (totalCharacters > MaxConversationHistoryCharacters && messages.length > 2) { + // Find the next valid user message to start from + const indexToTrim = this.findIndexToTrim(messages) + if (indexToTrim !== undefined && indexToTrim > 0) { + this.logger.debug( + `Removing the first ${indexToTrim} elements in the history due to character count limit` + ) + messages.splice(0, indexToTrim) + } else { + this.logger.debug('Could not find a valid point to trim, reset history to reduce character count') + return [] + } + totalCharacters = this.calculateCharacterCount(messages) + } + return messages + } + + private calculateCharacterCount(allMessages: Message[]): number { + let count = 0 + for (const message of allMessages) { + // Count characters of all message text + count += message.body.length + + // Count characters in tool uses + if (message.toolUses) { + try { + for (const toolUse of message.toolUses) { + count += JSON.stringify(toolUse).length + } + } catch (e) { + this.logger.error(`Error counting toolUses: ${String(e)}`) + } + } + // Count characters in tool results + if (message.userInputMessageContext?.toolResults) { + try { + for (const toolResul of message.userInputMessageContext.toolResults) { + count += JSON.stringify(toolResul).length + } + } catch (e) { + this.logger.error(`Error counting toolResults: ${String(e)}`) + } + } + } + this.logger.debug(`Current history characters: ${count}`) + return count + } + + 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 ensureValidMessageSequence(messages: Message[]): void { + // Make sure the first stored message is from the user (type === 'prompt'), else drop + while (messages.length > 0 && messages[0].type === ('answer' as ChatItemType)) { + messages.shift() + this.logger.debug('Dropped first message since it is not from user') + } + + // Make sure the last stored message is from the assistant (type === 'answer'), else drop + if (messages.length > 0 && messages[messages.length - 1].type === ('prompt' as ChatItemType)) { + messages.pop() + this.logger.debug('Dropped trailing user message') + } + } + + private handleToolUses(messages: Message[], newUserMessage: ChatMessage): void { + if (messages.length === 0) { + if (newUserMessage.userInputMessage?.userInputMessageContext?.toolResults) { + this.logger.debug('No history message found, but new user message has tool results.') + newUserMessage.userInputMessage.userInputMessageContext.toolResults = undefined + // tool results are empty, so content must not be empty + newUserMessage.userInputMessage.content = 'Conversation history was too large, so it was cleared.' + } + return + } + + const lastMsg = messages[messages.length - 1] + if (lastMsg.toolUses && lastMsg.toolUses.length > 0) { + const toolResults = newUserMessage.userInputMessage?.userInputMessageContext?.toolResults + if (!toolResults || toolResults.length === 0) { + this.logger.debug( + `No tools results in last user message following a tool use message from assisstant, marking as canceled` + ) + if (newUserMessage.userInputMessage?.userInputMessageContext) { + newUserMessage.userInputMessage.userInputMessageContext.toolResults = lastMsg.toolUses.map( + (toolUse) => ({ + toolUseId: toolUse.toolUseId, + content: [ + { + type: 'Text', + text: 'Tool use was cancelled by the user', + }, + ], + status: ToolResultStatus.ERROR, + }) + ) + } + } + } + } }