Skip to content

Commit 1286d3d

Browse files
ashishrp-awsashishrp-aws
committed
Setting up the trigger Id mapping to each agentic loop
1 parent 8fb3eff commit 1286d3d

File tree

8 files changed

+102
-59
lines changed

8 files changed

+102
-59
lines changed

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -252,7 +252,7 @@ export class Connector extends BaseConnector {
252252
}
253253

254254
if (messageData.type === 'customFormActionMessage') {
255-
this.onCustomFormAction(messageData.tabID, messageData.messageId, messageData.action)
255+
this.onCustomFormAction(messageData.tabID, messageData.messageId, messageData.action, messageData.triggerId)
256256
return
257257
}
258258

@@ -300,7 +300,8 @@ export class Connector extends BaseConnector {
300300
id: string
301301
text?: string | undefined
302302
formItemValues?: Record<string, string> | undefined
303-
}
303+
},
304+
triggerId: string
304305
) {
305306
if (action === undefined) {
306307
return
@@ -316,6 +317,7 @@ export class Connector extends BaseConnector {
316317
formSelectedValues: action.formItemValues,
317318
tabType: this.getTabType(),
318319
tabID: tabId,
320+
triggerId: triggerId,
319321
})
320322

321323
if (

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -743,7 +743,7 @@ export class Connector {
743743
tabType: 'cwc',
744744
})
745745
} else {
746-
this.cwChatConnector.onCustomFormAction(tabId, messageId ?? '', action)
746+
this.cwChatConnector.onCustomFormAction(tabId, messageId ?? '', action, messageId ?? '')
747747
}
748748
break
749749
case 'agentWalkthrough': {

‎packages/core/src/codewhisperer/client/codewhisperer.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@ export interface CodeWhispererConfig {
3131
}
3232

3333
export const defaultServiceConfig: CodeWhispererConfig = {
34-
region: 'us-west-2',
35-
endpoint: 'https://rts.alpha-us-west-2.codewhisperer.ai.aws.dev/',
34+
region: 'us-east-1',
35+
endpoint: 'https://codewhisperer.us-east-1.amazonaws.com/',
3636
}
3737

3838
export function getCodewhispererConfig(): CodeWhispererConfig {

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

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,6 @@ import { FsWriteParams } from '../../tools/fsWrite'
102102
import { tempDirPath } from '../../../shared/filesystemUtilities'
103103
import { Database } from '../../../shared/db/chatDb/chatDb'
104104
import { TabBarController } from './tabBarController'
105-
import { sleep } from '../../../shared'
106105

107106
export interface ChatControllerMessagePublishers {
108107
readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -409,23 +408,20 @@ export class ChatController {
409408

410409
private async processStopResponseMessage(message: StopResponseMessage) {
411410
const session = this.sessionStorage.getSession(message.tabID)
412-
const wasInAgenticLoop = session.agenticLoopInProgress
413411
session.tokenSource.cancel()
414412
console.log('process stop has been triggered')
415413
session.setAgenticLoopInProgress(false)
416414
session.setToolUseWithError(undefined)
417415

418-
// Mark any active triggers as completed when stopping the response
416+
// Mark any active triggers as cancelled when stopping the response
419417
const triggerEvents = this.triggerEventsStorage.getTriggerEventsByTabID(message.tabID)
420418
if (triggerEvents && triggerEvents.length > 0) {
421419
const conversationTracker = ConversationTracker.getInstance()
422420
triggerEvents.forEach((event) => {
423-
conversationTracker.markTriggerCompleted(event.id)
421+
conversationTracker.cancelTrigger(event.id)
424422
})
425423
}
426424

427-
wasInAgenticLoop && (await sleep(1000))
428-
429425
this.messenger.sendEmptyMessage(message.tabID, '', undefined)
430426
// this.chatHistoryStorage.getTabHistory(message.tabID).clearRecentHistory()
431427
}
@@ -705,7 +701,14 @@ export class ChatController {
705701
this.editorContextExtractor
706702
.extractContextForTrigger('ChatMessage')
707703
.then(async (context) => {
708-
const triggerID = message.triggerId ?? randomUUID()
704+
console.log('message in controller for tool use:', message)
705+
const triggerID = message.triggerId
706+
707+
// Check if this trigger has already been cancelled
708+
if (this.isTriggerCancelled(triggerID)) {
709+
return
710+
}
711+
709712
this.triggerEventsStorage.addTriggerEvent({
710713
id: triggerID,
711714
tabID: message.tabID,
@@ -717,7 +720,7 @@ export class ChatController {
717720
const session = this.sessionStorage.getSession(tabID)
718721

719722
// Check if the session has been cancelled before proceeding
720-
if (session.tokenSource.token.isCancellationRequested) {
723+
if (this.isTriggerCancelled(triggerID)) {
721724
getLogger().debug(`Tool execution cancelled for tabID: ${tabID}`)
722725
return
723726
}
@@ -758,8 +761,7 @@ export class ChatController {
758761
toolUse,
759762
{ requiresAcceptance: false },
760763
undefined,
761-
undefined,
762-
cancellationToken
764+
undefined
763765
)
764766

765767
if (tool.type === ToolType.FsWrite && toolUse.toolUseId) {
@@ -768,7 +770,7 @@ export class ChatController {
768770
}
769771

770772
// Check again if cancelled before invoking the tool
771-
if (cancellationToken.isCancellationRequested) {
773+
if (this.isTriggerCancelled(triggerID)) {
772774
getLogger().debug(`Tool execution cancelled before invoke for tabID: ${tabID}`)
773775
return
774776
}
@@ -1249,7 +1251,6 @@ export class ChatController {
12491251
// Register the trigger ID with the token for cancellation tracking
12501252
const conversationTracker = ConversationTracker.getInstance()
12511253
conversationTracker.registerTrigger(triggerID, session.tokenSource, message.tabID)
1252-
console.log('conversation tracker:', conversationTracker.getTokenForTrigger(triggerID))
12531254

12541255
this.triggerEventsStorage.addTriggerEvent({
12551256
id: triggerID,
@@ -1556,6 +1557,9 @@ export class ChatController {
15561557
this.messenger.sendInitalStream(tabID, triggerID)
15571558
this.messenger.sendAsyncEventProgress(tabID, true, '')
15581559
this.telemetryHelper.setConversationStreamStartTime(tabID)
1560+
if (this.isTriggerCancelled(triggerID)) {
1561+
return
1562+
}
15591563
if (isSsoConnection(AuthUtil.instance.conn)) {
15601564
const { $metadata, generateAssistantResponseResponse } = await session.chatSso(request)
15611565
response = {
@@ -1572,7 +1576,7 @@ export class ChatController {
15721576
this.telemetryHelper.recordEnterFocusConversation(triggerEvent.tabID)
15731577
this.telemetryHelper.recordStartConversation(triggerEvent, triggerPayload)
15741578

1575-
if (currentMessage && session.sessionIdentifier) {
1579+
if (currentMessage && session.sessionIdentifier && !this.isTriggerCancelled(triggerID)) {
15761580
chatHistory.appendUserMessage(currentMessage)
15771581
this.chatHistoryDb.addMessage(tabID, 'cwc', session.sessionIdentifier, {
15781582
body: triggerPayload.message,
@@ -1588,6 +1592,9 @@ export class ChatController {
15881592
response.$metadata.requestId
15891593
} metadata: ${inspect(response.$metadata, { depth: 12 })}`
15901594
)
1595+
if (this.isTriggerCancelled(triggerID)) {
1596+
return
1597+
}
15911598
await this.messenger.sendAIResponse(response, session, tabID, triggerID, triggerPayload, chatHistory)
15921599
} catch (e: any) {
15931600
this.telemetryHelper.recordMessageResponseError(triggerPayload, tabID, getHttpStatusCode(e) ?? 0)
@@ -1633,4 +1640,17 @@ export class ChatController {
16331640
return { relativeFilePath: filePath, lineRanges: mergedRanges }
16341641
})
16351642
}
1643+
1644+
/**
1645+
* Check if a trigger has been cancelled and should not proceed
1646+
* @param triggerId The trigger ID to check
1647+
* @returns true if the trigger is cancelled and should not proceed
1648+
*/
1649+
private isTriggerCancelled(triggerId: string): boolean {
1650+
if (!triggerId) {
1651+
return false
1652+
}
1653+
const conversationTracker = ConversationTracker.getInstance()
1654+
return conversationTracker.isTriggerCancelled(triggerId)
1655+
}
16361656
}

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

Lines changed: 57 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,6 @@ import { FsWriteParams } from '../../../tools/fsWrite'
6767
import { AsyncEventProgressMessage } from '../../../../amazonq/commons/connector/connectorMessages'
6868
import { localize } from '../../../../shared/utilities/vsCodeUtils'
6969
import { getDiffLinesFromChanges } from '../../../../shared/utilities/diffUtils'
70-
import { CancellationToken } from 'vscode'
7170
import { ConversationTracker } from '../../../storages/conversationTracker'
7271
import { waitUntilWithCancellation } from '../../../../shared/utilities/timeoutUtils'
7372

@@ -107,7 +106,7 @@ export class Messenger {
107106

108107
public sendInitalStream(tabID: string, triggerID: string) {
109108
// Check if the conversation has been cancelled
110-
if (ConversationTracker.getInstance().isTriggerCancelled(triggerID)) {
109+
if (this.isTriggerCancelled(triggerID)) {
111110
getLogger().debug(`Initial stream sending cancelled for tabID: ${tabID}, triggerID: ${triggerID}`)
112111
return
113112
}
@@ -140,7 +139,7 @@ export class Messenger {
140139
mergedRelevantDocuments: DocumentReference[] | undefined
141140
) {
142141
// Check if the conversation has been cancelled
143-
if (ConversationTracker.getInstance().isTriggerCancelled(triggerID)) {
142+
if (this.isTriggerCancelled(triggerID)) {
144143
getLogger().debug(`Context message sending cancelled for tabID: ${tabID}, triggerID: ${triggerID}`)
145144
return
146145
}
@@ -280,6 +279,9 @@ export class Messenger {
280279
}
281280

282281
if (cwChatEvent.toolUseEvent?.stop) {
282+
if (this.isTriggerCancelled(triggerID)) {
283+
return
284+
}
283285
toolUse.input = JSON.parse(toolUseInput)
284286
toolUse.toolUseId = cwChatEvent.toolUseEvent.toolUseId ?? ''
285287
toolUse.name = cwChatEvent.toolUseEvent.name ?? ''
@@ -313,6 +315,9 @@ export class Messenger {
313315
if (!validation.requiresAcceptance) {
314316
// Need separate id for read tool and safe bash command execution as 'run-shell-command' id is required to state in cwChatConnector.ts which will impact generic tool execution.
315317
if (tool.type === ToolType.ExecuteBash) {
318+
if (this.isTriggerCancelled(triggerID)) {
319+
return
320+
}
316321
this.dispatcher.sendCustomFormActionMessage(
317322
new CustomFormActionMessage(
318323
tabID,
@@ -323,6 +328,9 @@ export class Messenger {
323328
)
324329
)
325330
} else {
331+
if (this.isTriggerCancelled(triggerID)) {
332+
return
333+
}
326334
this.dispatcher.sendCustomFormActionMessage(
327335
new CustomFormActionMessage(
328336
tabID,
@@ -353,6 +361,11 @@ export class Messenger {
353361
if (codeBlockLanguage === 'plaintext') {
354362
codeBlockLanguage = extractCodeBlockLanguage(message)
355363
}
364+
// Check if this trigger has been cancelled
365+
if (this.isTriggerCancelled(triggerID)) {
366+
return
367+
}
368+
356369
this.dispatcher.sendChatMessage(
357370
new ChatMessage(
358371
{
@@ -485,6 +498,11 @@ export class Messenger {
485498
)
486499
}
487500

501+
// Check if this trigger has been cancelled before sending final message
502+
if (this.isTriggerCancelled(triggerID)) {
503+
return
504+
}
505+
488506
this.dispatcher.sendChatMessage(
489507
new ChatMessage(
490508
{
@@ -503,16 +521,18 @@ export class Messenger {
503521
)
504522
)
505523

506-
chatHistoryManager.pushAssistantMessage({
507-
assistantResponseMessage: {
508-
messageId: messageID,
509-
content: message,
510-
references: codeReference,
511-
...(toolUse &&
512-
toolUse.input !== undefined &&
513-
toolUse.input !== '' && { toolUses: [{ ...toolUse }] }),
514-
},
515-
})
524+
if (!this.isTriggerCancelled(triggerID)) {
525+
chatHistoryManager.pushAssistantMessage({
526+
assistantResponseMessage: {
527+
messageId: messageID,
528+
content: message,
529+
references: codeReference,
530+
...(toolUse &&
531+
toolUse.input !== undefined &&
532+
toolUse.input !== '' && { toolUses: [{ ...toolUse }] }),
533+
},
534+
})
535+
}
516536
if (!eventCounts.has('toolUseEvent')) {
517537
session.setAgenticLoopInProgress(false)
518538
session.setContext(undefined)
@@ -561,11 +581,10 @@ export class Messenger {
561581
triggerID: string,
562582
toolUse: ToolUse | undefined,
563583
validation: CommandValidation,
564-
changeList?: Change[],
565-
cancellationToken?: CancellationToken
584+
changeList?: Change[]
566585
) {
567586
// Check if the conversation has been cancelled before sending any tool logs
568-
if (ConversationTracker.getInstance().isTriggerCancelled(triggerID)) {
587+
if (this.isTriggerCancelled(triggerID)) {
569588
getLogger().debug(`Tool log sending cancelled for tabID: ${tabID}, triggerID: ${triggerID}`)
570589
return
571590
}
@@ -605,6 +624,9 @@ export class Messenger {
605624
// eslint-disable-next-line unicorn/no-null
606625
codeBlockActions = { 'insert-to-cursor': null, copy: null }
607626
} else if (toolUse?.name === ToolType.FsWrite) {
627+
if (this.isTriggerCancelled(triggerID)) {
628+
return
629+
}
608630
const input = toolUse.input as unknown as FsWriteParams
609631
const fileName = path.basename(input.path)
610632
const changes = getDiffLinesFromChanges(changeList)
@@ -643,6 +665,10 @@ export class Messenger {
643665
codeBlockActions = { 'insert-to-cursor': null, copy: null }
644666
}
645667

668+
if (this.isTriggerCancelled(triggerID)) {
669+
return
670+
}
671+
646672
this.dispatcher.sendChatMessage(
647673
new ChatMessage(
648674
{
@@ -680,7 +706,7 @@ export class Messenger {
680706

681707
public sendStaticTextResponse(type: StaticTextResponseType, triggerID: string, tabID: string) {
682708
// Check if the conversation has been cancelled
683-
if (ConversationTracker.getInstance().isTriggerCancelled(triggerID)) {
709+
if (this.isTriggerCancelled(triggerID)) {
684710
getLogger().debug(`Static text response sending cancelled for tabID: ${tabID}, triggerID: ${triggerID}`)
685711
return
686712
}
@@ -888,4 +914,18 @@ export class Messenger {
888914
)
889915
)
890916
}
917+
918+
/**
919+
* Check if a trigger has been cancelled and should not proceed
920+
* @param triggerId The trigger ID to check
921+
* @returns true if the trigger is cancelled and should not proceed
922+
*/
923+
private isTriggerCancelled(triggerId: string): boolean {
924+
if (!triggerId) {
925+
return false
926+
}
927+
928+
const conversationTracker = ConversationTracker.getInstance()
929+
return conversationTracker.isTriggerCancelled(triggerId)
930+
}
891931
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ export class ConversationTracker {
108108
*/
109109
public isTriggerCancelled(triggerID: string): boolean {
110110
if (!triggerID) {
111-
return false
111+
return true
112112
}
113113

114114
const tokenSource = this.triggerToToken.get(triggerID)

0 commit comments

Comments
 (0)