Trying another approach for read and list tool execution

Author: laileni-aws

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

@@ -787,7 +787,7 @@ export class ChatController {
                                 chatStream,
                                 ConversationTracker.getInstance().getTokenForTrigger(triggerID)
                             )
-                            ToolUtils.validateOutput(output)
+                            ToolUtils.validateOutput(output, tool.type)
 
                             let status: ToolResultStatus = ToolResultStatus.SUCCESS
                             if (output.output.success === false) {

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

@@ -355,7 +355,11 @@ export class Messenger {
                                     isReadOrList,
                                     changeList
                                 )
-                                await ToolUtils.queueDescription(tool, chatStream)
+                                await ToolUtils.queueDescription(
+                                    tool,
+                                    chatStream,
+                                    chatStream.validation.requiresAcceptance
+                                )
                                 if (session.messageIdToUpdate === undefined && tool.type === ToolType.FsRead) {
                                     // Store the first messageId in a chain of tool uses
                                     session.setMessageIdToUpdate(toolUse.toolUseId)

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

@@ -9,7 +9,6 @@ import { Writable } from 'stream'
 import { InvokeOutput, OutputKind, sanitizePath, CommandValidation, fsReadToolResponseSize } from './toolShared'
 import { isInDirectory } from '../../shared/filesystemUtilities'
 import path from 'path'
-import { ChatStream } from './chatStream'
 
 export interface FsReadParams {
     path: string
@@ -49,8 +48,8 @@ export class FsRead {
         this.logger.debug(`Validation succeeded for path: ${this.fsPath}`)
     }
 
-    public queueDescription(updates: ChatStream): void {
-        if (updates.validation.requiresAcceptance) {
+    public queueDescription(updates: Writable, requiresAcceptance: boolean): void {
+        if (requiresAcceptance) {
             const fileName = path.basename(this.fsPath)
             const fileUri = vscode.Uri.file(this.fsPath)
             updates.write(`Reading file: [${fileName}](${fileUri}), `)

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

@@ -10,7 +10,6 @@ import { Writable } from 'stream'
 import path from 'path'
 import { InvokeOutput, OutputKind, sanitizePath, CommandValidation } from './toolShared'
 import { isInDirectory } from '../../shared/filesystemUtilities'
-import { ChatStream } from './chatStream'
 
 export interface ListDirectoryParams {
     path: string
@@ -50,8 +49,8 @@ export class ListDirectory {
         }
     }
 
-    public queueDescription(updates: ChatStream): void {
-        if (updates.validation.requiresAcceptance) {
+    public queueDescription(updates: Writable, requiresAcceptance: boolean): void {
+        if (requiresAcceptance) {
             const fileName = path.basename(this.fsPath)
             if (this.maxDepth === undefined) {
                 updates.write(`Analyzing directories recursively: ${fileName}`)

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

@@ -15,7 +15,6 @@ import {
     fsReadToolResponseSize,
 } from './toolShared'
 import { ListDirectory, ListDirectoryParams } from './listDirectory'
-import { ChatStream } from './chatStream'
 
 export enum ToolType {
     FsRead = 'fsRead',
@@ -96,10 +95,10 @@ export class ToolUtils {
         }
     }
 
-    static async queueDescription(tool: Tool, updates: ChatStream): Promise<void> {
+    static async queueDescription(tool: Tool, updates: Writable, requiresAcceptance: boolean): Promise<void> {
         switch (tool.type) {
             case ToolType.FsRead:
-                tool.tool.queueDescription(updates)
+                tool.tool.queueDescription(updates, requiresAcceptance)
                 break
             case ToolType.FsWrite:
                 await tool.tool.queueDescription(updates)
@@ -108,7 +107,7 @@ export class ToolUtils {
                 tool.tool.queueDescription(updates)
                 break
             case ToolType.ListDirectory:
-                tool.tool.queueDescription(updates)
+                tool.tool.queueDescription(updates, requiresAcceptance)
                 break
         }
     }

packages/core/src/test/codewhispererChat/tools/toolShared.test.ts

@@ -22,7 +22,6 @@ import { ToolUse } from '@amzn/codewhisperer-streaming'
 import path from 'path'
 import fs from '../../../shared/fs/fs'
 import { ListDirectory } from '../../../codewhispererChat/tools/listDirectory'
-import { ChatStream } from '../../../codewhispererChat/tools/chatStream'
 
 describe('ToolUtils', function () {
     let sandbox: sinon.SinonSandbox
@@ -31,7 +30,6 @@ describe('ToolUtils', function () {
     let mockExecuteBash: sinon.SinonStubbedInstance<ExecuteBash>
     let mockListDirectory: sinon.SinonStubbedInstance<ListDirectory>
     let mockWritable: sinon.SinonStubbedInstance<Writable>
-    let mockChatStream: sinon.SinonStubbedInstance<ChatStream>
 
     beforeEach(function () {
         sandbox = sinon.createSandbox()
@@ -42,7 +40,6 @@ describe('ToolUtils', function () {
         mockWritable = {
             write: sandbox.stub(),
         } as unknown as sinon.SinonStubbedInstance<Writable>
-        mockChatStream = sandbox.createStubInstance(ChatStream)
         ;(mockFsRead.requiresAcceptance as sinon.SinonStub).returns({ requiresAcceptance: false })
         ;(mockListDirectory.requiresAcceptance as sinon.SinonStub).returns({ requiresAcceptance: false })
     })
@@ -237,30 +234,30 @@ describe('ToolUtils', function () {
         // TODO: Adding "void" to the following tests for the current implementation but in the next followup PR I will fix this issue.
         it('delegates to FsRead tool queueDescription method', function () {
             const tool: Tool = { type: ToolType.FsRead, tool: mockFsRead as unknown as FsRead }
-            void ToolUtils.queueDescription(tool, mockWritable as unknown as ChatStream)
+            void ToolUtils.queueDescription(tool, mockWritable as unknown as Writable, false)
 
-            assert(mockFsRead.queueDescription.calledOnceWith(mockChatStream))
+            assert(mockFsRead.queueDescription.calledOnceWith(mockWritable, false))
         })
 
         it('delegates to FsWrite tool queueDescription method', function () {
             const tool: Tool = { type: ToolType.FsWrite, tool: mockFsWrite as unknown as FsWrite }
-            void ToolUtils.queueDescription(tool, mockWritable as unknown as ChatStream)
+            void ToolUtils.queueDescription(tool, mockWritable as unknown as Writable, false)
 
             assert(mockFsWrite.queueDescription.calledOnceWith(mockWritable))
         })
 
         it('delegates to ExecuteBash tool queueDescription method', function () {
             const tool: Tool = { type: ToolType.ExecuteBash, tool: mockExecuteBash as unknown as ExecuteBash }
-            void ToolUtils.queueDescription(tool, mockWritable as unknown as ChatStream)
+            void ToolUtils.queueDescription(tool, mockWritable as unknown as Writable, false)
 
             assert(mockExecuteBash.queueDescription.calledOnceWith(mockWritable))
         })
 
         it('delegates to ListDirectory tool queueDescription method', function () {
             const tool: Tool = { type: ToolType.ListDirectory, tool: mockListDirectory as unknown as ListDirectory }
-            void ToolUtils.queueDescription(tool, mockWritable as unknown as ChatStream)
+            void ToolUtils.queueDescription(tool, mockWritable as unknown as Writable, false)
 
-            assert(mockListDirectory.queueDescription.calledOnceWith(mockChatStream))
+            assert(mockListDirectory.queueDescription.calledOnceWith(mockWritable, false))
         })
     })