1use std::sync::Arc;
41use tokio::sync::RwLock;
42
43use rmcp::schemars;
44use schemars::JsonSchema;
45use serde::{Deserialize, Serialize};
46
47use crate::cache::{CrateCache, storage::CacheStorage};
48use crate::search::config::{
49 DEFAULT_FUZZY_DISTANCE, DEFAULT_SEARCH_LIMIT, MAX_FUZZY_DISTANCE, MAX_SEARCH_LIMIT,
50};
51use crate::search::outputs::{SearchErrorOutput, SearchItemsFuzzyOutput};
52use crate::search::{FuzzySearchOptions, FuzzySearcher, SearchIndexer, SearchResult};
53
54#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
55pub struct SearchItemsFuzzyParams {
56 #[schemars(description = "The name of the crate")]
57 pub crate_name: String,
58 #[schemars(description = "The version of the crate")]
59 pub version: String,
60 #[schemars(description = "The search query")]
61 pub query: String,
62 #[schemars(description = "Enable fuzzy matching for typo tolerance")]
63 pub fuzzy_enabled: Option<bool>,
64 #[schemars(description = "Edit distance for fuzzy matching (0-2)")]
65 pub fuzzy_distance: Option<u8>,
66 #[schemars(description = "Maximum number of results to return")]
67 pub limit: Option<usize>,
68 #[schemars(description = "Filter by item kind")]
69 pub kind_filter: Option<String>,
70 #[schemars(
71 description = "For workspace crates, specify the member path (e.g., 'crates/rmcp')"
72 )]
73 pub member: Option<String>,
74}
75
76#[derive(Debug, Clone)]
77pub struct SearchTools {
78 cache: Arc<RwLock<CrateCache>>,
79}
80
81impl SearchTools {
82 pub fn new(cache: Arc<RwLock<CrateCache>>) -> Self {
83 Self { cache }
84 }
85
86 async fn has_search_index(
88 &self,
89 crate_name: &str,
90 version: &str,
91 member: Option<&str>,
92 ) -> bool {
93 let cache = self.cache.read().await;
94 cache.storage.has_search_index(crate_name, version, member)
95 }
96
97 async fn perform_search(
99 &self,
100 params: SearchItemsFuzzyParams,
101 storage: CacheStorage,
102 ) -> Result<Vec<SearchResult>, anyhow::Error> {
103 let indexer = SearchIndexer::new_for_crate(
105 ¶ms.crate_name,
106 ¶ms.version,
107 &storage,
108 params.member.as_deref(),
109 )?;
110
111 let fuzzy_searcher = FuzzySearcher::from_indexer(&indexer)?;
113
114 let fuzzy_distance = params.fuzzy_distance.unwrap_or(DEFAULT_FUZZY_DISTANCE);
116 if fuzzy_distance > MAX_FUZZY_DISTANCE {
117 return Err(anyhow::anyhow!(
118 "Fuzzy distance must be between 0 and {}",
119 MAX_FUZZY_DISTANCE
120 ));
121 }
122
123 let limit = params.limit.unwrap_or(DEFAULT_SEARCH_LIMIT);
125 if limit > MAX_SEARCH_LIMIT {
126 return Err(anyhow::anyhow!(
127 "Limit must not exceed {}",
128 MAX_SEARCH_LIMIT
129 ));
130 }
131
132 let options = FuzzySearchOptions {
134 fuzzy_enabled: params.fuzzy_enabled.unwrap_or(true),
135 fuzzy_distance,
136 limit,
137 kind_filter: params.kind_filter.clone(),
138 crate_filter: Some(params.crate_name.clone()),
139 member_filter: params.member.clone(),
140 };
141
142 fuzzy_searcher.search(¶ms.query, &options)
144 }
145
146 pub async fn search_items_fuzzy(
148 &self,
149 params: SearchItemsFuzzyParams,
150 ) -> Result<SearchItemsFuzzyOutput, SearchErrorOutput> {
151 let query = params.query.clone();
152 let fuzzy_enabled = params.fuzzy_enabled.unwrap_or(true);
153 let crate_name = params.crate_name.clone();
154 let version = params.version.clone();
155 let member = params.member.clone();
156 let result = async {
157 {
159 let cache = self.cache.read().await;
160 let has_docs = cache.has_docs(
161 ¶ms.crate_name,
162 ¶ms.version,
163 params.member.as_deref(),
164 );
165
166 if has_docs
167 && self
168 .has_search_index(
169 ¶ms.crate_name,
170 ¶ms.version,
171 params.member.as_deref(),
172 )
173 .await
174 {
175 let storage = cache.storage.clone();
177 drop(cache); return self.perform_search(params, storage).await;
180 }
181 }
182
183 {
185 let cache = self.cache.write().await;
186 let has_docs = cache.has_docs(
188 ¶ms.crate_name,
189 ¶ms.version,
190 params.member.as_deref(),
191 );
192
193 if !has_docs {
194 cache
195 .ensure_crate_or_member_docs(
196 ¶ms.crate_name,
197 ¶ms.version,
198 params.member.as_deref(),
199 )
200 .await?;
201 }
202 }
203
204 let cache = self.cache.read().await;
206 let storage = cache.storage.clone();
207 drop(cache);
208
209 if !self
211 .has_search_index(
212 ¶ms.crate_name,
213 ¶ms.version,
214 params.member.as_deref(),
215 )
216 .await
217 {
218 let cache = self.cache.write().await;
220 cache
221 .create_search_index(
222 ¶ms.crate_name,
223 ¶ms.version,
224 params.member.as_deref(),
225 )
226 .await?;
227 }
228
229 self.perform_search(params, storage).await
230 }
231 .await;
232
233 match result {
234 Ok(results) => {
235 let total = results.len();
236 Ok(SearchItemsFuzzyOutput {
237 results: results
238 .into_iter()
239 .map(|r| crate::search::outputs::SearchResult {
240 score: r.score,
241 item_id: r.item_id,
242 name: r.name,
243 path: r.path,
244 kind: r.kind,
245 crate_name: r.crate_name,
246 version: r.version,
247 visibility: r.visibility,
248 doc_preview: None, member: r.member,
250 })
251 .collect(),
252 query,
253 total_results: total,
254 fuzzy_enabled,
255 crate_name,
256 version,
257 member,
258 })
259 }
260 Err(e) => Err(SearchErrorOutput::new(format!("Search failed: {e}"))),
261 }
262 }
263}