Merge pull request aws#7019 from jguoamz/fsReadFix

fix(chat): Update FsRead toolspec, return truncated output is file is too large

Author: mr-lee

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

@@ -739,12 +739,17 @@ export class ChatController {
                 const toolUseError = toolUseWithError.error
                 const toolResults: ToolResult[] = []
 
+                let response = ''
                 if (toolUseError) {
                     toolResults.push({
                         content: [{ text: toolUseError.message }],
                         toolUseId: toolUse.toolUseId,
                         status: ToolResultStatus.ERROR,
                     })
+                    if (toolUseError instanceof SyntaxError) {
+                        response =
+                            "Your toolUse input isn't valid. Please check the syntax and make sure the input is complete. If the input is large, break it down into multiple tool uses with smaller input."
+                    }
                 } else {
                     const result = ToolUtils.tryFromToolUse(toolUse)
                     if ('type' in result) {
@@ -784,14 +789,19 @@ export class ChatController {
                             )
                             ToolUtils.validateOutput(output)
 
+                            let status: ToolResultStatus = ToolResultStatus.SUCCESS
+                            if (output.output.success === false) {
+                                status = ToolResultStatus.ERROR
+                            }
+
                             toolResults.push({
                                 content: [
                                     output.output.kind === OutputKind.Text
                                         ? { text: output.output.content }
                                         : { json: output.output.content },
                                 ],
                                 toolUseId: toolUse.toolUseId,
-                                status: ToolResultStatus.SUCCESS,
+                                status,
                             })
                         } catch (e: any) {
                             toolResults.push({
@@ -808,7 +818,7 @@ export class ChatController {
 
                 await this.generateResponse(
                     {
-                        message: '',
+                        message: response,
                         trigger: ChatTriggerType.ChatMessage,
                         query: undefined,
                         codeSelection: context?.focusAreaContext?.selectionInsideExtendedCodeBlock,

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

@@ -289,10 +289,12 @@ export class Messenger {
                             try {
                                 toolUse.input = JSON.parse(toolUseInput)
                             } catch (error: any) {
-                                getLogger().error(`JSON parse error for toolUseInput: ${toolUseInput}`)
+                                getLogger().error(
+                                    `JSON parse error: ${error.message} for toolUseInput: ${toolUseInput}`
+                                )
                                 // set toolUse.input to be empty valid json object
                                 toolUse.input = {}
-                                error.message = `Tool input has invalid JSON format: ${error.message}`
+                                error.message = `tooluse input is invalid: ${toolUseInput}`
                                 // throw it out to allow the error to be handled in the catch block
                                 throw error
                             }

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

@@ -343,6 +343,7 @@ export class ExecuteBash {
                 const exitStatus = result.exitCode ?? 0
                 const stdout = stdoutBuffer.join('\n')
                 const stderr = stderrBuffer.join('\n')
+                const success = exitStatus === 0 && !stderr
                 const [stdoutTrunc, stdoutSuffix] = ExecuteBash.truncateSafelyWithSuffix(
                     stdout,
                     maxBashToolResponseSize / 3
@@ -362,6 +363,7 @@ export class ExecuteBash {
                     output: {
                         kind: OutputKind.Json,
                         content: outputJson,
+                        success,
                     },
                 })
             } catch (err: any) {

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

@@ -6,7 +6,7 @@ import * as vscode from 'vscode'
 import { getLogger } from '../../shared/logger/logger'
 import fs from '../../shared/fs/fs'
 import { Writable } from 'stream'
-import { InvokeOutput, OutputKind, sanitizePath, CommandValidation } from './toolShared'
+import { InvokeOutput, OutputKind, sanitizePath, CommandValidation, fsReadToolResponseSize } from './toolShared'
 import { isInDirectory } from '../../shared/filesystemUtilities'
 
 export interface FsReadParams {
@@ -118,11 +118,24 @@ export class FsRead {
     }
 
     private createOutput(content: string): InvokeOutput {
+        if (content.length > fsReadToolResponseSize) {
+            this.logger.info(
+                `The file is too large, truncating output to the first ${fsReadToolResponseSize} characters.`
+            )
+            content = this.truncateContent(content)
+        }
         return {
             output: {
                 kind: OutputKind.Text,
                 content: content,
             },
         }
     }
+
+    private truncateContent(content: string): string {
+        if (content.length > fsReadToolResponseSize) {
+            return content.substring(0, fsReadToolResponseSize)
+        }
+        return content
+    }
 }

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

@@ -6,7 +6,9 @@
 import path from 'path'
 import fs from '../../shared/fs/fs'
 
-export const maxToolResponseSize = 200_000
+export const defaultMaxToolResponseSize = 100_000
+export const listDirectoryToolResponseSize = 30_000
+export const fsReadToolResponseSize = 200_000
 
 export enum OutputKind {
     Text = 'text',
@@ -17,6 +19,7 @@ export interface InvokeOutput {
     output: {
         kind: OutputKind
         content: string | any
+        success?: boolean
     }
 }
 

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

@@ -7,7 +7,12 @@ import { FsRead, FsReadParams } from './fsRead'
 import { FsWrite, FsWriteParams } from './fsWrite'
 import { CommandValidation, ExecuteBash, ExecuteBashParams } from './executeBash'
 import { ToolResult, ToolResultContentBlock, ToolResultStatus, ToolUse } from '@amzn/codewhisperer-streaming'
-import { InvokeOutput, maxToolResponseSize } from './toolShared'
+import {
+    InvokeOutput,
+    defaultMaxToolResponseSize,
+    listDirectoryToolResponseSize,
+    fsReadToolResponseSize,
+} from './toolShared'
 import { ListDirectory, ListDirectoryParams } from './listDirectory'
 import * as vscode from 'vscode'
 
@@ -73,9 +78,20 @@ export class ToolUtils {
         }
     }
 
-    static validateOutput(output: InvokeOutput): void {
+    static validateOutput(output: InvokeOutput, toolType: ToolType): void {
+        let maxToolResponseSize = defaultMaxToolResponseSize
+        switch (toolType) {
+            case ToolType.FsRead:
+                maxToolResponseSize = fsReadToolResponseSize
+                break
+            case ToolType.ListDirectory:
+                maxToolResponseSize = listDirectoryToolResponseSize
+                break
+            default:
+                break
+        }
         if (output.output.content.length > maxToolResponseSize) {
-            throw Error(`Tool output exceeds maximum character limit of ${maxToolResponseSize}`)
+            throw Error(`${toolType} output exceeds maximum character limit of ${maxToolResponseSize}`)
         }
     }
 

packages/core/src/codewhispererChat/tools/tool_index.json

@@ -1,7 +1,7 @@
 {
     "fsRead": {
         "name": "fsRead",
-        "description": "A tool for reading a file.\n * This tool returns the contents of a file, and the optional `readRange` determines what range of lines will be read from the specified file.",
+        "description": "A tool for reading a file.\n * This tool returns the contents of a file, and the optional `readRange` determines what range of lines will be read from the specified file.\n * If the file exceeds 200K characters, this tool will only read the first 200K characters of the file.",
         "inputSchema": {
             "type": "object",
             "properties": {
@@ -10,7 +10,7 @@
                     "type": "string"
                 },
                 "readRange": {
-                    "description": "Optional parameter when reading files.\n * If none is given, the full file is shown. If provided, the file will be shown in the indicated line number range, e.g. [11, 12] will show lines 11 and 12. Indexing at 1 to start. Setting `[startLine, -1]` shows all lines from `startLine` to the end of the file. If the whole file is too large, try reading 4000 lines at once, for example: after reading [1, 4000], read [4000, 8000] next and repeat. You should read atleast 250 lines per invocation of the tool. In some cases, if reading a range of lines results in too many invocations instead attempt to read 4000 lines.",
+                    "description": "Optional parameter when reading files.\n * If none is given, the full file is shown. If provided, the file will be shown in the indicated line number range, e.g. [11, 12] will show lines 11 and 12. Indexing at 1 to start.\n * Setting `[startLine, -1]` shows all lines from `startLine` to the end of the file.",
                     "items": {
                         "type": "integer"
                     },

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

@@ -8,6 +8,7 @@ import { TestFolder } from '../../testUtil'
 import path from 'path'
 import * as vscode from 'vscode'
 import sinon from 'sinon'
+import { fsReadToolResponseSize } from '../../../codewhispererChat/tools/toolShared'
 
 describe('FsRead Tool', () => {
     let testFolder: TestFolder
@@ -37,6 +38,20 @@ describe('FsRead Tool', () => {
         assert.strictEqual(result.output.content, fileContent, 'File content should match exactly')
     })
 
+    it('truncate output if too large', async () => {
+        const fileContent = 'A'.repeat(fsReadToolResponseSize + 10)
+        const filePath = await testFolder.write('largeFile.txt', fileContent)
+        const fsRead = new FsRead({ path: filePath })
+        await fsRead.validate()
+        const result = await fsRead.invoke(process.stdout)
+        assert.strictEqual(result.output.kind, 'text', 'Output kind should be "text"')
+        assert.strictEqual(
+            result.output.content.length,
+            fsReadToolResponseSize,
+            'Output should be truncated to the max size'
+        )
+    })
+
     it('reads partial lines of a file', async () => {
         const fileContent = 'A\nB\nC\nD\nE\nF'
         const filePath = await testFolder.write('partialFile.txt', fileContent)

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

@@ -6,7 +6,14 @@
 import assert from 'assert'
 import * as sinon from 'sinon'
 import { Writable } from 'stream'
-import { sanitizePath, OutputKind, InvokeOutput } from '../../../codewhispererChat/tools/toolShared'
+import {
+    sanitizePath,
+    OutputKind,
+    InvokeOutput,
+    listDirectoryToolResponseSize,
+    fsReadToolResponseSize,
+    defaultMaxToolResponseSize,
+} from '../../../codewhispererChat/tools/toolShared'
 import { ToolUtils, Tool, ToolType } from '../../../codewhispererChat/tools/toolUtils'
 import { FsRead } from '../../../codewhispererChat/tools/fsRead'
 import { FsWrite } from '../../../codewhispererChat/tools/fsWrite'
@@ -158,25 +165,67 @@ describe('ToolUtils', function () {
     })
 
     describe('validateOutput', function () {
-        it('does not throw error if output is within size limit', function () {
+        it('does not throw error if output is within size limit for fsRead', function () {
             const output: InvokeOutput = {
                 output: {
                     kind: OutputKind.Text,
-                    content: 'a'.repeat(150_000),
+                    content: 'a'.repeat(fsReadToolResponseSize - 1),
                 },
             }
-            assert.doesNotThrow(() => ToolUtils.validateOutput(output))
+            assert.doesNotThrow(() => ToolUtils.validateOutput(output, ToolType.FsRead))
         })
-        it('throws error if output exceeds max size', function () {
+        it('throws error if output exceeds max size for fsRead', function () {
             const output: InvokeOutput = {
                 output: {
                     kind: OutputKind.Text,
-                    content: 'a'.repeat(200_001), // 200,001 characters
+                    content: 'a'.repeat(fsReadToolResponseSize + 1),
                 },
             }
-            assert.throws(() => ToolUtils.validateOutput(output), {
+            assert.throws(() => ToolUtils.validateOutput(output, ToolType.FsRead), {
                 name: 'Error',
-                message: 'Tool output exceeds maximum character limit of 200000',
+                message: `fsRead output exceeds maximum character limit of ${fsReadToolResponseSize}`,
+            })
+        })
+        it('does not throw error if output is within size limit for listDirectory', function () {
+            const output: InvokeOutput = {
+                output: {
+                    kind: OutputKind.Text,
+                    content: 'a'.repeat(listDirectoryToolResponseSize - 1),
+                },
+            }
+            assert.doesNotThrow(() => ToolUtils.validateOutput(output, ToolType.ListDirectory))
+        })
+        it('throws error if output exceeds max size for listDirectory', function () {
+            const output: InvokeOutput = {
+                output: {
+                    kind: OutputKind.Text,
+                    content: 'a'.repeat(listDirectoryToolResponseSize + 1),
+                },
+            }
+            assert.throws(() => ToolUtils.validateOutput(output, ToolType.ListDirectory), {
+                name: 'Error',
+                message: `listDirectory output exceeds maximum character limit of ${listDirectoryToolResponseSize}`,
+            })
+        })
+        it('does not throw error if output is within size limit for fsWrite', function () {
+            const output: InvokeOutput = {
+                output: {
+                    kind: OutputKind.Text,
+                    content: 'a'.repeat(defaultMaxToolResponseSize - 1),
+                },
+            }
+            assert.doesNotThrow(() => ToolUtils.validateOutput(output, ToolType.FsWrite))
+        })
+        it('does not throw error if output is within size limit for fsWrite', function () {
+            const output: InvokeOutput = {
+                output: {
+                    kind: OutputKind.Text,
+                    content: 'a'.repeat(defaultMaxToolResponseSize + 1),
+                },
+            }
+            assert.throws(() => ToolUtils.validateOutput(output, ToolType.FsWrite), {
+                name: 'Error',
+                message: `fsWrite output exceeds maximum character limit of ${defaultMaxToolResponseSize}`,
             })
         })
     })