rig_tools/
search.rs

1//! The `search` tool, and [`find`], the search it runs, as data.
2
3use std::fmt;
4use std::path::{Path, PathBuf};
5use std::sync::Arc;
6
7use regex::{Regex, RegexBuilder};
8use rig_core::completion::ToolAnnotations;
9use rig_core::tool::{
10    ContextValue, FileRead, FilesRead, Scope, Tool, ToolContext, ToolExecutionError, Touches,
11    args_schema, content_hash,
12};
13use schemars::JsonSchema;
14use serde::{Deserialize, Serialize, Serializer};
15
16use crate::fs::listed;
17use crate::workspace::{Listing, Local, Workspace};
18use crate::{MAX_FILE_BYTES, blocking, clip};
19
20/// Most matches a search returns.
21pub const MAX_MATCHES: usize = 200;
22/// Most files a search for file names ([`SearchArgs::files`]) returns.
23pub const MAX_FILES: usize = 1000;
24const MAX_LINE: usize = 300;
25
26/// Searches file contents with a regular expression, honouring
27/// `.gitignore`, or lists the files with a match, or every file. Shows each
28/// path as `read` takes it from the same base: below the base when the
29/// searched path is relative or not given, absolute when it is absolute.
30/// Publishes what it found ([`SearchHits`]) and the files it found matches
31/// in ([`FilesRead`]), by their [resolved](Workspace) paths.
32#[derive(Clone, Debug)]
33pub struct Search {
34    /// Where the files are searched.
35    pub workspace: Arc<dyn Workspace>,
36    /// The directory searched, and a relative path is searched from;
37    /// `None` is the workspace's working directory. Hits of a relative
38    /// path are shown, and published, by their paths below it.
39    pub base: Option<PathBuf>,
40    /// Whether a path that resolves outside `base` is refused, before
41    /// anything is searched.
42    pub confine: bool,
43}
44
45impl Default for Search {
46    /// Searches this machine's files, from this process's working directory.
47    fn default() -> Self {
48        Self {
49            workspace: Arc::new(Local),
50            base: None,
51            confine: false,
52        }
53    }
54}
55
56/// Arguments of [`Search`].
57#[derive(Clone, Debug, PartialEq, Eq, Deserialize, JsonSchema)]
58#[serde(deny_unknown_fields)]
59pub struct SearchArgs {
60    /// The regular expression; with `files`, empty lists every file.
61    pub pattern: String,
62    /// The directory or file to search; default the current directory.
63    pub path: Option<String>,
64    /// Only files matching this glob, such as `*.rs`.
65    pub glob: Option<String>,
66    /// Also hidden files and directories (named with a leading `.`), such
67    /// as `.github/`; files git ignores stay out.
68    #[serde(default)]
69    pub hidden: bool,
70    /// Only the paths of the files with a match, one per line.
71    #[serde(default)]
72    pub files: bool,
73}
74
75impl SearchArgs {
76    /// Search the current directory for `pattern`.
77    pub fn new(pattern: impl Into<String>) -> Self {
78        Self {
79            pattern: pattern.into(),
80            path: None,
81            glob: None,
82            hidden: false,
83            files: false,
84        }
85    }
86
87    /// The same search, returning only the paths of the files with a
88    /// match; with an empty pattern, every file.
89    pub fn files(mut self) -> Self {
90        self.files = true;
91        self
92    }
93
94    /// The same search, in `path`.
95    pub fn in_path(mut self, path: impl Into<String>) -> Self {
96        self.path = Some(path.into());
97        self
98    }
99
100    /// The same search, only in files matching `glob`.
101    pub fn glob(mut self, glob: impl Into<String>) -> Self {
102        self.glob = Some(glob.into());
103        self
104    }
105
106    /// The same search, in hidden files too.
107    pub fn hidden(mut self) -> Self {
108        self.hidden = true;
109        self
110    }
111
112    /// Which files it searches.
113    pub fn listing(&self) -> Listing {
114        match self.hidden {
115            true => Listing::Hidden,
116            false => Listing::Visible,
117        }
118        .for_glob(self.glob.as_deref())
119    }
120}
121
122/// One line a search matched.
123#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
124pub struct Hit {
125    /// The file, as the walk from the searched path reached it, without a
126    /// leading `./`.
127    #[serde(serialize_with = "lossy")]
128    pub path: PathBuf,
129    /// The line's number, from 1.
130    pub line: usize,
131    /// The line, cut at 300 characters.
132    pub text: String,
133}
134
135/// `path:line: text`, as [`Search`] shows a hit.
136impl fmt::Display for Hit {
137    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
138        write!(f, "{}:{}: {}", self.path.display(), self.line, self.text)
139    }
140}
141
142/// `path` as text, any bytes that are not UTF-8 replaced.
143fn lossy<S: Serializer>(path: &Path, serializer: S) -> Result<S::Ok, S::Error> {
144    serializer.serialize_str(&path.to_string_lossy())
145}
146
147/// What a search found.
148#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
149pub struct Found {
150    /// The matched lines, file by file in walk order; a search for file
151    /// names ([`SearchArgs::files`]) has none.
152    pub hits: Vec<Hit>,
153    /// The files with a match, of a search for file names, in walk order,
154    /// as the walk from the searched path reached them, without a leading
155    /// `./`.
156    #[serde(default, skip_serializing_if = "Vec::is_empty")]
157    pub names: Vec<PathBuf>,
158    /// Whether it stopped at its [`MAX_MATCHES`] matches, or
159    /// [`MAX_FILES`] files, so there may be more.
160    pub stopped: bool,
161    /// The files with hits, each once, in walk order, by their
162    /// [resolved](Workspace) paths, with the hash of the bytes
163    /// searched.
164    #[serde(default)]
165    pub files: Vec<FileRead>,
166}
167
168/// What a [`Search`] call found, as it published it beside the text.
169#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
170#[serde(transparent)]
171pub struct SearchHits(pub Found);
172
173impl ContextValue for SearchHits {
174    const KEY: &'static str = "search_hits";
175}
176
177impl Tool for Search {
178    const NAME: &'static str = "search";
179    type Args = SearchArgs;
180    type Output = String;
181    type Error = ToolExecutionError;
182
183    fn description(&self) -> String {
184        format!(
185            "Search file contents with a regular expression (Rust regex syntax). Skips files \
186             ignored by git, target directories and, unless hidden is set or the glob names \
187             one, hidden files. Returns `path:line: text`, at most {MAX_MATCHES} matches; with \
188             files, only the paths of the files with a match, one per line, at most \
189             {MAX_FILES}, and with files and an empty pattern every file. Each path is one \
190             `read` takes: relative to the current directory when path is relative or not \
191             given, absolute when path is absolute. Use it to find code, not `grep` or `rg` in \
192             `shell`, and with files and an empty pattern to list files, not `ls` or `find`."
193        )
194    }
195
196    fn parameters(&self) -> serde_json::Value {
197        args_schema::<SearchArgs>()
198    }
199
200    fn annotations(&self) -> ToolAnnotations {
201        ToolAnnotations::default()
202            .with_read_only(true)
203            .with_open_world(false)
204    }
205
206    fn footprint(&self, args: &SearchArgs) -> Touches {
207        Touches::Reads(Scope::Paths(vec![
208            self.workspace
209                .resolve(self.base.as_deref(), args.path.as_deref().unwrap_or(".")),
210        ]))
211    }
212
213    async fn call(
214        &self,
215        context: &mut ToolContext,
216        mut args: SearchArgs,
217    ) -> Result<String, ToolExecutionError> {
218        let path = args.path.as_deref().unwrap_or(".");
219        self.workspace
220            .place(self.base.as_deref(), self.confine, path)?;
221        let base = self.base.clone().filter(|_| Path::new(path).is_relative());
222        if let Some(base) = &base {
223            let joined = self.workspace.join(base, Path::new(path));
224            args.path = Some(joined.to_string_lossy().into_owned());
225        }
226        let workspace = self.workspace.clone();
227        let mut found = blocking(move || find(&*workspace, &args)).await?;
228        if let Some(base) = &base {
229            let paths = found.hits.iter_mut().map(|hit| &mut hit.path);
230            for path in paths.chain(&mut found.names) {
231                if let Ok(below) = path.strip_prefix(base) {
232                    *path = below.to_path_buf();
233                }
234            }
235        }
236        if !found.files.is_empty() {
237            context.insert_result(FilesRead(found.files.clone()))?;
238        }
239        let text = render(&found);
240        context.insert_result(SearchHits(found))?;
241        Ok(text)
242    }
243}
244
245/// `found` as [`Search`] shows it: a `path:line: text` line per hit, or a
246/// path per file of a search for file names, and a last line when it
247/// stopped at [`MAX_MATCHES`] or [`MAX_FILES`].
248pub fn render(found: &Found) -> String {
249    let (mut lines, stop): (Vec<String>, _) = match found.names.is_empty() {
250        true => (
251            found.hits.iter().map(Hit::to_string).collect(),
252            format!("[stopped at {MAX_MATCHES} matches]"),
253        ),
254        false => (
255            found
256                .names
257                .iter()
258                .map(|name| name.display().to_string())
259                .collect(),
260            format!("[stopped at {MAX_FILES} files; narrow the search with path or glob]"),
261        ),
262    };
263    if lines.is_empty() {
264        return "No matches.".to_owned();
265    }
266    if found.stopped {
267        lines.push(stop);
268    }
269    lines.join("\n")
270}
271
272/// Searches as [`Search`] does, in `workspace`, and returns the matches as
273/// data: the [files](Workspace::files) under the searched path that its
274/// [listing](SearchArgs::listing) keeps, each
275/// of at most [`MAX_FILE_BYTES`] and read as UTF-8; a search for file
276/// names with an empty pattern reads none and keeps them all. Blocks: run
277/// it on a thread of its own ([`blocking`]).
278///
279/// # Errors
280///
281/// When the pattern or the glob is malformed, the path does not exist, or
282/// the walk fails.
283pub fn find(workspace: &dyn Workspace, args: &SearchArgs) -> Result<Found, ToolExecutionError> {
284    let pattern = Regex::new(&args.pattern)
285        .map_err(|error| ToolExecutionError::invalid_args(format!("bad pattern: {error}")))?;
286    let whole = whole_text(&args.pattern);
287    let root = args.path.as_deref().unwrap_or(".");
288    let files = listed(
289        workspace,
290        Path::new(root),
291        args.glob.as_deref(),
292        args.listing(),
293    )?;
294    let mut found = Found::default();
295    if args.files && args.pattern.is_empty() {
296        for file in files {
297            found.names.push(walked(&file).to_path_buf());
298            if found.names.len() == MAX_FILES {
299                found.stopped = true;
300                break;
301            }
302        }
303        return Ok(found);
304    }
305    let read = files.chunks(100).flat_map(|chunk| {
306        let read = workspace.read_files(chunk, MAX_FILE_BYTES);
307        chunk.iter().zip(read)
308    });
309    for (file, bytes) in read {
310        let Some(text) = bytes.and_then(|bytes| String::from_utf8(bytes).ok()) else {
311            continue;
312        };
313        if let Some(whole) = &whole
314            && !lone_cr(&text)
315            && !whole.is_match(&text)
316        {
317            continue;
318        }
319        let path = walked(file);
320        let mut lines = text
321            .lines()
322            .enumerate()
323            .filter(|(_, line)| pattern.is_match(line))
324            .peekable();
325        if lines.peek().is_none() {
326            continue;
327        }
328        found.files.push(FileRead::new(
329            workspace.canonical(path).to_string_lossy(),
330            content_hash(text.as_bytes()),
331        ));
332        if args.files {
333            found.names.push(path.to_path_buf());
334            if found.names.len() == MAX_FILES {
335                found.stopped = true;
336                return Ok(found);
337            }
338            continue;
339        }
340        for (number, line) in lines {
341            found.hits.push(Hit {
342                path: path.to_path_buf(),
343                line: number + 1,
344                text: clip(line, MAX_LINE).to_owned(),
345            });
346            if found.hits.len() == MAX_MATCHES {
347                found.stopped = true;
348                return Ok(found);
349            }
350        }
351    }
352    Ok(found)
353}
354
355/// `pattern` over a whole text, matching it wherever it matches one of its
356/// lines, so a file it does not match has no line to search: `^` and `$`
357/// match at each line end, `\n` or `\r\n`. `None` when the pattern anchors
358/// to the ends of the text (`\A`, `\z`) or may turn a flag off, such as
359/// `(?-m)`, so its lines' ends would not count.
360fn whole_text(pattern: &str) -> Option<Regex> {
361    let groups = pattern.split("(?").skip(1);
362    let mut flags = groups.filter_map(|group| group.split([')', ':']).next());
363    if pattern.contains("\\A") || pattern.contains("\\z") || flags.any(|it| it.contains('-')) {
364        return None;
365    }
366    RegexBuilder::new(pattern)
367        .multi_line(true)
368        .crlf(true)
369        .build()
370        .ok()
371}
372
373/// Whether `text` has a `\r` that does not end a line: [`whole_text`]'s
374/// `.` does not match it, though a line's does.
375fn lone_cr(text: &str) -> bool {
376    text.contains('\r') && text.matches('\r').count() != text.matches("\r\n").count()
377}
378
379/// `path` as the walk reached it, without a leading `./`.
380fn walked(path: &Path) -> &Path {
381    path.strip_prefix("./").unwrap_or(path)
382}
383
384#[cfg(test)]
385mod tests;
rig_tools/search.rs

rig_tools/
search.rs
