Skip to content

Commit c44a94a

Browse files
ashishrp-awsashishrp-aws
committed
Fix for chat history for multuple tabs and user tool ignore.
1 parent d9620cf commit c44a94a

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 } from '../../tools/toolUtils'
9493
import { ChatStream } from '../../tools/chatStream'
94+
import { ChatHistoryStorage } from '../../storages/chatHistoryStorage'
9595

9696
export interface ChatControllerMessagePublishers {
9797
readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -153,7 +153,7 @@ export class ChatController {
153153
private readonly userIntentRecognizer: UserIntentRecognizer
154154
private readonly telemetryHelper: CWCTelemetryHelper
155155
private userPromptsWatcher: vscode.FileSystemWatcher | undefined
156-
private readonly chatHistoryManager: ChatHistoryManager
156+
private readonly chatHistoryStorage: ChatHistoryStorage
157157

158158
public constructor(
159159
private readonly chatControllerMessageListeners: ChatControllerMessageListeners,
@@ -171,7 +171,7 @@ export class ChatController {
171171
this.editorContentController = new EditorContentController()
172172
this.promptGenerator = new PromptsGenerator()
173173
this.userIntentRecognizer = new UserIntentRecognizer()
174-
this.chatHistoryManager = new ChatHistoryManager()
174+
this.chatHistoryStorage = new ChatHistoryStorage()
175175

176176
onDidChangeAmazonQVisibility((visible) => {
177177
if (visible) {
@@ -434,7 +434,7 @@ export class ChatController {
434434

435435
private async processTabCloseMessage(message: TabClosedMessage) {
436436
this.sessionStorage.deleteSession(message.tabID)
437-
this.chatHistoryManager.clear()
437+
this.chatHistoryStorage.deleteHistory(message.tabID)
438438
this.triggerEventsStorage.removeTabEvents(message.tabID)
439439
// this.telemetryHelper.recordCloseChat(message.tabID)
440440
}
@@ -747,7 +747,6 @@ export class ChatController {
747747
getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)
748748

749749
this.sessionStorage.deleteSession(tabID)
750-
this.chatHistoryManager.clear()
751750
}
752751

753752
private async processContextMenuCommand(command: EditorContextCommand) {
@@ -821,7 +820,6 @@ export class ChatController {
821820
codeQuery: context?.focusAreaContext?.names,
822821
userIntent: this.userIntentRecognizer.getFromContextMenuCommand(command),
823822
customization: getSelectedCustomization(),
824-
chatHistory: this.chatHistoryManager.getHistory(),
825823
additionalContents: [],
826824
relevantTextDocuments: [],
827825
documentReferences: [],
@@ -869,7 +867,7 @@ export class ChatController {
869867
switch (message.command) {
870868
case 'clear':
871869
this.sessionStorage.deleteSession(message.tabID)
872-
this.chatHistoryManager.clear()
870+
this.chatHistoryStorage.getHistory(message.tabID).clear()
873871
this.triggerEventsStorage.removeTabEvents(message.tabID)
874872
recordTelemetryChatRunCommand('clear')
875873
return
@@ -908,7 +906,7 @@ export class ChatController {
908906
codeQuery: lastTriggerEvent.context?.focusAreaContext?.names,
909907
userIntent: message.userIntent,
910908
customization: getSelectedCustomization(),
911-
chatHistory: this.chatHistoryManager.getHistory(),
909+
chatHistory: this.chatHistoryStorage.getHistory(message.tabID).getHistory(),
912910
contextLengths: {
913911
...defaultContextLengths,
914912
},
@@ -996,7 +994,7 @@ export class ChatController {
996994
customization: getSelectedCustomization(),
997995
toolResults: toolResults,
998996
origin: Origin.IDE,
999-
chatHistory: this.chatHistoryManager.getHistory(),
997+
chatHistory: this.chatHistoryStorage.getHistory(tabID).getHistory(),
1000998
context: [],
1001999
relevantTextDocuments: [],
10021000
additionalContents: [],
@@ -1042,7 +1040,7 @@ export class ChatController {
10421040
codeQuery: context?.focusAreaContext?.names,
10431041
userIntent: this.userIntentRecognizer.getFromPromptChatMessage(message),
10441042
customization: getSelectedCustomization(),
1045-
chatHistory: this.chatHistoryManager.getHistory(),
1043+
chatHistory: this.chatHistoryStorage.getHistory(message.tabID).getHistory(),
10461044
origin: Origin.IDE,
10471045
context: message.context ?? [],
10481046
relevantTextDocuments: [],
@@ -1269,16 +1267,28 @@ export class ChatController {
12691267

12701268
triggerPayload.contextLengths.userInputContextLength = triggerPayload.message.length
12711269
triggerPayload.contextLengths.focusFileContextLength = triggerPayload.fileText.length
1272-
const request = triggerPayloadToChatRequest(triggerPayload)
1273-
if (
1274-
this.chatHistoryManager.getConversationId() !== undefined &&
1275-
this.chatHistoryManager.getConversationId() !== ''
1276-
) {
1277-
request.conversationState.conversationId = this.chatHistoryManager.getConversationId()
1278-
} else {
1279-
this.chatHistoryManager.setConversationId(randomUUID())
1280-
request.conversationState.conversationId = this.chatHistoryManager.getConversationId()
1270+
1271+
const chatHistory = this.chatHistoryStorage.getHistory(tabID)
1272+
const newUserMessage = {
1273+
userInputMessage: {
1274+
content: triggerPayload.message,
1275+
userIntent: triggerPayload.userIntent,
1276+
...(triggerPayload.origin && { origin: triggerPayload.origin }),
1277+
userInputMessageContext: {
1278+
tools: tools,
1279+
...(triggerPayload.toolResults && { toolResults: triggerPayload.toolResults }),
1280+
},
1281+
},
1282+
}
1283+
const fixedHistoryMessage = chatHistory.fixHistory(newUserMessage)
1284+
if (fixedHistoryMessage.userInputMessage?.userInputMessageContext) {
1285+
triggerPayload.toolResults = fixedHistoryMessage.userInputMessage.userInputMessageContext.toolResults
12811286
}
1287+
const request = triggerPayloadToChatRequest(triggerPayload)
1288+
const conversationId = chatHistory.getConversationId() || randomUUID()
1289+
chatHistory.setConversationId(conversationId)
1290+
request.conversationState.conversationId = conversationId
1291+
12821292
triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments)
12831293

12841294
// Update context transparency after it's truncated dynamically to show users only the context sent.
@@ -1328,32 +1338,14 @@ export class ChatController {
13281338
}
13291339
this.telemetryHelper.recordEnterFocusConversation(triggerEvent.tabID)
13301340
this.telemetryHelper.recordStartConversation(triggerEvent, triggerPayload)
1331-
1332-
this.chatHistoryManager.appendUserMessage({
1333-
userInputMessage: {
1334-
content: triggerPayload.message,
1335-
userIntent: triggerPayload.userIntent,
1336-
...(triggerPayload.origin && { origin: triggerPayload.origin }),
1337-
userInputMessageContext: {
1338-
tools: tools,
1339-
...(triggerPayload.toolResults && { toolResults: triggerPayload.toolResults }),
1340-
},
1341-
},
1342-
})
1341+
chatHistory.appendUserMessage(fixedHistoryMessage)
13431342

13441343
getLogger().info(
13451344
`response to tab: ${tabID} conversationID: ${session.sessionIdentifier} requestID: ${
13461345
response.$metadata.requestId
13471346
} metadata: ${inspect(response.$metadata, { depth: 12 })}`
13481347
)
1349-
await this.messenger.sendAIResponse(
1350-
response,
1351-
session,
1352-
tabID,
1353-
triggerID,
1354-
triggerPayload,
1355-
this.chatHistoryManager
1356-
)
1348+
await this.messenger.sendAIResponse(response, session, tabID, triggerID, triggerPayload, chatHistory)
13571349
} catch (e: any) {
13581350
this.telemetryHelper.recordMessageResponseError(triggerPayload, tabID, getHttpStatusCode(e) ?? 0)
13591351
// 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)