molo_coding/coding/
search.rs

1use crate::harness::{ExecutionPolicy, NetworkPolicy, SandboxPolicy};
2use crate::{RunContext, RunMetadata};
3use async_trait::async_trait;
4use regex::Regex;
5use serde::{Deserialize, Serialize};
6use std::sync::Arc;
7use std::time::Duration;
8
9use super::command::{CommandExecutor, CommandOutputLimit, CommandRequest};
10use super::workspace::{FileBody, FileReadOptions, ListFilesQuery, Workspace, WorkspacePath};
11
12/// Repository search mode.
13#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
14#[non_exhaustive]
15pub enum SearchMode {
16    /// Treat the query as a literal substring.
17    #[default]
18    Literal,
19    /// Treat the query as a regular expression.
20    Regex,
21}
22
23/// Repository search request.
24#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
25pub struct RepoSearchRequest {
26    /// Query string or regex pattern.
27    pub query: String,
28    /// Paths to search. Empty means workspace root.
29    pub paths: Vec<WorkspacePath>,
30    /// Search mode.
31    pub mode: SearchMode,
32    /// Maximum matches returned.
33    pub max_matches: usize,
34    /// Context lines around each match.
35    pub context_lines: usize,
36    /// Whether hidden paths are included.
37    pub include_hidden: bool,
38    /// Whether simple `.gitignore` rules are respected.
39    pub respect_gitignore: bool,
40}
41
42impl RepoSearchRequest {
43    /// Constructs a literal search request.
44    pub fn literal(query: impl Into<String>) -> Self {
45        Self {
46            query: query.into(),
47            paths: Vec::new(),
48            mode: SearchMode::Literal,
49            max_matches: 100,
50            context_lines: 0,
51            include_hidden: false,
52            respect_gitignore: true,
53        }
54    }
55}
56
57/// One repository search match.
58#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
59pub struct SearchMatch {
60    /// Matched file.
61    pub path: WorkspacePath,
62    /// 1-based line number.
63    pub line: usize,
64    /// 1-based column number when known.
65    pub column: Option<usize>,
66    /// Matched line text.
67    pub line_text: String,
68    /// Context lines before the match.
69    pub before: Vec<String>,
70    /// Context lines after the match.
71    pub after: Vec<String>,
72}
73
74/// Repository search results.
75#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
76pub struct RepoSearchResults {
77    /// Matches in deterministic order.
78    pub matches: Vec<SearchMatch>,
79    /// Whether matching stopped at the output budget.
80    pub truncated: bool,
81    /// Warnings such as missing external search binaries.
82    pub warnings: Vec<String>,
83    /// Host-owned metadata.
84    pub metadata: RunMetadata,
85}
86
87/// Search implementation for repositories.
88#[async_trait]
89pub trait RepoSearcher: Send + Sync {
90    /// Searches repository content.
91    async fn search(
92        &self,
93        request: RepoSearchRequest,
94        context: &RunContext,
95    ) -> Result<RepoSearchResults, SearchError>;
96}
97
98/// In-process fallback searcher based on [`Workspace`] reads.
99#[derive(Debug, Clone)]
100pub struct WorkspaceSearcher<W> {
101    workspace: W,
102    max_file_bytes: usize,
103}
104
105impl<W> WorkspaceSearcher<W> {
106    /// Constructs a workspace searcher.
107    pub fn new(workspace: W) -> Self {
108        Self {
109            workspace,
110            max_file_bytes: 256 * 1024,
111        }
112    }
113
114    /// Sets the per-file read budget.
115    pub fn with_max_file_bytes(mut self, max_file_bytes: usize) -> Self {
116        self.max_file_bytes = max_file_bytes;
117        self
118    }
119}
120
121#[async_trait]
122impl<W> RepoSearcher for WorkspaceSearcher<W>
123where
124    W: Workspace,
125{
126    async fn search(
127        &self,
128        request: RepoSearchRequest,
129        _context: &RunContext,
130    ) -> Result<RepoSearchResults, SearchError> {
131        if request.query.is_empty() {
132            return Err(SearchError::InvalidQuery {
133                message: "search query must not be empty".to_string(),
134            });
135        }
136        let matcher = Matcher::new(&request.query, request.mode)?;
137        let paths = if request.paths.is_empty() {
138            vec![WorkspacePath::root()]
139        } else {
140            request.paths.clone()
141        };
142        let mut matches = Vec::new();
143        let mut truncated = false;
144        for path in paths {
145            let entries = self
146                .workspace
147                .list_files(ListFilesQuery {
148                    path,
149                    recursive: true,
150                    max_entries: Some(20_000),
151                    include_hidden: request.include_hidden,
152                    respect_gitignore: request.respect_gitignore,
153                })
154                .await
155                .map_err(|error| SearchError::Workspace {
156                    message: error.to_string(),
157                })?;
158            for entry in entries {
159                if matches.len() >= request.max_matches {
160                    truncated = true;
161                    break;
162                }
163                if entry.kind != super::workspace::ResolvedPathKind::File {
164                    continue;
165                }
166                let content = self
167                    .workspace
168                    .read_file(
169                        &entry.path,
170                        FileReadOptions {
171                            max_bytes: Some(self.max_file_bytes),
172                            include_binary: false,
173                        },
174                    )
175                    .await
176                    .map_err(|error| SearchError::Workspace {
177                        message: error.to_string(),
178                    })?;
179                let FileBody::Text { text, .. } = content.body else {
180                    continue;
181                };
182                let file_matches = search_text(&entry.path, &text, &matcher, request.context_lines);
183                for mat in file_matches {
184                    if matches.len() >= request.max_matches {
185                        truncated = true;
186                        break;
187                    }
188                    matches.push(mat);
189                }
190            }
191        }
192        Ok(RepoSearchResults {
193            matches,
194            truncated,
195            warnings: Vec::new(),
196            metadata: RunMetadata::new(),
197        })
198    }
199}
200
201/// Repository searcher that invokes `rg` through a [`CommandExecutor`].
202#[derive(Clone)]
203pub struct RipgrepSearcher<C> {
204    commands: Arc<C>,
205    timeout: Duration,
206}
207
208impl<C> std::fmt::Debug for RipgrepSearcher<C> {
209    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
210        f.debug_struct("RipgrepSearcher")
211            .field("timeout", &self.timeout)
212            .finish_non_exhaustive()
213    }
214}
215
216impl<C> RipgrepSearcher<C> {
217    /// Constructs a ripgrep searcher from a command executor.
218    pub fn new(commands: C) -> Self {
219        Self {
220            commands: Arc::new(commands),
221            timeout: Duration::from_secs(10),
222        }
223    }
224
225    /// Sets the ripgrep command timeout.
226    pub fn with_timeout(mut self, timeout: Duration) -> Self {
227        self.timeout = timeout;
228        self
229    }
230}
231
232#[async_trait]
233impl<C> RepoSearcher for RipgrepSearcher<C>
234where
235    C: CommandExecutor + 'static,
236{
237    async fn search(
238        &self,
239        request: RepoSearchRequest,
240        context: &RunContext,
241    ) -> Result<RepoSearchResults, SearchError> {
242        if request.mode == SearchMode::Regex {
243            Regex::new(&request.query).map_err(|error| SearchError::InvalidQuery {
244                message: format!("invalid regex: {error}"),
245            })?;
246        }
247        let mut argv = vec![
248            "rg".to_string(),
249            "--line-number".to_string(),
250            "--column".to_string(),
251            "--no-heading".to_string(),
252            "--color".to_string(),
253            "never".to_string(),
254        ];
255        if request.mode == SearchMode::Literal {
256            argv.push("--fixed-strings".to_string());
257        }
258        if request.include_hidden {
259            argv.push("--hidden".to_string());
260        }
261        if !request.respect_gitignore {
262            argv.push("--no-ignore".to_string());
263        }
264        argv.push(request.query.clone());
265        for path in &request.paths {
266            argv.push(path.display());
267        }
268        let mut command = CommandRequest::new(argv);
269        command.timeout = Some(self.timeout);
270        command.output_limit = CommandOutputLimit {
271            stdout_bytes: 512 * 1024,
272            stderr_bytes: 64 * 1024,
273        };
274        let output = self
275            .commands
276            .execute(
277                command,
278                &ExecutionPolicy::new(SandboxPolicy::ReadOnly, NetworkPolicy::Deny)
279                    .with_timeout(Some(self.timeout)),
280                context,
281            )
282            .await
283            .map_err(|error| SearchError::Command {
284                message: error.to_string(),
285            })?;
286        let mut warnings = Vec::new();
287        if !output.stderr.text.trim().is_empty() {
288            warnings.push(output.stderr.text);
289        }
290        let mut matches = Vec::new();
291        for line in output.stdout.text.lines() {
292            if matches.len() >= request.max_matches {
293                break;
294            }
295            if let Some(mat) = parse_rg_line(line) {
296                matches.push(mat);
297            }
298        }
299        let hit_match_limit = matches.len() >= request.max_matches;
300        Ok(RepoSearchResults {
301            matches,
302            truncated: output.truncated || output.stdout.truncated || hit_match_limit,
303            warnings,
304            metadata: output.metadata,
305        })
306    }
307}
308
309enum Matcher {
310    Literal(String),
311    Regex(Regex),
312}
313
314impl Matcher {
315    fn new(query: &str, mode: SearchMode) -> Result<Self, SearchError> {
316        match mode {
317            SearchMode::Literal => Ok(Self::Literal(query.to_string())),
318            SearchMode::Regex => {
319                Regex::new(query)
320                    .map(Self::Regex)
321                    .map_err(|error| SearchError::InvalidQuery {
322                        message: format!("invalid regex: {error}"),
323                    })
324            }
325        }
326    }
327
328    fn find(&self, line: &str) -> Option<usize> {
329        match self {
330            Self::Literal(query) => line.find(query),
331            Self::Regex(regex) => regex.find(line).map(|mat| mat.start()),
332        }
333    }
334}
335
336fn search_text(
337    path: &WorkspacePath,
338    text: &str,
339    matcher: &Matcher,
340    context_lines: usize,
341) -> Vec<SearchMatch> {
342    let lines: Vec<_> = text.lines().collect();
343    let mut matches = Vec::new();
344    for (index, line) in lines.iter().enumerate() {
345        let Some(column) = matcher.find(line) else {
346            continue;
347        };
348        let before_start = index.saturating_sub(context_lines);
349        let after_end = (index + context_lines + 1).min(lines.len());
350        matches.push(SearchMatch {
351            path: path.clone(),
352            line: index + 1,
353            column: Some(column + 1),
354            line_text: (*line).to_string(),
355            before: lines[before_start..index]
356                .iter()
357                .map(|line| (*line).to_string())
358                .collect(),
359            after: lines[index + 1..after_end]
360                .iter()
361                .map(|line| (*line).to_string())
362                .collect(),
363        });
364    }
365    matches
366}
367
368fn parse_rg_line(line: &str) -> Option<SearchMatch> {
369    let mut parts = line.splitn(4, ':');
370    let path = WorkspacePath::parse(parts.next()?).ok()?;
371    let line_number = parts.next()?.parse().ok()?;
372    let column = parts.next()?.parse().ok();
373    let line_text = parts.next().unwrap_or_default().to_string();
374    Some(SearchMatch {
375        path,
376        line: line_number,
377        column,
378        line_text,
379        before: Vec::new(),
380        after: Vec::new(),
381    })
382}
383
384/// Search errors.
385#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error, Serialize, Deserialize)]
386#[non_exhaustive]
387pub enum SearchError {
388    /// Query is invalid.
389    #[error("invalid search query: {message}")]
390    InvalidQuery {
391        /// Model-safe explanation.
392        message: String,
393    },
394    /// Workspace operation failed.
395    #[error("workspace search error: {message}")]
396    Workspace {
397        /// Model-safe explanation.
398        message: String,
399    },
400    /// Command search failed.
401    #[error("search command error: {message}")]
402    Command {
403        /// Model-safe explanation.
404        message: String,
405    },
406    /// Search implementation is unsupported.
407    #[error("unsupported search operation: {message}")]
408    Unsupported {
409        /// Model-safe explanation.
410        message: String,
411    },
412}
413
414#[cfg(test)]
415mod tests {
416    use super::*;
417    use crate::coding::LocalWorkspace;
418
419    fn temp_dir(tag: &str) -> std::path::PathBuf {
420        let dir =
421            std::env::temp_dir().join(format!("molo-search-test-{}-{tag}", std::process::id()));
422        let _ = std::fs::remove_dir_all(&dir);
423        std::fs::create_dir_all(&dir).unwrap();
424        dir
425    }
426
427    #[tokio::test]
428    async fn workspace_search_returns_structured_matches() {
429        let root = temp_dir("workspace");
430        std::fs::write(root.join("a.txt"), "one\ntwo\nthree").unwrap();
431        let searcher = WorkspaceSearcher::new(LocalWorkspace::new(&root).unwrap());
432        let results = searcher
433            .search(
434                RepoSearchRequest {
435                    query: "two".to_string(),
436                    paths: Vec::new(),
437                    mode: SearchMode::Literal,
438                    max_matches: 10,
439                    context_lines: 1,
440                    include_hidden: false,
441                    respect_gitignore: true,
442                },
443                &RunContext::new("search"),
444            )
445            .await
446            .unwrap();
447        assert_eq!(results.matches.len(), 1);
448        assert_eq!(results.matches[0].line, 2);
449        assert_eq!(results.matches[0].before, vec!["one"]);
450        assert_eq!(results.matches[0].after, vec!["three"]);
451        let _ = std::fs::remove_dir_all(root);
452    }
453}
molo_coding/coding/search.rs

molo_coding/coding/
search.rs
