Skip to content

Commit 22a132d

Browse files
mr-leemr-lee
authored
Merge pull request aws#6888 from jguoamz/listDirectory
feat(chat): Add ListDirectory tool
2 parents d9620cf + 1e747fc commit 22a132d

File tree

9 files changed

+244
-75
lines changed

9 files changed

+244
-75
lines changed

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

Lines changed: 7 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
*/
55
import * as vscode from 'vscode'
66
import { getLogger } from '../../shared/logger/logger'
7-
import { readDirectoryRecursively } from '../../shared/utilities/workspaceUtils'
87
import fs from '../../shared/fs/fs'
98
import { InvokeOutput, maxToolResponseSize, OutputKind, sanitizePath } from './toolShared'
109
import { Writable } from 'stream'
@@ -18,7 +17,6 @@ export interface FsReadParams {
1817
export class FsRead {
1918
private fsPath: string
2019
private readonly readRange?: number[]
21-
private isFile?: boolean // true for file, false for directory
2220
private readonly logger = getLogger('fsRead')
2321

2422
constructor(params: FsReadParams) {
@@ -36,17 +34,16 @@ export class FsRead {
3634
this.fsPath = sanitized
3735

3836
const fileUri = vscode.Uri.file(this.fsPath)
39-
let exists: boolean
37+
let fileExists: boolean
4038
try {
41-
exists = await fs.exists(fileUri)
42-
if (!exists) {
39+
fileExists = await fs.existsFile(fileUri)
40+
if (!fileExists) {
4341
throw new Error(`Path: "${this.fsPath}" does not exist or cannot be accessed.`)
4442
}
4543
} catch (err) {
4644
throw new Error(`Path: "${this.fsPath}" does not exist or cannot be accessed. (${err})`)
4745
}
4846

49-
this.isFile = await fs.existsFile(fileUri)
5047
this.logger.debug(`Validation succeeded for path: ${this.fsPath}`)
5148
}
5249

@@ -61,20 +58,12 @@ export class FsRead {
6158
try {
6259
const fileUri = vscode.Uri.file(this.fsPath)
6360

64-
if (this.isFile) {
65-
const fileContents = await this.readFile(fileUri)
66-
this.logger.info(`Read file: ${this.fsPath}, size: ${fileContents.length}`)
67-
return this.handleFileRange(fileContents)
68-
} else if (!this.isFile) {
69-
const maxDepth = this.getDirectoryDepth() ?? 0
70-
const listing = await readDirectoryRecursively(fileUri, maxDepth)
71-
return this.createOutput(listing.join('\n'))
72-
} else {
73-
throw new Error(`"${this.fsPath}" is neither a standard file nor directory.`)
74-
}
61+
const fileContents = await this.readFile(fileUri)
62+
this.logger.info(`Read file: ${this.fsPath}, size: ${fileContents.length}`)
63+
return this.handleFileRange(fileContents)
7564
} catch (error: any) {
7665
this.logger.error(`Failed to read "${this.fsPath}": ${error.message || error}`)
77-
throw new Error(`[fs_read] Failed to read "${this.fsPath}": ${error.message || error}`)
66+
throw new Error(`Failed to read "${this.fsPath}": ${error.message || error}`)
7867
}
7968
}
8069

@@ -118,13 +107,6 @@ export class FsRead {
118107
return [finalStart, finalEnd]
119108
}
120109

121-
private getDirectoryDepth(): number | undefined {
122-
if (!this.readRange || this.readRange.length === 0) {
123-
return 0
124-
}
125-
return this.readRange[0]
126-
}
127-
128110
private enforceMaxSize(content: string): string {
129111
const byteCount = Buffer.byteLength(content, 'utf8')
130112
if (byteCount > maxToolResponseSize) {

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
import * as vscode from 'vscode'
6+
import { getLogger } from '../../shared/logger/logger'
7+
import { readDirectoryRecursively } from '../../shared/utilities/workspaceUtils'
8+
import fs from '../../shared/fs/fs'
9+
import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'
10+
import { Writable } from 'stream'
11+
import path from 'path'
12+
13+
export interface ListDirectoryParams {
14+
path: string
15+
}
16+
17+
export class ListDirectory {
18+
private fsPath: string
19+
private readonly logger = getLogger('listDirectory')
20+
21+
constructor(params: ListDirectoryParams) {
22+
this.fsPath = params.path
23+
}
24+
25+
public async validate(): Promise<void> {
26+
if (!this.fsPath || this.fsPath.trim().length === 0) {
27+
throw new Error('Path cannot be empty.')
28+
}
29+
30+
const sanitized = sanitizePath(this.fsPath)
31+
this.fsPath = sanitized
32+
33+
const pathUri = vscode.Uri.file(this.fsPath)
34+
let pathExists: boolean
35+
try {
36+
pathExists = await fs.existsDir(pathUri)
37+
if (!pathExists) {
38+
throw new Error(`Path: "${this.fsPath}" does not exist or cannot be accessed.`)
39+
}
40+
} catch (err) {
41+
throw new Error(`Path: "${this.fsPath}" does not exist or cannot be accessed. (${err})`)
42+
}
43+
}
44+
45+
public queueDescription(updates: Writable): void {
46+
const fileName = path.basename(this.fsPath)
47+
updates.write(`Listing directory on filePath: ${fileName}`)
48+
updates.end()
49+
}
50+
51+
public async invoke(updates: Writable): Promise<InvokeOutput> {
52+
try {
53+
const fileUri = vscode.Uri.file(this.fsPath)
54+
const listing = await readDirectoryRecursively(fileUri, 0)
55+
return this.createOutput(listing.join('\n'))
56+
} catch (error: any) {
57+
this.logger.error(`Failed to list directory "${this.fsPath}": ${error.message || error}`)
58+
throw new Error(`Failed to list directory "${this.fsPath}": ${error.message || error}`)
59+
}
60+
}
61+
62+
private createOutput(content: string): InvokeOutput {
63+
return {
64+
output: {
65+
kind: OutputKind.Text,
66+
content: content,
67+
},
68+
}
69+
}
70+
}

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,17 +8,20 @@ import { FsWrite, FsWriteParams } from './fsWrite'
88
import { ExecuteBash, ExecuteBashParams } from './executeBash'
99
import { ToolResult, ToolResultContentBlock, ToolResultStatus, ToolUse } from '@amzn/codewhisperer-streaming'
1010
import { InvokeOutput } from './toolShared'
11+
import { ListDirectory, ListDirectoryParams } from './listDirectory'
1112

1213
export enum ToolType {
1314
FsRead = 'fsRead',
1415
FsWrite = 'fsWrite',
1516
ExecuteBash = 'executeBash',
17+
ListDirectory = 'listDirectory',
1618
}
1719

1820
export type Tool =
1921
| { type: ToolType.FsRead; tool: FsRead }
2022
| { type: ToolType.FsWrite; tool: FsWrite }
2123
| { type: ToolType.ExecuteBash; tool: ExecuteBash }
24+
| { type: ToolType.ListDirectory; tool: ListDirectory }
2225

2326
export class ToolUtils {
2427
static displayName(tool: Tool): string {
@@ -29,6 +32,8 @@ export class ToolUtils {
2932
return 'Write to filesystem'
3033
case ToolType.ExecuteBash:
3134
return 'Execute shell command'
35+
case ToolType.ListDirectory:
36+
return 'List directory from filesystem'
3237
}
3338
}
3439

@@ -40,6 +45,8 @@ export class ToolUtils {
4045
return true
4146
case ToolType.ExecuteBash:
4247
return tool.tool.requiresAcceptance()
48+
case ToolType.ListDirectory:
49+
return false
4350
}
4451
}
4552

@@ -51,6 +58,8 @@ export class ToolUtils {
5158
return tool.tool.invoke(updates)
5259
case ToolType.ExecuteBash:
5360
return tool.tool.invoke(updates)
61+
case ToolType.ListDirectory:
62+
return tool.tool.invoke(updates)
5463
}
5564
}
5665

@@ -65,6 +74,9 @@ export class ToolUtils {
6574
case ToolType.ExecuteBash:
6675
tool.tool.queueDescription(updates)
6776
break
77+
case ToolType.ListDirectory:
78+
tool.tool.queueDescription(updates)
79+
break
6880
}
6981
}
7082

@@ -76,6 +88,8 @@ export class ToolUtils {
7688
return tool.tool.validate()
7789
case ToolType.ExecuteBash:
7890
return tool.tool.validate()
91+
case ToolType.ListDirectory:
92+
return tool.tool.validate()
7993
}
8094
}
8195

@@ -108,6 +122,11 @@ export class ToolUtils {
108122
type: ToolType.ExecuteBash,
109123
tool: new ExecuteBash(value.input as unknown as ExecuteBashParams),
110124
}
125+
case ToolType.ListDirectory:
126+
return {
127+
type: ToolType.ListDirectory,
128+
tool: new ListDirectory(value.input as unknown as ListDirectoryParams),
129+
}
111130
default:
112131
return {
113132
toolUseId: value.toolUseId,

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

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,16 @@
11
{
22
"fsRead": {
33
"name": "fsRead",
4-
"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`.",
4+
"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.",
55
"inputSchema": {
66
"type": "object",
77
"properties": {
88
"path": {
9-
"description": "Absolute path to file or directory, e.g. `/repo/file.py` or `/repo`.",
9+
"description": "Absolute path to a file, e.g. `/repo/file.py`.",
1010
"type": "string"
1111
},
1212
"readRange": {
13-
"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.",
13+
"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.",
1414
"items": {
1515
"type": "integer"
1616
},
@@ -22,7 +22,7 @@
2222
},
2323
"fsWrite": {
2424
"name": "fsWrite",
25-
"description": "A tool for creating and editing files\n * The `create` command will override the file at `path` if it already exists as a file, and otherwise create a new file\n * The `append` command will add content to the end of an existing file, automatically adding a newline if the file doesn't end with one. The file must exist.\n Notes for using the `strReplace` command:\n * The `oldStr` parameter should match EXACTLY one or more consecutive lines from the original file. Be mindful of whitespaces!\n * If the `oldStr` parameter is not unique in the file, the replacement will not be performed. Make sure to include enough context in `oldStr` to make it unique\n * The `newStr` parameter should contain the edited lines that should replace the `oldStr`. The `insert` command will insert `newStr` after `insertLine` and place it on its own line.",
25+
"description": "A tool for creating and editing a file.\n * The `create` command will override the file at `path` if it already exists as a file, and otherwise create a new file\n * The `append` command will add content to the end of an existing file, automatically adding a newline if the file doesn't end with one. The file must exist.\n Notes for using the `strReplace` command:\n * The `oldStr` parameter should match EXACTLY one or more consecutive lines from the original file. Be mindful of whitespaces!\n * If the `oldStr` parameter is not unique in the file, the replacement will not be performed. Make sure to include enough context in `oldStr` to make it unique\n * The `newStr` parameter should contain the edited lines that should replace the `oldStr`. The `insert` command will insert `newStr` after `insertLine` and place it on its own line.",
2626
"inputSchema": {
2727
"type": "object",
2828
"properties": {
@@ -72,5 +72,23 @@
7272
},
7373
"required": ["command", "cwd"]
7474
}
75+
},
76+
"listDirectory": {
77+
"name": "listDirectory",
78+
"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.",
79+
"inputSchema": {
80+
"type": "object",
81+
"properties": {
82+
"explanation": {
83+
"description": "One sentence explanation as to why this tool is being used, and how it contributes to the goal.",
84+
"type": "string"
85+
},
86+
"path": {
87+
"type": "string",
88+
"description": "Absolute path to a directory, e.g., `/repo`."
89+
}
90+
},
91+
"required": ["path"]
92+
}
7593
}
7694
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ export type LogTopic =
1919
| 'fsRead'
2020
| 'fsWrite'
2121
| 'executeBash'
22+
| 'listDirectory'
2223
| 'chatStream'
2324
| 'unknown'
2425

‎packages/core/src/shared/utilities/workspaceUtils.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -673,14 +673,14 @@ export async function findStringInDirectory(searchStr: string, dirPath: string)
673673
}
674674

675675
/**
676-
* Returns a one-character tag for a directory ('d'), symlink ('l'), or file ('-').
676+
* Returns a prefix for a directory ('[DIR]'), symlink ('[LINK]'), or file ('[FILE]').
677677
*/
678678
export function formatListing(name: string, fileType: vscode.FileType, fullPath: string): string {
679-
let typeChar = '-'
679+
let typeChar = '[FILE]'
680680
if (fileType === vscode.FileType.Directory) {
681-
typeChar = 'd'
681+
typeChar = '[DIR]'
682682
} else if (fileType === vscode.FileType.SymbolicLink) {
683-
typeChar = 'l'
683+
typeChar = '[LINK]'
684684
}
685685
return `${typeChar} ${fullPath}`
686686
}

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

Lines changed: 0 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -43,23 +43,6 @@ describe('FsRead Tool', () => {
4343
assert.strictEqual(result.output.content, 'B\nC\nD')
4444
})
4545

46-
it('lists directory contents up to depth = 1', async () => {
47-
await testFolder.mkdir('subfolder')
48-
await testFolder.write('fileA.txt', 'fileA content')
49-
await testFolder.write(path.join('subfolder', 'fileB.md'), '# fileB')
50-
51-
const fsRead = new FsRead({ path: testFolder.path, readRange: [1] })
52-
await fsRead.validate()
53-
const result = await fsRead.invoke(process.stdout)
54-
55-
const lines = result.output.content.split('\n')
56-
const hasFileA = lines.some((line: string | string[]) => line.includes('- ') && line.includes('fileA.txt'))
57-
const hasSubfolder = lines.some((line: string | string[]) => line.includes('d ') && line.includes('subfolder'))
58-
59-
assert.ok(hasFileA, 'Should list fileA.txt in the directory output')
60-
assert.ok(hasSubfolder, 'Should list the subfolder in the directory output')
61-
})
62-
6346
it('throws error if path does not exist', async () => {
6447
const missingPath = path.join(testFolder.path, 'no_such_file.txt')
6548
const fsRead = new FsRead({ path: missingPath })
@@ -94,29 +77,4 @@ describe('FsRead Tool', () => {
9477
assert.strictEqual(result.output.kind, 'text')
9578
assert.strictEqual(result.output.content, '')
9679
})
97-
98-
it('expands ~ path', async () => {
99-
const fsRead = new FsRead({ path: '~' })
100-
await fsRead.validate()
101-
const result = await fsRead.invoke(process.stdout)
102-
103-
assert.strictEqual(result.output.kind, 'text')
104-
assert.ok(result.output.content.length > 0)
105-
})
106-
107-
it('resolves relative path', async () => {
108-
await testFolder.mkdir('relTest')
109-
const filePath = path.join('relTest', 'relFile.txt')
110-
const content = 'Hello from a relative file!'
111-
await testFolder.write(filePath, content)
112-
113-
const relativePath = path.relative(process.cwd(), path.join(testFolder.path, filePath))
114-
115-
const fsRead = new FsRead({ path: relativePath })
116-
await fsRead.validate()
117-
const result = await fsRead.invoke(process.stdout)
118-
119-
assert.strictEqual(result.output.kind, 'text')
120-
assert.strictEqual(result.output.content, content)
121-
})
12280
})

0 commit comments

Comments
 (0)