Skip to content

Commit 91bfef8

Browse files
authored
feat: add file search tool (#1103)
1 parent 40098dd commit 91bfef8

File tree

4 files changed

+368
-5
lines changed

4 files changed

+368
-5
lines changed

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ import { ListDirectory, ListDirectoryParams } from './tools/listDirectory'
101101
import { FsWrite, FsWriteParams, getDiffChanges } from './tools/fsWrite'
102102
import { ExecuteBash, ExecuteBashOutput, ExecuteBashParams } from './tools/executeBash'
103103
import { ExplanatoryParams, InvokeOutput, ToolApprovalException } from './tools/toolShared'
104+
import { FileSearch, FileSearchParams } from './tools/fileSearch'
104105
import { GrepSearch, SanitizedRipgrepOutput } from './tools/grepSearch'
105106

106107
type ChatHandlers = Omit<
@@ -573,6 +574,7 @@ export class AgenticChatController implements ChatHandlers {
573574
switch (toolUse.name) {
574575
case 'fsRead':
575576
case 'listDirectory':
577+
case 'fileSearch':
576578
case 'grepSearch':
577579
case 'fsWrite':
578580
case 'executeBash': {
@@ -581,6 +583,7 @@ export class AgenticChatController implements ChatHandlers {
581583
listDirectory: { Tool: ListDirectory },
582584
fsWrite: { Tool: FsWrite },
583585
executeBash: { Tool: ExecuteBash },
586+
fileSearch: { Tool: FileSearch },
584587
grepSearch: { Tool: GrepSearch },
585588
}
586589

@@ -615,8 +618,8 @@ export class AgenticChatController implements ChatHandlers {
615618
await chatResultStream.writeResultBlock({ messageId: loadingMessageId, type: 'answer' })
616619
this.#features.chat.sendChatUpdate({ tabId, state: { inProgress: true } })
617620

618-
if (['fsRead', 'listDirectory'].includes(toolUse.name)) {
619-
const initialListDirResult = this.#processReadOrList(toolUse, chatResultStream)
621+
if (['fsRead', 'listDirectory', 'fileSearch'].includes(toolUse.name)) {
622+
const initialListDirResult = this.#processReadOrListOrSearch(toolUse, chatResultStream)
620623
if (initialListDirResult) {
621624
await chatResultStream.writeResultBlock(initialListDirResult)
622625
}
@@ -658,6 +661,8 @@ export class AgenticChatController implements ChatHandlers {
658661
switch (toolUse.name) {
659662
case 'fsRead':
660663
case 'listDirectory':
664+
case 'fileSearch':
665+
// no need to write tool result for listDir,fsRead,fileSearch into chat stream
661666
case 'executeBash':
662667
// no need to write tool result for listDir and fsRead into chat stream
663668
// executeBash will stream the output instead of waiting until the end
@@ -898,7 +903,7 @@ export class AgenticChatController implements ChatHandlers {
898903
}
899904
}
900905

901-
#processReadOrList(toolUse: ToolUse, chatResultStream: AgenticChatResultStream): ChatMessage | undefined {
906+
#processReadOrListOrSearch(toolUse: ToolUse, chatResultStream: AgenticChatResultStream): ChatMessage | undefined {
902907
let messageIdToUpdate = toolUse.toolUseId!
903908
const currentId = chatResultStream.getMessageIdToUpdateForTool(toolUse.name!)
904909

@@ -908,7 +913,7 @@ export class AgenticChatController implements ChatHandlers {
908913
chatResultStream.setMessageIdToUpdateForTool(toolUse.name!, messageIdToUpdate)
909914
}
910915

911-
const currentPath = (toolUse.input as unknown as FsReadParams | ListDirectoryParams)?.path
916+
const currentPath = (toolUse.input as unknown as FsReadParams | ListDirectoryParams | FileSearchParams)?.path
912917
if (!currentPath) return
913918
const existingPaths = chatResultStream.getMessageOperation(messageIdToUpdate)?.filePaths || []
914919
// Check if path already exists in the list
@@ -932,7 +937,9 @@ export class AgenticChatController implements ChatHandlers {
932937
title =
933938
toolUse.name === 'fsRead'
934939
? `${itemCount} file${itemCount > 1 ? 's' : ''} read`
935-
: `${itemCount} ${itemCount === 1 ? 'directory' : 'directories'} listed`
940+
: toolUse.name === 'fileSearch'
941+
? `${itemCount} ${itemCount === 1 ? 'directory' : 'directories'} searched`
942+
: `${itemCount} ${itemCount === 1 ? 'directory' : 'directories'} listed`
936943
}
937944
const fileDetails: Record<string, FileDetails> = {}
938945
for (const item of filePathsPushed) {
Lines changed: 211 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,211 @@
1+
import * as assert from 'assert'
2+
import { FileSearch } from './fileSearch'
3+
import { testFolder } from '@aws/lsp-core'
4+
import * as path from 'path'
5+
import * as fs from 'fs/promises'
6+
import { TestFeatures } from '@aws/language-server-runtimes/testing'
7+
import { Features } from '@aws/language-server-runtimes/server-interface/server'
8+
9+
describe('FileSearch Tool', () => {
10+
let tempFolder: testFolder.TestFolder
11+
let testFeatures: TestFeatures
12+
13+
before(async () => {
14+
testFeatures = new TestFeatures()
15+
// @ts-ignore does not require all fs operations to be implemented
16+
testFeatures.workspace.fs = {
17+
exists: path =>
18+
fs
19+
.access(path)
20+
.then(() => true)
21+
.catch(() => false),
22+
readdir: path => fs.readdir(path, { withFileTypes: true }),
23+
} as Features['workspace']['fs']
24+
tempFolder = await testFolder.TestFolder.create()
25+
})
26+
27+
after(async () => {
28+
await tempFolder.delete()
29+
})
30+
31+
it('invalidates empty path', async () => {
32+
const fileSearch = new FileSearch(testFeatures)
33+
await assert.rejects(
34+
fileSearch.validate({ path: '', pattern: '.*' }),
35+
/Path cannot be empty/i,
36+
'Expected an error about empty path'
37+
)
38+
})
39+
40+
it('invalidates invalid regex pattern', async () => {
41+
const fileSearch = new FileSearch(testFeatures)
42+
await assert.rejects(
43+
fileSearch.validate({ path: tempFolder.path, pattern: '[' }),
44+
/Invalid regex pattern/i,
45+
'Expected an error about invalid regex pattern'
46+
)
47+
})
48+
49+
it('invalidates negative maxDepth', async () => {
50+
const fileSearch = new FileSearch(testFeatures)
51+
await assert.rejects(
52+
fileSearch.validate({ path: '~', pattern: '.*', maxDepth: -1 }),
53+
/MaxDepth cannot be negative/i,
54+
'Expected an error about negative maxDepth'
55+
)
56+
})
57+
58+
it('searches for files matching pattern', async () => {
59+
await tempFolder.write('fileA.txt', 'fileA content')
60+
await tempFolder.write('fileB.md', '# fileB content')
61+
await tempFolder.write('fileC.js', 'console.log("fileC");')
62+
63+
const fileSearch = new FileSearch(testFeatures)
64+
const result = await fileSearch.invoke({
65+
path: tempFolder.path,
66+
pattern: '\\.txt$',
67+
maxDepth: 0,
68+
})
69+
70+
assert.strictEqual(result.output.kind, 'text')
71+
const lines = result.output.content.split('\n')
72+
const hasFileA = lines.some(line => line.includes('[F] ') && line.includes('fileA.txt'))
73+
const hasFileB = lines.some(line => line.includes('[F] ') && line.includes('fileB.md'))
74+
75+
assert.ok(hasFileA, 'Should find fileA.txt matching the pattern')
76+
assert.ok(!hasFileB, 'Should not find fileB.md as it does not match the pattern')
77+
})
78+
79+
it('searches recursively in subdirectories', async () => {
80+
const subfolder = await tempFolder.nest('subfolder')
81+
await tempFolder.write('fileA.txt', 'fileA content')
82+
await subfolder.write('fileB.txt', 'fileB content')
83+
await subfolder.write('fileC.md', '# fileC content')
84+
85+
const fileSearch = new FileSearch(testFeatures)
86+
const result = await fileSearch.invoke({
87+
path: tempFolder.path,
88+
pattern: '\\.txt$',
89+
})
90+
91+
assert.strictEqual(result.output.kind, 'text')
92+
const lines = result.output.content.split('\n')
93+
const hasFileA = lines.some(line => line.includes('[F] ') && line.includes('fileA.txt'))
94+
const hasFileB = lines.some(line => line.includes('[F] ') && line.includes('fileB.txt'))
95+
const hasFileC = lines.some(line => line.includes('[F] ') && line.includes('fileC.md'))
96+
97+
assert.ok(hasFileA, 'Should find fileA.txt in root directory')
98+
assert.ok(hasFileB, 'Should find fileB.txt in subfolder')
99+
assert.ok(!hasFileC, 'Should not find fileC.md as it does not match the pattern')
100+
})
101+
102+
it('respects maxDepth parameter', async () => {
103+
const subfolder1 = await tempFolder.nest('subfolder1')
104+
const subfolder2 = await subfolder1.nest('subfolder2')
105+
106+
await tempFolder.write('root.txt', 'root content')
107+
await subfolder1.write('level1.txt', 'level1 content')
108+
await subfolder2.write('level2.txt', 'level2 content')
109+
110+
const fileSearch = new FileSearch(testFeatures)
111+
const result = await fileSearch.invoke({
112+
path: tempFolder.path,
113+
pattern: '\\.txt$',
114+
maxDepth: 1,
115+
})
116+
117+
assert.strictEqual(result.output.kind, 'text')
118+
const lines = result.output.content.split('\n')
119+
const hasRootFile = lines.some(line => line.includes('[F] ') && line.includes('root.txt'))
120+
const hasLevel1File = lines.some(line => line.includes('[F] ') && line.includes('level1.txt'))
121+
const hasLevel2File = lines.some(line => line.includes('[F] ') && line.includes('level2.txt'))
122+
123+
assert.ok(hasRootFile, 'Should find root.txt in root directory')
124+
assert.ok(hasLevel1File, 'Should find level1.txt in subfolder1')
125+
assert.ok(!hasLevel2File, 'Should not find level2.txt as it exceeds maxDepth')
126+
})
127+
128+
it('performs case-insensitive search by default', async () => {
129+
await tempFolder.write('FileUpper.txt', 'upper case filename')
130+
await tempFolder.write('fileLower.txt', 'lower case filename')
131+
132+
const fileSearch = new FileSearch(testFeatures)
133+
const result = await fileSearch.invoke({
134+
path: tempFolder.path,
135+
pattern: 'file',
136+
maxDepth: 0,
137+
})
138+
139+
assert.strictEqual(result.output.kind, 'text')
140+
const lines = result.output.content.split('\n')
141+
const hasUpperFile = lines.some(line => line.includes('[F] ') && line.includes('FileUpper.txt'))
142+
const hasLowerFile = lines.some(line => line.includes('[F] ') && line.includes('fileLower.txt'))
143+
144+
assert.ok(hasUpperFile, 'Should find FileUpper.txt with case-insensitive search')
145+
assert.ok(hasLowerFile, 'Should find fileLower.txt with case-insensitive search')
146+
})
147+
148+
it('performs case-sensitive search when specified', async () => {
149+
await tempFolder.write('FileUpper.txt', 'upper case filename')
150+
await tempFolder.write('fileLower.txt', 'lower case filename')
151+
152+
const fileSearch = new FileSearch(testFeatures)
153+
const result = await fileSearch.invoke({
154+
path: tempFolder.path,
155+
pattern: 'file',
156+
maxDepth: 0,
157+
caseSensitive: true,
158+
})
159+
160+
assert.strictEqual(result.output.kind, 'text')
161+
const lines = result.output.content.split('\n')
162+
const hasUpperFile = lines.some(line => line.includes('[F] ') && line.includes('FileUpper.txt'))
163+
const hasLowerFile = lines.some(line => line.includes('[F] ') && line.includes('fileLower.txt'))
164+
165+
assert.ok(!hasUpperFile, 'Should not find FileUpper.txt with case-sensitive search')
166+
assert.ok(hasLowerFile, 'Should find fileLower.txt with case-sensitive search')
167+
})
168+
169+
it('ignores excluded directories', async () => {
170+
const nodeModules = await tempFolder.nest('node_modules')
171+
await tempFolder.write('regular.txt', 'regular content')
172+
await nodeModules.write('excluded.txt', 'excluded content')
173+
174+
const fileSearch = new FileSearch(testFeatures)
175+
const result = await fileSearch.invoke({
176+
path: tempFolder.path,
177+
pattern: '\\.txt$',
178+
})
179+
180+
assert.strictEqual(result.output.kind, 'text')
181+
const lines = result.output.content.split('\n')
182+
const hasRegularFile = lines.some(line => line.includes('[F] ') && line.includes('regular.txt'))
183+
const hasExcludedFile = lines.some(line => line.includes('[F] ') && line.includes('excluded.txt'))
184+
185+
assert.ok(hasRegularFile, 'Should find regular.txt in root directory')
186+
assert.ok(!hasExcludedFile, 'Should not find excluded.txt in node_modules directory')
187+
})
188+
189+
it('throws error if path does not exist', async () => {
190+
const missingPath = path.join(tempFolder.path, 'no_such_directory')
191+
const fileSearch = new FileSearch(testFeatures)
192+
193+
await assert.rejects(
194+
fileSearch.invoke({ path: missingPath, pattern: '.*' }),
195+
/Failed to search directory/i,
196+
'Expected an error about non-existent path'
197+
)
198+
})
199+
200+
it('expands ~ path', async () => {
201+
const fileSearch = new FileSearch(testFeatures)
202+
const result = await fileSearch.invoke({
203+
path: '~',
204+
pattern: '.*',
205+
maxDepth: 0,
206+
})
207+
208+
assert.strictEqual(result.output.kind, 'text')
209+
assert.ok(result.output.content.length > 0)
210+
})
211+
})

0 commit comments

Comments
 (0)