From 8d97e37985c1681fda9b6dffe263076443a945a6 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Sun, 30 Mar 2025 22:53:57 +0700
Subject: [PATCH] refactor: clean up

---
 src-tauri/src/core/cmd.rs | 53 +++++++++++++++++++++++++++++----------
 src-tauri/src/core/mcp.rs | 14 +++++++----
 2 files changed, 49 insertions(+), 18 deletions(-)

diff --git a/src-tauri/src/core/cmd.rs b/src-tauri/src/core/cmd.rs
index ef7f3af78..62f53bf56 100644
--- a/src-tauri/src/core/cmd.rs
+++ b/src-tauri/src/core/cmd.rs
@@ -277,6 +277,20 @@ pub async fn stop_server() -> Result<(), String> {
     Ok(())
 }
 
+/// Retrieves all available tools from all MCP servers
+///
+/// # Arguments
+/// * `state` - Application state containing MCP server connections
+///
+/// # Returns
+/// * `Result<Vec<Tool>, String>` - A vector of all tools if successful, or an error message if failed
+///
+/// This function:
+/// 1. Locks the MCP servers mutex to access server connections
+/// 2. Iterates through all connected servers
+/// 3. Gets the list of tools from each server
+/// 4. Combines all tools into a single vector
+/// 5. Returns the combined list of all available tools
 #[tauri::command]
 pub async fn get_tools(state: State<'_, AppState>) -> Result<Vec<Tool>, String> {
     let servers = state.mcp_servers.lock().await;
@@ -294,6 +308,21 @@ pub async fn get_tools(state: State<'_, AppState>) -> Result<Vec<Tool>, String>
     Ok(all_tools)
 }
 
+/// Calls a tool on an MCP server by name with optional arguments
+///
+/// # Arguments
+/// * `state` - Application state containing MCP server connections
+/// * `tool_name` - Name of the tool to call
+/// * `arguments` - Optional map of argument names to values
+///
+/// # Returns
+/// * `Result<CallToolResult, String>` - Result of the tool call if successful, or error message if failed
+///
+/// This function:
+/// 1. Locks the MCP servers mutex to access server connections
+/// 2. Searches through all servers for one containing the named tool
+/// 3. When found, calls the tool on that server with the provided arguments
+/// 4. Returns error if no server has the requested tool
 #[tauri::command]
 pub async fn call_tool(
     state: State<'_, AppState>,
@@ -302,22 +331,20 @@ pub async fn call_tool(
 ) -> Result<CallToolResult, String> {
     let servers = state.mcp_servers.lock().await;
 
+    // Iterate through servers and find the first one that contains the tool
     for (_, service) in servers.iter() {
-        if let Ok(tool) = service.list_all_tools().await {
-            for t in tool {
-                if t.name == tool_name {
-                    let result = service
-                        .call_tool(CallToolRequestParam {
-                            name: tool_name.into(),
-                            arguments,
-                        })
-                        .await
-                        .map_err(|e| e.to_string())?;
-                    return Ok(result);
-                }
+        if let Ok(tools) = service.list_all_tools().await {
+            if tools.iter().any(|t| t.name == tool_name) {
+                return service
+                    .call_tool(CallToolRequestParam {
+                        name: tool_name.into(),
+                        arguments,
+                    })
+                    .await
+                    .map_err(|e| e.to_string());
             }
         }
     }
 
-    return Err(format!("Tool {} not found", tool_name));
+    Err(format!("Tool {} not found", tool_name))
 }
diff --git a/src-tauri/src/core/mcp.rs b/src-tauri/src/core/mcp.rs
index 2196ecb01..654dca2e1 100644
--- a/src-tauri/src/core/mcp.rs
+++ b/src-tauri/src/core/mcp.rs
@@ -3,6 +3,15 @@ use std::{collections::HashMap, sync::Arc};
 use rmcp::{service::RunningService, transport::TokioChildProcess, RoleClient, ServiceExt};
 use tokio::{process::Command, sync::Mutex};
 
+/// Runs MCP commands by reading configuration from a JSON file and initializing servers
+///
+/// # Arguments
+/// * `app_path` - Path to the application directory containing mcp_config.json
+/// * `servers_state` - Shared state containing running MCP services
+///
+/// # Returns
+/// * `Ok(())` if servers were initialized successfully
+/// * `Err(String)` if there was an error reading config or starting servers
 pub async fn run_mcp_commands(
     app_path: String,
     servers_state: Arc<Mutex<HashMap<String, RunningService<RoleClient, ()>>>>,
@@ -26,7 +35,6 @@ pub async fn run_mcp_commands(
         println!("MCP Servers: {servers:#?}");
         if let Some(server_map) = servers.as_object() {
             for (name, config) in server_map {
-                println!("Server Name: {}", name);
                 if let Some(config_obj) = config.as_object() {
                     if let (Some(command), Some(args)) = (
                         config_obj.get("command").and_then(|v| v.as_str()),
@@ -59,10 +67,6 @@ pub async fn run_mcp_commands(
         // Initialize
         let _server_info = service.peer_info();
         println!("Connected to server: {_server_info:#?}");
-        // List tools
-        let _tools = service.list_all_tools().await.map_err(|e| e.to_string())?;
-
-        println!("Tools: {_tools:#?}");
     }
     Ok(())
 }