rust_docs_mcp/docs/
tools.rs

1use std::sync::Arc;
2use tokio::sync::RwLock;
3
4use rmcp::schemars;
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7
8use crate::cache::CrateCache;
9use crate::docs::{
10    DocQuery,
11    outputs::{
12        DetailedItem, DocsErrorOutput, GetItemDetailsOutput, GetItemDocsOutput,
13        GetItemSourceOutput, ItemInfo, ItemPreview, ListCrateItemsOutput, PaginationInfo,
14        SearchItemsOutput, SearchItemsPreviewOutput, SourceInfo, SourceLocation,
15    },
16};
17
18/// Maximum size for response in bytes (roughly 25k tokens * 4 bytes/token)
19const MAX_RESPONSE_SIZE: usize = 100_000;
20
21#[derive(Debug, Serialize, Deserialize, JsonSchema)]
22pub struct ListItemsParams {
23    #[schemars(description = "The name of the crate")]
24    pub crate_name: String,
25    #[schemars(description = "The version of the crate")]
26    pub version: String,
27    #[schemars(description = "Optional filter by item kind (e.g., 'function', 'struct', 'enum')")]
28    pub kind_filter: Option<String>,
29    #[schemars(description = "Maximum number of items to return (default: 100)")]
30    pub limit: Option<i64>,
31    #[schemars(description = "Starting position for pagination (default: 0)")]
32    pub offset: Option<i64>,
33    #[schemars(
34        description = "For workspace crates, specify the member path (e.g., 'crates/rmcp')"
35    )]
36    pub member: Option<String>,
37}
38
39#[derive(Debug, Serialize, Deserialize, JsonSchema)]
40pub struct SearchItemsParams {
41    #[schemars(description = "The name of the crate")]
42    pub crate_name: String,
43    #[schemars(description = "The version of the crate")]
44    pub version: String,
45    #[schemars(
46        description = "The pattern to search for in item names. Note: passing '*' will not return any items - use specific Rust symbols or generalize over common names (e.g., 'new', 'parse', 'Error') to get meaningful results"
47    )]
48    pub pattern: String,
49    #[schemars(description = "Maximum number of items to return (default: 100)")]
50    pub limit: Option<i64>,
51    #[schemars(description = "Starting position for pagination (default: 0)")]
52    pub offset: Option<i64>,
53    #[schemars(description = "Optional filter by item kind (e.g., 'function', 'struct', 'enum')")]
54    pub kind_filter: Option<String>,
55    #[schemars(description = "Optional filter by module path prefix")]
56    pub path_filter: Option<String>,
57    #[schemars(
58        description = "For workspace crates, specify the member path (e.g., 'crates/rmcp')"
59    )]
60    pub member: Option<String>,
61}
62
63#[derive(Debug, Serialize, Deserialize, JsonSchema)]
64pub struct SearchItemsPreviewParams {
65    #[schemars(description = "The name of the crate")]
66    pub crate_name: String,
67    #[schemars(description = "The version of the crate")]
68    pub version: String,
69    #[schemars(
70        description = "The pattern to search for in item names. Note: passing '*' will not return any items - use specific Rust symbols or generalize over common names (e.g., 'new', 'parse', 'Error') to get meaningful results"
71    )]
72    pub pattern: String,
73    #[schemars(description = "Maximum number of items to return (default: 100)")]
74    pub limit: Option<i64>,
75    #[schemars(description = "Starting position for pagination (default: 0)")]
76    pub offset: Option<i64>,
77    #[schemars(description = "Optional filter by item kind (e.g., 'function', 'struct', 'enum')")]
78    pub kind_filter: Option<String>,
79    #[schemars(description = "Optional filter by module path prefix")]
80    pub path_filter: Option<String>,
81    #[schemars(
82        description = "For workspace crates, specify the member path (e.g., 'crates/rmcp')"
83    )]
84    pub member: Option<String>,
85}
86
87#[derive(Debug, Serialize, Deserialize, JsonSchema)]
88pub struct GetItemDetailsParams {
89    #[schemars(description = "The name of the crate")]
90    pub crate_name: String,
91    #[schemars(description = "The version of the crate")]
92    pub version: String,
93    #[schemars(description = "The numeric ID of the item")]
94    pub item_id: i32,
95    #[schemars(
96        description = "For workspace crates, specify the member path (e.g., 'crates/rmcp')"
97    )]
98    pub member: Option<String>,
99}
100
101#[derive(Debug, Serialize, Deserialize, JsonSchema)]
102pub struct GetItemDocsParams {
103    #[schemars(description = "The name of the crate")]
104    pub crate_name: String,
105    #[schemars(description = "The version of the crate")]
106    pub version: String,
107    #[schemars(description = "The numeric ID of the item")]
108    pub item_id: i32,
109    #[schemars(
110        description = "For workspace crates, specify the member path (e.g., 'crates/rmcp')"
111    )]
112    pub member: Option<String>,
113}
114
115#[derive(Debug, Serialize, Deserialize, JsonSchema)]
116pub struct GetItemSourceParams {
117    #[schemars(description = "The name of the crate")]
118    pub crate_name: String,
119    #[schemars(description = "The version of the crate")]
120    pub version: String,
121    #[schemars(description = "The numeric ID of the item")]
122    pub item_id: i32,
123    #[schemars(
124        description = "Number of context lines to include before and after the item (default: 3)"
125    )]
126    pub context_lines: Option<i64>,
127    #[schemars(
128        description = "For workspace crates, specify the member path (e.g., 'crates/rmcp')"
129    )]
130    pub member: Option<String>,
131}
132
133#[derive(Debug, Clone)]
134pub struct DocsTools {
135    cache: Arc<RwLock<CrateCache>>,
136}
137
138impl DocsTools {
139    pub fn new(cache: Arc<RwLock<CrateCache>>) -> Self {
140        Self { cache }
141    }
142
143    /// Helper to check if a response might exceed size limits
144    fn estimate_response_size<T: Serialize>(data: &T) -> usize {
145        serde_json::to_string(data).map(|s| s.len()).unwrap_or(0)
146    }
147
148    pub async fn list_crate_items(
149        &self,
150        params: ListItemsParams,
151    ) -> Result<ListCrateItemsOutput, DocsErrorOutput> {
152        let cache = self.cache.write().await;
153        match cache
154            .ensure_crate_or_member_docs(
155                &params.crate_name,
156                &params.version,
157                params.member.as_deref(),
158            )
159            .await
160        {
161            Ok(crate_data) => {
162                let query = DocQuery::new(crate_data);
163                let items = query.list_items(params.kind_filter.as_deref());
164
165                let total_count = items.len();
166                let limit = params.limit.unwrap_or(100).max(0) as usize;
167                let offset = params.offset.unwrap_or(0).max(0) as usize;
168
169                // Apply pagination
170                let paginated_items: Vec<_> = items
171                    .into_iter()
172                    .skip(offset)
173                    .take(limit)
174                    .map(|item| ItemInfo {
175                        id: item.id.to_string(),
176                        name: item.name.clone(),
177                        kind: item.kind.clone(),
178                        path: item.path.clone(),
179                        docs: item.docs.clone(),
180                        visibility: item.visibility.clone(),
181                    })
182                    .collect();
183
184                Ok(ListCrateItemsOutput {
185                    items: paginated_items,
186                    pagination: PaginationInfo {
187                        total: total_count,
188                        limit,
189                        offset,
190                        has_more: offset + limit < total_count,
191                    },
192                })
193            }
194            Err(e) => Err(DocsErrorOutput::new(format!(
195                "Failed to get crate docs: {e}"
196            ))),
197        }
198    }
199
200    pub async fn search_items(
201        &self,
202        params: SearchItemsParams,
203    ) -> Result<SearchItemsOutput, DocsErrorOutput> {
204        let cache = self.cache.write().await;
205        match cache
206            .ensure_crate_or_member_docs(
207                &params.crate_name,
208                &params.version,
209                params.member.as_deref(),
210            )
211            .await
212        {
213            Ok(crate_data) => {
214                let query = DocQuery::new(crate_data);
215                let mut items = query.search_items(&params.pattern);
216
217                // Apply kind filter if provided
218                if let Some(kind_filter) = &params.kind_filter {
219                    items.retain(|item| item.kind == *kind_filter);
220                }
221
222                // Apply path filter if provided
223                if let Some(path_filter) = &params.path_filter {
224                    items.retain(|item| {
225                        let item_path = item.path.join("::");
226                        item_path.starts_with(path_filter)
227                    });
228                }
229
230                let total_count = items.len();
231                let limit = params.limit.unwrap_or(100).max(0) as usize;
232                let offset = params.offset.unwrap_or(0).max(0) as usize;
233
234                // Apply pagination
235                let mut paginated_items: Vec<_> =
236                    items.into_iter().skip(offset).take(limit).collect();
237
238                // Check response size and truncate if necessary
239                let mut actual_limit = limit;
240                let mut truncated = false;
241
242                loop {
243                    let test_response = serde_json::json!({
244                        "items": &paginated_items,
245                        "pagination": {
246                            "total": total_count,
247                            "limit": actual_limit,
248                            "offset": offset,
249                            "has_more": offset + paginated_items.len() < total_count
250                        }
251                    });
252
253                    if Self::estimate_response_size(&test_response) <= MAX_RESPONSE_SIZE {
254                        break;
255                    }
256
257                    // Reduce items by half if too large
258                    let new_len = paginated_items.len() / 2;
259                    if new_len == 0 {
260                        break; // Can't reduce further
261                    }
262                    paginated_items.truncate(new_len);
263                    actual_limit = new_len;
264                    truncated = true;
265                }
266
267                let warning = if truncated {
268                    Some("Response was truncated to stay within size limits. Use smaller limit or preview mode.".to_string())
269                } else {
270                    None
271                };
272
273                Ok(SearchItemsOutput {
274                    items: paginated_items
275                        .into_iter()
276                        .map(|item| ItemInfo {
277                            id: item.id.to_string(),
278                            name: item.name.clone(),
279                            kind: item.kind.clone(),
280                            path: item.path.clone(),
281                            docs: item.docs.clone(),
282                            visibility: item.visibility.clone(),
283                        })
284                        .collect(),
285                    pagination: PaginationInfo {
286                        total: total_count,
287                        limit: actual_limit,
288                        offset,
289                        has_more: offset + actual_limit < total_count,
290                    },
291                    warning,
292                })
293            }
294            Err(e) => Err(DocsErrorOutput::new(format!(
295                "Failed to get crate docs: {e}"
296            ))),
297        }
298    }
299
300    pub async fn search_items_preview(
301        &self,
302        params: SearchItemsPreviewParams,
303    ) -> Result<SearchItemsPreviewOutput, DocsErrorOutput> {
304        let cache = self.cache.write().await;
305        match cache
306            .ensure_crate_or_member_docs(
307                &params.crate_name,
308                &params.version,
309                params.member.as_deref(),
310            )
311            .await
312        {
313            Ok(crate_data) => {
314                let query = DocQuery::new(crate_data);
315                let mut items = query.search_items(&params.pattern);
316
317                // Apply kind filter if provided
318                if let Some(kind_filter) = &params.kind_filter {
319                    items.retain(|item| item.kind == *kind_filter);
320                }
321
322                // Apply path filter if provided
323                if let Some(path_filter) = &params.path_filter {
324                    items.retain(|item| {
325                        let item_path = item.path.join("::");
326                        item_path.starts_with(path_filter)
327                    });
328                }
329
330                let total_count = items.len();
331                let limit = params.limit.unwrap_or(100).max(0) as usize;
332                let offset = params.offset.unwrap_or(0).max(0) as usize;
333
334                // Apply pagination and create preview items
335                let preview_items: Vec<_> = items
336                    .into_iter()
337                    .skip(offset)
338                    .take(limit)
339                    .map(|item| {
340                        serde_json::json!({
341                            "id": item.id,
342                            "name": item.name,
343                            "kind": item.kind,
344                            "path": item.path,
345                        })
346                    })
347                    .collect();
348
349                Ok(SearchItemsPreviewOutput {
350                    items: preview_items
351                        .into_iter()
352                        .map(|item| ItemPreview {
353                            id: item["id"].as_str().unwrap_or("").to_string(),
354                            name: item["name"].as_str().unwrap_or("").to_string(),
355                            kind: item["kind"].as_str().unwrap_or("").to_string(),
356                            path: item["path"]
357                                .as_array()
358                                .map(|arr| {
359                                    arr.iter()
360                                        .filter_map(|v| v.as_str().map(String::from))
361                                        .collect()
362                                })
363                                .unwrap_or_default(),
364                        })
365                        .collect(),
366                    pagination: PaginationInfo {
367                        total: total_count,
368                        limit,
369                        offset,
370                        has_more: offset + limit < total_count,
371                    },
372                })
373            }
374            Err(e) => Err(DocsErrorOutput::new(format!(
375                "Failed to get crate docs: {e}"
376            ))),
377        }
378    }
379
380    pub async fn get_item_details(&self, params: GetItemDetailsParams) -> GetItemDetailsOutput {
381        let cache = self.cache.write().await;
382        match cache
383            .ensure_crate_or_member_docs(
384                &params.crate_name,
385                &params.version,
386                params.member.as_deref(),
387            )
388            .await
389        {
390            Ok(crate_data) => {
391                let query = DocQuery::new(crate_data);
392                match query.get_item_details(params.item_id.max(0) as u32) {
393                    Ok(details) => {
394                        // Convert the details to our output format
395                        GetItemDetailsOutput::Success(Box::new(DetailedItem {
396                            info: ItemInfo {
397                                id: details.info.id.clone(),
398                                name: details.info.name.clone(),
399                                kind: details.info.kind.clone(),
400                                path: details.info.path.clone(),
401                                docs: details.info.docs.clone(),
402                                visibility: details.info.visibility.clone(),
403                            },
404                            signature: details.signature.clone(),
405                            generics: details.generics.clone(),
406                            fields: details.fields.map(|fields| {
407                                fields
408                                    .into_iter()
409                                    .map(|f| ItemInfo {
410                                        id: f.id,
411                                        name: f.name,
412                                        kind: f.kind,
413                                        path: f.path,
414                                        docs: f.docs,
415                                        visibility: f.visibility,
416                                    })
417                                    .collect()
418                            }),
419                            variants: details.variants.map(|variants| {
420                                variants
421                                    .into_iter()
422                                    .map(|v| ItemInfo {
423                                        id: v.id,
424                                        name: v.name,
425                                        kind: v.kind,
426                                        path: v.path,
427                                        docs: v.docs,
428                                        visibility: v.visibility,
429                                    })
430                                    .collect()
431                            }),
432                            methods: details.methods.map(|methods| {
433                                methods
434                                    .into_iter()
435                                    .map(|m| ItemInfo {
436                                        id: m.id,
437                                        name: m.name,
438                                        kind: m.kind,
439                                        path: m.path,
440                                        docs: m.docs,
441                                        visibility: m.visibility,
442                                    })
443                                    .collect()
444                            }),
445                            source_location: details.source_location.map(|loc| SourceLocation {
446                                filename: loc.filename,
447                                line_start: loc.line_start,
448                                column_start: loc.column_start,
449                                line_end: loc.line_end,
450                                column_end: loc.column_end,
451                            }),
452                        }))
453                    }
454                    Err(e) => GetItemDetailsOutput::Error {
455                        error: format!("Item not found: {e}"),
456                    },
457                }
458            }
459            Err(e) => GetItemDetailsOutput::Error {
460                error: format!("Failed to get crate docs: {e}"),
461            },
462        }
463    }
464
465    pub async fn get_item_docs(
466        &self,
467        params: GetItemDocsParams,
468    ) -> Result<GetItemDocsOutput, DocsErrorOutput> {
469        let cache = self.cache.write().await;
470        match cache
471            .ensure_crate_or_member_docs(
472                &params.crate_name,
473                &params.version,
474                params.member.as_deref(),
475            )
476            .await
477        {
478            Ok(crate_data) => {
479                let query = DocQuery::new(crate_data);
480                match query.get_item_docs(params.item_id.max(0) as u32) {
481                    Ok(docs) => {
482                        let message = if docs.is_none() {
483                            Some("No documentation available for this item".to_string())
484                        } else {
485                            None
486                        };
487                        Ok(GetItemDocsOutput {
488                            documentation: docs,
489                            message,
490                        })
491                    }
492                    Err(e) => Err(DocsErrorOutput::new(format!("Failed to get docs: {e}"))),
493                }
494            }
495            Err(e) => Err(DocsErrorOutput::new(format!(
496                "Failed to get crate docs: {e}"
497            ))),
498        }
499    }
500
501    pub async fn get_item_source(&self, params: GetItemSourceParams) -> GetItemSourceOutput {
502        let cache = self.cache.write().await;
503        let source_base_path = match cache.get_source_path(&params.crate_name, &params.version) {
504            Ok(path) => path,
505            Err(e) => {
506                return GetItemSourceOutput::Error {
507                    error: format!("Failed to get source path: {e}"),
508                };
509            }
510        };
511
512        match cache
513            .ensure_crate_or_member_docs(
514                &params.crate_name,
515                &params.version,
516                params.member.as_deref(),
517            )
518            .await
519        {
520            Ok(crate_data) => {
521                let query = DocQuery::new(crate_data);
522                let context_lines = params.context_lines.unwrap_or(3).max(0) as usize;
523
524                match query.get_item_source(
525                    params.item_id.max(0) as u32,
526                    &source_base_path,
527                    context_lines,
528                ) {
529                    Ok(source_info) => GetItemSourceOutput::Success(SourceInfo {
530                        location: SourceLocation {
531                            filename: source_info.location.filename,
532                            line_start: source_info.location.line_start,
533                            column_start: source_info.location.column_start,
534                            line_end: source_info.location.line_end,
535                            column_end: source_info.location.column_end,
536                        },
537                        code: source_info.code,
538                        context_lines: source_info.context_lines,
539                    }),
540                    Err(e) => GetItemSourceOutput::Error {
541                        error: format!("Failed to get source: {e}"),
542                    },
543                }
544            }
545            Err(e) => GetItemSourceOutput::Error {
546                error: format!("Failed to get crate docs: {e}"),
547            },
548        }
549    }
550}