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

Author: jguoamz

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

@@ -746,7 +746,7 @@ export class ChatController {
                                 session.setFsWriteBackup(toolUse.toolUseId, backup)
                             }
                             const output = await ToolUtils.invoke(tool, chatStream)
-                            ToolUtils.validateOutput(output)
+                            ToolUtils.validateOutput(output, tool.type)
 
                             toolResults.push({
                                 content: [

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',

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'
 
 export enum ToolType {
@@ -63,9 +68,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}`,
             })
         })
     })