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/// Empty params struct for list_cached_crates tool
113#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
114pub struct ListCachedCratesParams {}
115
116#[derive(Debug, Clone)]
117pub struct CacheTools {
118    cache: Arc<RwLock<CrateCache>>,
119}
120
121impl CacheTools {
122    pub fn new(cache: Arc<RwLock<CrateCache>>) -> Self {
123        Self { cache }
124    }
125
126    pub async fn cache_crate_from_cratesio(
127        &self,
128        params: CacheCrateFromCratesIOParams,
129    ) -> CacheCrateOutput {
130        let cache = self.cache.write().await;
131        let source = CrateSource::CratesIO(params);
132        let json_response = cache.cache_crate_with_source(source).await;
133        serde_json::from_str(&json_response).unwrap_or_else(|_| CacheCrateOutput::Error {
134            error: "Failed to parse cache response".to_string(),
135        })
136    }
137
138    pub async fn cache_crate_from_github(
139        &self,
140        params: CacheCrateFromGitHubParams,
141    ) -> CacheCrateOutput {
142        // Validate that only one of branch or tag is provided
143        match (&params.branch, &params.tag) {
144            (Some(_), Some(_)) => {
145                return CacheCrateOutput::Error {
146                    error: "Only one of 'branch' or 'tag' can be specified, not both".to_string(),
147                };
148            }
149            (None, None) => {
150                return CacheCrateOutput::Error {
151                    error: "Either 'branch' or 'tag' must be specified".to_string(),
152                };
153            }
154            _ => {} // Valid: exactly one is provided
155        }
156
157        let cache = self.cache.write().await;
158        let source = CrateSource::GitHub(params);
159        let json_response = cache.cache_crate_with_source(source).await;
160        serde_json::from_str(&json_response).unwrap_or_else(|_| CacheCrateOutput::Error {
161            error: "Failed to parse cache response".to_string(),
162        })
163    }
164
165    pub async fn cache_crate_from_local(
166        &self,
167        params: CacheCrateFromLocalParams,
168    ) -> CacheCrateOutput {
169        let cache = self.cache.write().await;
170        let source = CrateSource::LocalPath(params);
171        let json_response = cache.cache_crate_with_source(source).await;
172        serde_json::from_str(&json_response).unwrap_or_else(|_| CacheCrateOutput::Error {
173            error: "Failed to parse cache response".to_string(),
174        })
175    }
176
177    pub async fn remove_crate(
178        &self,
179        params: RemoveCrateParams,
180    ) -> Result<RemoveCrateOutput, ErrorOutput> {
181        let cache = self.cache.write().await;
182        match cache
183            .remove_crate(&params.crate_name, &params.version)
184            .await
185        {
186            Ok(_) => Ok(RemoveCrateOutput {
187                status: "success".to_string(),
188                message: format!(
189                    "Successfully removed {}-{}",
190                    params.crate_name, params.version
191                ),
192                crate_name: params.crate_name,
193                version: params.version,
194            }),
195            Err(e) => Err(ErrorOutput::new(format!("Failed to remove crate: {e}"))),
196        }
197    }
198
199    pub async fn list_cached_crates(&self) -> Result<ListCachedCratesOutput, ErrorOutput> {
200        let cache = self.cache.read().await;
201        match cache.list_all_cached_crates().await {
202            Ok(mut crates) => {
203                // Sort by name and version for consistent output
204                crates.sort_by(|a, b| {
205                    a.name.cmp(&b.name).then_with(|| b.version.cmp(&a.version)) // Newer versions first
206                });
207
208                // Calculate total size
209                let total_size_bytes: u64 = crates.iter().map(|c| c.size_bytes).sum();
210
211                // Group by crate name for better organization
212                let mut grouped: std::collections::HashMap<String, Vec<VersionInfo>> =
213                    std::collections::HashMap::new();
214                for crate_meta in crates {
215                    let crate_name = crate_meta.name.clone();
216                    let version = crate_meta.version.clone();
217
218                    // Get workspace members for this crate version
219                    let members = match cache.storage.list_workspace_members(&crate_name, &version)
220                    {
221                        Ok(members) if !members.is_empty() => Some(members),
222                        _ => None,
223                    };
224
225                    let version_info = VersionInfo {
226                        version: crate_meta.version,
227                        cached_at: crate_meta.cached_at.to_string(),
228                        doc_generated: crate_meta.doc_generated,
229                        size_bytes: crate_meta.size_bytes,
230                        size_human: format_bytes(crate_meta.size_bytes),
231                        members,
232                    };
233
234                    grouped.entry(crate_name).or_default().push(version_info);
235                }
236
237                Ok(ListCachedCratesOutput {
238                    crates: grouped.clone(),
239                    total_crates: grouped.len(),
240                    total_versions: grouped.values().map(|v| v.len()).sum::<usize>(),
241                    total_size: SizeInfo {
242                        bytes: total_size_bytes,
243                        human: format_bytes(total_size_bytes),
244                    },
245                })
246            }
247            Err(e) => Err(ErrorOutput::new(format!(
248                "Failed to list cached crates: {e}"
249            ))),
250        }
251    }
252
253    pub async fn list_crate_versions(
254        &self,
255        params: ListCrateVersionsParams,
256    ) -> Result<ListCrateVersionsOutput, ErrorOutput> {
257        let cache = self.cache.read().await;
258
259        // Get all cached metadata for this crate
260        match cache.storage.list_cached_crates() {
261            Ok(all_crates) => {
262                // Filter to just this crate's versions
263                let mut versions: Vec<VersionInfo> = all_crates
264                    .into_iter()
265                    .filter(|meta| meta.name == params.crate_name)
266                    .map(|meta| {
267                        // Get workspace members if any
268                        let members = match cache
269                            .storage
270                            .list_workspace_members(&meta.name, &meta.version)
271                        {
272                            Ok(members) if !members.is_empty() => Some(members),
273                            _ => None,
274                        };
275
276                        VersionInfo {
277                            version: meta.version,
278                            cached_at: meta.cached_at.to_string(),
279                            doc_generated: meta.doc_generated,
280                            size_bytes: meta.size_bytes,
281                            size_human: format_bytes(meta.size_bytes),
282                            members,
283                        }
284                    })
285                    .collect();
286
287                // Sort versions (newest first)
288                versions.sort_by(|a, b| b.version.cmp(&a.version));
289
290                Ok(ListCrateVersionsOutput {
291                    crate_name: params.crate_name.clone(),
292                    versions: versions.clone(),
293                    count: versions.len(),
294                })
295            }
296            Err(e) => Err(ErrorOutput::new(format!(
297                "Failed to get cached versions: {e}"
298            ))),
299        }
300    }
301
302    pub async fn get_crates_metadata(
303        &self,
304        params: GetCratesMetadataParams,
305    ) -> GetCratesMetadataOutput {
306        let cache = self.cache.read().await;
307        let mut metadata_list = Vec::new();
308        let mut total_cached = 0;
309        let total_queried = params.queries.len();
310
311        for query in params.queries {
312            let crate_name = &query.crate_name;
313            let version = &query.version;
314
315            // Check if main crate is cached
316            if cache.storage.is_cached(crate_name, version) {
317                total_cached += 1;
318
319                let main_metadata = match cache.storage.load_metadata(crate_name, version, None) {
320                    Ok(metadata) => {
321                        // Check if docs are analyzed
322                        let analyzed = cache.storage.has_docs(crate_name, version, None);
323
324                        CrateMetadata {
325                            crate_name: crate_name.clone(),
326                            version: version.clone(),
327                            cached: true,
328                            analyzed,
329                            cache_size_bytes: Some(metadata.size_bytes),
330                            cache_size_human: Some(format_bytes(metadata.size_bytes)),
331                            member: None,
332                            workspace_members: None,
333                        }
334                    }
335                    Err(_) => CrateMetadata {
336                        crate_name: crate_name.clone(),
337                        version: version.clone(),
338                        cached: true,
339                        analyzed: false,
340                        cache_size_bytes: None,
341                        cache_size_human: None,
342                        member: None,
343                        workspace_members: None,
344                    },
345                };
346                metadata_list.push(main_metadata);
347            } else {
348                metadata_list.push(CrateMetadata {
349                    crate_name: crate_name.clone(),
350                    version: version.clone(),
351                    cached: false,
352                    analyzed: false,
353                    cache_size_bytes: None,
354                    cache_size_human: None,
355                    member: None,
356                    workspace_members: None,
357                });
358            }
359
360            // Check requested members if any
361            if let Some(members) = query.members {
362                for member_path in members {
363                    if cache
364                        .storage
365                        .is_member_cached(crate_name, version, &member_path)
366                    {
367                        total_cached += 1;
368
369                        let member_metadata = match cache.storage.load_metadata(
370                            crate_name,
371                            version,
372                            Some(&member_path),
373                        ) {
374                            Ok(metadata) => {
375                                let analyzed =
376                                    cache
377                                        .storage
378                                        .has_docs(crate_name, version, Some(&member_path));
379
380                                CrateMetadata {
381                                    crate_name: crate_name.clone(),
382                                    version: version.clone(),
383                                    cached: true,
384                                    analyzed,
385                                    cache_size_bytes: Some(metadata.size_bytes),
386                                    cache_size_human: Some(format_bytes(metadata.size_bytes)),
387                                    member: Some(member_path),
388                                    workspace_members: None,
389                                }
390                            }
391                            Err(_) => CrateMetadata {
392                                crate_name: crate_name.clone(),
393                                version: version.clone(),
394                                cached: true,
395                                analyzed: false,
396                                cache_size_bytes: None,
397                                cache_size_human: None,
398                                member: Some(member_path),
399                                workspace_members: None,
400                            },
401                        };
402                        metadata_list.push(member_metadata);
403                    } else {
404                        metadata_list.push(CrateMetadata {
405                            crate_name: crate_name.clone(),
406                            version: version.clone(),
407                            cached: false,
408                            analyzed: false,
409                            cache_size_bytes: None,
410                            cache_size_human: None,
411                            member: Some(member_path),
412                            workspace_members: None,
413                        });
414                    }
415                }
416            }
417        }
418
419        GetCratesMetadataOutput {
420            metadata: metadata_list,
421            total_queried,
422            total_cached,
423        }
424    }
425}