Skip to content

Commit 3a5d085

Browse files
tsmithszbywang56
authored andcommitted
fix(amazonq): Address multiple history bugs (aws#7059)
## Problem - `fixHistory` only processes the first conversation associated with the tab, but there can be multiple. Results in an incomplete history and messes up the history sent in the request - We currently delete the entire history for a tab if an exception is thrown. This is not ideal and leads to customer confusion when they are not able to continue using their chat history after an error. - Right-click actions like `Explain this code` fail when preceding message history contains toolResults. Mynah does not recognize this enum - Fake messages for tooluse failures are being stored the chat history. ## Solution - Collect message across all conversation associated with the current tab and consolidate them into a single conversation - Remove clearTab() call when exception is thrown - Do not include history for requests to Mynah - Remove clearRecentHistory as this is no longer used - Do not store these fake messages to the chatHistory ## Testing - Manually tested fixes: - Chat history still remains after exception is thrown - Mynah no longer throws exception for right-click action following tool execution ![Screenshot 2025-04-15 at 4 51 12 PM](https://github.com/user-attachments/assets/58817f2f-0fe0-452b-bb08-457bf95d729b) --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
1 parent 2dc850b commit 3a5d085

File tree

6 files changed

+51
-56
lines changed

6 files changed

+51
-56
lines changed

packages/core/src/codewhispererChat/controllers/chat/controller.ts

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -757,6 +757,7 @@ export class ChatController {
757757
const toolResults: ToolResult[] = []
758758

759759
let response = ''
760+
let shouldDisplayMessage = true
760761
if (toolUseError) {
761762
toolResults.push({
762763
content: [{ text: toolUseError.message }],
@@ -766,6 +767,7 @@ export class ChatController {
766767
if (toolUseError instanceof SyntaxError) {
767768
response =
768769
"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."
770+
shouldDisplayMessage = false
769771
}
770772
} else {
771773
const result = ToolUtils.tryFromToolUse(toolUse)
@@ -857,6 +859,7 @@ export class ChatController {
857859
contextLengths: {
858860
...defaultContextLengths,
859861
},
862+
shouldDisplayMessage: shouldDisplayMessage,
860863
},
861864
triggerID
862865
)
@@ -1151,9 +1154,6 @@ export class ChatController {
11511154

11521155
this.messenger.sendErrorMessage(errorMessage, tabID, requestID, statusCode)
11531156
getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)
1154-
1155-
this.sessionStorage.deleteSession(tabID)
1156-
this.chatHistoryDb.clearTab(tabID)
11571157
}
11581158

11591159
private async processContextMenuCommand(command: EditorContextCommand) {
@@ -1330,6 +1330,7 @@ export class ChatController {
13301330
this.processException(e, message.tabID)
13311331
}
13321332
}
1333+
13331334
private initialCleanUp(session: ChatSession) {
13341335
// Create a fresh token for this new conversation
13351336
session.createNewTokenSource()
@@ -1608,11 +1609,12 @@ export class ChatController {
16081609

16091610
const currentMessage = request.conversationState.currentMessage
16101611
if (currentMessage) {
1611-
this.chatHistoryDb.fixHistory(tabID, currentMessage)
1612+
this.chatHistoryDb.fixHistory(tabID, currentMessage, session.sessionIdentifier)
16121613
}
1613-
request.conversationState.history = this.chatHistoryDb
1614-
.getMessages(tabID)
1615-
.map((chat) => messageToChatMessage(chat))
1614+
// Do not include chatHistory for requests going to Mynah
1615+
request.conversationState.history = request.conversationState.currentMessage?.userInputMessage?.userIntent
1616+
? []
1617+
: this.chatHistoryDb.getMessages(tabID).map((chat) => messageToChatMessage(chat))
16161618
request.conversationState.conversationId = session.sessionIdentifier
16171619

16181620
triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments)
@@ -1682,6 +1684,7 @@ export class ChatController {
16821684
userIntent: currentMessage.userInputMessage?.userIntent,
16831685
origin: currentMessage.userInputMessage?.origin,
16841686
userInputMessageContext: currentMessage.userInputMessage?.userInputMessageContext,
1687+
shouldDisplayMessage: triggerPayload.shouldDisplayMessage ?? true,
16851688
})
16861689
}
16871690

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -560,6 +560,7 @@ export class Messenger {
560560
toolUse && toolUse.input !== undefined && toolUse.input !== ''
561561
? [{ ...toolUse }]
562562
: undefined,
563+
shouldDisplayMessage: triggerPayload.shouldDisplayMessage ?? true,
563564
})
564565
}
565566
if (

packages/core/src/codewhispererChat/controllers/chat/model.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,7 @@ export interface TriggerPayload {
238238
origin?: Origin
239239
pairProgrammingModeOn?: boolean
240240
history?: Message[]
241+
shouldDisplayMessage?: boolean
241242
}
242243

243244
export type ContextLengths = {

packages/core/src/codewhispererChat/controllers/chat/tabBarController.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,9 @@ export class TabBarController {
9191
selectedTab.historyId,
9292
selectedTab.tabType,
9393
selectedTab.conversations.flatMap((conv: Conversation) =>
94-
conv.messages.map((message) => messageToChatItem(message))
94+
conv.messages
95+
.filter((message) => message.shouldDisplayMessage !== false)
96+
.map((message) => messageToChatItem(message))
9597
),
9698
exportTab
9799
)

packages/core/src/shared/db/chatDb/chatDb.ts

Lines changed: 35 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ import { ChatMessage, ToolResultStatus } from '@amzn/codewhisperer-streaming'
2424

2525
// Maximum number of characters to keep in history
2626
const MaxConversationHistoryCharacters = 600_000
27+
// Maximum number of messages to keep in history
28+
const MaxConversationHistoryMessages = 250
2729

2830
/**
2931
* A singleton database class that manages chat history persistence using LokiJS.
@@ -144,18 +146,6 @@ export class Database {
144146
return undefined
145147
}
146148

147-
getActiveConversation(historyId: string): Conversation | undefined {
148-
const tabCollection = this.db.getCollection<Tab>(TabCollection)
149-
const tabData = tabCollection.findOne({ historyId })
150-
151-
if (!tabData?.conversations.length) {
152-
this.logger.debug('No active conversations found')
153-
return undefined
154-
}
155-
156-
return tabData.conversations[0]
157-
}
158-
159149
clearTab(tabId: string) {
160150
if (this.initialized) {
161151
const tabCollection = this.db.getCollection<Tab>(TabCollection)
@@ -170,34 +160,6 @@ export class Database {
170160
}
171161
}
172162

173-
// Removes the most recent message(s) from the chat history for a given tab
174-
clearRecentHistory(tabId: string): void {
175-
if (this.initialized) {
176-
const historyId = this.historyIdMapping.get(tabId)
177-
this.logger.info(`Clearing recent history: tabId=${tabId}, historyId=${historyId || 'undefined'}`)
178-
if (historyId) {
179-
const tabCollection = this.db.getCollection<Tab>(TabCollection)
180-
const tabData = tabCollection.findOne({ historyId })
181-
if (tabData) {
182-
const activeConversation = tabData.conversations[0]
183-
const allMessages = this.getMessages(tabId)
184-
const lastMessage = allMessages[allMessages.length - 1]
185-
this.logger.debug(`Last message type: ${lastMessage.type}`)
186-
187-
if (lastMessage.type === ('prompt' as ChatItemType)) {
188-
allMessages.pop()
189-
this.logger.info(`Removed last user message`)
190-
} else {
191-
allMessages.splice(-2)
192-
this.logger.info(`Removed last assistant message and user message`)
193-
}
194-
activeConversation.messages = allMessages
195-
tabCollection.update(tabData)
196-
}
197-
}
198-
}
199-
}
200-
201163
updateTabOpenState(tabId: string, isOpen: boolean) {
202164
if (this.initialized) {
203165
const tabCollection = this.db.getCollection<Tab>(TabCollection)
@@ -348,13 +310,14 @@ export class Database {
348310

349311
/**
350312
* Fixes the history to maintain the following invariants:
351-
* 1. The history character length is <= MAX_CONVERSATION_HISTORY_CHARACTERS. Oldest messages are dropped.
352-
* 2. The first message is from the user. Oldest messages are dropped if needed.
353-
* 3. The last message is from the assistant. The last message is dropped if it is from the user.
354-
* 4. If the last message is from the assistant and it contains tool uses, and a next user
313+
* 1. The history contains at most MaxConversationHistoryMessages messages. Oldest messages are dropped.
314+
* 2. The history character length is <= MaxConversationHistoryCharacters. Oldest messages are dropped.
315+
* 3. The first message is from the user. Oldest messages are dropped if needed.
316+
* 4. The last message is from the assistant. The last message is dropped if it is from the user.
317+
* 5. If the last message is from the assistant and it contains tool uses, and a next user
355318
* message is set without tool results, then the user message will have cancelled tool results.
356319
*/
357-
fixHistory(tabId: string, newUserMessage: ChatMessage): void {
320+
fixHistory(tabId: string, newUserMessage: ChatMessage, conversationId: string): void {
358321
if (!this.initialized) {
359322
return
360323
}
@@ -371,10 +334,12 @@ export class Database {
371334
return
372335
}
373336

374-
const activeConversation = tabData.conversations[0]
375-
let allMessages = activeConversation.messages
337+
let allMessages = tabData.conversations.flatMap((conversation: Conversation) => conversation.messages)
376338
this.logger.info(`Found ${allMessages.length} messages in conversation`)
377339

340+
// Make sure we don't exceed MaxConversationHistoryMessages
341+
allMessages = this.trimHistoryToMaxLength(allMessages)
342+
378343
// Drop empty assistant partial if it’s the last message
379344
this.handleEmptyAssistantMessage(allMessages)
380345

@@ -387,11 +352,33 @@ export class Database {
387352
// 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.
388353
this.handleToolUses(allMessages, newUserMessage)
389354

390-
activeConversation.messages = allMessages
355+
tabData.conversations = [
356+
{
357+
conversationId: conversationId,
358+
clientType: ClientType.VSCode,
359+
messages: allMessages,
360+
},
361+
]
362+
tabData.updatedAt = new Date()
391363
tabCollection.update(tabData)
392364
this.logger.info(`Updated tab data in collection`)
393365
}
394366

367+
private trimHistoryToMaxLength(messages: Message[]): Message[] {
368+
while (messages.length > MaxConversationHistoryMessages) {
369+
// Find the next valid user message to start from
370+
const indexToTrim = this.findIndexToTrim(messages)
371+
if (indexToTrim !== undefined && indexToTrim > 0) {
372+
this.logger.debug(`Removing the first ${indexToTrim} elements to maintain valid history length`)
373+
messages.splice(0, indexToTrim)
374+
} else {
375+
this.logger.debug('Could not find a valid point to trim, reset history to reduce history size')
376+
return []
377+
}
378+
}
379+
return messages
380+
}
381+
395382
private handleEmptyAssistantMessage(messages: Message[]): void {
396383
if (messages.length === 0) {
397384
return

packages/core/src/shared/db/chatDb/util.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ export type Message = {
5656
origin?: Origin
5757
userInputMessageContext?: UserInputMessageContext
5858
toolUses?: ToolUse[]
59+
shouldDisplayMessage?: boolean
5960
}
6061

6162
/**

0 commit comments

Comments
 (0)