From 9ed98614fee32a7dd98635493b36305e3d49c101 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Mon, 11 Aug 2025 19:42:59 +0700
Subject: [PATCH] fix: factory reset process got blocked

---
 src-tauri/src/core/cmd.rs   |  5 +++--
 src-tauri/src/core/mcp.rs   | 20 ++++++++++++++++++++
 src-tauri/src/core/setup.rs | 24 ------------------------
 src-tauri/src/lib.rs        |  5 ++++-
 4 files changed, 27 insertions(+), 27 deletions(-)

diff --git a/src-tauri/src/core/cmd.rs b/src-tauri/src/core/cmd.rs
index 37937cfe0..dbfb41fc3 100644
--- a/src-tauri/src/core/cmd.rs
+++ b/src-tauri/src/core/cmd.rs
@@ -2,7 +2,7 @@ use serde::{Deserialize, Serialize};
 use std::{fs, io, path::PathBuf};
 use tauri::{AppHandle, Manager, Runtime, State};
 
-use crate::core::utils::extensions::inference_llamacpp_extension::cleanup::cleanup_processes;
+use crate::core::{mcp::clean_up_mcp_servers, utils::extensions::inference_llamacpp_extension::cleanup::cleanup_processes};
 
 use super::{server, setup, state::AppState};
 
@@ -125,6 +125,7 @@ pub fn factory_reset(app_handle: tauri::AppHandle, state: State<'_, AppState>) {
     log::info!("Factory reset, removing data folder: {:?}", data_folder);
 
     tauri::async_runtime::block_on(async {
+        clean_up_mcp_servers(state.clone()).await;
         cleanup_processes(state).await;
 
         if data_folder.exists() {
@@ -138,7 +139,7 @@ pub fn factory_reset(app_handle: tauri::AppHandle, state: State<'_, AppState>) {
         let _ = fs::create_dir_all(&data_folder).map_err(|e| e.to_string());
 
         // Reset the configuration
-        let mut default_config = AppConfiguration::default();
+        let mut default_config: AppConfiguration = AppConfiguration::default();
         default_config.data_folder = default_data_folder_path(app_handle.clone());
         let _ = update_app_configuration(app_handle.clone(), default_config);
 
diff --git a/src-tauri/src/core/mcp.rs b/src-tauri/src/core/mcp.rs
index 7485fd17e..93497772c 100644
--- a/src-tauri/src/core/mcp.rs
+++ b/src-tauri/src/core/mcp.rs
@@ -751,6 +751,26 @@ pub async fn reset_mcp_restart_count(state: State<'_, AppState>, server_name: St
     Ok(())
 }
 
+pub async fn clean_up_mcp_servers(
+    state: State<'_, AppState>,
+) {
+    log::info!("Cleaning up MCP servers");
+    
+    // Stop all running MCP servers
+    let _ = stop_mcp_servers(state.mcp_servers.clone()).await;
+    
+    // Clear active servers and restart counts
+    {
+        let mut active_servers = state.mcp_active_servers.lock().await;
+        active_servers.clear();
+    }
+    {
+        let mut restart_counts = state.mcp_restart_counts.lock().await;
+        restart_counts.clear();
+    }
+    log::info!("MCP servers cleaned up successfully");
+}
+
 pub async fn stop_mcp_servers(
     servers_state: Arc<Mutex<HashMap<String, RunningService<RoleClient, ()>>>>,
 ) -> Result<(), String> {
diff --git a/src-tauri/src/core/setup.rs b/src-tauri/src/core/setup.rs
index 940122235..0f32f4c7e 100644
--- a/src-tauri/src/core/setup.rs
+++ b/src-tauri/src/core/setup.rs
@@ -199,30 +199,6 @@ pub fn setup_mcp(app: &App) {
     let state = app.state::<AppState>();
     let servers = state.mcp_servers.clone();
     let app_handle: tauri::AppHandle = app.handle().clone();
-    // Setup kill-mcp-servers event listener (similar to kill-sidecar)
-    let app_handle_for_kill = app_handle.clone();
-    app_handle.listen("kill-mcp-servers", move |_event| {
-        let app_handle = app_handle_for_kill.clone();
-        tauri::async_runtime::spawn(async move {
-            log::info!("Received kill-mcp-servers event - cleaning up MCP servers");
-            let app_state = app_handle.state::<AppState>();
-            // Stop all running MCP servers
-            if let Err(e) = super::mcp::stop_mcp_servers(app_state.mcp_servers.clone()).await {
-                log::error!("Failed to stop MCP servers: {}", e);
-                return;
-            }
-            // Clear active servers and restart counts
-            {
-                let mut active_servers = app_state.mcp_active_servers.lock().await;
-                active_servers.clear();
-            }
-            {
-                let mut restart_counts = app_state.mcp_restart_counts.lock().await;
-                restart_counts.clear();
-            }
-            log::info!("MCP servers cleaned up successfully");
-        });
-    });
     tauri::async_runtime::spawn(async move {
         if let Err(e) = run_mcp_commands(&app_handle, servers).await {
             log::error!("Failed to run mcp commands: {}", e);
diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs
index e449fc739..b0935206b 100644
--- a/src-tauri/src/lib.rs
+++ b/src-tauri/src/lib.rs
@@ -10,6 +10,8 @@ use std::{collections::HashMap, sync::Arc};
 use tauri::{Emitter, Manager, RunEvent};
 use tokio::sync::Mutex;
 
+use crate::core::mcp::clean_up_mcp_servers;
+
 #[cfg_attr(mobile, tauri::mobile_entry_point)]
 pub fn run() {
     let mut builder = tauri::Builder::default();
@@ -143,10 +145,10 @@ pub fn run() {
         .on_window_event(|window, event| match event {
             tauri::WindowEvent::CloseRequested { .. } => {
                 if window.label() == "main" {
-                    window.emit("kill-mcp-servers", ()).unwrap();
                     let state = window.app_handle().state::<AppState>();
 
                     tauri::async_runtime::block_on(async {
+                        clean_up_mcp_servers(state.clone()).await;
                         cleanup_processes(state).await;
                     });
                 }
@@ -173,6 +175,7 @@ pub fn run() {
                     }
 
                     // Quick cleanup with shorter timeout
+                    clean_up_mcp_servers(state.clone()).await;
                     cleanup_processes(state).await;
                 });
             });