feat(chat): Add ListDirectory tool

Author: jguoamz

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

@@ -4,7 +4,6 @@
  */
 import * as vscode from 'vscode'
 import { getLogger } from '../../shared/logger/logger'
-import { readDirectoryRecursively } from '../../shared/utilities/workspaceUtils'
 import fs from '../../shared/fs/fs'
 import { InvokeOutput, maxToolResponseSize, OutputKind, sanitizePath } from './toolShared'
 import { Writable } from 'stream'
@@ -18,7 +17,6 @@ export interface FsReadParams {
 export class FsRead {
     private fsPath: string
     private readonly readRange?: number[]
-    private isFile?: boolean // true for file, false for directory
     private readonly logger = getLogger('fsRead')
 
     constructor(params: FsReadParams) {
@@ -46,7 +44,6 @@ export class FsRead {
             throw new Error(`Path: "${this.fsPath}" does not exist or cannot be accessed. (${err})`)
         }
 
-        this.isFile = await fs.existsFile(fileUri)
         this.logger.debug(`Validation succeeded for path: ${this.fsPath}`)
     }
 
@@ -61,17 +58,9 @@ export class FsRead {
         try {
             const fileUri = vscode.Uri.file(this.fsPath)
 
-            if (this.isFile) {
-                const fileContents = await this.readFile(fileUri)
-                this.logger.info(`Read file: ${this.fsPath}, size: ${fileContents.length}`)
-                return this.handleFileRange(fileContents)
-            } else if (!this.isFile) {
-                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.`)
-            }
+            const fileContents = await this.readFile(fileUri)
+            this.logger.info(`Read file: ${this.fsPath}, size: ${fileContents.length}`)
+            return this.handleFileRange(fileContents)
         } 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}`)
@@ -118,13 +107,6 @@ export class FsRead {
         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) {

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

@@ -0,0 +1,73 @@
+/*!
+ * 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'
+import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'
+import { Writable } from 'stream'
+import path from 'path'
+
+export interface ListDirectoryParams {
+    path: string
+}
+
+export class ListDirectory {
+    private fsPath: string
+    private readonly logger = getLogger('listDirectory')
+
+    constructor(params: ListDirectoryParams) {
+        this.fsPath = params.path
+    }
+
+    public async validate(): Promise<void> {
+        this.logger.debug(`Validating fsPath: ${this.fsPath}`)
+        if (!this.fsPath || this.fsPath.trim().length === 0) {
+            throw new Error('Path cannot be empty.')
+        }
+
+        const sanitized = sanitizePath(this.fsPath)
+        this.fsPath = sanitized
+
+        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})`)
+        }
+
+        this.logger.debug(`Validation succeeded for path: ${this.fsPath}`)
+    }
+
+    public queueDescription(updates: Writable): void {
+        const fileName = path.basename(this.fsPath)
+        updates.write(`Listing directory on: [${fileName}]`)
+        updates.end()
+    }
+
+    public async invoke(updates: Writable): Promise<InvokeOutput> {
+        try {
+            const fileUri = vscode.Uri.file(this.fsPath)
+            const listing = await readDirectoryRecursively(fileUri, 0)
+            return this.createOutput(listing.join('\n'))
+        } catch (error: any) {
+            this.logger.error(`Failed to list directory "${this.fsPath}": ${error.message || error}`)
+            throw new Error(`[fs_read] Failed to list directory "${this.fsPath}": ${error.message || error}`)
+        }
+    }
+
+    private createOutput(content: string): InvokeOutput {
+        return {
+            output: {
+                kind: OutputKind.Text,
+                content: content,
+            },
+        }
+    }
+}

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

@@ -8,17 +8,20 @@ import { FsWrite, FsWriteParams } from './fsWrite'
 import { ExecuteBash, ExecuteBashParams } from './executeBash'
 import { ToolResult, ToolResultContentBlock, ToolResultStatus, ToolUse } from '@amzn/codewhisperer-streaming'
 import { InvokeOutput } from './toolShared'
+import { ListDirectory, ListDirectoryParams } from './listDirectory'
 
 export enum ToolType {
     FsRead = 'fsRead',
     FsWrite = 'fsWrite',
     ExecuteBash = 'executeBash',
+    ListDirectory = 'listDirectory',
 }
 
 export type Tool =
     | { type: ToolType.FsRead; tool: FsRead }
     | { type: ToolType.FsWrite; tool: FsWrite }
     | { type: ToolType.ExecuteBash; tool: ExecuteBash }
+    | { type: ToolType.ListDirectory; tool: ListDirectory }
 
 export class ToolUtils {
     static displayName(tool: Tool): string {
@@ -29,6 +32,8 @@ export class ToolUtils {
                 return 'Write to filesystem'
             case ToolType.ExecuteBash:
                 return 'Execute shell command'
+            case ToolType.ListDirectory:
+                return 'List directory from filesystem'
         }
     }
 
@@ -40,6 +45,8 @@ export class ToolUtils {
                 return true
             case ToolType.ExecuteBash:
                 return tool.tool.requiresAcceptance()
+            case ToolType.ListDirectory:
+                return false
         }
     }
 
@@ -51,6 +58,8 @@ export class ToolUtils {
                 return tool.tool.invoke(updates)
             case ToolType.ExecuteBash:
                 return tool.tool.invoke(updates)
+            case ToolType.ListDirectory:
+                return tool.tool.invoke(updates)
         }
     }
 
@@ -65,6 +74,9 @@ export class ToolUtils {
             case ToolType.ExecuteBash:
                 tool.tool.queueDescription(updates)
                 break
+            case ToolType.ListDirectory:
+                tool.tool.queueDescription(updates)
+                break
         }
     }
 
@@ -76,6 +88,8 @@ export class ToolUtils {
                 return tool.tool.validate()
             case ToolType.ExecuteBash:
                 return tool.tool.validate()
+            case ToolType.ListDirectory:
+                return tool.tool.validate()
         }
     }
 
@@ -108,6 +122,11 @@ export class ToolUtils {
                         type: ToolType.ExecuteBash,
                         tool: new ExecuteBash(value.input as unknown as ExecuteBashParams),
                     }
+                case ToolType.ListDirectory:
+                    return {
+                        type: ToolType.ListDirectory,
+                        tool: new ListDirectory(value.input as unknown as ListDirectoryParams),
+                    }
                 default:
                     return {
                         toolUseId: value.toolUseId,

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

@@ -1,16 +1,16 @@
 {
     "fsRead": {
         "name": "fsRead",
-        "description": "A tool for reading files (e.g. `cat -n`), or listing files/directories (e.g. `ls -la` or `find . -maxdepth 2). The behavior of this tool is determined by the `path` parameter pointing to a file or directory.\n* If `path` is a file, this tool returns the result of running `cat -n`, and the optional `readRange` determines what range of lines will be read from the specified file.\n* If `path` is a directory, this tool returns the listed files and directories of the specified path, as if running `ls -la`. If the `readRange` parameter is provided, the tool acts like the `find . -maxdepth <readRange>`, where `readRange` is the number of subdirectories deep to search, e.g. [2] will run `find . -maxdepth 2`.",
+        "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.",
         "inputSchema": {
             "type": "object",
             "properties": {
                 "path": {
-                    "description": "Absolute path to file or directory, e.g. `/repo/file.py` or `/repo`.",
+                    "description": "Absolute path to a file, e.g. `/repo/file.py`.",
                     "type": "string"
                 },
                 "readRange": {
-                    "description": "Optional parameter when reading either files or directories.\n* When `path` is a file, 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.\n* When `path` is a directory, if none is given, the results of `ls -l` are given. If provided, the current directory and indicated number of subdirectories will be shown, e.g. [2] will show the current directory and directories two levels deep.",
+                    "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.",
                     "items": {
                         "type": "integer"
                     },
@@ -72,5 +72,23 @@
             },
             "required": ["command", "cwd"]
         }
+    },
+    "listDirectory": {
+        "name": "listDirectory",
+        "description": "List the contents of a directory. Use this tool for discovery, before using more targeted tools like fsRead. Useful to try to understand the file structure before diving deeper into specific files. Can be used to explore the codebase.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "explanation": {
+                    "description": "One sentence explanation as to why this tool is being used, and how it contributes to the goal.",
+                    "type": "string"
+                },
+                "path": {
+                    "type": "string",
+                    "description": "Absolute path to file or directory, e.g. `/repo/file.py` or `/repo`."
+                }
+            },
+            "required": ["path"]
+        }
     }
 }

packages/core/src/shared/logger/logger.ts

@@ -19,6 +19,7 @@ export type LogTopic =
     | 'fsRead'
     | 'fsWrite'
     | 'executeBash'
+    | 'listDirectory'
     | 'chatStream'
     | 'unknown'
 

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

@@ -43,23 +43,6 @@ describe('FsRead Tool', () => {
         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] })
-        await fsRead.validate()
-        const result = await fsRead.invoke(process.stdout)
-
-        const lines = result.output.content.split('\n')
-        const hasFileA = lines.some((line: string | string[]) => line.includes('- ') && line.includes('fileA.txt'))
-        const hasSubfolder = lines.some((line: string | string[]) => 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 })
@@ -95,15 +78,6 @@ describe('FsRead Tool', () => {
         assert.strictEqual(result.output.content, '')
     })
 
-    it('expands ~ path', async () => {
-        const fsRead = new FsRead({ path: '~' })
-        await fsRead.validate()
-        const result = await fsRead.invoke(process.stdout)
-
-        assert.strictEqual(result.output.kind, 'text')
-        assert.ok(result.output.content.length > 0)
-    })
-
     it('resolves relative path', async () => {
         await testFolder.mkdir('relTest')
         const filePath = path.join('relTest', 'relFile.txt')

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

@@ -0,0 +1,74 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import assert from 'assert'
+import { ListDirectory } from '../../../codewhispererChat/tools/listDirectory'
+import { TestFolder } from '../../testUtil'
+import path from 'path'
+
+describe('ListDirectory Tool', () => {
+    let testFolder: TestFolder
+
+    before(async () => {
+        testFolder = await TestFolder.create()
+    })
+
+    it('throws if path is empty', async () => {
+        const listDirectory = new ListDirectory({ path: '' })
+        await assert.rejects(listDirectory.validate(), /Path cannot be empty/i, 'Expected an error about empty path')
+    })
+
+    it('lists directory contents', async () => {
+        await testFolder.mkdir('subfolder')
+        await testFolder.write('fileA.txt', 'fileA content')
+        await testFolder.write(path.join('subfolder', 'fileB.md'), '# fileB')
+
+        const listDirectory = new ListDirectory({ path: testFolder.path })
+        await listDirectory.validate()
+        const result = await listDirectory.invoke(process.stdout)
+
+        const lines = result.output.content.split('\n')
+        const hasFileA = lines.some((line: string | string[]) => line.includes('- ') && line.includes('fileA.txt'))
+        const hasSubfolder = lines.some((line: string | string[]) => 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 listDirectory = new ListDirectory({ path: missingPath })
+
+        await assert.rejects(
+            listDirectory.validate(),
+            /does not exist or cannot be accessed/i,
+            'Expected an error indicating the path does not exist'
+        )
+    })
+
+    it('expands ~ path', async () => {
+        const listDirectory = new ListDirectory({ path: '~' })
+        await listDirectory.validate()
+        const result = await listDirectory.invoke(process.stdout)
+
+        assert.strictEqual(result.output.kind, 'text')
+        assert.ok(result.output.content.length > 0)
+    })
+
+    it('resolves relative path', async () => {
+        await testFolder.mkdir('relTest')
+        const filePath = path.join('relTest', 'relFile.txt')
+        const content = 'Hello from a relative file!'
+        await testFolder.write(filePath, content)
+
+        const relativePath = path.relative(process.cwd(), path.join(testFolder.path, filePath))
+
+        const listDirectory = new ListDirectory({ path: relativePath })
+        await listDirectory.validate()
+        const result = await listDirectory.invoke(process.stdout)
+
+        assert.strictEqual(result.output.kind, 'text')
+        assert.strictEqual(result.output.content, content)
+    })
+})