rust_docs_mcp/cache/
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::{
9    CrateCache,
10    downloader::CrateSource,
11    outputs::{
12        CacheCrateOutput, CrateMetadata, ErrorOutput, GetCratesMetadataOutput,
13        ListCachedCratesOutput, ListCrateVersionsOutput, RemoveCrateOutput, SizeInfo, VersionInfo,
14    },
15    utils::format_bytes,
16};
17
18#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
19pub struct CacheCrateFromCratesIOParams {
20    #[schemars(description = "The name of the crate")]
21    pub crate_name: String,
22    #[schemars(description = "The version of the crate")]
23    pub version: String,
24    #[schemars(
25        description = "Optional list of workspace members to cache. If the crate is a workspace and this is not provided, the tool will return a list of available members. Specify member paths relative to the workspace root (e.g., [\"crates/rmcp\", \"crates/rmcp-macros\"])."
26    )]
27    pub members: Option<Vec<String>>,
28    #[schemars(
29        description = "Force re-download and re-cache the crate even if it already exists. Defaults to false. The existing cache is preserved until the update succeeds."
30    )]
31    pub update: Option<bool>,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
35pub struct CacheCrateFromGitHubParams {
36    #[schemars(description = "The name of the crate")]
37    pub crate_name: String,
38    #[schemars(description = "GitHub repository URL (e.g., https://github.com/user/repo)")]
39    pub github_url: String,
40    #[schemars(
41        description = "Branch to use (e.g., 'main', 'develop'). Only one of branch or tag can be specified."
42    )]
43    pub branch: Option<String>,
44    #[schemars(
45        description = "Tag to use (e.g., 'v1.0.0', '0.2.1'). Only one of branch or tag can be specified."
46    )]
47    pub tag: Option<String>,
48    #[schemars(
49        description = "Optional list of workspace members to cache. If the crate is a workspace and this is not provided, the tool will return a list of available members. Specify member paths relative to the workspace root (e.g., [\"crates/rmcp\", \"crates/rmcp-macros\"])."
50    )]
51    pub members: Option<Vec<String>>,
52    #[schemars(
53        description = "Force re-download and re-cache the crate even if it already exists. Defaults to false. The existing cache is preserved until the update succeeds."
54    )]
55    pub update: Option<bool>,
56}
57
58#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
59pub struct CacheCrateFromLocalParams {
60    #[schemars(description = "The name of the crate")]
61    pub crate_name: String,
62    #[schemars(
63        description = "Optional version to use for caching. If not provided, the version from the local crate's Cargo.toml will be used. If provided, it will be validated against the actual version."
64    )]
65    pub version: Option<String>,
66    #[schemars(
67        description = "Local file system path. Supports absolute paths (/path), home paths (~/path), and relative paths (./path, ../path)"
68    )]
69    pub path: String,
70    #[schemars(
71        description = "Optional list of workspace members to cache. If the crate is a workspace and this is not provided, the tool will return a list of available members. Specify member paths relative to the workspace root (e.g., [\"crates/rmcp\", \"crates/rmcp-macros\"])."
72    )]
73    pub members: Option<Vec<String>>,
74    #[schemars(
75        description = "Force re-download and re-cache the crate even if it already exists. Defaults to false. The existing cache is preserved until the update succeeds."
76    )]
77    pub update: Option<bool>,
78}
79
80#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
81pub struct CrateMetadataQuery {
82    #[schemars(description = "The name of the crate")]
83    pub crate_name: String,
84    #[schemars(description = "The version of the crate")]
85    pub version: String,
86    #[schemars(
87        description = "Optional list of workspace members to query (e.g., ['crates/rmcp', 'crates/rmcp-macros'])"
88    )]
89    pub members: Option<Vec<String>>,
90}
91
92#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
93pub struct GetCratesMetadataParams {
94    #[schemars(description = "List of crates and their members to query metadata for")]
95    pub queries: Vec<CrateMetadataQuery>,
96}
97
98#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
99pub struct RemoveCrateParams {
100    #[schemars(description = "The name of the crate")]
101    pub crate_name: String,
102    #[schemars(description = "The version of the crate")]
103    pub version: String,
104}
105
106#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
107pub struct ListCrateVersionsParams {
108    #[schemars(description = "The name of the crate")]
109    pub crate_name: String,
110}
111
112#[derive(Debug, Clone)]
113pub struct CacheTools {
114    cache: Arc<RwLock<CrateCache>>,
115}
116
117impl CacheTools {
118    pub fn new(cache: Arc<RwLock<CrateCache>>) -> Self {
119        Self { cache }
120    }
121
122    pub async fn cache_crate_from_cratesio(
123        &self,
124        params: CacheCrateFromCratesIOParams,
125    ) -> CacheCrateOutput {
126        let cache = self.cache.write().await;
127        let source = CrateSource::CratesIO(params);
128        let json_response = cache.cache_crate_with_source(source).await;
129        serde_json::from_str(&json_response).unwrap_or_else(|_| CacheCrateOutput::Error {
130            error: "Failed to parse cache response".to_string(),
131        })
132    }
133
134    pub async fn cache_crate_from_github(
135        &self,
136        params: CacheCrateFromGitHubParams,
137    ) -> CacheCrateOutput {
138        // Validate that only one of branch or tag is provided
139        match (&params.branch, &params.tag) {
140            (Some(_), Some(_)) => {
141                return CacheCrateOutput::Error {
142                    error: "Only one of 'branch' or 'tag' can be specified, not both".to_string(),
143                };
144            }
145            (None, None) => {
146                return CacheCrateOutput::Error {
147                    error: "Either 'branch' or 'tag' must be specified".to_string(),
148                };
149            }
150            _ => {} // Valid: exactly one is provided
151        }
152
153        let cache = self.cache.write().await;
154        let source = CrateSource::GitHub(params);
155        let json_response = cache.cache_crate_with_source(source).await;
156        serde_json::from_str(&json_response).unwrap_or_else(|_| CacheCrateOutput::Error {
157            error: "Failed to parse cache response".to_string(),
158        })
159    }
160
161    pub async fn cache_crate_from_local(
162        &self,
163        params: CacheCrateFromLocalParams,
164    ) -> CacheCrateOutput {
165        let cache = self.cache.write().await;
166        let source = CrateSource::LocalPath(params);
167        let json_response = cache.cache_crate_with_source(source).await;
168        serde_json::from_str(&json_response).unwrap_or_else(|_| CacheCrateOutput::Error {
169            error: "Failed to parse cache response".to_string(),
170        })
171    }
172
173    pub async fn remove_crate(
174        &self,
175        params: RemoveCrateParams,
176    ) -> Result<RemoveCrateOutput, ErrorOutput> {
177        let cache = self.cache.write().await;
178        match cache
179            .remove_crate(&params.crate_name, &params.version)
180            .await
181        {
182            Ok(_) => Ok(RemoveCrateOutput {
183                status: "success".to_string(),
184                message: format!(
185                    "Successfully removed {}-{}",
186                    params.crate_name, params.version
187                ),
188                crate_name: params.crate_name,
189                version: params.version,
190            }),
191            Err(e) => Err(ErrorOutput::new(format!("Failed to remove crate: {e}"))),
192        }
193    }
194
195    pub async fn list_cached_crates(&self) -> Result<ListCachedCratesOutput, ErrorOutput> {
196        let cache = self.cache.read().await;
197        match cache.list_all_cached_crates().await {
198            Ok(mut crates) => {
199                // Sort by name and version for consistent output
200                crates.sort_by(|a, b| {
201                    a.name.cmp(&b.name).then_with(|| b.version.cmp(&a.version)) // Newer versions first
202                });
203
204                // Calculate total size
205                let total_size_bytes: u64 = crates.iter().map(|c| c.size_bytes).sum();
206
207                // Group by crate name for better organization
208                let mut grouped: std::collections::HashMap<String, Vec<VersionInfo>> =
209                    std::collections::HashMap::new();
210                for crate_meta in crates {
211                    let crate_name = crate_meta.name.clone();
212                    let version = crate_meta.version.clone();
213
214                    // Get workspace members for this crate version
215                    let members = match cache.storage.list_workspace_members(&crate_name, &version)
216                    {
217                        Ok(members) if !members.is_empty() => Some(members),
218                        _ => None,
219                    };
220
221                    let version_info = VersionInfo {
222                        version: crate_meta.version,
223                        cached_at: crate_meta.cached_at.to_string(),
224                        doc_generated: crate_meta.doc_generated,
225                        size_bytes: crate_meta.size_bytes,
226                        size_human: format_bytes(crate_meta.size_bytes),
227                        members,
228                    };
229
230                    grouped.entry(crate_name).or_default().push(version_info);
231                }
232
233                Ok(ListCachedCratesOutput {
234                    crates: grouped.clone(),
235                    total_crates: grouped.len(),
236                    total_versions: grouped.values().map(|v| v.len()).sum::<usize>(),
237                    total_size: SizeInfo {
238                        bytes: total_size_bytes,
239                        human: format_bytes(total_size_bytes),
240                    },
241                })
242            }
243            Err(e) => Err(ErrorOutput::new(format!(
244                "Failed to list cached crates: {e}"
245            ))),
246        }
247    }
248
249    pub async fn list_crate_versions(
250        &self,
251        params: ListCrateVersionsParams,
252    ) -> Result<ListCrateVersionsOutput, ErrorOutput> {
253        let cache = self.cache.read().await;
254
255        // Get all cached metadata for this crate
256        match cache.storage.list_cached_crates() {
257            Ok(all_crates) => {
258                // Filter to just this crate's versions
259                let mut versions: Vec<VersionInfo> = all_crates
260                    .into_iter()
261                    .filter(|meta| meta.name == params.crate_name)
262                    .map(|meta| {
263                        // Get workspace members if any
264                        let members = match cache
265                            .storage
266                            .list_workspace_members(&meta.name, &meta.version)
267                        {
268                            Ok(members) if !members.is_empty() => Some(members),
269                            _ => None,
270                        };
271
272                        VersionInfo {
273                            version: meta.version,
274                            cached_at: meta.cached_at.to_string(),
275                            doc_generated: meta.doc_generated,
276                            size_bytes: meta.size_bytes,
277                            size_human: format_bytes(meta.size_bytes),
278                            members,
279                        }
280                    })
281                    .collect();
282
283                // Sort versions (newest first)
284                versions.sort_by(|a, b| b.version.cmp(&a.version));
285
286                Ok(ListCrateVersionsOutput {
287                    crate_name: params.crate_name.clone(),
288                    versions: versions.clone(),
289                    count: versions.len(),
290                })
291            }
292            Err(e) => Err(ErrorOutput::new(format!(
293                "Failed to get cached versions: {e}"
294            ))),
295        }
296    }
297
298    pub async fn get_crates_metadata(
299        &self,
300        params: GetCratesMetadataParams,
301    ) -> GetCratesMetadataOutput {
302        let cache = self.cache.read().await;
303        let mut metadata_list = Vec::new();
304        let mut total_cached = 0;
305        let total_queried = params.queries.len();
306
307        for query in params.queries {
308            let crate_name = &query.crate_name;
309            let version = &query.version;
310
311            // Check if main crate is cached
312            if cache.storage.is_cached(crate_name, version) {
313                total_cached += 1;
314
315                let main_metadata = match cache.storage.load_metadata(crate_name, version, None) {
316                    Ok(metadata) => {
317                        // Check if docs are analyzed
318                        let analyzed = cache.storage.has_docs(crate_name, version, None);
319
320                        CrateMetadata {
321                            crate_name: crate_name.clone(),
322                            version: version.clone(),
323                            cached: true,
324                            analyzed,
325                            cache_size_bytes: Some(metadata.size_bytes),
326                            cache_size_human: Some(format_bytes(metadata.size_bytes)),
327                            member: None,
328                            workspace_members: None,
329                        }
330                    }
331                    Err(_) => CrateMetadata {
332                        crate_name: crate_name.clone(),
333                        version: version.clone(),
334                        cached: true,
335                        analyzed: false,
336                        cache_size_bytes: None,
337                        cache_size_human: None,
338                        member: None,
339                        workspace_members: None,
340                    },
341                };
342                metadata_list.push(main_metadata);
343            } else {
344                metadata_list.push(CrateMetadata {
345                    crate_name: crate_name.clone(),
346                    version: version.clone(),
347                    cached: false,
348                    analyzed: false,
349                    cache_size_bytes: None,
350                    cache_size_human: None,
351                    member: None,
352                    workspace_members: None,
353                });
354            }
355
356            // Check requested members if any
357            if let Some(members) = query.members {
358                for member_path in members {
359                    if cache
360                        .storage
361                        .is_member_cached(crate_name, version, &member_path)
362                    {
363                        total_cached += 1;
364
365                        let member_metadata = match cache.storage.load_metadata(
366                            crate_name,
367                            version,
368                            Some(&member_path),
369                        ) {
370                            Ok(metadata) => {
371                                let analyzed =
372                                    cache
373                                        .storage
374                                        .has_docs(crate_name, version, Some(&member_path));
375
376                                CrateMetadata {
377                                    crate_name: crate_name.clone(),
378                                    version: version.clone(),
379                                    cached: true,
380                                    analyzed,
381                                    cache_size_bytes: Some(metadata.size_bytes),
382                                    cache_size_human: Some(format_bytes(metadata.size_bytes)),
383                                    member: Some(member_path),
384                                    workspace_members: None,
385                                }
386                            }
387                            Err(_) => CrateMetadata {
388                                crate_name: crate_name.clone(),
389                                version: version.clone(),
390                                cached: true,
391                                analyzed: false,
392                                cache_size_bytes: None,
393                                cache_size_human: None,
394                                member: Some(member_path),
395                                workspace_members: None,
396                            },
397                        };
398                        metadata_list.push(member_metadata);
399                    } else {
400                        metadata_list.push(CrateMetadata {
401                            crate_name: crate_name.clone(),
402                            version: version.clone(),
403                            cached: false,
404                            analyzed: false,
405                            cache_size_bytes: None,
406                            cache_size_human: None,
407                            member: Some(member_path),
408                            workspace_members: None,
409                        });
410                    }
411                }
412            }
413        }
414
415        GetCratesMetadataOutput {
416            metadata: metadata_list,
417            total_queried,
418            total_cached,
419        }
420    }
421}
rust_docs_mcp/cache/tools.rs

rust_docs_mcp/cache/
tools.rs
