fix: update listDirectory tool to output in tree-like format to reduce toolSize (#1260)

Author: jguoamz

core/aws-lsp-core/src/util/workspaceUtils.test.ts

@@ -3,7 +3,13 @@ import * as fs from 'fs/promises'
 import * as assert from 'assert'
 import * as path from 'path'
 import { TestFolder } from '../test/testFolder'
-import { readDirectoryRecursively, getEntryPath, isParentFolder, isInWorkspace } from './workspaceUtils'
+import {
+    readDirectoryRecursively,
+    readDirectoryWithTreeOutput,
+    getEntryPath,
+    isParentFolder,
+    isInWorkspace,
+} from './workspaceUtils'
 import { TestFeatures } from '@aws/language-server-runtimes/testing'
 
 describe('workspaceUtils', function () {
@@ -209,4 +215,160 @@ describe('workspaceUtils', function () {
             assert.deepStrictEqual(result, [subdir1.path, file1, subdir11.path, file3, subdir2.path, file2].sort())
         })
     })
+
+    describe('readDirectoryWithTreeOuput', function () {
+        let tempFolder: TestFolder
+        let testFeatures: TestFeatures
+
+        before(async function () {
+            tempFolder = await TestFolder.create()
+            testFeatures = new TestFeatures()
+            // Taken from https://github.com/aws/language-server-runtimes/blob/674c02696c150838b4bc93543fb0009c5982e7ad/runtimes/runtimes/standalone.ts#L216
+            testFeatures.workspace.fs.readdir = path => fs.readdir(path, { withFileTypes: true })
+            testFeatures.workspace.fs.exists = path =>
+                fs.access(path).then(
+                    () => true,
+                    () => false
+                )
+        })
+
+        afterEach(async function () {
+            await tempFolder.clear()
+        })
+
+        after(async function () {
+            await tempFolder.delete()
+        })
+
+        it('recurses into subdirectories', async function () {
+            const subdir1 = await tempFolder.nest('subdir1')
+            await subdir1.write('file1', 'this is content')
+            const subdir2 = await tempFolder.nest('subdir2')
+            await subdir2.write('file2', 'this is other content')
+
+            const subdir11 = await subdir1.nest('subdir11')
+            await subdir11.write('file3', 'this is also content')
+            const subdir12 = await subdir1.nest('subdir12')
+            await subdir12.write('file4', 'this is even more content')
+            await subdir12.write('file5', 'and this is it')
+
+            const expected =
+                '|-- subdir1/\n' +
+                '|   |-- subdir11/\n' +
+                '|   |   `-- file3\n' +
+                '|   |-- subdir12/\n' +
+                '|   |   |-- file4\n' +
+                '|   |   `-- file5\n' +
+                '|   `-- file1\n' +
+                '`-- subdir2/\n' +
+                '    `-- file2\n'
+
+            const result = await readDirectoryWithTreeOutput(testFeatures, tempFolder.path)
+            assert.ok(result.includes(expected))
+        })
+
+        it('respects maxDepth parameter', async function () {
+            const subdir1 = await tempFolder.nest('subdir1')
+            await subdir1.write('file1', 'this is content')
+            const subdir2 = await subdir1.nest('subdir2')
+            await subdir2.write('file2', 'this is other content')
+            const subdir3 = await subdir2.nest('subdir3')
+            await subdir3.write('file3', 'this is also content')
+
+            const testDepth = async (maxDepth: number, expected: string) => {
+                const result = await readDirectoryWithTreeOutput(testFeatures, tempFolder.path, {
+                    maxDepth,
+                })
+                assert.ok(result.includes(expected))
+            }
+
+            await testDepth(0, '`-- subdir1/\n')
+            await testDepth(1, '`-- subdir1/\n    |-- subdir2/\n    `-- file1\n')
+            await testDepth(
+                2,
+                '`-- subdir1/\n    |-- subdir2/\n    |   |-- subdir3/\n    |   `-- file2\n    `-- file1\n'
+            )
+            await testDepth(
+                3,
+                '`-- subdir1/\n    |-- subdir2/\n    |   |-- subdir3/\n    |   |   `-- file3\n    |   `-- file2\n    `-- file1\n'
+            )
+        })
+
+        // This test doesn't work on windows since it modifies file permissions
+        if (process.platform !== 'win32') {
+            it('respects the failOnError flag', async function () {
+                const subdir1 = await tempFolder.nest('subdir1')
+                await subdir1.write('file1', 'this is content')
+
+                // Temporarily make the file unreadable.
+                await fs.chmod(subdir1.path, 0)
+                await assert.rejects(
+                    readDirectoryWithTreeOutput(testFeatures, tempFolder.path, {
+                        failOnError: true,
+                    })
+                )
+
+                const result = await readDirectoryWithTreeOutput(testFeatures, tempFolder.path)
+                await fs.chmod(subdir1.path, 0o755)
+                assert.ok(result.includes('`-- subdir1/\n'))
+            })
+        }
+
+        it('always fail if directory does not exist', async function () {
+            await assert.rejects(readDirectoryWithTreeOutput(testFeatures, path.join(tempFolder.path, 'notReal')))
+        })
+
+        it('ignores files in the exclude entries', async function () {
+            const subdir1 = await tempFolder.nest('subdir1')
+            await subdir1.write('file1', 'this is content')
+            const subdir2 = await tempFolder.nest('subdir2')
+            await subdir2.write('file2-bad', 'this is other content')
+
+            const subdir11 = await subdir1.nest('subdir11')
+            await subdir11.write('file3', 'this is also content')
+            const subdir12 = await subdir1.nest('subdir12')
+            await subdir12.write('file4-bad', 'this is even more content')
+            await subdir12.write('file5', 'and this is it')
+
+            const result = await readDirectoryWithTreeOutput(testFeatures, tempFolder.path, {
+                excludeFiles: ['file4-bad', 'file2-bad'],
+            })
+
+            const expected =
+                '|-- subdir1/\n' +
+                '|   |-- subdir11/\n' +
+                '|   |   `-- file3\n' +
+                '|   |-- subdir12/\n' +
+                '|   |   `-- file5\n' +
+                '|   `-- file1\n' +
+                '`-- subdir2/\n'
+            assert.ok(result.includes(expected))
+        })
+
+        it('ignores directories in the exclude entries', async function () {
+            const subdir1 = await tempFolder.nest('subdir1')
+            await subdir1.write('file1', 'this is content')
+            const subdir2 = await tempFolder.nest('subdir2')
+            await subdir2.write('file2', 'this is other content')
+
+            const subdir11 = await subdir1.nest('subdir11')
+            await subdir11.write('file3', 'this is also content')
+            const subdir12 = await subdir1.nest('subdir12-bad')
+            await subdir12.write('file4-bad', 'this is even more content')
+            await subdir12.write('file5-bad', 'and this is it')
+            await subdir12.write('file6-bad', 'and this is it')
+
+            const result = await readDirectoryWithTreeOutput(testFeatures, tempFolder.path, {
+                excludeDirs: ['subdir12-bad'],
+            })
+            const expected =
+                '|-- subdir1/\n' +
+                '|   |-- subdir11/\n' +
+                '|   |   `-- file3\n' +
+                '|   `-- file1\n' +
+                '`-- subdir2/\n' +
+                '    `-- file2\n'
+            assert.ok(result.includes(expected))
+        })
+    })
 })

core/aws-lsp-core/src/util/workspaceUtils.ts

@@ -109,3 +109,124 @@ export function isParentFolder(parentPath: string, childPath: string): boolean {
 export function isInWorkspace(workspaceFolderPaths: string[], filepath: string) {
     return workspaceFolderPaths.some(wsFolder => isParentFolder(wsFolder, filepath) || wsFolder === filepath)
 }
+
+/**
+ * Output the result in tree-like format, this will save tokens as the output contains much less characters
+ * project-root/
+ * |-- src/
+ * |   |-- components/
+ * |   |   |-- Button.js
+ * |   |   |-- Card.js
+ * |   |   `-- index.js
+ * |   |-- utils/
+ * |   |   |-- helpers.js
+ * |   |   `-- formatters.js
+ * |   `-- index.js
+ * |-- public/
+ * |   |-- images/
+ * |   |   |-- logo.png
+ * |   |   `-- banner.jpg
+ * |   `-- index.html
+ * |-- package.json
+ * `-- README.md
+ */
+export async function readDirectoryWithTreeOutput(
+    features: Pick<Features, 'workspace' | 'logging'> & Partial<Features>,
+    folderPath: string,
+    options?: {
+        maxDepth?: number
+        excludeDirs?: string[]
+        excludeFiles?: string[]
+        failOnError?: boolean
+    },
+    token?: CancellationToken
+): Promise<string> {
+    const dirExists = await features.workspace.fs.exists(folderPath)
+    const excludeFiles = options?.excludeFiles ?? []
+    const excludeDirs = options?.excludeDirs ?? []
+    if (!dirExists) {
+        throw new Error(`Directory does not exist: ${folderPath}`)
+    }
+
+    features.logging.info(
+        `Reading directory in tree like format: ${folderPath} to max depth: ${options?.maxDepth === undefined ? 'unlimited' : options.maxDepth}`
+    )
+
+    // Initialize result with the root directory
+    let result = `${folderPath}/\n`
+
+    // Recursive function to build the tree
+    const processDirectory = async (dirPath: string, depth: number, prefix: string): Promise<string> => {
+        if (token?.isCancellationRequested) {
+            features.logging.info('cancelled readDirectoryRecursively')
+            throw new CancellationError('user')
+        }
+
+        if (options?.maxDepth !== undefined && depth > options?.maxDepth) {
+            features.logging.info(`Skipping directory: ${dirPath} (depth ${depth} > max ${options.maxDepth})`)
+            return ''
+        }
+
+        let entries: Awaited<ReturnType<typeof features.workspace.fs.readdir>>
+        try {
+            entries = await features.workspace.fs.readdir(dirPath)
+        } catch (err) {
+            if (options?.failOnError) {
+                throw err
+            }
+            const errMsg = `Failed to read: ${dirPath} (${err})`
+            features.logging.warn(errMsg)
+            return ''
+        }
+
+        // Sort entries: directories first, then files, both alphabetically
+        entries.sort((a, b) => {
+            const aIsDir = a.isDirectory()
+            const bIsDir = b.isDirectory()
+            if (aIsDir && !bIsDir) return -1
+            if (!aIsDir && bIsDir) return 1
+            return a.name.localeCompare(b.name)
+        })
+
+        let treeOutput = ''
+
+        for (let i = 0; i < entries.length; i++) {
+            const entry = entries[i]
+            const isLast = i === entries.length - 1
+
+            if (entry.isDirectory()) {
+                if (excludeDirs.includes(entry.name)) {
+                    features.logging.log(`Skipping directory ${entry.name} due to match in [${excludeDirs.join(', ')}]`)
+                    continue
+                }
+            } else {
+                if (excludeFiles.includes(entry.name)) {
+                    features.logging.log(`Skipping file ${entry.name} due to match in [${excludeFiles.join(', ')}]`)
+                    continue
+                }
+            }
+
+            // Format this entry with proper tree characters
+            const entryType = entry.isDirectory() ? '/' : entry.isSymbolicLink() ? '@' : ''
+
+            // Use '`--' for the last item, '|--' for others
+            const connector = isLast ? '`-- ' : '|-- '
+            treeOutput += `${prefix}${connector}${entry.name}${entryType}\n`
+
+            // Process subdirectories with proper indentation for the next level
+            if (entry.isDirectory() && (options?.maxDepth === undefined || depth < options?.maxDepth)) {
+                const childPath = getEntryPath(entry)
+                // For the next level, add vertical line for non-last items, spaces for last items
+                const childPrefix = prefix + (isLast ? '    ' : '|   ')
+                treeOutput += await processDirectory(childPath, depth + 1, childPrefix)
+            }
+        }
+
+        return treeOutput
+    }
+
+    // Start processing from the root directory
+    result += await processDirectory(folderPath, 0, '')
+
+    return result
+}

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/listDirectory.test.ts

@@ -1,5 +1,4 @@
 import * as assert from 'assert'
-import { Writable } from 'stream'
 import { ListDirectory } from './listDirectory'
 import { testFolder } from '@aws/lsp-core'
 import * as path from 'path'
@@ -55,11 +54,8 @@ describe('ListDirectory Tool', () => {
         const result = await listDirectory.invoke({ path: tempFolder.path, maxDepth: 0 })
 
         assert.strictEqual(result.output.kind, 'text')
-        const lines = result.output.content.split('\n')
-        const hasFileA = lines.some((line: string | string[]) => line.includes('[F] ') && line.includes('fileA.txt'))
-        const hasSubfolder = lines.some(
-            (line: string | string[]) => line.includes('[D] ') && line.includes('subfolder')
-        )
+        const hasFileA = result.output.content.includes('`-- fileA.txt')
+        const hasSubfolder = result.output.content.includes('subfolder/\n')
 
         assert.ok(hasFileA, 'Should list fileA.txt in the directory output')
         assert.ok(hasSubfolder, 'Should list the subfolder in the directory output')
@@ -74,12 +70,9 @@ describe('ListDirectory Tool', () => {
         const result = await listDirectory.invoke({ path: tempFolder.path })
 
         assert.strictEqual(result.output.kind, 'text')
-        const lines = result.output.content.split('\n')
-        const hasFileA = lines.some((line: string | string[]) => line.includes('[F] ') && line.includes('fileA.txt'))
-        const hasSubfolder = lines.some(
-            (line: string | string[]) => line.includes('[D] ') && line.includes('subfolder')
-        )
-        const hasFileB = lines.some((line: string | string[]) => line.includes('[F] ') && line.includes('fileB.md'))
+        const hasFileA = result.output.content.includes('`-- fileA.txt')
+        const hasSubfolder = result.output.content.includes('subfolder/\n')
+        const hasFileB = result.output.content.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')
@@ -95,11 +88,8 @@ describe('ListDirectory Tool', () => {
         const result = await listDirectory.invoke({ path: tempFolder.path })
 
         assert.strictEqual(result.output.kind, 'text')
-        const lines = result.output.content.split('\n')
-        const hasNodeModules = lines.some(
-            (line: string | string[]) => line.includes('[D] ') && line.includes('node_modules')
-        )
-        const hasFileC = lines.some((line: string | string[]) => line.includes('[F] ') && line.includes('fileC.md'))
+        const hasNodeModules = result.output.content.includes('node_modules/\n')
+        const hasFileC = result.output.content.includes('`-- fileC.md')
 
         assert.ok(!hasNodeModules, 'Should not list node_modules in the directory output')
         assert.ok(!hasFileC, 'Should not list fileC.md under node_modules in the directory output')
@@ -113,8 +103,7 @@ describe('ListDirectory Tool', () => {
         await listDirectory.validate({ path: tempFolder.path })
         const result = await listDirectory.invoke({ path: tempFolder.path })
         assert.strictEqual(result.output.kind, 'text')
-        const lines = result.output.content.split('\n')
-        const hasOutput = lines.some((line: string) => line.includes('output.md'))
+        const hasOutput = result.output.content.includes('`-- output.md')
 
         assert.ok(hasOutput, 'Should list output.md under foo in the directory output')
     })