feat(chat): Add fsRead tool for Amazon Q Agentic Chat

tsmithsz · tsmithsz · commit b883c67d3d02 · 2025-03-21T15:22:39.000-07:00
diff --git a/packages/core/src/codewhispererChat/tools/fsRead.ts b/packages/core/src/codewhispererChat/tools/fsRead.ts
@@ -0,0 +1,138 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import * as vscode from 'vscode'
+import { getLogger } from '../../shared/logger/logger'
+import { readDirectoryRecursively } from '../../shared/utilities/workspaceUtils'
+import fs from '../../shared/fs/fs'
+
+const maxToolResponseSize = 30720 // 30KB
+
+enum OutputKind {
+    Text = 'text',
+    Json = 'json',
+}
+
+export interface InvokeOutput {
+    output: {
+        kind: OutputKind
+        content: string
+    }
+}
+
+export interface FsReadParams {
+    path: string
+    readRange?: number[]
+}
+
+export class FsRead {
+    private readonly fsPath: string
+    private readonly readRange?: number[]
+    private readonly logger = getLogger('fsRead')
+
+    constructor(params: FsReadParams) {
+        this.fsPath = params.path
+        this.readRange = params.readRange
+    }
+
+    public async invoke(): Promise<InvokeOutput> {
+        try {
+            const fileUri = vscode.Uri.file(this.fsPath)
+            let exists: boolean
+            try {
+                exists = await fs.exists(fileUri)
+                if (!exists) {
+                    throw new Error(`Path: "${this.fsPath}" does not exist or cannot be accessed.`)
+                }
+            } catch (err) {
+                throw new Error(`Path: "${this.fsPath}" does not exist or cannot be accessed. (${err})`)
+            }
+
+            const isFile = await fs.existsFile(fileUri)
+            const isDirectory = await fs.existsDir(fileUri)
+
+            if (isFile) {
+                const fileContents = await this.readFile(fileUri)
+                this.logger.info(`Read file: ${this.fsPath}, size: ${fileContents.length}`)
+                return this.handleFileRange(fileContents)
+            } else if (isDirectory) {
+                const maxDepth = this.getDirectoryDepth() ?? 0
+                const listing = await readDirectoryRecursively(fileUri, maxDepth)
+                return this.createOutput(listing.join('\n'))
+            } else {
+                throw new Error(`"${this.fsPath}" is neither a standard file nor directory.`)
+            }
+        } catch (error: any) {
+            this.logger.error(`Failed to read "${this.fsPath}": ${error.message || error}`)
+            throw new Error(`[fs_read] Failed to read "${this.fsPath}": ${error.message || error}`)
+        }
+    }
+
+    private async readFile(fileUri: vscode.Uri): Promise<string> {
+        this.logger.info(`Reading file: ${fileUri.fsPath}`)
+        return await fs.readFileText(fileUri)
+    }
+
+    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))
+        }
+
+        const lines = fullText.split('\n')
+        const [start, end] = this.parseLineRange(lines.length, this.readRange)
+        if (start > end) {
+            this.logger.error(`Invalid range: ${this.readRange.join('-')}`)
+            return this.createOutput('')
+        }
+
+        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))
+    }
+
+    private parseLineRange(lineCount: number, range: number[]): [number, number] {
+        const startIdx = range[0]
+        let endIdx = range.length >= 2 ? range[1] : undefined
+
+        if (endIdx === undefined) {
+            endIdx = -1
+        }
+
+        const convert = (i: number): number => {
+            return i < 0 ? lineCount + i : i - 1
+        }
+
+        const finalStart = Math.max(0, Math.min(lineCount - 1, convert(startIdx)))
+        const finalEnd = Math.max(0, Math.min(lineCount - 1, convert(endIdx)))
+        return [finalStart, finalEnd]
+    }
+
+    private getDirectoryDepth(): number | undefined {
+        if (!this.readRange || this.readRange.length === 0) {
+            return 0
+        }
+        return this.readRange[0]
+    }
+
+    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: {
+                kind: OutputKind.Text,
+                content: content,
+            },
+        }
+    }
+}
diff --git a/packages/core/src/shared/logger/logger.ts b/packages/core/src/shared/logger/logger.ts
@@ -14,6 +14,7 @@ export type LogTopic =
     | 'unknown'
     | 'chat'
     | 'stepfunctions'
+    | 'fsRead'
 
 class ErrorLog {
     constructor(
diff --git a/packages/core/src/shared/utilities/workspaceUtils.ts b/packages/core/src/shared/utilities/workspaceUtils.ts
@@ -671,3 +671,69 @@ export async function findStringInDirectory(searchStr: string, dirPath: string)
     })
     return spawnResult
 }
+
+/**
+ * Returns a one-character tag for a directory ('d'), symlink ('l'), or file ('-').
+ */
+export function formatListing(name: string, fileType: vscode.FileType, fullPath: string): string {
+    let typeChar = '-'
+    if (fileType === vscode.FileType.Directory) {
+        typeChar = 'd'
+    } else if (fileType === vscode.FileType.SymbolicLink) {
+        typeChar = 'l'
+    }
+    return `${typeChar} ${fullPath}`
+}
+
+/**
+ * Recursively lists directories using a BFS approach, returning lines like:
+ * d /absolute/path/to/folder
+ * - /absolute/path/to/file.txt
+ *
+ * You can either pass a custom callback or rely on the default `formatListing`.
+ *
+ * @param dirUri The folder to begin traversing
+ * @param maxDepth Maximum depth to descend (0 => just this folder)
+ * @param customFormatCallback Optional. If given, it will override the default line-formatting
+ */
+export async function readDirectoryRecursively(
+    dirUri: vscode.Uri,
+    maxDepth: number,
+    customFormatCallback?: (name: string, fileType: vscode.FileType, fullPath: string) => string
+): Promise<string[]> {
+    const logger = getLogger()
+    logger.info(`Reading directory: ${dirUri.fsPath} to max depth: ${maxDepth}`)
+
+    const queue: Array<{ uri: vscode.Uri; depth: number }> = [{ uri: dirUri, depth: 0 }]
+    const results: string[] = []
+
+    const formatter = customFormatCallback ?? formatListing
+
+    while (queue.length > 0) {
+        const { uri, depth } = queue.shift()!
+        if (depth > maxDepth) {
+            logger.info(`Skipping directory: ${uri.fsPath} (depth ${depth} > max ${maxDepth})`)
+            continue
+        }
+
+        let entries: [string, vscode.FileType][]
+        try {
+            entries = await fs.readdir(uri)
+        } catch (err) {
+            logger.error(`Cannot read directory: ${uri.fsPath} (${err})`)
+            results.push(`Cannot read directory: ${uri.fsPath} (${err})`)
+            continue
+        }
+
+        for (const [name, fileType] of entries) {
+            const childUri = vscode.Uri.joinPath(uri, name)
+            results.push(formatter(name, fileType, childUri.fsPath))
+
+            if (fileType === vscode.FileType.Directory && depth < maxDepth) {
+                queue.push({ uri: childUri, depth: depth + 1 })
+            }
+        }
+    }
+
+    return results
+}
diff --git a/packages/core/src/test/codewhispererChat/tools/fsRead.test.ts b/packages/core/src/test/codewhispererChat/tools/fsRead.test.ts
@@ -0,0 +1,87 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import assert from 'assert'
+import { FsRead } from '../../../codewhispererChat/tools/fsRead'
+import { TestFolder } from '../../testUtil'
+import path from 'path'
+
+describe('FsRead Tool', () => {
+    let testFolder: TestFolder
+
+    before(async () => {
+        testFolder = await TestFolder.create()
+    })
+
+    it('reads entire file', async () => {
+        const fileContent = 'Line 1\nLine 2\nLine 3'
+        const filePath = await testFolder.write('fullFile.txt', fileContent)
+
+        const fsRead = new FsRead({ path: filePath })
+        const result = await fsRead.invoke()
+
+        assert.strictEqual(result.output.kind, 'text', 'Output kind should be "text"')
+        assert.strictEqual(result.output.content, fileContent, 'File content should match exactly')
+    })
+
+    it('reads partial lines of a file', async () => {
+        const fileContent = 'A\nB\nC\nD\nE\nF'
+        const filePath = await testFolder.write('partialFile.txt', fileContent)
+
+        const fsRead = new FsRead({ path: filePath, readRange: [2, 4] })
+        const result = await fsRead.invoke()
+
+        assert.strictEqual(result.output.kind, 'text')
+        assert.strictEqual(result.output.content, 'B\nC\nD')
+    })
+
+    it('lists directory contents up to depth = 1', async () => {
+        await testFolder.mkdir('subfolder')
+        await testFolder.write('fileA.txt', 'fileA content')
+        await testFolder.write(path.join('subfolder', 'fileB.md'), '# fileB')
+
+        const fsRead = new FsRead({ path: testFolder.path, readRange: [1] })
+        const result = await fsRead.invoke()
+
+        const lines = result.output.content.split('\n')
+        const hasFileA = lines.some((line) => line.includes('- ') && line.includes('fileA.txt'))
+        const hasSubfolder = lines.some((line) => line.includes('d ') && line.includes('subfolder'))
+
+        assert.ok(hasFileA, 'Should list fileA.txt in the directory output')
+        assert.ok(hasSubfolder, 'Should list the subfolder in the directory output')
+    })
+
+    it('throws error if path does not exist', async () => {
+        const missingPath = path.join(testFolder.path, 'no_such_file.txt')
+        const fsRead = new FsRead({ path: missingPath })
+
+        await assert.rejects(
+            fsRead.invoke(),
+            /does not exist or cannot be accessed/i,
+            'Expected an error indicating the path does not exist'
+        )
+    })
+
+    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 assert.rejects(
+            fsRead.invoke(),
+            /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] })
+
+        const result = await fsRead.invoke()
+        assert.strictEqual(result.output.kind, 'text')
+        assert.strictEqual(result.output.content, '')
+    })
+})
