Fix for Agentic loop follopwup

Author: ashishrp-aws

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

@@ -73,6 +73,8 @@ export class ChatSession {
     async chatSso(chatRequest: GenerateAssistantResponseRequest): Promise<GenerateAssistantResponseCommandOutput> {
         const client = await createCodeWhispererChatStreamingClient()
 
+        // eslint-disable-next-line aws-toolkits/no-console-log
+        console.log(chatRequest)
         const response = await client.generateAssistantResponse(chatRequest)
         if (!response.generateAssistantResponseResponse) {
             throw new ToolkitError(

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

@@ -78,6 +78,7 @@ import {
     aditionalContentNameLimit,
     additionalContentInnerContextLimit,
     contextMaxLength,
+    tools,
 } from '../../constants'
 import { ChatSession } from '../../clients/chat/v0/chat'
 import { ChatHistoryManager } from '../../storages/chatHistory'
@@ -722,13 +723,6 @@ export class ChatController {
                     command,
                 })
 
-                this.chatHistoryManager.appendUserMessage({
-                    userInputMessage: {
-                        content: prompt,
-                        userIntent: this.userIntentRecognizer.getFromContextMenuCommand(command),
-                    },
-                })
-
                 return this.generateResponse(
                     {
                         message: prompt,
@@ -808,13 +802,6 @@ export class ChatController {
                 context: lastTriggerEvent.context,
             })
 
-            this.chatHistoryManager.appendUserMessage({
-                userInputMessage: {
-                    content: message.message,
-                    userIntent: message.userIntent,
-                },
-            })
-
             return this.generateResponse(
                 {
                     message: message.message,
@@ -870,7 +857,7 @@ export class ChatController {
                         //     result = await executeBash.invoke(process.stdout)
                         //     break
                         // }
-                        case 'fs_read': {
+                        case 'fsRead': {
                             const fsRead = new FsRead(toolUse.input as unknown as FsReadParams)
                             await fsRead.validate()
                             result = await fsRead.invoke()
@@ -902,17 +889,9 @@ export class ChatController {
                     toolResults.push({ content: [{ text: e.message }], toolUseId: toolUse.toolUseId, status: 'error' })
                 }
 
-                this.chatHistoryManager.appendUserMessage({
-                    userInputMessage: {
-                        content: 'Tool Results',
-                        userIntent: undefined,
-                        origin: Origin.IDE,
-                    },
-                })
-
                 await this.generateResponse(
                     {
-                        message: 'Tool Results',
+                        message: '',
                         trigger: ChatTriggerType.ChatMessage,
                         query: undefined,
                         codeSelection: context?.focusAreaContext?.selectionInsideExtendedCodeBlock,
@@ -926,6 +905,7 @@ export class ChatController {
                         context: undefined,
                         toolResults: toolResults,
                         origin: Origin.IDE,
+                        chatHistory: this.chatHistoryManager.getHistory(),
                     },
                     triggerID
                 )
@@ -947,13 +927,6 @@ export class ChatController {
                     type: 'chat_message',
                     context,
                 })
-                this.chatHistoryManager.appendUserMessage({
-                    userInputMessage: {
-                        content: message.message,
-                        userIntent: message.userIntent,
-                        origin: Origin.IDE,
-                    },
-                })
                 await this.generateResponse(
                     {
                         message: message.message,
@@ -1293,6 +1266,18 @@ export class ChatController {
             this.telemetryHelper.recordEnterFocusConversation(triggerEvent.tabID)
             this.telemetryHelper.recordStartConversation(triggerEvent, triggerPayload)
 
+            this.chatHistoryManager.appendUserMessage({
+                userInputMessage: {
+                    content: triggerPayload.message,
+                    userIntent: triggerPayload.userIntent,
+                    ...(triggerPayload.origin && { origin: triggerPayload.origin }),
+                    userInputMessageContext: {
+                        tools: tools,
+                        ...(triggerPayload.toolResults && { toolResults: triggerPayload.toolResults }),
+                    },
+                },
+            })
+
             getLogger().info(
                 `response to tab: ${tabID} conversationID: ${session.sessionIdentifier} requestID: ${
                     response.$metadata.requestId

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

@@ -9,6 +9,7 @@ import {
     AuthNeededException,
     CodeReference,
     ContextCommandData,
+    CustomFormActionMessage,
     EditorContextCommandMessage,
     OpenSettingsMessage,
     QuickActionMessage,
@@ -228,13 +229,14 @@ export class Messenger {
                                 tabID
                             )
                         )
+
+                        this.dispatcher.sendCustomFormActionMessage(
+                            new CustomFormActionMessage(tabID, {
+                                id: 'confirm-tool-use',
+                            })
+                        )
                         // TODO: setup permission action
                         // if (!isConfirmationRequired) {
-                        //     this.dispatcher.sendCustomFormActionMessage(
-                        //         new CustomFormActionMessage(tabID, {
-                        //             id: 'confirm-tool-use',
-                        //         })
-                        //     )
                         // }
                     }
 
@@ -388,7 +390,9 @@ export class Messenger {
                         messageId: messageID,
                         content: message,
                         references: codeReference,
-                        toolUses: [{ ...toolUse }],
+                        ...(toolUse &&
+                            toolUse.input !== undefined &&
+                            toolUse.input !== '' && { toolUses: [{ ...toolUse }] }),
                     },
                 })
 

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

@@ -69,7 +69,7 @@ export class ChatHistoryManager {
         if (!newMessage.userInputMessage?.content || newMessage.userInputMessage?.content.trim() === '') {
             this.logger.warn('input must not be empty when adding new messages')
         }
-        this.history.push(this.lastUserMessage)
+        this.history.push(this.formatChatHistoryMessage(this.lastUserMessage))
     }
 
     /**
@@ -196,4 +196,19 @@ export class ChatHistoryManager {
             this.lastUserMessage.userInputMessage = msg
         }
     }
+
+    private formatChatHistoryMessage(message: ChatMessage): ChatMessage {
+        if (message.userInputMessage !== undefined) {
+            return {
+                userInputMessage: {
+                    ...message.userInputMessage,
+                    userInputMessageContext: {
+                        ...message.userInputMessage.userInputMessageContext,
+                        tools: undefined,
+                    },
+                },
+            }
+        }
+        return message
+    }
 }