diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts index a9711c3ab02..573a26fc5e2 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/controller.ts @@ -744,6 +744,7 @@ export class ChatController { const toolResults: ToolResult[] = [] let response = '' + let shouldDisplayMessage = true if (toolUseError) { toolResults.push({ content: [{ text: toolUseError.message }], @@ -753,6 +754,7 @@ export class ChatController { if (toolUseError instanceof SyntaxError) { response = "Your toolUse input isn't valid. Please check the syntax and make sure the input is complete. If the input is large, break it down into multiple tool uses with smaller input." + shouldDisplayMessage = false } } else { const result = ToolUtils.tryFromToolUse(toolUse) @@ -844,6 +846,7 @@ export class ChatController { contextLengths: { ...defaultContextLengths, }, + shouldDisplayMessage: shouldDisplayMessage, }, triggerID ) @@ -1138,9 +1141,6 @@ export class ChatController { this.messenger.sendErrorMessage(errorMessage, tabID, requestID, statusCode) getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`) - - this.sessionStorage.deleteSession(tabID) - this.chatHistoryDb.clearTab(tabID) } private async processContextMenuCommand(command: EditorContextCommand) { @@ -1317,6 +1317,7 @@ export class ChatController { this.processException(e, message.tabID) } } + private initialCleanUp(session: ChatSession) { // Create a fresh token for this new conversation session.createNewTokenSource() @@ -1595,11 +1596,12 @@ export class ChatController { const currentMessage = request.conversationState.currentMessage if (currentMessage) { - this.chatHistoryDb.fixHistory(tabID, currentMessage) + this.chatHistoryDb.fixHistory(tabID, currentMessage, session.sessionIdentifier) } - request.conversationState.history = this.chatHistoryDb - .getMessages(tabID) - .map((chat) => messageToChatMessage(chat)) + // Do not include chatHistory for requests going to Mynah + request.conversationState.history = request.conversationState.currentMessage?.userInputMessage?.userIntent + ? [] + : this.chatHistoryDb.getMessages(tabID).map((chat) => messageToChatMessage(chat)) request.conversationState.conversationId = session.sessionIdentifier triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments) @@ -1669,6 +1671,7 @@ export class ChatController { userIntent: currentMessage.userInputMessage?.userIntent, origin: currentMessage.userInputMessage?.origin, userInputMessageContext: currentMessage.userInputMessage?.userInputMessageContext, + shouldDisplayMessage: triggerPayload.shouldDisplayMessage ?? true, }) } diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts index 34e2ddc093c..92026d5cbe3 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts @@ -558,6 +558,7 @@ export class Messenger { toolUse && toolUse.input !== undefined && toolUse.input !== '' ? [{ ...toolUse }] : undefined, + shouldDisplayMessage: triggerPayload.shouldDisplayMessage ?? true, }) } if ( diff --git a/packages/core/src/codewhispererChat/controllers/chat/model.ts b/packages/core/src/codewhispererChat/controllers/chat/model.ts index fcbe5fd4302..7ba0dfef5ac 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/model.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/model.ts @@ -238,6 +238,7 @@ export interface TriggerPayload { origin?: Origin pairProgrammingModeOn?: boolean history?: Message[] + shouldDisplayMessage?: boolean } export type ContextLengths = { diff --git a/packages/core/src/codewhispererChat/controllers/chat/tabBarController.ts b/packages/core/src/codewhispererChat/controllers/chat/tabBarController.ts index 13088e096f8..1be498a6c80 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/tabBarController.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/tabBarController.ts @@ -88,7 +88,9 @@ export class TabBarController { selectedTab.historyId, selectedTab.tabType, selectedTab.conversations.flatMap((conv: Conversation) => - conv.messages.map((message) => messageToChatItem(message)) + conv.messages + .filter((message) => message.shouldDisplayMessage !== false) + .map((message) => messageToChatItem(message)) ), exportTab ) diff --git a/packages/core/src/shared/db/chatDb/chatDb.ts b/packages/core/src/shared/db/chatDb/chatDb.ts index 2b949b23869..1cafa5990c9 100644 --- a/packages/core/src/shared/db/chatDb/chatDb.ts +++ b/packages/core/src/shared/db/chatDb/chatDb.ts @@ -24,6 +24,8 @@ import { ChatMessage, ToolResultStatus } from '@amzn/codewhisperer-streaming' // Maximum number of characters to keep in history const MaxConversationHistoryCharacters = 600_000 +// Maximum number of messages to keep in history +const MaxConversationHistoryMessages = 250 /** * A singleton database class that manages chat history persistence using LokiJS. @@ -144,18 +146,6 @@ 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) @@ -170,34 +160,6 @@ 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) @@ -348,13 +310,14 @@ export class Database { /** * 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 + * 1. The history contains at most MaxConversationHistoryMessages messages. Oldest messages are dropped. + * 2. The history character length is <= MaxConversationHistoryCharacters. Oldest messages are dropped. + * 3. The first message is from the user. Oldest messages are dropped if needed. + * 4. The last message is from the assistant. The last message is dropped if it is from the user. + * 5. 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 { + fixHistory(tabId: string, newUserMessage: ChatMessage, conversationId: string): void { if (!this.initialized) { return } @@ -371,10 +334,12 @@ export class Database { return } - const activeConversation = tabData.conversations[0] - let allMessages = activeConversation.messages + let allMessages = tabData.conversations.flatMap((conversation: Conversation) => conversation.messages) this.logger.info(`Found ${allMessages.length} messages in conversation`) + // Make sure we don't exceed MaxConversationHistoryMessages + allMessages = this.trimHistoryToMaxLength(allMessages) + // Drop empty assistant partial if it’s the last message this.handleEmptyAssistantMessage(allMessages) @@ -387,11 +352,33 @@ export class Database { // 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 + tabData.conversations = [ + { + conversationId: conversationId, + clientType: ClientType.VSCode, + messages: allMessages, + }, + ] + tabData.updatedAt = new Date() tabCollection.update(tabData) this.logger.info(`Updated tab data in collection`) } + private trimHistoryToMaxLength(messages: Message[]): Message[] { + while (messages.length > MaxConversationHistoryMessages) { + // 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 to maintain valid history length`) + messages.splice(0, indexToTrim) + } else { + this.logger.debug('Could not find a valid point to trim, reset history to reduce history size') + return [] + } + } + return messages + } + private handleEmptyAssistantMessage(messages: Message[]): void { if (messages.length === 0) { return diff --git a/packages/core/src/shared/db/chatDb/util.ts b/packages/core/src/shared/db/chatDb/util.ts index fc681b2b5a5..8995db90e8a 100644 --- a/packages/core/src/shared/db/chatDb/util.ts +++ b/packages/core/src/shared/db/chatDb/util.ts @@ -56,6 +56,7 @@ export type Message = { origin?: Origin userInputMessageContext?: UserInputMessageContext toolUses?: ToolUse[] + shouldDisplayMessage?: boolean } /**