1use std::path::{Path, PathBuf};
2use mcp_attr::Result;
3use mcp_attr::server::{mcp_server, McpServer};
4use mcp_attr::schema::{GetPromptResult, CallToolResult};
5use std::sync::Mutex;
6use std::collections::HashMap;
7use serde::Deserialize;
8use schemars::JsonSchema;
9use reqwest;
10use crate::mcp::crates_io::{CratesIoClient, RequestOptions, FetchResponse};
11use crate::mcp::function_signatures;
12use crate::mcp::patch::{parse_hunks, find_candidates, rebuild_hunks, rebuild_patch};
13use std::fs;
14use std::process::Command;
15use crate::mcp::prompts::{CODE_CHANGE_WORKFLOW, MCP_TOOLS_GUIDE};
16
17
18pub mod mcp;
19
20
21#[derive(Deserialize, JsonSchema)]
24struct SearchCratesArgs {
25 query: String,
26 page: Option<u32>,
27 per_page: Option<u32>,
28}
29
30#[derive(Deserialize, JsonSchema)]
31struct GetCrateArgs {
32 crate_name: String,
33}
34
35#[derive(Deserialize, JsonSchema)]
36struct GetCrateVersionsArgs {
37 crate_name: String,
38}
39
40#[derive(Deserialize, JsonSchema)]
41struct GetCrateDependenciesArgs {
42 crate_name: String,
43 version: String,
44}
45
46#[derive(Deserialize, JsonSchema)]
47struct ListFunctionSignaturesArgs {
48 file_path: Option<String>,
50}
51
52#[derive(Deserialize, JsonSchema)]
53struct LookupCrateDocsArgs {
54 #[serde(rename = "crateName")]
55 crate_name: Option<String>,
56}
57
58
59pub struct ServerData {
60 pub current_working_dir: PathBuf,
61 pub http_client: reqwest::Client,
62}
63
64pub struct CorrodeMcpServer(pub Mutex<ServerData>);
65
66#[mcp_server]
67impl McpServer for CorrodeMcpServer {
68 #[prompt]
70 async fn search_crates(
71 &self,
72 query: String,
74 _page: Option<String>, _per_page: Option<String>, ) -> Result<GetPromptResult> { let prompt_text = format!("Search crates.io for '{}'. Summarize the top results.", query);
81 Ok(GetPromptResult::from(prompt_text))
83 }
84
85 #[prompt]
87 async fn cd(
88 &self,
89 target_directory: String,
91 ) -> Result<GetPromptResult> {
92 let prompt_text = format!("Please enter the full path to the project directory you want to change to, starting from: {}", target_directory);
93 Ok(GetPromptResult::from(prompt_text))
94 }
95
96 #[prompt]
98 async fn code_change_workflow(
99 &self,
100 _aspect: Option<String>,
102 ) -> Result<GetPromptResult> {
103 let workflow = CODE_CHANGE_WORKFLOW;
104
105 Ok(GetPromptResult::from(workflow))
107 }
108
109 #[prompt]
111 async fn mcp_tools_guide(
112 &self,
113 _tool: Option<String>,
115 ) -> Result<GetPromptResult> {
116 let guide = MCP_TOOLS_GUIDE;
117
118 Ok(GetPromptResult::from(guide))
124 }
125
126
127 #[tool]
196 async fn execute_bash(&self, command: String) -> Result<CallToolResult> {
197 let mut result = String::new();
198
199 let commands: Vec<&str> = if command.contains("&&") {
201 command.split("&&").collect()
202 } else if command.contains(';') {
203 command.split(';').collect()
204 } else {
205 vec![&command]
206 };
207
208 let mut server_state = self.0.lock().unwrap();
210
211 for cmd in commands {
212 let cmd = cmd.trim();
213 let current_dir_path = server_state.current_working_dir.clone();
214
215 if let Some(new_dir) = handle_cd_command(¤t_dir_path, cmd) {
217 if new_dir.exists() && new_dir.is_dir() {
219 server_state.current_working_dir = new_dir.clone();
221 result.push_str(&format!("Changed directory to: {}\n", new_dir.display()));
222 } else {
223 let error_message = format!(
225 "Directory change failed:\n- Command: '{}'\n- Target: {}\n- Current directory: {}\n- Error: The specified directory does not exist or is not accessible",
226 cmd,
227 new_dir.display(),
228 current_dir_path.display()
229 );
230
231 result.push_str(&format!("{}\n", error_message));
232
233 mcp_attr::bail!("{}", error_message);
236 }
237
238 if cmd == "cd" || (cmd.starts_with("cd ") && !cmd.contains("&&") && !cmd.contains(';')) {
240 continue;
241 }
242 }
243
244 let output = Command::new("bash")
247 .arg("-l") .current_dir(¤t_dir_path) .arg("-c")
250 .arg(cmd) .output();
252
253 match output {
254 Ok(output) => {
255 let stdout = String::from_utf8_lossy(&output.stdout).to_string();
256 let stderr = String::from_utf8_lossy(&output.stderr).to_string();
257
258 let cmd_result = format!("$ {}\n", cmd);
259 result.push_str(&cmd_result);
260
261 let exit_status = output.status.code().unwrap_or(-1);
262 let cmd_is_error = !output.status.success();
263 result.push_str(&format!("Exit code: {}\n", exit_status));
266
267 if !stdout.is_empty() {
268 result.push_str(&format!("\nStandard output:\n{}", stdout));
269 }
270
271 if !stderr.is_empty() {
272 result.push_str(&format!("\nStandard error:\n{}\n", stderr));
273 }
274
275 if cmd_is_error {
277 let error_message = format!("Command '{}' failed with exit code {}.\n\nSTDOUT:\n{}\n\nSTDERR:\n{}",
279 cmd,
280 exit_status,
281 stdout,
282 stderr
283 );
284
285 result.push_str(&format!("{}", error_message));
287 }
288 },
289 Err(e) => {
290 let error_details = format!(
292 "Failed to execute command '{}':\n- Error: {}\n- Working Directory: {}\n- Note: This typically happens when the command or shell is not found, or due to permissions issues",
293 cmd,
294 e,
295 current_dir_path.display()
296 );
297
298 result.push_str(&format!("{}\n", error_details));
299
300 mcp_attr::bail!("{}", error_details);
302 }
303 }
304 }
305
306 drop(server_state);
308
309 Ok(CallToolResult::from(result))
312 }
313
314 #[tool]
318 async fn patch_file(&self,
333 file_name: String,
335 patch: String) -> Result<CallToolResult> {
337 let current_dir = self.0.lock().unwrap().current_working_dir.clone();
339 let file_path_buf = resolve_path(¤t_dir, &file_name);
340 let display_path = file_path_buf.display().to_string();
341
342 let mut old_content = match fs::read_to_string(&file_path_buf) {
344 Ok(content) => content,
345 Err(e) => mcp_attr::bail!("Failed to read file {}: {}", display_path, e),
346 };
347
348 if !old_content.ends_with('\n') {
350 old_content.push('\n');
351 }
352
353 let old_hunks = match parse_hunks(&patch) {
355 Ok(hunks) => hunks,
356 Err(e) => mcp_attr::bail!("Failed to parse patch: {}", e),
357 };
358
359 let candidates = find_candidates(&old_content, &old_hunks);
361
362 let new_hunks = rebuild_hunks(&candidates);
364
365 let updated_patch = match rebuild_patch(&patch, &new_hunks) {
367 Ok(patch) => patch,
368 Err(e) => mcp_attr::bail!("Failed to render fixed patch: {}", e),
369 };
370
371 let diffy_patch = match diffy::Patch::from_str(&updated_patch) {
373 Ok(patch) => patch,
374 Err(e) => mcp_attr::bail!("Failed to parse patch: {}", e),
375 };
376
377 let patched = match diffy::apply(&old_content, &diffy_patch) {
379 Ok(patched) => patched,
380 Err(e) => mcp_attr::bail!("Failed to apply patch: {}", e),
381 };
382
383 match fs::write(&file_path_buf, &patched) {
385 Ok(_) => {
386 if new_hunks.len() != old_hunks.len() {
387 let failed = old_hunks
388 .iter()
389 .filter(|h| !new_hunks.iter().any(|h2| h2.body == h.body))
390 .collect::<Vec<_>>();
391
392 return Ok(CallToolResult::from(format!(
393 "Failed to apply all hunks. {} hunks failed to apply.\n\nThe following hunks failed to apply as their context lines could not be matched to the file, no changes were applied:\n\n---\n{}\n---\n\nMake sure all lines are correct. Are you also sure that the changes have not been applied already?",
394 failed.len(),
395 failed.iter().map(|h| h.body.as_str()).collect::<Vec<_>>().join("\n")
396 )));
397 }
398
399 Ok(CallToolResult::from(format!("Patch applied successfully to {}", display_path)))
400 },
401 Err(e) => mcp_attr::bail!("Error writing to file '{}': {}", display_path, e),
402 }
403 }
404
405 #[tool]
407 async fn write_file(&self, file_path: String, content: String) -> Result<CallToolResult> {
408 let current_dir = self.0.lock().unwrap().current_working_dir.clone();
409 let file_path_buf = resolve_path(¤t_dir, &file_path);
410 let display_path = file_path_buf.display().to_string();
411
412 if let Some(parent) = file_path_buf.parent() {
413 if !parent.exists() {
414 if let Err(e) = fs::create_dir_all(parent) {
415 mcp_attr::bail!("Error creating directory structure for '{}': {}", display_path, e); }
417 }
418 }
419
420 match fs::write(&file_path_buf, &content) {
421 Ok(_) => Ok(CallToolResult::from(format!("Successfully wrote to file: {}", display_path))), Err(e) => mcp_attr::bail!("Error writing to file '{}': {}", display_path, e), }
424 }
425
426 #[tool]
429 async fn check_code(&self) -> Result<CallToolResult> {
430 let current_dir = self.0.lock().unwrap().current_working_dir.clone();
431 let cargo_toml_path = current_dir.join("Cargo.toml");
432
433 if !cargo_toml_path.exists() {
434 mcp_attr::bail!("No Cargo.toml found in '{}'. This doesn't appear to be a Rust project.", current_dir.display()); }
436
437 self.execute_bash("cargo check".to_string()).await }
439
440 #[tool]
445 async fn read_file(&self, file_path: String) -> Result<CallToolResult> {
446 let current_dir = self.0.lock().unwrap().current_working_dir.clone();
447 let file_path_buf = resolve_path(¤t_dir, &file_path);
448 let display_path = file_path_buf.display().to_string();
449
450 match fs::read_to_string(&file_path_buf) {
451 Ok(content) => {
452 Ok(CallToolResult::from(content)) },
455 Err(e) => mcp_attr::bail!("Error reading file '{}': {}", display_path, e), }
457 }
458 #[tool]
465 async fn tool_search_crates(&self, args: SearchCratesArgs) -> Result<String> {
466 let mut query_params = HashMap::new();
467 query_params.insert("q".to_string(), args.query.clone());
468
469 let crates_client = {
471 let server_data = self.0.lock().unwrap();
472 CratesIoClient::with_client(server_data.http_client.clone())
473 }; if let Some(page) = args.page {
476 query_params.insert("page".to_string(), page.to_string());
477 }
478 if let Some(per_page) = args.per_page {
479 query_params.insert("per_page".to_string(), per_page.to_string());
480 }
481 let options = RequestOptions { params: Some(query_params), ..Default::default() };
482
483 match crates_client.get("crates", Some(options)).await {
484 Ok(response) => match response {
485 FetchResponse::Json { data, status, .. } => {
486 let json_string = match serde_json::to_string_pretty(&data) {
487 Ok(s) => s,
488 Err(e) => mcp_attr::bail!("Error serializing JSON response: {}", e),
489 };
490 Ok(format!("Status: {}\n\n{}", status, json_string))
491 },
492 FetchResponse::Text { data, status, .. } => {
493 Ok(format!("Status: {}\n{}", status, data))
494 }
495 },
496 Err(e) => mcp_attr::bail!("Error searching crates: {}", e),
497 }
498 }
499
500 #[tool]
502 async fn get_crate(&self, args: GetCrateArgs) -> Result<String> {
503 let (crates_client, path) = {
505 let server_data = self.0.lock().unwrap();
506 let client = CratesIoClient::with_client(server_data.http_client.clone());
507 let path_str = format!("crates/{}", args.crate_name);
508 (client, path_str)
509 };
510
511 match crates_client.get(&path, None).await {
512 Ok(response) => match response {
513 FetchResponse::Json { data, status, .. } => {
514 let json_string = match serde_json::to_string_pretty(&data) {
515 Ok(s) => s,
516 Err(e) => mcp_attr::bail!("Error serializing JSON response: {}", e),
517 };
518 Ok(format!("Status: {}\n\n{}", status, json_string))
519 },
520 FetchResponse::Text { data, status, .. } => {
521 Ok(format!("Status: {}\n{}", status, data))
522 }
523 },
524 Err(e) => mcp_attr::bail!("Error getting crate details: {}", e),
525 }
526 }
527
528 #[tool]
530 async fn get_crate_versions(&self, args: GetCrateVersionsArgs) -> Result<String> {
531 let (crates_client, path) = {
533 let server_data = self.0.lock().unwrap();
534 let client = CratesIoClient::with_client(server_data.http_client.clone());
535 let path_str = format!("crates/{}/versions", args.crate_name);
536 (client, path_str)
537 };
538
539 match crates_client.get(&path, None).await {
540 Ok(response) => match response {
541 FetchResponse::Json { data, status, .. } => {
542 let json_string = serde_json::to_string_pretty(&data)?;
543 Ok(format!("Status: {}\n\n{}", status, json_string))
544 },
545 FetchResponse::Text { data, status, .. } => {
546 Ok(format!("Status: {}\n{}", status, data) )
547 }
548 },
549 Err(e) => mcp_attr::bail!("Error getting crate versions: {}", e),
550 }
551 }
552
553 #[tool]
555 async fn get_crate_dependencies(&self, args: GetCrateDependenciesArgs) -> Result<String> {
556 let (crates_client, path) = {
558 let server_data = self.0.lock().unwrap();
559 let client = CratesIoClient::with_client(server_data.http_client.clone());
560 let path_str = format!("crates/{}/{}/dependencies", args.crate_name, args.version);
561 (client, path_str)
562 };
563
564 match crates_client.get(&path, None).await {
565 Ok(response) => match response {
566 FetchResponse::Json { data, status, .. } => {
567 let json_string = serde_json::to_string_pretty(&data)?;
568
569 Ok(format!("Status: {}\n\n{}", status, json_string))
570 },
571 FetchResponse::Text { data, status, .. } => {
572 Ok(format!("Status: {}\n{}", status, data))
573 }
574 },
575 Err(e) => mcp_attr::bail!("Error getting crate dependencies: {}", e),
576 }
577 }
578
579 #[tool]
581 async fn lookup_crate_docs(&self, args: LookupCrateDocsArgs) -> Result<CallToolResult> {
582 let crate_name = args.crate_name.unwrap_or_else(|| "tokio".to_string());
583 let url = format!("https://docs.rs/{}/latest/{}/", crate_name, crate_name.replace('-', "_"));
584
585 let client = {
587 let server_state = self.0.lock().unwrap();
588 server_state.http_client.clone()
589 };
590
591 match client.get(&url).send().await {
592 Ok(response) => {
593 if !response.status().is_success() {
594 let error_text = format!("Error: Could not fetch documentation from {}. HTTP status: {}", url, response.status());
595 mcp_attr::bail_public!(mcp_attr::ErrorCode::INTERNAL_ERROR, "{}", error_text);
596 }
597
598 match response.text().await {
599 Ok(html_content) => {
600 let html_result = html2text::from_read(html_content.as_bytes(), 130);
602 if let Err(e) = &html_result {
603 mcp_attr::bail!("Error converting HTML to text: {}", e);
604 }
605 let text_content = html_result.unwrap();
606
607 const MAX_LENGTH: usize = 8000;
609 let truncated_text = if text_content.chars().count() > MAX_LENGTH {
610 format!("{}\n\n[Content truncated. Full documentation available at {}]",
611 text_content.chars().take(MAX_LENGTH).collect::<String>(), url)
612 } else {
613 text_content
614 };
615 Ok(CallToolResult::from(truncated_text))
616 }
617 Err(e) => {
618 mcp_attr::bail!("Error reading documentation content: {}", e)
619 }
620 }
621 }
622 Err(e) => {
623 mcp_attr::bail!("Error fetching documentation from {}: {}", url, e)
624 }
625 }
626 }
627
628 #[tool]
630 async fn list_function_signatures(&self, args: Option<ListFunctionSignaturesArgs>) -> Result<CallToolResult> {
631 let current_dir = self.0.lock().unwrap().current_working_dir.clone();
632
633 let mut result_string = format!("Current working directory: {}\n\n", current_dir.display());
635
636 let signatures = if let Some(args) = args {
637 if let Some(file_path) = args.file_path {
638 let file_path_buf = resolve_path(¤t_dir, &file_path);
639 result_string.push_str(&format!("Checking specific file: {}\n\n", file_path_buf.display()));
640
641 if !file_path_buf.exists() {
642 return Ok(CallToolResult::from(format!(
643 "Error: File '{}' does not exist.",
644 file_path_buf.display()
645 )));
646 }
647
648 function_signatures::extract_function_signatures(&file_path_buf, None)
649 } else {
650 result_string.push_str("Scanning entire project directory\n\n");
651 function_signatures::extract_project_signatures(¤t_dir)
652 }
653 } else {
654 result_string.push_str("Scanning entire project directory\n\n");
655 function_signatures::extract_project_signatures(¤t_dir)
656 };
657
658 if signatures.is_empty() {
659 result_string.push_str("No function signatures found.");
660 return Ok(CallToolResult::from(result_string));
661 }
662
663 result_string.push_str(&format!("Found {} function signatures:\n\n", signatures.len()));
665
666 for sig in signatures {
667 let formatted_line = format!(
669 "{}:{}: {}\n",
670 sig.file_path,
671 sig.line_number,
672 sig.signature.trim() );
674 result_string.push_str(&formatted_line);
675 }
676
677 Ok(CallToolResult::from(result_string))
678 }
679
680}
681pub fn resolve_path(current_dir: &Path, file_path: &str) -> PathBuf {
684 if file_path.starts_with('/') {
685 PathBuf::from(file_path)
687 } else if file_path.starts_with("~/") || file_path == "~" {
688 let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
690 home.join(file_path.trim_start_matches("~/"))
691 } else {
692 current_dir.join(file_path)
694 }
695}
696
697pub fn handle_cd_command(current_dir: &Path, command: &str) -> Option<PathBuf> {
700 let command = command.trim();
701
702 if command == "cd" || command.starts_with("cd ") {
704 let parts: Vec<&str> = command.splitn(2, ' ').collect();
705 if parts.len() == 1 {
706 return Some(dirs::home_dir().unwrap_or_else(|| PathBuf::from(".")));
708 } else if parts.len() == 2 {
709 let dir = parts[1].trim();
710 let new_path = resolve_path(current_dir, dir);
712 return Some(new_path);
713 }
714 }
715 None
716}
717