Skip to content

Commit 62f6ac5

Browse files
authored
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 ed50d42 commit 62f6ac5

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
@@ -744,6 +744,7 @@ export class ChatController {
744744
const toolResults: ToolResult[] = []
745745

746746
let response = ''
747+
let shouldDisplayMessage = true
747748
if (toolUseError) {
748749
toolResults.push({
749750
content: [{ text: toolUseError.message }],
@@ -753,6 +754,7 @@ export class ChatController {
753754
if (toolUseError instanceof SyntaxError) {
754755
response =
755756
"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."
757+
shouldDisplayMessage = false
756758
}
757759
} else {
758760
const result = ToolUtils.tryFromToolUse(toolUse)
@@ -844,6 +846,7 @@ export class ChatController {
844846
contextLengths: {
845847
...defaultContextLengths,
846848
},
849+
shouldDisplayMessage: shouldDisplayMessage,
847850
},
848851
triggerID
849852
)
@@ -1138,9 +1141,6 @@ export class ChatController {
11381141

11391142
this.messenger.sendErrorMessage(errorMessage, tabID, requestID, statusCode)
11401143
getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)
1141-
1142-
this.sessionStorage.deleteSession(tabID)
1143-
this.chatHistoryDb.clearTab(tabID)
11441144
}
11451145

11461146
private async processContextMenuCommand(command: EditorContextCommand) {
@@ -1317,6 +1317,7 @@ export class ChatController {
13171317
this.processException(e, message.tabID)
13181318
}
13191319
}
1320+
13201321
private initialCleanUp(session: ChatSession) {
13211322
// Create a fresh token for this new conversation
13221323
session.createNewTokenSource()
@@ -1595,11 +1596,12 @@ export class ChatController {
15951596

15961597
const currentMessage = request.conversationState.currentMessage
15971598
if (currentMessage) {
1598-
this.chatHistoryDb.fixHistory(tabID, currentMessage)
1599+
this.chatHistoryDb.fixHistory(tabID, currentMessage, session.sessionIdentifier)
15991600
}
1600-
request.conversationState.history = this.chatHistoryDb
1601-
.getMessages(tabID)
1602-
.map((chat) => messageToChatMessage(chat))
1601+
// Do not include chatHistory for requests going to Mynah
1602+
request.conversationState.history = request.conversationState.currentMessage?.userInputMessage?.userIntent
1603+
? []
1604+
: this.chatHistoryDb.getMessages(tabID).map((chat) => messageToChatMessage(chat))
16031605
request.conversationState.conversationId = session.sessionIdentifier
16041606

16051607
triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments)
@@ -1669,6 +1671,7 @@ export class ChatController {
16691671
userIntent: currentMessage.userInputMessage?.userIntent,
16701672
origin: currentMessage.userInputMessage?.origin,
16711673
userInputMessageContext: currentMessage.userInputMessage?.userInputMessageContext,
1674+
shouldDisplayMessage: triggerPayload.shouldDisplayMessage ?? true,
16721675
})
16731676
}
16741677

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
@@ -88,7 +88,9 @@ export class TabBarController {
8888
selectedTab.historyId,
8989
selectedTab.tabType,
9090
selectedTab.conversations.flatMap((conv: Conversation) =>
91-
conv.messages.map((message) => messageToChatItem(message))
91+
conv.messages
92+
.filter((message) => message.shouldDisplayMessage !== false)
93+
.map((message) => messageToChatItem(message))
9294
),
9395
exportTab
9496
)

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)