Skip to content

Commit 2e327b1

Browse files
author
linyuan.yang
committed
grep glob 加上超时
1 parent 52cc9e7 commit 2e327b1

2 files changed

Lines changed: 53 additions & 22 deletions

File tree

‎packages/sbot/src/Tools/FileSystem/content/grep.ts‎

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ const MAX_LINE_LENGTH = 2000;
1515
const DEFAULT_MAX_MATCHES = 100;
1616
const MAX_FILE_SIZE = 1 * 1024 * 1024; // 1MB - skip large files in Node.js fallback
1717
const MAX_COLUMNS = 4096;
18+
const DEFAULT_TIMEOUT_SEC = 30;
1819

1920
// ─── 类型 ─────────────────────────────────────────────────────────────────────
2021
interface MatchLine { lineNum: number; text: string; }
@@ -32,6 +33,7 @@ function searchWithRg(
3233
includeHidden: boolean,
3334
fileGlob: string | undefined,
3435
maxMatches: number,
36+
timeoutMs: number,
3537
): Promise<SearchResult> {
3638
return new Promise((resolve, reject) => {
3739
const args = ['--json', `--max-columns=${MAX_COLUMNS}`, ...RG_EXCLUDE_ARGS];
@@ -45,17 +47,21 @@ function searchWithRg(
4547
const fileOrder: string[] = [];
4648
let totalMatches = 0;
4749
let reachedLimit = false;
50+
let timedOut = false;
4851
let killed = false;
4952
let buffer = '';
5053
let stderr = '';
5154

52-
const stop = (limit: boolean) => {
55+
const stop = (limit: boolean, timeout = false) => {
5356
if (killed) return;
5457
killed = true;
5558
if (limit) reachedLimit = true;
59+
if (timeout) timedOut = true;
5660
try { proc.kill('SIGTERM'); } catch { /* ignore */ }
5761
};
5862

63+
const timer = setTimeout(() => stop(true, true), timeoutMs);
64+
5965
const processLine = (line: string): boolean => {
6066
if (!line) return true;
6167
let parsed: any;
@@ -91,13 +97,15 @@ function searchWithRg(
9197
if (buffer.length > 10 * MAX_COLUMNS) buffer = '';
9298
});
9399
proc.stderr.on('data', (d: Buffer) => { stderr += d.toString(); });
94-
proc.on('error', e => { reject(e); });
100+
proc.on('error', e => { clearTimeout(timer); reject(e); });
95101
proc.on('close', code => {
102+
clearTimeout(timer);
96103
if (!killed && buffer) processLine(buffer.replace(/\r$/, ''));
97104
// code 1 = 无匹配；被 kill 时 code/signal 不可靠，按 totalMatches 判断
98105
if (!killed && code !== 0 && code !== 1 && totalMatches === 0) {
99106
return reject(new Error(`ripgrep: ${stderr.trim() || `exit ${code}`}`));
100107
}
108+
if (timedOut) logger.warn(`ripgrep timed out after ${timeoutMs}ms; returning ${totalMatches} matches`);
101109
const results = fileOrder.map(fp => byFile.get(fp)!);
102110
results.sort((a, b) => b.mtime - a.mtime);
103111
resolve({ results, reachedLimit });
@@ -125,15 +133,20 @@ async function searchWithNodeJs(
125133
includeHidden: boolean,
126134
maxFileSize: number,
127135
maxMatches: number,
136+
timeoutMs: number,
128137
): Promise<SearchResult> {
129138
const searchRegex = useRegex ? new RegExp(pattern) : null;
139+
const deadline = Date.now() + timeoutMs;
140+
const expired = () => Date.now() >= deadline;
130141

131142
const allFiles: Array<{ path: string; mtime: number }> = [];
132143
async function walk(d: string): Promise<void> {
144+
if (expired()) return;
133145
let entries;
134146
try { entries = await fsAsync.readdir(d, { withFileTypes: true }); }
135147
catch (e: any) { logger.warn(`Cannot access ${d}: ${e.message}`); return; }
136148
for (const entry of entries) {
149+
if (expired()) return;
137150
if (!includeHidden && entry.name.startsWith('.')) continue;
138151
const full = path.join(d, entry.name);
139152
if (entry.isDirectory()) {
@@ -155,6 +168,7 @@ async function searchWithNodeJs(
155168
let reachedLimit = false;
156169

157170
outer: for (const { path: fp, mtime } of allFiles) {
171+
if (expired()) { reachedLimit = true; break; }
158172
if (await isBinaryAsync(fp)) continue;
159173
try {
160174
const fileMatches: MatchLine[] = [];
@@ -176,6 +190,8 @@ async function searchWithNodeJs(
176190
} catch { /* skip unreadable */ }
177191
}
178192

193+
if (expired() && !reachedLimit) reachedLimit = true;
194+
if (Date.now() >= deadline) logger.warn(`Node.js grep timed out after ${timeoutMs}ms; returning ${totalMatches} matches`);
179195
return { results, reachedLimit };
180196
}
181197

@@ -211,17 +227,19 @@ export function createGrepFilesTool(): StructuredToolInterface {
211227
useRegex: z.boolean().optional().default(false).describe('Treat pattern as a regex, default false (literal search)'),
212228
includeHidden: z.boolean().optional().default(false).describe('Include hidden files (starting with .), default false'),
213229
maxMatches: z.number().optional().default(DEFAULT_MAX_MATCHES).describe('Maximum number of matching lines across all files, default 100'),
230+
timeoutSec: z.number().positive().optional().default(DEFAULT_TIMEOUT_SEC).describe(`Search timeout in seconds; on timeout returns partial results marked truncated. Default ${DEFAULT_TIMEOUT_SEC}`),
214231
}) as any,
215-
func: async ({ path: searchPath, pattern, glob, useRegex = false, includeHidden = false, maxMatches = DEFAULT_MAX_MATCHES }: any): Promise<MCPToolResult> => {
232+
func: async ({ path: searchPath, pattern, glob, useRegex = false, includeHidden = false, maxMatches = DEFAULT_MAX_MATCHES, timeoutSec = DEFAULT_TIMEOUT_SEC }: any): Promise<MCPToolResult> => {
216233
try {
217234
const abs = checkDir(searchPath);
235+
const timeoutMs = Math.round(timeoutSec * 1000);
218236
let result: SearchResult;
219237

220238
if (await checkRg()) {
221-
result = await searchWithRg(abs, pattern, useRegex, includeHidden, glob, maxMatches);
239+
result = await searchWithRg(abs, pattern, useRegex, includeHidden, glob, maxMatches, timeoutMs);
222240
} else {
223241
const fileRegex = globToRegex(glob ?? '*');
224-
result = await searchWithNodeJs(abs, pattern, fileRegex, useRegex, includeHidden, MAX_FILE_SIZE, maxMatches);
242+
result = await searchWithNodeJs(abs, pattern, fileRegex, useRegex, includeHidden, MAX_FILE_SIZE, maxMatches, timeoutMs);
225243
}
226244

227245
if (result.results.length === 0) return createSuccessResult(createTextContent('No matches found'));

‎packages/sbot/src/Tools/FileSystem/operations/glob.ts‎

Lines changed: 30 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@ import { loadPrompt } from '../../../Core/PromptLoader';
1111
const logger = LoggerService.getLogger('Tools/FileSystem/operations/glob.ts');
1212

1313
const LIMIT = 100;
14+
const DEFAULT_TIMEOUT_SEC = 30;
15+
16+
interface GlobSearchResult { files: Array<{ path: string; mtime: number }>; truncated: boolean; }
1417

1518
// rg 默认排除目录（gitignore 语法）
1619
const RG_EXCLUDE_ARGS = [...EXCLUDE_DIRS].map(d => `--glob=!${d}`);
@@ -43,7 +46,7 @@ function globToRegex(pattern: string): RegExp {
4346
}
4447

4548
// ─── ripgrep 搜索（流式 + 达到上限即终止）──────────────────────────────────
46-
function searchWithRg(dir: string, pattern: string, includeHidden: boolean): Promise<Array<{ path: string; mtime: number }>> {
49+
function searchWithRg(dir: string, pattern: string, includeHidden: boolean, timeoutMs: number): Promise<GlobSearchResult> {
4750
return new Promise((resolve, reject) => {
4851
const args = ['--files', ...RG_EXCLUDE_ARGS, `--iglob=${pattern}`];
4952
if (includeHidden) args.push('--hidden');
@@ -52,30 +55,35 @@ function searchWithRg(dir: string, pattern: string, includeHidden: boolean): Pro
5255
const proc = spawn('rg', args);
5356
const files: Array<{ path: string; mtime: number }> = [];
5457
let killed = false;
58+
let timedOut = false;
5559
let buffer = '';
5660
let stderr = '';
5761
let pending = 0;
5862
let submitted = 0;
5963
let closed = false;
6064
let exitCode: number | null = null;
6165

62-
const stop = () => {
66+
const stop = (timeout = false) => {
6367
if (killed) return;
6468
killed = true;
69+
if (timeout) timedOut = true;
6570
try { proc.kill('SIGTERM'); } catch { /* ignore */ }
6671
};
6772

73+
const timer = setTimeout(() => stop(true), timeoutMs);
74+
6875
const tryResolve = () => {
6976
if (closed && pending === 0) {
7077
if (!killed && exitCode !== 0 && exitCode !== 1 && files.length === 0) {
7178
return reject(new Error(`ripgrep: ${stderr.trim() || `exit ${exitCode}`}`));
7279
}
73-
resolve(files);
80+
if (timedOut) logger.warn(`ripgrep glob timed out after ${timeoutMs}ms; returning ${files.length} files`);
81+
resolve({ files, truncated: timedOut });
7482
}
7583
};
7684

7785
const processLine = (line: string): boolean => {
78-
if (!line) return true;
86+
if (!line || killed) return !killed;
7987
if (submitted >= LIMIT) return false;
8088
submitted++;
8189
pending++;
@@ -97,8 +105,9 @@ function searchWithRg(dir: string, pattern: string, includeHidden: boolean): Pro
97105
}
98106
});
99107
proc.stderr.on('data', (d: Buffer) => { stderr += d.toString(); });
100-
proc.on('error', e => reject(e));
108+
proc.on('error', e => { clearTimeout(timer); reject(e); });
101109
proc.on('close', code => {
110+
clearTimeout(timer);
102111
if (!killed && buffer) processLine(buffer.replace(/\r$/, ''));
103112
closed = true;
104113
exitCode = code;
@@ -108,17 +117,20 @@ function searchWithRg(dir: string, pattern: string, includeHidden: boolean): Pro
108117
}
109118

110119
// ─── Node.js fallback（全异步，避免阻塞事件循环）────────────────────────────
111-
async function searchWithNodeJs(dir: string, pattern: string, includeHidden: boolean): Promise<Array<{ path: string; mtime: number }>> {
120+
async function searchWithNodeJs(dir: string, pattern: string, includeHidden: boolean, timeoutMs: number): Promise<GlobSearchResult> {
112121
const useFullPath = hasPathPattern(pattern);
113122
const regex = globToRegex(pattern);
114123
const results: Array<{ path: string; mtime: number }> = [];
124+
const deadline = Date.now() + timeoutMs;
125+
const expired = () => Date.now() >= deadline;
115126

116127
async function walk(d: string): Promise<boolean> {
128+
if (expired()) return false;
117129
let entries;
118130
try { entries = await fsAsync.readdir(d, { withFileTypes: true }); }
119131
catch (e: any) { logger.warn(`Cannot access ${d}: ${e.message}`); return true; }
120132
for (const entry of entries) {
121-
if (results.length >= LIMIT) return false;
133+
if (results.length >= LIMIT || expired()) return false;
122134
if (!includeHidden && entry.name.startsWith('.')) continue;
123135
const full = path.join(d, entry.name);
124136
if (entry.isDirectory()) {
@@ -138,7 +150,9 @@ async function searchWithNodeJs(dir: string, pattern: string, includeHidden: boo
138150
return true;
139151
}
140152
await walk(dir);
141-
return results;
153+
const timedOut = expired();
154+
if (timedOut) logger.warn(`Node.js glob timed out after ${timeoutMs}ms; returning ${results.length} files`);
155+
return { files: results, truncated: timedOut };
142156
}
143157

144158
// ─── Tool 定义 ────────────────────────────────────────────────────────────────
@@ -152,21 +166,20 @@ export function createGlobTool(): StructuredToolInterface {
152166
pattern: z.string().describe('Glob pattern, e.g. **/*.ts, src/**/*.test.js, *.json'),
153167
path: z.string().describe('Absolute path of the directory to search'),
154168
includeHidden: z.boolean().optional().default(false).describe('Include hidden files, default false'),
169+
timeoutSec: z.number().positive().optional().default(DEFAULT_TIMEOUT_SEC).describe(`Search timeout in seconds; on timeout returns partial results marked truncated. Default ${DEFAULT_TIMEOUT_SEC}`),
155170
}) as any,
156-
func: async ({ pattern, path: searchPath, includeHidden = false }: any): Promise<MCPToolResult> => {
171+
func: async ({ pattern, path: searchPath, includeHidden = false, timeoutSec = DEFAULT_TIMEOUT_SEC }: any): Promise<MCPToolResult> => {
157172
try {
158173
const abs = checkDir(searchPath);
159-
let files: Array<{ path: string; mtime: number }>;
160-
161-
if (await checkRg()) {
162-
files = await searchWithRg(abs, pattern, includeHidden);
163-
} else {
164-
files = await searchWithNodeJs(abs, pattern, includeHidden);
165-
}
174+
const timeoutMs = Math.round(timeoutSec * 1000);
175+
const search = await checkRg()
176+
? await searchWithRg(abs, pattern, includeHidden, timeoutMs)
177+
: await searchWithNodeJs(abs, pattern, includeHidden, timeoutMs);
178+
let files = search.files;
179+
let truncated = search.truncated;
166180

167181
files.sort((a, b) => b.mtime - a.mtime);
168182

169-
let truncated = false;
170183
if (files.length > LIMIT) {
171184
files = files.slice(0, LIMIT);
172185
truncated = true;

0 commit comments

Comments
 (0)