use std::{cmp::Reverse, path::Path}; use crate::session::{Agent, Session}; use nucleo_matcher::{ Config, Matcher, Utf32Str, pattern::{CaseMatching, Normalization, Pattern}, }; pub struct SearchIndex { pub sessions: Vec, text: Vec, } struct SearchText { title: String, messages: Vec, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum MatchKind { Recent, Title, Message, } #[derive(Clone, Debug)] pub struct Hit { pub index: usize, pub kind: MatchKind, pub message: Option, score: u32, tier: u8, } impl SearchIndex { pub fn new(sessions: Vec) -> Self { let text = sessions .iter() .map(|s| SearchText { title: s.title.to_lowercase(), messages: s.messages.iter().map(|s| s.to_lowercase()).collect(), }) .collect(); Self { sessions, text } } pub fn search(&self, query: &str, agent: Option, cwd: Option<&Path>) -> Vec { let query = query.trim().to_lowercase(); let words: Vec<_> = query.split_whitespace().collect(); let pattern = Pattern::new( &query, CaseMatching::Ignore, Normalization::Smart, nucleo_matcher::pattern::AtomKind::Fuzzy, ); let mut matcher = Matcher::new(Config::DEFAULT); let mut buffer = vec![]; let mut hits = vec![]; for (index, (session, text)) in self.sessions.iter().zip(&self.text).enumerate() { if agent.is_some_and(|a| a != session.agent) || cwd.is_some_and(|p| p != session.cwd) { continue; } let mut hit = Hit { index, kind: MatchKind::Recent, message: None, score: 0, tier: 0, }; if !query.is_empty() { if text.title.contains(&query) { hit.kind = MatchKind::Title; hit.tier = 4; hit.score = if text.title == query { 2 } else if text.title.starts_with(&query) { 1 } else { 0 }; } else if words.iter().all(|w| text.title.contains(w)) { hit.kind = MatchKind::Title; hit.tier = 3; } else if let Some(score) = pattern.score(Utf32Str::new(&session.title, &mut buffer), &mut matcher) { hit.kind = MatchKind::Title; hit.tier = 2; hit.score = score; } else { let mut found = None; for (i, message) in text.messages.iter().enumerate() { if message.contains(&query) { found = Some((i, 1)); break; } if found.is_none() && words.iter().all(|w| message.contains(w)) { found = Some((i, 0)); } } let Some((i, score)) = found else { continue }; hit.kind = MatchKind::Message; hit.tier = 1; hit.score = score; hit.message = Some(i); } } hits.push(hit); } hits.sort_by_key(|h| { ( Reverse(h.tier), Reverse(h.score), Reverse(self.sessions[h.index].modified), h.index, ) }); hits } } /// Context around the matching portion, not just the start of a long prompt. pub fn snippet(text: &str, query: &str, max_chars: usize) -> String { let lower = text.to_lowercase(); let query = query.to_lowercase(); let pos = query .split_whitespace() .filter_map(|w| lower.find(w)) .min() .unwrap_or(0); // Lowercasing can change byte lengths. Convert to a character offset first. let char_pos = lower[..pos].chars().count(); let start = char_pos.saturating_sub(70); let mut output = if start > 0 { "…".to_owned() } else { String::new() }; output.extend(text.chars().skip(start).take(max_chars)); if text.chars().count() > start + max_chars { output.push('…'); } output } #[cfg(test)] mod tests { use super::*; fn session(title: &str, message: &str, modified: u64) -> Session { Session { agent: Agent::Pi, id: title.into(), path: "/x".into(), cwd: "/work".into(), title: title.into(), modified, messages: vec![message.into()], } } #[test] fn title_hits_always_beat_body_hits_and_recency_breaks_ties() { let index = SearchIndex::new(vec![ session("Other", "authentication failure", 999), session("Authentication", "", 1), session("Auth guide", "", 2), ]); let hits = index.search("auth", None, None); assert_eq!(hits.iter().map(|h| h.index).collect::>(), [2, 1, 0]); assert_eq!(hits[2].kind, MatchKind::Message); } #[test] fn fuzzy_titles_and_unordered_message_tokens() { let index = SearchIndex::new(vec![ session("Authentication refactor", "", 1), session("Other", "cache invalidation fix", 2), ]); assert_eq!(index.search("athref", None, None)[0].index, 0); let hits = index.search("fix cache", None, None); assert_eq!(hits[0].index, 1); assert_eq!(hits[0].message, Some(0)); assert!(index.search("no match at all", None, None).is_empty()); } #[test] fn scope_is_exact_and_filters_compose() { let index = SearchIndex::new(vec![session("title", "text", 1)]); assert!(index.search("", Some(Agent::Claude), None).is_empty()); assert!(index.search("", None, Some(Path::new("/wor"))).is_empty()); assert_eq!( index .search("", Some(Agent::Pi), Some(Path::new("/work"))) .len(), 1 ); } #[test] fn unicode_and_match_context() { let index = SearchIndex::new(vec![session("École 日本語", "the café is open", 1)]); assert_eq!(index.search("日本", None, None).len(), 1); assert_eq!(index.search("CAFÉ", None, None).len(), 1); let text = format!("{}needle", "x".repeat(1000)); assert!(snippet(&text, "needle", 100).contains("needle")); } }