fix(chat): Update FsRead toolspec, set different output limit for tools

Author: jguoamz

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

@@ -736,7 +736,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/toolShared.ts

@@ -6,7 +6,10 @@
 import path from 'path'
 import fs from '../../shared/fs/fs'
 
-export const maxToolResponseSize = 200_000
+export const defaultMaxToolResponseSize = 100_000
+export const listDirectoryToolResponseSize = 50_000
+// TODO: temporary fix to allow large file read, reduce the limit once falcon context is working properly
+export const fsReadToolResponseSize = 300_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

@@ -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.",
+                    "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. If the whole file is too large, try reading 5000 lines at once, for example, use [1, 5000] first and then [5000, 10000] next and so on. You should read AT LEAST 1000 lines per invocation of the tool.",
                     "items": {
                         "type": "integer"
                     },

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

@@ -158,27 +158,57 @@ 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(299_000),
                 },
             }
-            assert.doesNotThrow(() => ToolUtils.validateOutput(output))
+            assert.doesNotThrow(() => ToolUtils.validateOutput(output, ToolType.FsRead), 'No error should be thrown')
         })
-        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(300_001), // 200,001 characters
                 },
             }
-            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 exceeds maximum character limit of 300000',
             })
         })
+        it('does not throw error if output is within size limit for listDirectory', function () {
+            const output: InvokeOutput = {
+                output: {
+                    kind: OutputKind.Text,
+                    content: 'a'.repeat(49_000),
+                },
+            }
+            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(50_001),
+                },
+            }
+            assert.throws(() => ToolUtils.validateOutput(output, ToolType.ListDirectory), {
+                name: 'Error',
+                message: 'fsRead exceeds maximum character limit of 50000',
+            })
+        })
+        it('does not throw error if output is within size limit for fsWrite', function () {
+            const output: InvokeOutput = {
+                output: {
+                    kind: OutputKind.Text,
+                    content: 'a'.repeat(199_999),
+                },
+            }
+            assert.doesNotThrow(() => ToolUtils.validateOutput(output, ToolType.FsWrite))
+        })
     })
 
     describe('queueDescription', function () {