making updates optional for executeBash tool

Author: laileni-aws

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

@@ -90,9 +90,9 @@ import { ChatSession } from '../../clients/chat/v0/chat'
 import { ChatHistoryManager } from '../../storages/chatHistory'
 import { amazonQTabSuffix } from '../../../shared/constants'
 import { OutputKind } from '../../tools/toolShared'
-import { ToolUtils, Tool } from '../../tools/toolUtils'
+import { ToolUtils, Tool, ToolType } from '../../tools/toolUtils'
 import { ChatStream } from '../../tools/chatStream'
-import { FsWriteParams } from '../../tools/fsWrite'
+import { FsWrite, FsWriteParams } from '../../tools/fsWrite'
 import { tempDirPath } from '../../../shared/filesystemUtilities'
 
 export interface ChatControllerMessagePublishers {
@@ -661,10 +661,6 @@ export class ChatController {
                             toolUseId: toolUse.toolUseId,
                             status: ToolResultStatus.SUCCESS,
                         })
-                        // Close the diff view if User accept the generated code changes.
-                        if (vscode.window.tabGroups.activeTabGroup.activeTab?.label.includes(amazonQTabSuffix)) {
-                            await vscode.commands.executeCommand('workbench.action.closeActiveEditor')
-                        }
                     } catch (e: any) {
                         toolResults.push({
                             content: [{ text: e.message }],
@@ -677,6 +673,13 @@ export class ChatController {
                     toolResults.push(toolResult)
                 }
 
+                if (toolUse.name === ToolType.FsWrite) {
+                    await vscode.commands.executeCommand(
+                        'vscode.open',
+                        vscode.Uri.file((toolUse.input as unknown as FsWriteParams).path)
+                    )
+                }
+
                 await this.generateResponse(
                     {
                         message: '',
@@ -710,6 +713,13 @@ export class ChatController {
             })
     }
 
+    private async closeDiffView() {
+        // Close the diff view if User reject the generated code changes.
+        if (vscode.window.tabGroups.activeTabGroup.activeTab?.label.includes(amazonQTabSuffix)) {
+            await vscode.commands.executeCommand('workbench.action.closeActiveEditor')
+        }
+    }
+
     private async processCustomFormAction(message: CustomFormActionMessage) {
         switch (message.action.id) {
             case 'submit-create-prompt':
@@ -718,15 +728,11 @@ export class ChatController {
             case 'accept-code-diff':
             case 'confirm-tool-use':
             case 'generic-tool-execution':
+                await this.closeDiffView()
                 await this.processToolUseMessage(message)
                 break
             case 'reject-code-diff':
-                // Close the diff view if User reject the generated code changes.
-                if (vscode.window.tabGroups.activeTabGroup.activeTab?.label.includes(amazonQTabSuffix)) {
-                    await vscode.commands.executeCommand('workbench.action.closeActiveEditor')
-                }
-                // TODO: Do session cleanUp.
-                getLogger().info('Generated response is rejected')
+                await this.closeDiffView()
                 break
             default:
                 getLogger().warn(`Unhandled action: ${message.action.id}`)
@@ -771,13 +777,9 @@ export class ChatController {
             }
             const input = clonedToolUse.input as unknown as FsWriteParams
             input.path = tempFilePath
-            const result = ToolUtils.tryFromToolUse(clonedToolUse)
-            if (!('type' in result)) {
-                return
-            }
-            const tool: Tool = result
-            await ToolUtils.validate(tool)
-            await ToolUtils.invoke(tool)
+
+            const fsWrite = new FsWrite(input)
+            await fsWrite.invoke()
 
             // Check if fileExists=false, If yes, return instead of showing broken diff experience.
             if (!tempFilePath) {

packages/core/src/codewhispererChat/tools/executeBash.ts

@@ -67,17 +67,6 @@ export class ExecuteBash {
 
     public async invoke(updates?: Writable): Promise<InvokeOutput> {
         this.logger.info(`Invoking bash command: "${this.command}" in cwd: "${this.workingDirectory}"`)
-        if (!updates) {
-            // updates passed as undefined only in the code diff is clicked to avoid showing Q Streaming spinner component to the user.
-            return new Promise(async (resolve) => {
-                resolve({
-                    output: {
-                        kind: OutputKind.Json,
-                        content: '',
-                    },
-                })
-            })
-        }
 
         return new Promise(async (resolve, reject) => {
             this.logger.debug(`Spawning process with command: bash -c "${this.command}" (cwd=${this.workingDirectory})`)
@@ -137,9 +126,9 @@ export class ExecuteBash {
         })
     }
 
-    private static handleChunk(chunk: string, buffer: string[], updates: Writable) {
+    private static handleChunk(chunk: string, buffer: string[], updates?: Writable) {
         try {
-            updates.write(chunk)
+            updates?.write(chunk)
             const lines = chunk.split(/\r?\n/)
             for (const line of lines) {
                 buffer.push(line)

packages/core/src/codewhispererChat/tools/toolUtils.ts

@@ -46,11 +46,11 @@ export class ToolUtils {
     static async invoke(tool: Tool, updates?: Writable): Promise<InvokeOutput> {
         switch (tool.type) {
             case ToolType.FsRead:
-                return tool.tool.invoke(updates)
+                return tool.tool.invoke(updates ?? undefined)
             case ToolType.FsWrite:
-                return tool.tool.invoke(updates)
+                return tool.tool.invoke(updates ?? undefined)
             case ToolType.ExecuteBash:
-                return tool.tool.invoke(updates)
+                return tool.tool.invoke(updates ?? undefined)
         }
     }