fix(chat): Add maxDepth parameter to ListDirectory tool (aws#6919)

jguoamz · web-flow · commit a8672fd1940e · 2025-04-03T14:58:26.000-07:00
## Problem - ListDirectory tool will get called too many times if the directory has many subdirectories ## Solution - Add maxDepth parameter to ListDirectory tool to allow recursively listing directories - Minor fix on the tool_spec --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
diff --git a/packages/core/src/codewhispererChat/tools/listDirectory.ts b/packages/core/src/codewhispererChat/tools/listDirectory.ts
@@ -12,20 +12,26 @@ import path from 'path'
 
 export interface ListDirectoryParams {
     path: string
+    maxDepth?: number
 }
 
 export class ListDirectory {
     private fsPath: string
+    private maxDepth?: number
     private readonly logger = getLogger('listDirectory')
 
     constructor(params: ListDirectoryParams) {
         this.fsPath = params.path
+        this.maxDepth = params.maxDepth
     }
 
     public async validate(): Promise<void> {
         if (!this.fsPath || this.fsPath.trim().length === 0) {
             throw new Error('Path cannot be empty.')
         }
+        if (this.maxDepth !== undefined && this.maxDepth < 0) {
+            throw new Error('MaxDepth cannot be negative.')
+        }
 
         const sanitized = sanitizePath(this.fsPath)
         this.fsPath = sanitized
@@ -44,14 +50,21 @@ export class ListDirectory {
 
     public queueDescription(updates: Writable): void {
         const fileName = path.basename(this.fsPath)
-        updates.write(`Listing directory: ${fileName}`)
+        if (this.maxDepth === undefined) {
+            updates.write(`Listing directory recursively: ${fileName}`)
+        } else if (this.maxDepth === 0) {
+            updates.write(`Listing directory: ${fileName}`)
+        } else {
+            const level = this.maxDepth > 1 ? 'levels' : 'level'
+            updates.write(`Listing directory: ${fileName} limited to ${this.maxDepth} subfolder ${level}`)
+        }
         updates.end()
     }
 
     public async invoke(updates?: Writable): Promise<InvokeOutput> {
         try {
             const fileUri = vscode.Uri.file(this.fsPath)
-            const listing = await readDirectoryRecursively(fileUri, 0)
+            const listing = await readDirectoryRecursively(fileUri, this.maxDepth)
             return this.createOutput(listing.join('\n'))
         } catch (error: any) {
             this.logger.error(`Failed to list directory "${this.fsPath}": ${error.message || error}`)
diff --git a/packages/core/src/codewhispererChat/tools/tool_index.json b/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.",
         "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.",
+                    "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"
                     },
@@ -75,7 +75,7 @@
     },
     "listDirectory": {
         "name": "listDirectory",
-        "description": "List the contents of a directory.\n * Use this tool for discovery, before using more targeted tools like fsRead.\n *Useful to try to understand the file structure before diving deeper into specific files.\n *Can be used to explore the codebase.\n *Results clearly distinguish between files, directories or symlinks with [FILE], [DIR] and [LINK] prefixes.",
+        "description": "List the contents of a directory and its subdirectories.\n * Use this tool for discovery, before using more targeted tools like fsRead.\n *Useful to try to understand the file structure before diving deeper into specific files.\n *Can be used to explore the codebase.\n *Results clearly distinguish between files, directories or symlinks with [FILE], [DIR] and [LINK] prefixes.",
         "inputSchema": {
             "type": "object",
             "properties": {
@@ -86,6 +86,10 @@
                 "path": {
                     "type": "string",
                     "description": "Absolute path to a directory, e.g., `/repo`."
+                },
+                "maxDepth": {
+                    "type": "integer",
+                    "description": "Maximum depth to traverse when listing directories. Use `0` to list only the specified directory, `1` to include immediate subdirectories, etc. If it's not provided, it will list all subdirectories recursively."
                 }
             },
             "required": ["path"]
diff --git a/packages/core/src/shared/utilities/workspaceUtils.ts b/packages/core/src/shared/utilities/workspaceUtils.ts
@@ -693,16 +693,17 @@ export function formatListing(name: string, fileType: vscode.FileType, fullPath:
  * 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 maxDepth Maximum depth to descend (0 => just this folder, if it's missing => recursively)
  * @param customFormatCallback Optional. If given, it will override the default line-formatting
  */
 export async function readDirectoryRecursively(
     dirUri: vscode.Uri,
-    maxDepth: number,
+    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 depthDescription = maxDepth === undefined ? 'unlimited' : maxDepth
+    logger.info(`Reading directory: ${dirUri.fsPath} to max depth: ${depthDescription}`)
 
     const queue: Array<{ uri: vscode.Uri; depth: number }> = [{ uri: dirUri, depth: 0 }]
     const results: string[] = []
@@ -711,7 +712,7 @@ export async function readDirectoryRecursively(
 
     while (queue.length > 0) {
         const { uri, depth } = queue.shift()!
-        if (depth > maxDepth) {
+        if (maxDepth !== undefined && depth > maxDepth) {
             logger.info(`Skipping directory: ${uri.fsPath} (depth ${depth} > max ${maxDepth})`)
             continue
         }
@@ -729,7 +730,7 @@ export async function readDirectoryRecursively(
             const childUri = vscode.Uri.joinPath(uri, name)
             results.push(formatter(name, fileType, childUri.fsPath))
 
-            if (fileType === vscode.FileType.Directory && depth < maxDepth) {
+            if (fileType === vscode.FileType.Directory && (maxDepth === undefined || depth < maxDepth)) {
                 queue.push({ uri: childUri, depth: depth + 1 })
             }
         }
diff --git a/packages/core/src/test/codewhispererChat/tools/listDirectory.test.ts b/packages/core/src/test/codewhispererChat/tools/listDirectory.test.ts
@@ -15,13 +15,40 @@ describe('ListDirectory Tool', () => {
     })
 
     it('throws if path is empty', async () => {
-        const listDirectory = new ListDirectory({ path: '' })
+        const listDirectory = new ListDirectory({ path: '', maxDepth: 0 })
         await assert.rejects(listDirectory.validate(), /Path cannot be empty/i, 'Expected an error about empty path')
     })
 
+    it('throws if maxDepth is negative', async () => {
+        const listDirectory = new ListDirectory({ path: '~', maxDepth: -1 })
+        await assert.rejects(
+            listDirectory.validate(),
+            /MaxDepth cannot be negative/i,
+            'Expected an error about negative maxDepth'
+        )
+    })
+
     it('lists directory contents', async () => {
         await testFolder.mkdir('subfolder')
         await testFolder.write('fileA.txt', 'fileA content')
+
+        const listDirectory = new ListDirectory({ path: testFolder.path, maxDepth: 0 })
+        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('[FILE] ') && line.includes('fileA.txt'))
+        const hasSubfolder = lines.some(
+            (line: string | string[]) => line.includes('[DIR] ') && 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('lists directory contents recursively', 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 })
@@ -33,14 +60,16 @@ describe('ListDirectory Tool', () => {
         const hasSubfolder = lines.some(
             (line: string | string[]) => line.includes('[DIR] ') && line.includes('subfolder')
         )
+        const hasFileB = lines.some((line: string | string[]) => line.includes('[FILE] ') && line.includes('fileB.md'))
 
         assert.ok(hasFileA, 'Should list fileA.txt in the directory output')
         assert.ok(hasSubfolder, 'Should list the subfolder in the directory output')
+        assert.ok(hasFileB, 'Should list fileB.md in 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 })
+        const listDirectory = new ListDirectory({ path: missingPath, maxDepth: 0 })
 
         await assert.rejects(
             listDirectory.validate(),
@@ -50,7 +79,7 @@ describe('ListDirectory Tool', () => {
     })
 
     it('expands ~ path', async () => {
-        const listDirectory = new ListDirectory({ path: '~' })
+        const listDirectory = new ListDirectory({ path: '~', maxDepth: 0 })
         await listDirectory.validate()
         const result = await listDirectory.invoke(process.stdout)
 
