fix: migrate new mcp server config (#6651)

louis-jan · web-flow · commit 54d17c9c7292 · 2025-09-30T00:07:57.000+07:00
diff --git a/src-tauri/src/core/mcp/helpers.rs b/src-tauri/src/core/mcp/helpers.rs
@@ -934,3 +934,47 @@ pub async fn should_restart_server(
         }
     }
 }
+
+// Add a new server configuration to the MCP config file
+pub fn add_server_config<R: Runtime>(
+    app_handle: tauri::AppHandle<R>,
+    server_key: String,
+    server_value: Value,
+) -> Result<(), String> {
+    add_server_config_with_path(app_handle, server_key, server_value, None)
+}
+
+// Add a new server configuration to the MCP config file with custom path support
+pub fn add_server_config_with_path<R: Runtime>(
+    app_handle: tauri::AppHandle<R>,
+    server_key: String,
+    server_value: Value,
+    config_filename: Option<&str>,
+) -> Result<(), String> {
+    let config_filename = config_filename.unwrap_or("mcp_config.json");
+    let config_path = get_jan_data_folder_path(app_handle).join(config_filename);
+
+    let mut config: Value = serde_json::from_str(
+        &std::fs::read_to_string(&config_path)
+            .map_err(|e| format!("Failed to read config file: {e}"))?,
+    )
+    .map_err(|e| format!("Failed to parse config: {e}"))?;
+
+    config
+        .as_object_mut()
+        .ok_or("Config root is not an object")?
+        .entry("mcpServers")
+        .or_insert_with(|| Value::Object(serde_json::Map::new()))
+        .as_object_mut()
+        .ok_or("mcpServers is not an object")?
+        .insert(server_key, server_value);
+
+    std::fs::write(
+        &config_path,
+        serde_json::to_string_pretty(&config)
+            .map_err(|e| format!("Failed to serialize config: {e}"))?,
+    )
+    .map_err(|e| format!("Failed to write config file: {e}"))?;
+
+    Ok(())
+}
diff --git a/src-tauri/src/core/mcp/tests.rs b/src-tauri/src/core/mcp/tests.rs
@@ -1,4 +1,4 @@
-use super::helpers::run_mcp_commands;
+use super::helpers::{add_server_config, add_server_config_with_path, run_mcp_commands};
 use crate::core::app::commands::get_jan_data_folder_path;
 use crate::core::state::SharedMcpServers;
 use std::collections::HashMap;
@@ -37,3 +37,140 @@ async fn test_run_mcp_commands() {
     // Clean up the mock config file
     std::fs::remove_file(&config_path).expect("Failed to remove config file");
 }
+
+#[test]
+fn test_add_server_config_new_file() {
+    let app = mock_app();
+    let app_path = get_jan_data_folder_path(app.handle().clone());
+    let config_path = app_path.join("mcp_config_test_new.json");
+
+    // Ensure the directory exists
+    if let Some(parent) = config_path.parent() {
+        std::fs::create_dir_all(parent).expect("Failed to create parent directory");
+    }
+
+    // Create initial config file with empty mcpServers
+    let mut file = File::create(&config_path).expect("Failed to create config file");
+    file.write_all(b"{\"mcpServers\":{}}")
+        .expect("Failed to write to config file");
+    drop(file);
+
+    // Test adding a new server config
+    let server_value = serde_json::json!({
+        "command": "npx",
+        "args": ["-y", "test-server"],
+        "env": { "TEST_API_KEY": "test_key" },
+        "active": false
+    });
+
+    let result = add_server_config_with_path(
+        app.handle().clone(),
+        "test_server".to_string(),
+        server_value.clone(),
+        Some("mcp_config_test_new.json"),
+    );
+
+    assert!(result.is_ok(), "Failed to add server config: {:?}", result);
+
+    // Verify the config was added correctly
+    let config_content = std::fs::read_to_string(&config_path)
+        .expect("Failed to read config file");
+    let config: serde_json::Value = serde_json::from_str(&config_content)
+        .expect("Failed to parse config");
+
+    assert!(config["mcpServers"]["test_server"].is_object());
+    assert_eq!(config["mcpServers"]["test_server"]["command"], "npx");
+    assert_eq!(config["mcpServers"]["test_server"]["args"][0], "-y");
+    assert_eq!(config["mcpServers"]["test_server"]["args"][1], "test-server");
+
+    // Clean up
+    std::fs::remove_file(&config_path).expect("Failed to remove config file");
+}
+
+#[test]
+fn test_add_server_config_existing_servers() {
+    let app = mock_app();
+    let app_path = get_jan_data_folder_path(app.handle().clone());
+    let config_path = app_path.join("mcp_config_test_existing.json");
+
+    // Ensure the directory exists
+    if let Some(parent) = config_path.parent() {
+        std::fs::create_dir_all(parent).expect("Failed to create parent directory");
+    }
+
+    // Create config file with existing server
+    let initial_config = serde_json::json!({
+        "mcpServers": {
+            "existing_server": {
+                "command": "existing_command",
+                "args": ["arg1"],
+                "active": true
+            }
+        }
+    });
+
+    let mut file = File::create(&config_path).expect("Failed to create config file");
+    file.write_all(serde_json::to_string_pretty(&initial_config).unwrap().as_bytes())
+        .expect("Failed to write to config file");
+    drop(file);
+
+    // Add new server
+    let new_server_value = serde_json::json!({
+        "command": "new_command",
+        "args": ["new_arg"],
+        "active": false
+    });
+
+    let result = add_server_config_with_path(
+        app.handle().clone(),
+        "new_server".to_string(),
+        new_server_value,
+        Some("mcp_config_test_existing.json"),
+    );
+
+    assert!(result.is_ok(), "Failed to add server config: {:?}", result);
+
+    // Verify both servers exist
+    let config_content = std::fs::read_to_string(&config_path)
+        .expect("Failed to read config file");
+    let config: serde_json::Value = serde_json::from_str(&config_content)
+        .expect("Failed to parse config");
+
+    // Check existing server is still there
+    assert!(config["mcpServers"]["existing_server"].is_object());
+    assert_eq!(config["mcpServers"]["existing_server"]["command"], "existing_command");
+
+    // Check new server was added
+    assert!(config["mcpServers"]["new_server"].is_object());
+    assert_eq!(config["mcpServers"]["new_server"]["command"], "new_command");
+
+    // Clean up
+    std::fs::remove_file(&config_path).expect("Failed to remove config file");
+}
+
+#[test]
+fn test_add_server_config_missing_config_file() {
+    let app = mock_app();
+    let app_path = get_jan_data_folder_path(app.handle().clone());
+    let config_path = app_path.join("nonexistent_config.json");
+
+    // Ensure the file doesn't exist
+    if config_path.exists() {
+        std::fs::remove_file(&config_path).ok();
+    }
+
+    let server_value = serde_json::json!({
+        "command": "test",
+        "args": [],
+        "active": false
+    });
+
+    let result = add_server_config(
+        app.handle().clone(),
+        "test".to_string(),
+        server_value,
+    );
+
+    assert!(result.is_err(), "Expected error when config file doesn't exist");
+    assert!(result.unwrap_err().contains("Failed to read config file"));
+}
diff --git a/src-tauri/src/core/setup.rs b/src-tauri/src/core/setup.rs
@@ -3,39 +3,23 @@ use std::{
     fs::{self, File},
     io::Read,
     path::PathBuf,
+    sync::Arc,
 };
 use tar::Archive;
 use tauri::{
     menu::{Menu, MenuItem, PredefinedMenuItem},
     tray::{MouseButton, MouseButtonState, TrayIcon, TrayIconBuilder, TrayIconEvent},
-    App, Emitter, Manager,
+    App, Emitter, Manager, Wry,
 };
-use tauri_plugin_store::StoreExt;
-// use tokio::sync::Mutex;
-// use tokio::time::{sleep, Duration}; // Using tokio::sync::Mutex
-//                                     // MCP
+use tauri_plugin_store::Store;
+
+use crate::core::mcp::helpers::add_server_config;
 
-// MCP
 use super::{
-    app::commands::get_jan_data_folder_path, extensions::commands::get_jan_extensions_path,
-    mcp::helpers::run_mcp_commands, state::AppState,
+    extensions::commands::get_jan_extensions_path, mcp::helpers::run_mcp_commands, state::AppState,
 };
 
 pub fn install_extensions(app: tauri::AppHandle, force: bool) -> Result<(), String> {
-    let mut store_path = get_jan_data_folder_path(app.clone());
-    store_path.push("store.json");
-    let store = app.store(store_path).expect("Store not initialized");
-    let stored_version = store
-        .get("version")
-        .and_then(|v| v.as_str().map(String::from))
-        .unwrap_or_default();
-
-    let app_version = app
-        .config()
-        .version
-        .clone()
-        .unwrap_or_else(|| "".to_string());
-
     let extensions_path = get_jan_extensions_path(app.clone());
     let pre_install_path = app
         .path()
@@ -50,13 +34,8 @@ pub fn install_extensions(app: tauri::AppHandle, force: bool) -> Result<(), Stri
     if std::env::var("IS_CLEAN").is_ok() {
         clean_up = true;
     }
-    log::info!(
-        "Installing extensions. Clean up: {}, Stored version: {}, App version: {}",
-        clean_up,
-        stored_version,
-        app_version
-    );
-    if !clean_up && stored_version == app_version && extensions_path.exists() {
+    log::info!("Installing extensions. Clean up: {}", clean_up);
+    if !clean_up && extensions_path.exists() {
         return Ok(());
     }
 
@@ -160,10 +139,36 @@ pub fn install_extensions(app: tauri::AppHandle, force: bool) -> Result<(), Stri
     )
     .map_err(|e| e.to_string())?;
 
-    // Store the new app version
-    store.set("version", serde_json::json!(app_version));
-    store.save().expect("Failed to save store");
+    Ok(())
+}
 
+// Migrate MCP servers configuration
+pub fn migrate_mcp_servers(
+    app_handle: tauri::AppHandle,
+    store: Arc<Store<Wry>>,
+) -> Result<(), String> {
+    let mcp_version = store
+        .get("mcp_version")
+        .and_then(|v| v.as_i64())
+        .unwrap_or_else(|| 0);
+    if mcp_version < 1 {
+        log::info!("Migrating MCP schema version 1");
+        let result = add_server_config(
+            app_handle,
+            "exa".to_string(),
+            serde_json::json!({
+                  "command": "npx",
+                  "args": ["-y", "exa-mcp-server"],
+                  "env": { "EXA_API_KEY": "YOUR_EXA_API_KEY_HERE" },
+                  "active": false
+            }),
+        );
+        if let Err(e) = result {
+            log::error!("Failed to add server config: {}", e);
+        }
+    }
+    store.set("mcp_version", 1);
+    store.save().expect("Failed to save store");
     Ok(())
 }
 
diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs
@@ -10,6 +10,7 @@ use jan_utils::generate_app_token;
 use std::{collections::HashMap, sync::Arc};
 use tauri::{Emitter, Manager, RunEvent};
 use tauri_plugin_llamacpp::cleanup_llama_processes;
+use tauri_plugin_store::StoreExt;
 use tokio::sync::Mutex;
 
 use crate::core::setup::setup_tray;
@@ -151,11 +152,40 @@ pub fn run() {
             )?;
             app.handle()
                 .plugin(tauri_plugin_updater::Builder::new().build())?;
-            // Install extensions
-            if let Err(e) = setup::install_extensions(app.handle().clone(), false) {
+
+            // Start migration
+            let mut store_path = get_jan_data_folder_path(app.handle().clone());
+            store_path.push("store.json");
+            let store = app
+                .handle()
+                .store(store_path)
+                .expect("Store not initialized");
+            let stored_version = store
+                .get("version")
+                .and_then(|v| v.as_str().map(String::from))
+                .unwrap_or_default();
+            let app_version = app
+                .config()
+                .version
+                .clone()
+                .unwrap_or_else(|| "".to_string());
+            // Migrate extensions
+            if let Err(e) =
+                setup::install_extensions(app.handle().clone(), stored_version != app_version)
+            {
                 log::error!("Failed to install extensions: {}", e);
             }
 
+            // Migrate MCP servers
+            if let Err(e) = setup::migrate_mcp_servers(app.handle().clone(), store.clone()) {
+                log::error!("Failed to migrate MCP servers: {}", e);
+            }
+
+            // Store the new app version
+            store.set("version", serde_json::json!(app_version));
+            store.save().expect("Failed to save store");
+            // Migration completed
+
             if option_env!("ENABLE_SYSTEM_TRAY_ICON").unwrap_or("false") == "true" {
                 log::info!("Enabling system tray icon");
                 let _ = setup_tray(app);
