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, 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 match (¶ms.branch, ¶ms.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 _ => {} }
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(¶ms.crate_name, ¶ms.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 crates.sort_by(|a, b| {
205 a.name.cmp(&b.name).then_with(|| b.version.cmp(&a.version)) });
207
208 let total_size_bytes: u64 = crates.iter().map(|c| c.size_bytes).sum();
210
211 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 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 match cache.storage.list_cached_crates() {
261 Ok(all_crates) => {
262 let mut versions: Vec<VersionInfo> = all_crates
264 .into_iter()
265 .filter(|meta| meta.name == params.crate_name)
266 .map(|meta| {
267 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 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 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 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 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}