use wsb::prelude::*;
use grep::regex::RegexMatcher;
use grep::searcher::{sinks::UTF8, SearcherBuilder};
use ignore::WalkBuilder;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::sync::mpsc;
#[derive(Serialize, Clone, Debug)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum SearchHit {
Session {
session_id: String,
project: Option<String>,
worktree: Option<String>,
agent: String,
name: String,
ordinal: i64,
session_ended_at: Option<String>,
snippet: String,
},
File {
path: String,
line_number: u64,
snippet: String,
},
}
#[derive(Serialize, Clone, Debug)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum SearchEvent {
Hits {
source: String,
hits: Vec<SearchHit>,
},
Done {
source: String,
},
Error {
source: String,
message: String,
},
}
pub struct CancelOnDrop(pub Arc<AtomicBool>);
impl Drop for CancelOnDrop {
fn drop(&mut self) {
self.0.store(true, Ordering::Relaxed);
}
}
pub struct SessionSource {
db: Arc<SessionRepo>,
sort: Option<String>,
place: wsb::persist::global::models::PlaceFilter,
}
impl SessionSource {
pub fn new(
db: Arc<SessionRepo>,
sort: Option<String>,
place: wsb::persist::global::models::PlaceFilter,
) -> Self {
Self { db, sort, place }
}
pub async fn run(
self,
query: String,
tx: mpsc::Sender<SearchEvent>,
cancelled: Arc<AtomicBool>,
) {
if cancelled.load(Ordering::Relaxed) {
return;
}
let sort = self.sort.clone().unwrap_or_else(|| "relevance".to_string());
let place = self.place.clone();
let db = self.db.clone();
let q = query.clone();
let result = tokio::task::spawn_blocking(move || db.global_search(&q, 50, 0, &sort, &place))
.await;
if cancelled.load(Ordering::Relaxed) {
return;
}
match result {
Ok(Ok(rows)) => {
let hits: Vec<SearchHit> = rows
.into_iter()
.map(|r| SearchHit::Session {
session_id: r.session_id,
project: r.project,
worktree: r.worktree,
agent: r.agent,
name: r.name,
ordinal: r.ordinal,
session_ended_at: r.session_ended_at,
snippet: r.snippet,
})
.collect();
if !hits.is_empty() {
let _ = tx
.send(SearchEvent::Hits {
source: "sessions".into(),
hits,
})
.await;
}
let _ = tx
.send(SearchEvent::Done {
source: "sessions".into(),
})
.await;
}
Ok(Err(e)) => {
let _ = tx
.send(SearchEvent::Error {
source: "sessions".into(),
message: e.to_string(),
})
.await;
}
Err(e) => {
let _ = tx
.send(SearchEvent::Error {
source: "sessions".into(),
message: format!("join error: {}", e),
})
.await;
}
}
}
}
pub struct FileSource {
workspace_root: PathBuf,
}
impl FileSource {
pub fn new(workspace_root: PathBuf) -> Self {
Self { workspace_root }
}
pub fn run(self, query: String, tx: mpsc::Sender<SearchEvent>, cancelled: Arc<AtomicBool>) {
if cancelled.load(Ordering::Relaxed) {
return;
}
let pattern = regex::escape(&query);
let matcher = match RegexMatcher::new_line_matcher(&pattern) {
Ok(m) => m,
Err(e) => {
let _ = tx.blocking_send(SearchEvent::Error {
source: "files".into(),
message: format!("invalid search pattern: {}", e),
});
return;
}
};
let walker = WalkBuilder::new(&self.workspace_root)
.standard_filters(true)
.build();
let mut searcher = SearcherBuilder::new().line_number(true).build();
for entry in walker {
if cancelled.load(Ordering::Relaxed) {
break;
}
let entry = match entry {
Ok(e) => e,
Err(_) => continue,
};
if !entry.file_type().map(|ft| ft.is_file()).unwrap_or(false) {
continue;
}
let path = entry.path();
let rel_path = path
.strip_prefix(&self.workspace_root)
.unwrap_or(path)
.to_string_lossy()
.into_owned();
let mut file_hits: Vec<SearchHit> = Vec::new();
let _ = searcher.search_path(
&matcher,
path,
UTF8(|line_number, line| {
file_hits.push(SearchHit::File {
path: rel_path.clone(),
line_number,
snippet: line.trim_end_matches('\n').to_string(),
});
Ok(true)
}),
);
if !file_hits.is_empty() {
if tx
.blocking_send(SearchEvent::Hits {
source: "files".into(),
hits: file_hits,
})
.is_err()
{
return;
}
}
}
let _ = tx.blocking_send(SearchEvent::Done {
source: "files".into(),
});
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
use tokio::sync::mpsc;
fn run_file_source(
source: FileSource,
query: &str,
cancelled: Arc<AtomicBool>,
) -> Vec<SearchEvent> {
let (tx, mut rx) = mpsc::channel(256);
let q = query.to_string();
let handle = std::thread::spawn(move || source.run(q, tx, cancelled));
handle.join().unwrap();
let mut events = Vec::new();
while let Ok(e) = rx.try_recv() {
events.push(e);
}
events
}
#[test]
fn file_source_finds_literal_matches() {
let dir = TempDir::new().unwrap();
let file = dir.path().join("hello.txt");
std::fs::write(&file, "foo bar\nbaz qux\nfoo again\n").unwrap();
let source = FileSource::new(dir.path().to_path_buf());
let cancelled = Arc::new(AtomicBool::new(false));
let events = run_file_source(source, "foo", cancelled);
let hits: Vec<&SearchHit> = events
.iter()
.filter_map(|e| match e {
SearchEvent::Hits { hits, .. } => Some(hits.iter()),
_ => None,
})
.flatten()
.collect();
assert_eq!(hits.len(), 2, "expected two 'foo' matches");
let done_count = events
.iter()
.filter(|e| matches!(e, SearchEvent::Done { source } if source == "files"))
.count();
assert_eq!(done_count, 1, "expected exactly one Done event");
}
#[test]
fn file_source_escapes_regex_metacharacters() {
let dir = TempDir::new().unwrap();
let file = dir.path().join("dots.txt");
std::fs::write(&file, "1.2.3\nxyz\n").unwrap();
let source = FileSource::new(dir.path().to_path_buf());
let cancelled = Arc::new(AtomicBool::new(false));
let events = run_file_source(source, "1.2.3", cancelled);
let hits: Vec<&SearchHit> = events
.iter()
.filter_map(|e| match e {
SearchEvent::Hits { hits, .. } => Some(hits.iter()),
_ => None,
})
.flatten()
.collect();
assert_eq!(hits.len(), 1);
match &hits[0] {
SearchHit::File { snippet, .. } => assert!(snippet.contains("1.2.3")),
other => panic!("expected File hit, got {:?}", other),
}
}
#[test]
fn file_source_respects_cancellation() {
let dir = TempDir::new().unwrap();
for i in 0..20 {
let path = dir.path().join(format!("file{}.txt", i));
std::fs::write(&path, "needle\nneedle\nneedle\n").unwrap();
}
let cancelled = Arc::new(AtomicBool::new(true));
let source = FileSource::new(dir.path().to_path_buf());
let events = run_file_source(source, "needle", cancelled);
let hit_count: usize = events
.iter()
.filter(|e| matches!(e, SearchEvent::Hits { .. }))
.count();
assert_eq!(hit_count, 0, "pre-cancelled source must send no hits");
}
#[test]
fn file_source_sends_done_even_with_no_matches() {
let dir = TempDir::new().unwrap();
std::fs::write(dir.path().join("a.txt"), "nothing here\n").unwrap();
let source = FileSource::new(dir.path().to_path_buf());
let cancelled = Arc::new(AtomicBool::new(false));
let events = run_file_source(source, "UNLIKELY_TOKEN_XYZ", cancelled);
let done_count = events
.iter()
.filter(|e| matches!(e, SearchEvent::Done { .. }))
.count();
assert_eq!(
done_count, 1,
"Done must be sent even when no matches found"
);
let hit_count = events
.iter()
.filter(|e| matches!(e, SearchEvent::Hits { .. }))
.count();
assert_eq!(hit_count, 0);
}
#[test]
fn file_source_relative_paths() {
let dir = TempDir::new().unwrap();
let sub = dir.path().join("subdir");
std::fs::create_dir(&sub).unwrap();
std::fs::write(sub.join("code.rs"), "let x = target;\n").unwrap();
let source = FileSource::new(dir.path().to_path_buf());
let cancelled = Arc::new(AtomicBool::new(false));
let events = run_file_source(source, "target", cancelled);
let paths: Vec<String> = events
.iter()
.filter_map(|e| match e {
SearchEvent::Hits { hits, .. } => Some(hits.iter()),
_ => None,
})
.flatten()
.filter_map(|h| match h {
SearchHit::File { path, .. } => Some(path.clone()),
_ => None,
})
.collect();
assert_eq!(paths.len(), 1);
assert!(
!paths[0].starts_with('/'),
"path must be relative, got: {}",
paths[0]
);
assert!(paths[0].contains("subdir"), "path must include subdir");
}
}