Skip to content

Commit a8a4579

Browse files
ashishrp-awsashishrp-aws
committed
fix for chat history for multuple tabs and user tool ignore
1 parent 7a47e33 commit a8a4579

File tree

3 files changed

+156
-49
lines changed

3 files changed

+156
-49
lines changed

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

Lines changed: 31 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -87,11 +87,11 @@ import {
8787
defaultContextLengths,
8888
} from '../../constants'
8989
import { ChatSession } from '../../clients/chat/v0/chat'
90-
import { ChatHistoryManager } from '../../storages/chatHistory'
9190
import { amazonQTabSuffix } from '../../../shared/constants'
9291
import { OutputKind } from '../../tools/toolShared'
9392
import { ToolUtils, Tool, ToolType } from '../../tools/toolUtils'
9493
import { ChatStream } from '../../tools/chatStream'
94+
import { ChatHistoryStorage } from '../../storages/chatHistoryStorage'
9595
import { FsWrite, FsWriteParams } from '../../tools/fsWrite'
9696
import { tempDirPath } from '../../../shared/filesystemUtilities'
9797

@@ -155,7 +155,7 @@ export class ChatController {
155155
private readonly userIntentRecognizer: UserIntentRecognizer
156156
private readonly telemetryHelper: CWCTelemetryHelper
157157
private userPromptsWatcher: vscode.FileSystemWatcher | undefined
158-
private readonly chatHistoryManager: ChatHistoryManager
158+
private readonly chatHistoryStorage: ChatHistoryStorage
159159

160160
public constructor(
161161
private readonly chatControllerMessageListeners: ChatControllerMessageListeners,
@@ -173,7 +173,7 @@ export class ChatController {
173173
this.editorContentController = new EditorContentController()
174174
this.promptGenerator = new PromptsGenerator()
175175
this.userIntentRecognizer = new UserIntentRecognizer()
176-
this.chatHistoryManager = new ChatHistoryManager()
176+
this.chatHistoryStorage = new ChatHistoryStorage()
177177

178178
onDidChangeAmazonQVisibility((visible) => {
179179
if (visible) {
@@ -424,7 +424,7 @@ export class ChatController {
424424

425425
private async processTabCloseMessage(message: TabClosedMessage) {
426426
this.sessionStorage.deleteSession(message.tabID)
427-
this.chatHistoryManager.clear()
427+
this.chatHistoryStorage.deleteHistory(message.tabID)
428428
this.triggerEventsStorage.removeTabEvents(message.tabID)
429429
// this.telemetryHelper.recordCloseChat(message.tabID)
430430
}
@@ -710,7 +710,7 @@ export class ChatController {
710710
customization: getSelectedCustomization(),
711711
toolResults: toolResults,
712712
origin: Origin.IDE,
713-
chatHistory: this.chatHistoryManager.getHistory(),
713+
chatHistory: this.chatHistoryStorage.getHistory(tabID).getHistory(),
714714
context: session.context ?? [],
715715
relevantTextDocuments: [],
716716
additionalContents: [],
@@ -890,7 +890,6 @@ export class ChatController {
890890
getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)
891891

892892
this.sessionStorage.deleteSession(tabID)
893-
this.chatHistoryManager.clear()
894893
}
895894

896895
private async processContextMenuCommand(command: EditorContextCommand) {
@@ -964,7 +963,6 @@ export class ChatController {
964963
codeQuery: context?.focusAreaContext?.names,
965964
userIntent: this.userIntentRecognizer.getFromContextMenuCommand(command),
966965
customization: getSelectedCustomization(),
967-
chatHistory: this.chatHistoryManager.getHistory(),
968966
additionalContents: [],
969967
relevantTextDocuments: [],
970968
documentReferences: [],
@@ -1012,7 +1010,7 @@ export class ChatController {
10121010
switch (message.command) {
10131011
case 'clear':
10141012
this.sessionStorage.deleteSession(message.tabID)
1015-
this.chatHistoryManager.clear()
1013+
this.chatHistoryStorage.getHistory(message.tabID).clear()
10161014
this.triggerEventsStorage.removeTabEvents(message.tabID)
10171015
recordTelemetryChatRunCommand('clear')
10181016
return
@@ -1051,7 +1049,7 @@ export class ChatController {
10511049
codeQuery: lastTriggerEvent.context?.focusAreaContext?.names,
10521050
userIntent: message.userIntent,
10531051
customization: getSelectedCustomization(),
1054-
chatHistory: this.chatHistoryManager.getHistory(),
1052+
chatHistory: this.chatHistoryStorage.getHistory(message.tabID).getHistory(),
10551053
contextLengths: {
10561054
...defaultContextLengths,
10571055
},
@@ -1100,7 +1098,7 @@ export class ChatController {
11001098
codeQuery: context?.focusAreaContext?.names,
11011099
userIntent: this.userIntentRecognizer.getFromPromptChatMessage(message),
11021100
customization: getSelectedCustomization(),
1103-
chatHistory: this.chatHistoryManager.getHistory(),
1101+
chatHistory: this.chatHistoryStorage.getHistory(message.tabID).getHistory(),
11041102
origin: Origin.IDE,
11051103
context: message.context ?? [],
11061104
relevantTextDocuments: [],
@@ -1327,16 +1325,28 @@ export class ChatController {
13271325

13281326
triggerPayload.contextLengths.userInputContextLength = triggerPayload.message.length
13291327
triggerPayload.contextLengths.focusFileContextLength = triggerPayload.fileText.length
1330-
const request = triggerPayloadToChatRequest(triggerPayload)
1331-
if (
1332-
this.chatHistoryManager.getConversationId() !== undefined &&
1333-
this.chatHistoryManager.getConversationId() !== ''
1334-
) {
1335-
request.conversationState.conversationId = this.chatHistoryManager.getConversationId()
1336-
} else {
1337-
this.chatHistoryManager.setConversationId(randomUUID())
1338-
request.conversationState.conversationId = this.chatHistoryManager.getConversationId()
1328+
1329+
const chatHistory = this.chatHistoryStorage.getHistory(tabID)
1330+
const newUserMessage = {
1331+
userInputMessage: {
1332+
content: triggerPayload.message,
1333+
userIntent: triggerPayload.userIntent,
1334+
...(triggerPayload.origin && { origin: triggerPayload.origin }),
1335+
userInputMessageContext: {
1336+
tools: tools,
1337+
...(triggerPayload.toolResults && { toolResults: triggerPayload.toolResults }),
1338+
},
1339+
},
1340+
}
1341+
const fixedHistoryMessage = chatHistory.fixHistory(newUserMessage)
1342+
if (fixedHistoryMessage.userInputMessage?.userInputMessageContext) {
1343+
triggerPayload.toolResults = fixedHistoryMessage.userInputMessage.userInputMessageContext.toolResults
13391344
}
1345+
const request = triggerPayloadToChatRequest(triggerPayload)
1346+
const conversationId = chatHistory.getConversationId() || randomUUID()
1347+
chatHistory.setConversationId(conversationId)
1348+
request.conversationState.conversationId = conversationId
1349+
13401350
triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments)
13411351

13421352
// Update context transparency after it's truncated dynamically to show users only the context sent.
@@ -1386,32 +1396,14 @@ export class ChatController {
13861396
}
13871397
this.telemetryHelper.recordEnterFocusConversation(triggerEvent.tabID)
13881398
this.telemetryHelper.recordStartConversation(triggerEvent, triggerPayload)
1389-
1390-
this.chatHistoryManager.appendUserMessage({
1391-
userInputMessage: {
1392-
content: triggerPayload.message,
1393-
userIntent: triggerPayload.userIntent,
1394-
...(triggerPayload.origin && { origin: triggerPayload.origin }),
1395-
userInputMessageContext: {
1396-
tools: tools,
1397-
...(triggerPayload.toolResults && { toolResults: triggerPayload.toolResults }),
1398-
},
1399-
},
1400-
})
1399+
chatHistory.appendUserMessage(fixedHistoryMessage)
14011400

14021401
getLogger().info(
14031402
`response to tab: ${tabID} conversationID: ${session.sessionIdentifier} requestID: ${
14041403
response.$metadata.requestId
14051404
} metadata: ${inspect(response.$metadata, { depth: 12 })}`
14061405
)
1407-
await this.messenger.sendAIResponse(
1408-
response,
1409-
session,
1410-
tabID,
1411-
triggerID,
1412-
triggerPayload,
1413-
this.chatHistoryManager
1414-
)
1406+
await this.messenger.sendAIResponse(response, session, tabID, triggerID, triggerPayload, chatHistory)
14151407
} catch (e: any) {
14161408
this.telemetryHelper.recordMessageResponseError(triggerPayload, tabID, getHttpStatusCode(e) ?? 0)
14171409
// clears session, record telemetry before this call

‎packages/core/src/codewhispererChat/storages/chatHistory.ts‎

Lines changed: 82 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -23,14 +23,15 @@ const MaxConversationHistoryLength = 100
2323
*/
2424
export class ChatHistoryManager {
2525
private conversationId: string
26+
private tabId: string
2627
private history: ChatMessage[] = []
2728
private logger = getLogger()
2829
private lastUserMessage?: ChatMessage
2930
private tools: Tool[] = []
3031

31-
constructor() {
32+
constructor(tabId?: string) {
3233
this.conversationId = randomUUID()
33-
this.logger.info(`Generated new conversation id: ${this.conversationId}`)
34+
this.tabId = tabId ?? randomUUID()
3435
this.tools = tools
3536
}
3637

@@ -45,6 +46,20 @@ export class ChatHistoryManager {
4546
this.conversationId = conversationId
4647
}
4748

49+
/**
50+
* Get the tab ID
51+
*/
52+
public getTabId(): string {
53+
return this.tabId
54+
}
55+
56+
/**
57+
* Set the tab ID
58+
*/
59+
public setTabId(tabId: string) {
60+
this.tabId = tabId
61+
}
62+
4863
/**
4964
* Get the full chat history
5065
*/
@@ -65,7 +80,6 @@ export class ChatHistoryManager {
6580
*/
6681
public appendUserMessage(newMessage: ChatMessage): void {
6782
this.lastUserMessage = newMessage
68-
this.fixHistory()
6983
if (!newMessage.userInputMessage?.content || newMessage.userInputMessage?.content.trim() === '') {
7084
this.logger.warn('input must not be empty when adding new messages')
7185
}
@@ -90,7 +104,7 @@ export class ChatHistoryManager {
90104
* 4. If the last message is from the assistant and it contains tool uses, and a next user
91105
* message is set without tool results, then the user message will have cancelled tool results.
92106
*/
93-
public fixHistory(): void {
107+
public fixHistory(newUserMessage: ChatMessage): ChatMessage {
94108
// Trim the conversation history if it exceeds the maximum length
95109
if (this.history.length > MaxConversationHistoryLength) {
96110
// Find the second oldest user message without tool results
@@ -123,22 +137,22 @@ export class ChatHistoryManager {
123137
this.history.pop()
124138
}
125139

126-
// TODO: If the last message from the assistant contains tool uses, ensure the next user message contains tool results
140+
// If the last message from the assistant contains tool uses, ensure the next user message contains tool results
127141

128142
const lastHistoryMessage = this.history[this.history.length - 1]
129143

130144
if (
131145
lastHistoryMessage &&
132146
(lastHistoryMessage.assistantResponseMessage ||
133147
lastHistoryMessage.assistantResponseMessage !== undefined) &&
134-
this.lastUserMessage
148+
newUserMessage
135149
) {
136150
const toolUses = lastHistoryMessage.assistantResponseMessage.toolUses
137151

138152
if (toolUses && toolUses.length > 0) {
139-
if (this.lastUserMessage.userInputMessage) {
140-
if (this.lastUserMessage.userInputMessage.userInputMessageContext) {
141-
const ctx = this.lastUserMessage.userInputMessage.userInputMessageContext
153+
if (newUserMessage.userInputMessage) {
154+
if (newUserMessage.userInputMessage.userInputMessageContext) {
155+
const ctx = newUserMessage.userInputMessage.userInputMessageContext
142156

143157
if (!ctx.toolResults || ctx.toolResults.length === 0) {
144158
ctx.toolResults = toolUses.map((toolUse) => ({
@@ -164,16 +178,21 @@ export class ChatHistoryManager {
164178
status: ToolResultStatus.ERROR,
165179
}))
166180

167-
this.lastUserMessage.userInputMessage.userInputMessageContext = {
181+
newUserMessage.userInputMessage.userInputMessageContext = {
168182
shellState: undefined,
169183
envState: undefined,
170184
toolResults: toolResults,
171185
tools: this.tools.length === 0 ? undefined : [...this.tools],
172186
}
187+
188+
return newUserMessage
173189
}
174190
}
175191
}
176192
}
193+
194+
// Always return the message to fix the TypeScript error
195+
return newUserMessage
177196
}
178197

179198
/**
@@ -197,6 +216,59 @@ export class ChatHistoryManager {
197216
}
198217
}
199218

219+
/**
220+
* Checks if the latest message in history is an Assistant Message.
221+
* If it is and doesn't have toolUse, it will be removed.
222+
* If it has toolUse, an assistantResponse message with cancelled tool status will be added.
223+
*/
224+
public checkLatestAssistantMessage(): void {
225+
if (this.history.length === 0) {
226+
return
227+
}
228+
229+
const lastMessage = this.history[this.history.length - 1]
230+
231+
if (lastMessage.assistantResponseMessage) {
232+
const toolUses = lastMessage.assistantResponseMessage.toolUses
233+
234+
if (!toolUses || toolUses.length === 0) {
235+
// If there are no tool uses, remove the assistant message
236+
this.logger.debug('Removing assistant message without tool uses')
237+
this.history.pop()
238+
} else {
239+
// If there are tool uses, add cancelled tool results
240+
const toolResults = toolUses.map((toolUse) => ({
241+
toolUseId: toolUse.toolUseId,
242+
content: [
243+
{
244+
type: 'Text',
245+
text: 'Tool use was cancelled by the user',
246+
},
247+
],
248+
status: ToolResultStatus.ERROR,
249+
}))
250+
251+
// Create a new user message with cancelled tool results
252+
const userInputMessageContext: UserInputMessageContext = {
253+
shellState: undefined,
254+
envState: undefined,
255+
toolResults: toolResults,
256+
tools: this.tools.length === 0 ? undefined : [...this.tools],
257+
}
258+
259+
const userMessage: ChatMessage = {
260+
userInputMessage: {
261+
content: '',
262+
userInputMessageContext: userInputMessageContext,
263+
},
264+
}
265+
266+
this.history.push(this.formatChatHistoryMessage(userMessage))
267+
this.logger.debug('Added user message with cancelled tool results')
268+
}
269+
}
270+
}
271+
200272
private formatChatHistoryMessage(message: ChatMessage): ChatMessage {
201273
if (message.userInputMessage !== undefined) {
202274
return {

‎packages/core/src/codewhispererChat/storages/chatHistoryStorage.ts‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import { ChatHistoryManager } from './chatHistory'
7+
8+
/**
9+
* ChatHistoryStorage manages ChatHistoryManager instances for multiple tabs.
10+
* Each tab has its own ChatHistoryManager to maintain separate chat histories.
11+
*/
12+
export class ChatHistoryStorage {
13+
private histories: Map<string, ChatHistoryManager> = new Map()
14+
15+
/**
16+
* Gets the ChatHistoryManager for a specific tab.
17+
* If no history exists for the tab, creates a new one.
18+
*
19+
* @param tabId The ID of the tab
20+
* @returns The ChatHistoryManager for the specified tab
21+
*/
22+
public getHistory(tabId: string): ChatHistoryManager {
23+
const historyFromStorage = this.histories.get(tabId)
24+
if (historyFromStorage !== undefined) {
25+
return historyFromStorage
26+
}
27+
28+
// Create a new ChatHistoryManager with the tabId
29+
const newHistory = new ChatHistoryManager(tabId)
30+
this.histories.set(tabId, newHistory)
31+
32+
return newHistory
33+
}
34+
35+
/**
36+
* Deletes the ChatHistoryManager for a specific tab.
37+
*
38+
* @param tabId The ID of the tab
39+
*/
40+
public deleteHistory(tabId: string) {
41+
this.histories.delete(tabId)
42+
}
43+
}

0 commit comments

Comments
 (0)