rust_docs_mcp/search/
tools.rs

1//! # Search Tools Module
2//!
3//! Provides MCP tool integration for fuzzy search functionality.
4//!
5//! ## Key Components
6//! - [`SearchTools`] - MCP tool handler for search operations
7//! - [`SearchItemsFuzzyParams`] - Parameters for fuzzy search requests
8//!
9//! ## Features
10//! - Automatic crate indexing on first search
11//! - Fuzzy search with configurable edit distance
12//! - Result filtering by kind and crate
13//!
14//! ## Example
15//! ```no_run
16//! # use std::sync::Arc;
17//! # use tokio::sync::RwLock;
18//! # use rust_docs_mcp::cache::CrateCache;
19//! # use rust_docs_mcp::search::tools::{SearchTools, SearchItemsFuzzyParams};
20//! # async fn example() -> anyhow::Result<()> {
21//! let cache = Arc::new(RwLock::new(CrateCache::new(None)?));
22//! let tools = SearchTools::new(cache);
23//!
24//! let params = SearchItemsFuzzyParams {
25//!     crate_name: "serde".to_string(),
26//!     version: "1.0.0".to_string(),
27//!     query: "deserialize".to_string(),
28//!     fuzzy_enabled: Some(true),
29//!     fuzzy_distance: Some(1),
30//!     limit: Some(10),
31//!     kind_filter: None,
32//!     member: None,
33//! };
34//!
35//! let results = tools.search_items_fuzzy(params).await;
36//! # Ok(())
37//! # }
38//! ```
39
40use 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    /// Check if a crate has a search index
87    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    /// Perform the actual search without holding any locks
98    async fn perform_search(
99        &self,
100        params: SearchItemsFuzzyParams,
101        storage: CacheStorage,
102    ) -> Result<Vec<SearchResult>, anyhow::Error> {
103        // Create indexer for the specific crate
104        let indexer = SearchIndexer::new_for_crate(
105            &params.crate_name,
106            &params.version,
107            &storage,
108            params.member.as_deref(),
109        )?;
110
111        // Create fuzzy searcher
112        let fuzzy_searcher = FuzzySearcher::from_indexer(&indexer)?;
113
114        // Validate fuzzy distance
115        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        // Validate limit
124        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        // Build search options
133        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        // Perform search
143        fuzzy_searcher.search(&params.query, &options)
144    }
145
146    /// Perform fuzzy search on crate items
147    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            // First check with read lock if docs already exist
158            {
159                let cache = self.cache.read().await;
160                let has_docs = cache.has_docs(
161                    &params.crate_name,
162                    &params.version,
163                    params.member.as_deref(),
164                );
165
166                if has_docs
167                    && self
168                        .has_search_index(
169                            &params.crate_name,
170                            &params.version,
171                            params.member.as_deref(),
172                        )
173                        .await
174                {
175                    // Docs and index exist, proceed with search using read lock only
176                    let storage = cache.storage.clone();
177                    drop(cache); // Release read lock early
178
179                    return self.perform_search(params, storage).await;
180                }
181            }
182
183            // Need to generate docs/index, acquire write lock
184            {
185                let cache = self.cache.write().await;
186                // Double-check in case another task generated it
187                let has_docs = cache.has_docs(
188                    &params.crate_name,
189                    &params.version,
190                    params.member.as_deref(),
191                );
192
193                if !has_docs {
194                    cache
195                        .ensure_crate_or_member_docs(
196                            &params.crate_name,
197                            &params.version,
198                            params.member.as_deref(),
199                        )
200                        .await?;
201                }
202            }
203
204            // Now perform search with read lock
205            let cache = self.cache.read().await;
206            let storage = cache.storage.clone();
207            drop(cache);
208
209            // Check if search index exists after ensuring docs
210            if !self
211                .has_search_index(
212                    &params.crate_name,
213                    &params.version,
214                    params.member.as_deref(),
215                )
216                .await
217            {
218                // Docs exist but search index is missing - regenerate it
219                let cache = self.cache.write().await;
220                cache
221                    .create_search_index(
222                        &params.crate_name,
223                        &params.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, // fuzzy::SearchResult doesn't have doc_preview
249                            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}
rust_docs_mcp/search/tools.rs

rust_docs_mcp/search/
tools.rs
