Skip to content

Commit b43f2a0

Browse files
ashishrp-awsashishrp-aws
committed
Adding Stop button and stopping generation for new user prompt
1 parent 69de65c commit b43f2a0

File tree

8 files changed

+88
-123
lines changed

8 files changed

+88
-123
lines changed

‎packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts‎

Lines changed: 23 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,22 @@ export interface ConnectorProps extends BaseConnectorProps {
2727
description?: string
2828
) => void
2929
onChatAnswerUpdated?: (tabID: string, message: ChatItem) => void
30+
onAsyncEventProgress: (
31+
tabID: string,
32+
inProgress: boolean,
33+
message: string,
34+
messageId: string | undefined,
35+
enableStopAction: boolean,
36+
isPromptInputDisabled: boolean
37+
) => void
3038
}
3139

3240
export class Connector extends BaseConnector {
3341
private readonly onCWCContextCommandMessage
3442
private readonly onContextCommandDataReceived
3543
private readonly onShowCustomForm
3644
private readonly onChatAnswerUpdated
45+
private readonly onAsyncEventProgress
3746
private chatItems: Map<string, Map<string, ChatItem>> = new Map() // tabId -> messageId -> ChatItem
3847

3948
override getTabType(): TabType {
@@ -46,6 +55,7 @@ export class Connector extends BaseConnector {
4655
this.onContextCommandDataReceived = props.onContextCommandDataReceived
4756
this.onShowCustomForm = props.onShowCustomForm
4857
this.onChatAnswerUpdated = props.onChatAnswerUpdated
58+
this.onAsyncEventProgress = props.onAsyncEventProgress
4959
}
5060

5161
onSourceLinkClick = (tabID: string, messageId: string, link: string): void => {
@@ -168,35 +178,6 @@ export class Connector extends BaseConnector {
168178
}
169179
}
170180

171-
private processToolMessage = async (messageData: any): Promise<void> => {
172-
if (this.onChatAnswerUpdated === undefined) {
173-
return
174-
}
175-
const answer: CWCChatItem = {
176-
type: messageData.messageType,
177-
messageId: messageData.messageID ?? messageData.triggerID,
178-
body: messageData.message,
179-
followUp: messageData.followUps,
180-
canBeVoted: messageData.canBeVoted ?? false,
181-
codeReference: messageData.codeReference,
182-
userIntent: messageData.contextList,
183-
codeBlockLanguage: messageData.codeBlockLanguage,
184-
contextList: messageData.contextList,
185-
title: messageData.title,
186-
buttons: messageData.buttons,
187-
fileList: messageData.fileList,
188-
header: messageData.header ?? undefined,
189-
padding: messageData.padding ?? undefined,
190-
fullWidth: messageData.fullWidth ?? undefined,
191-
codeBlockActions: messageData.codeBlockActions ?? undefined,
192-
}
193-
if (answer.messageId) {
194-
this.storeChatItem(messageData.tabID, answer.messageId, answer)
195-
}
196-
this.onChatAnswerUpdated(messageData.tabID, answer)
197-
return
198-
}
199-
200181
private storeChatItem(tabId: string, messageId: string, item: ChatItem): void {
201182
if (!this.chatItems.has(tabId)) {
202183
this.chatItems.set(tabId, new Map())
@@ -253,11 +234,6 @@ export class Connector extends BaseConnector {
253234
return
254235
}
255236

256-
if (messageData.type === 'toolMessage') {
257-
await this.processToolMessage(messageData)
258-
return
259-
}
260-
261237
if (messageData.type === 'editorContextCommandMessage') {
262238
await this.processEditorContextCommandMessage(messageData)
263239
return
@@ -276,6 +252,19 @@ export class Connector extends BaseConnector {
276252
this.onCustomFormAction(messageData.tabID, messageData.messageId, messageData.action)
277253
return
278254
}
255+
256+
if (messageData.type === 'asyncEventProgressMessage') {
257+
const enableStopAction = true
258+
this.onAsyncEventProgress(
259+
messageData.tabID,
260+
messageData.inProgress,
261+
messageData.message ?? undefined,
262+
messageData.messageId ?? undefined,
263+
enableStopAction,
264+
false
265+
)
266+
return
267+
}
279268
// For other message types, call the base class handleMessageReceive
280269
await this.baseHandleMessageReceive(messageData)
281270
}

‎packages/core/src/amazonq/webview/ui/main.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -287,12 +287,13 @@ export const createMynahUI = (
287287
inProgress: boolean,
288288
message: string | undefined,
289289
messageId: string | undefined = undefined,
290-
enableStopAction: boolean = false
290+
enableStopAction: boolean = false,
291+
isPromptInputDisabled: boolean = true
291292
) => {
292293
if (inProgress) {
293294
mynahUI.updateStore(tabID, {
294295
loadingChat: true,
295-
promptInputDisabledState: true,
296+
promptInputDisabledState: isPromptInputDisabled,
296297
cancelButtonWhenLoading: enableStopAction,
297298
})
298299

‎packages/core/src/codewhispererChat/clients/chat/v0/chat.ts‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,12 @@ export class ChatSession {
2222
* _readFiles = list of files read from the project to gather context before generating response.
2323
* _showDiffOnFileWrite = Controls whether to show diff view (true) or file context view (false) to the user
2424
* _context = Additional context to be passed to the LLM for generating the response
25-
* _messageIdToUpdate = messageId of a chat message to be updated, used for reducing consecutive tool messages
2625
*/
2726
private _readFiles: string[] = []
2827
private _toolUse: ToolUse | undefined
2928
private _showDiffOnFileWrite: boolean = false
3029
private _context: PromptMessage['context']
3130
private _pairProgrammingModeOn: boolean = true
32-
private _messageIdToUpdate: string | undefined
3331

3432
contexts: Map<string, { first: number; second: number }[]> = new Map()
3533
// TODO: doesn't handle the edge case when two files share the same relativePath string but from different root
@@ -63,14 +61,6 @@ export class ChatSession {
6361
this._context = context
6462
}
6563

66-
public get messageIdToUpdate(): string | undefined {
67-
return this._messageIdToUpdate
68-
}
69-
70-
public setMessageIdToUpdate(messageId: string | undefined) {
71-
this._messageIdToUpdate = messageId
72-
}
73-
7464
public tokenSource!: vscode.CancellationTokenSource
7565

7666
constructor() {

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

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -373,6 +373,7 @@ export class ChatController {
373373
private async processStopResponseMessage(message: StopResponseMessage) {
374374
const session = this.sessionStorage.getSession(message.tabID)
375375
session.tokenSource.cancel()
376+
this.messenger.sendEmptyMessage(message.tabID, '', undefined)
376377
this.chatHistoryStorage.getTabHistory(message.tabID).clearRecentHistory()
377378
}
378379

@@ -716,6 +717,7 @@ export class ChatController {
716717
type: 'chat_message',
717718
context,
718719
})
720+
this.messenger.sendAsyncEventProgress(tabID, true, '')
719721
const session = this.sessionStorage.getSession(tabID)
720722
const toolUse = session.toolUse
721723
if (!toolUse || !toolUse.input) {
@@ -734,16 +736,9 @@ export class ChatController {
734736
try {
735737
await ToolUtils.validate(tool)
736738

737-
const chatStream = new ChatStream(
738-
this.messenger,
739-
tabID,
740-
triggerID,
741-
// Pass in a different toolUseId so that the output does not overwrite
742-
// any previous messages
743-
{ ...toolUse, toolUseId: `${toolUse.toolUseId}-output` },
744-
{ requiresAcceptance: false },
745-
undefined
746-
)
739+
const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse, {
740+
requiresAcceptance: false,
741+
})
747742
const output = await ToolUtils.invoke(tool, chatStream)
748743
if (output.output.content.length > maxToolOutputCharacterLength) {
749744
throw Error(
@@ -1185,6 +1180,8 @@ export class ChatController {
11851180
context,
11861181
})
11871182

1183+
this.messenger.sendAsyncEventProgress(message.tabID, true, '')
1184+
11881185
// Save the context for the agentic loop
11891186
session.setContext(message.context)
11901187

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

Lines changed: 33 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ import {
1414
OpenSettingsMessage,
1515
QuickActionMessage,
1616
ShowCustomFormMessage,
17-
ToolMessage,
1817
} from '../../../view/connector/connector'
1918
import { EditorContextCommandType } from '../../../commands/registerCommands'
2019
import { ChatResponseStream as qdevChatResponseStream } from '@amzn/amazon-q-developer-streaming-client'
@@ -49,6 +48,7 @@ import { extractErrorInfo } from '../../../../shared/utilities/messageUtil'
4948
import { noWriteTools, tools } from '../../../constants'
5049
import { Change } from 'diff'
5150
import { FsWriteParams } from '../../../tools/fsWrite'
51+
import { AsyncEventProgressMessage } from '../../../../amazonq/commons/connector/connectorMessages'
5252

5353
export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help'
5454

@@ -238,36 +238,10 @@ export class Messenger {
238238
session.setShowDiffOnFileWrite(true)
239239
changeList = await tool.tool.getDiffChanges()
240240
}
241-
if (
242-
tool.type === ToolType.FsWrite ||
243-
tool.type === ToolType.ExecuteBash ||
244-
eventCounts.has('assistantResponseEvent')
245-
) {
246-
// FsWrite and ExecuteBash should never replace older messages
247-
// If the current stream also has assistantResponseEvent then reset this as well.
248-
session.setMessageIdToUpdate(undefined)
249-
}
250241
const validation = ToolUtils.requiresAcceptance(tool)
251-
252-
const chatStream = new ChatStream(
253-
this,
254-
tabID,
255-
triggerID,
256-
toolUse,
257-
validation,
258-
session.messageIdToUpdate,
259-
changeList
260-
)
242+
const chatStream = new ChatStream(this, tabID, triggerID, toolUse, validation, changeList)
261243
await ToolUtils.queueDescription(tool, chatStream)
262244

263-
if (
264-
session.messageIdToUpdate === undefined &&
265-
(tool.type === ToolType.FsRead || tool.type === ToolType.ListDirectory)
266-
) {
267-
// Store the first messageId in a chain of tool uses
268-
session.setMessageIdToUpdate(toolUse.toolUseId)
269-
}
270-
271245
if (!validation.requiresAcceptance) {
272246
// Need separate id for read tool and safe bash command execution as 'confirm-tool-use' id is required to change button status from `Confirm` to `Confirmed` state in cwChatConnector.ts which will impact generic tool execution.
273247
this.dispatcher.sendCustomFormActionMessage(
@@ -469,33 +443,12 @@ export class Messenger {
469443
)
470444
}
471445

472-
public sendInitialToolMessage(tabID: string, triggerID: string, toolUseId: string | undefined) {
473-
this.dispatcher.sendChatMessage(
474-
new ChatMessage(
475-
{
476-
message: '',
477-
messageType: 'answer',
478-
followUps: undefined,
479-
followUpsHeader: undefined,
480-
relatedSuggestions: undefined,
481-
triggerID,
482-
messageID: toolUseId ?? 'toolUse',
483-
userIntent: undefined,
484-
codeBlockLanguage: undefined,
485-
contextList: undefined,
486-
},
487-
tabID
488-
)
489-
)
490-
}
491-
492446
public sendPartialToolLog(
493447
message: string,
494448
tabID: string,
495449
triggerID: string,
496450
toolUse: ToolUse | undefined,
497451
validation: CommandValidation,
498-
messageIdToUpdate: string | undefined,
499452
changeList?: Change[]
500453
) {
501454
const buttons: ChatItemButton[] = []
@@ -550,16 +503,16 @@ export class Messenger {
550503
})
551504
}
552505

553-
this.dispatcher.sendToolMessage(
554-
new ToolMessage(
506+
this.dispatcher.sendChatMessage(
507+
new ChatMessage(
555508
{
556509
message: message,
557510
messageType: 'answer-part',
558511
followUps: undefined,
559512
followUpsHeader: undefined,
560513
relatedSuggestions: undefined,
561514
triggerID,
562-
messageID: messageIdToUpdate ?? toolUse?.toolUseId ?? '',
515+
messageID: toolUse?.toolUseId ?? '',
563516
userIntent: undefined,
564517
codeBlockLanguage: undefined,
565518
contextList: undefined,
@@ -738,4 +691,32 @@ export class Messenger {
738691
new ShowCustomFormMessage(tabID, formItems, buttons, title, description)
739692
)
740693
}
694+
695+
public sendAsyncEventProgress(tabID: string, inProgress: boolean, message: string | undefined) {
696+
this.dispatcher.sendAsyncEventProgress(new AsyncEventProgressMessage(tabID, 'CWChat', inProgress, message))
697+
}
698+
699+
public sendEmptyMessage(
700+
tabID: string,
701+
triggerId: string,
702+
mergedRelevantDocuments: DocumentReference[] | undefined
703+
) {
704+
this.dispatcher.sendChatMessage(
705+
new ChatMessage(
706+
{
707+
message: '',
708+
messageType: 'answer',
709+
followUps: undefined,
710+
followUpsHeader: undefined,
711+
relatedSuggestions: undefined,
712+
triggerID: triggerId,
713+
messageID: '',
714+
userIntent: undefined,
715+
codeBlockLanguage: undefined,
716+
contextList: undefined,
717+
},
718+
tabID
719+
)
720+
)
721+
}
741722
}

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,10 @@ export class ChatHistoryManager {
8383
if (newMessage !== undefined && this.lastUserMessage !== undefined) {
8484
this.logger.warn('last Message should not be defined when pushing an assistant message')
8585
}
86+
// check if last message in histroy is assistant message and now replace it in that case
87+
if (this.history.length > 0 && this.history.at(-1)?.assistantResponseMessage) {
88+
this.history.pop()
89+
}
8690
this.history.push(newMessage)
8791
}
8892

@@ -101,6 +105,18 @@ export class ChatHistoryManager {
101105
}
102106

103107
private trimConversationHistory(): void {
108+
// make sure the UseInputMessage is the first stored message
109+
if (this.history.length === 1 && this.history[0].assistantResponseMessage) {
110+
this.history = []
111+
}
112+
113+
if (
114+
this.history.at(-1)?.assistantResponseMessage?.content === '' &&
115+
this.history.at(-1)?.assistantResponseMessage?.toolUses === undefined
116+
) {
117+
this.clearRecentHistory()
118+
}
119+
104120
if (this.history.length <= MaxConversationHistoryLength) {
105121
return
106122
}

‎packages/core/src/codewhispererChat/tools/chatStream.ts‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,17 +23,12 @@ export class ChatStream extends Writable {
2323
private readonly triggerID: string,
2424
private readonly toolUse: ToolUse | undefined,
2525
private readonly validation: CommandValidation,
26-
private readonly messageIdToUpdate: string | undefined,
2726
private readonly changeList?: Change[],
2827
private readonly logger = getLogger('chatStream')
2928
) {
3029
super()
3130
this.logger.debug(`ChatStream created for tabID: ${tabID}, triggerID: ${triggerID}`)
32-
if (!messageIdToUpdate) {
33-
// If messageIdToUpdate is undefined, we need to first create an empty message
34-
// with messageId so it can be updated later
35-
this.messenger.sendInitialToolMessage(tabID, triggerID, toolUse?.toolUseId)
36-
}
31+
this.messenger.sendInitalStream(tabID, triggerID, undefined)
3732
}
3833

3934
override _write(chunk: Buffer, encoding: BufferEncoding, callback: (error?: Error | null) => void): void {
@@ -46,7 +41,6 @@ export class ChatStream extends Writable {
4641
this.triggerID,
4742
this.toolUse,
4843
this.validation,
49-
this.messageIdToUpdate,
5044
this.changeList
5145
)
5246
callback()

0 commit comments

Comments
 (0)