From ed440d371af25d8a53a2fb84bb9c3a573e7070f5 Mon Sep 17 00:00:00 2001 From: Jason Guo Date: Mon, 7 Apr 2025 18:04:49 -0700 Subject: [PATCH] fix(chat): Remove enforceMaxSize from fsRead --- .../controllers/chat/controller.ts | 8 ++----- .../src/codewhispererChat/tools/fsRead.ts | 17 +++---------- .../src/codewhispererChat/tools/toolShared.ts | 3 +-- .../src/codewhispererChat/tools/toolUtils.ts | 8 ++++++- .../codewhispererChat/tools/fsRead.test.ts | 14 ----------- .../tools/toolShared.test.ts | 24 +++++++++++++++++++ 6 files changed, 37 insertions(+), 37 deletions(-) diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts index 32ab4968aed..a116e45fda9 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/controller.ts @@ -93,7 +93,7 @@ import { } from '../../constants' import { ChatSession } from '../../clients/chat/v0/chat' import { amazonQTabSuffix } from '../../../shared/constants' -import { maxToolOutputCharacterLength, OutputKind } from '../../tools/toolShared' +import { OutputKind } from '../../tools/toolShared' import { ToolUtils, Tool, ToolType } from '../../tools/toolUtils' import { ChatStream } from '../../tools/chatStream' import { ChatHistoryStorage } from '../../storages/chatHistoryStorage' @@ -723,11 +723,7 @@ export class ChatController { requiresAcceptance: false, }) const output = await ToolUtils.invoke(tool, chatStream) - if (output.output.content.length > maxToolOutputCharacterLength) { - throw Error( - `Tool output exceeds maximum character limit of ${maxToolOutputCharacterLength}` - ) - } + ToolUtils.validateOutput(output) toolResults.push({ content: [ diff --git a/packages/core/src/codewhispererChat/tools/fsRead.ts b/packages/core/src/codewhispererChat/tools/fsRead.ts index 07fd64d5293..05519641bd0 100644 --- a/packages/core/src/codewhispererChat/tools/fsRead.ts +++ b/packages/core/src/codewhispererChat/tools/fsRead.ts @@ -5,7 +5,7 @@ import * as vscode from 'vscode' import { getLogger } from '../../shared/logger/logger' import fs from '../../shared/fs/fs' -import { InvokeOutput, maxToolResponseSize, OutputKind, sanitizePath } from './toolShared' +import { InvokeOutput, OutputKind, sanitizePath } from './toolShared' import { Writable } from 'stream' import path from 'path' @@ -89,7 +89,7 @@ export class FsRead { private handleFileRange(fullText: string): InvokeOutput { if (!this.readRange || this.readRange.length === 0) { this.logger.info('No range provided. returning entire file.') - return this.createOutput(this.enforceMaxSize(fullText)) + return this.createOutput(fullText) } const lines = fullText.split('\n') @@ -101,7 +101,7 @@ export class FsRead { this.logger.info(`Reading file: ${this.fsPath}, lines ${start + 1}-${end + 1}`) const slice = lines.slice(start, end + 1).join('\n') - return this.createOutput(this.enforceMaxSize(slice)) + return this.createOutput(slice) } private parseLineRange(lineCount: number, range: number[]): [number, number] { @@ -121,17 +121,6 @@ export class FsRead { return [finalStart, finalEnd] } - private enforceMaxSize(content: string): string { - const byteCount = Buffer.byteLength(content, 'utf8') - if (byteCount > maxToolResponseSize) { - throw new Error( - `This tool only supports reading ${maxToolResponseSize} bytes at a time. - You tried to read ${byteCount} bytes. Try executing with fewer lines specified.` - ) - } - return content - } - private createOutput(content: string): InvokeOutput { return { output: { diff --git a/packages/core/src/codewhispererChat/tools/toolShared.ts b/packages/core/src/codewhispererChat/tools/toolShared.ts index d00881f4a34..0fc349bde3b 100644 --- a/packages/core/src/codewhispererChat/tools/toolShared.ts +++ b/packages/core/src/codewhispererChat/tools/toolShared.ts @@ -6,8 +6,7 @@ import path from 'path' import fs from '../../shared/fs/fs' -export const maxToolResponseSize = 30720 // 30KB -export const maxToolOutputCharacterLength = 800_000 +export const maxToolResponseSize = 800_000 export enum OutputKind { Text = 'text', diff --git a/packages/core/src/codewhispererChat/tools/toolUtils.ts b/packages/core/src/codewhispererChat/tools/toolUtils.ts index 93d51244ebf..cbaa2149217 100644 --- a/packages/core/src/codewhispererChat/tools/toolUtils.ts +++ b/packages/core/src/codewhispererChat/tools/toolUtils.ts @@ -7,7 +7,7 @@ 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 } from './toolShared' +import { InvokeOutput, maxToolResponseSize } from './toolShared' import { ListDirectory, ListDirectoryParams } from './listDirectory' export enum ToolType { @@ -63,6 +63,12 @@ export class ToolUtils { } } + static validateOutput(output: InvokeOutput): void { + if (output.output.content.length > maxToolResponseSize) { + throw Error(`Tool output exceeds maximum character limit of ${maxToolResponseSize}`) + } + } + static async queueDescription(tool: Tool, updates: Writable): Promise { switch (tool.type) { case ToolType.FsRead: diff --git a/packages/core/src/test/codewhispererChat/tools/fsRead.test.ts b/packages/core/src/test/codewhispererChat/tools/fsRead.test.ts index d7bf76456b1..985c1a86a05 100644 --- a/packages/core/src/test/codewhispererChat/tools/fsRead.test.ts +++ b/packages/core/src/test/codewhispererChat/tools/fsRead.test.ts @@ -54,20 +54,6 @@ describe('FsRead Tool', () => { ) }) - it('throws error if content exceeds 30KB', async function () { - const bigContent = 'x'.repeat(35_000) - const bigFilePath = await testFolder.write('bigFile.txt', bigContent) - - const fsRead = new FsRead({ path: bigFilePath }) - await fsRead.validate() - - await assert.rejects( - fsRead.invoke(process.stdout), - /This tool only supports reading \d+ bytes at a time/i, - 'Expected a size-limit error' - ) - }) - it('invalid line range', async () => { const filePath = await testFolder.write('rangeTest.txt', '1\n2\n3') const fsRead = new FsRead({ path: filePath, readRange: [3, 2] }) diff --git a/packages/core/src/test/codewhispererChat/tools/toolShared.test.ts b/packages/core/src/test/codewhispererChat/tools/toolShared.test.ts index b6dec845e3e..0c638370694 100644 --- a/packages/core/src/test/codewhispererChat/tools/toolShared.test.ts +++ b/packages/core/src/test/codewhispererChat/tools/toolShared.test.ts @@ -155,6 +155,30 @@ describe('ToolUtils', function () { }) }) + describe('validateOutput', function () { + it('does not throw error if output is within size limit', function () { + const output: InvokeOutput = { + output: { + kind: OutputKind.Text, + content: 'a'.repeat(700_000), + }, + } + assert.doesNotThrow(() => ToolUtils.validateOutput(output)) + }) + it('throws error if output exceeds max size', function () { + const output: InvokeOutput = { + output: { + kind: OutputKind.Text, + content: 'a'.repeat(900_000), // 900,000 characters + }, + } + assert.throws(() => ToolUtils.validateOutput(output), { + name: 'Error', + message: 'Tool output exceeds maximum character limit of 800000', + }) + }) + }) + describe('queueDescription', 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 () {