From 62ba503b8633ee44dd23bea402b3f0948fdcb04e Mon Sep 17 00:00:00 2001
From: Akarshan <akarshan@menlo.ai>
Date: Mon, 23 Jun 2025 11:11:51 +0530
Subject: [PATCH] chore: cleanup llama-server processes upon app exit

---
 src-tauri/src/lib.rs | 94 +++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 92 insertions(+), 2 deletions(-)

diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs
index e44ec043a..c56e092e7 100644
--- a/src-tauri/src/lib.rs
+++ b/src-tauri/src/lib.rs
@@ -1,16 +1,16 @@
 mod core;
-use reqwest::Client;
 use core::{
     cmd::get_jan_data_folder_path,
     setup::{self, setup_mcp},
     state::{generate_app_token, AppState},
     utils::download::DownloadManagerState,
 };
+use reqwest::Client;
 use std::{collections::HashMap, sync::Arc};
+use tauri::Manager;
 
 use tokio::sync::Mutex;
 
-
 #[cfg_attr(mobile, tauri::mobile_entry_point)]
 pub fn run() {
     let mut builder = tauri::Builder::default();
@@ -133,6 +133,96 @@ pub fn run() {
                 if window.label() == "main" {
                     window.emit("kill-mcp-servers", ()).unwrap();
                     clean_up();
+                    let state = window.app_handle().state::<AppState>();
+
+                    tauri::async_runtime::block_on(async {
+                        let mut map = state.llama_server_process.lock().await;
+                        let pids: Vec<String> = map.keys().cloned().collect();
+                        for pid in pids {
+                            if let Some(mut child) = map.remove(&pid) {
+                                #[cfg(unix)]
+                                {
+                                    use nix::sys::signal::{kill, Signal};
+                                    use nix::unistd::Pid;
+                                    use tokio::time::{timeout, Duration};
+
+                                    if let Some(raw_pid) = child.id() {
+                                        let raw_pid = raw_pid as i32;
+                                        log::info!(
+                                            "Sending SIGTERM to PID {} during shutdown",
+                                            raw_pid
+                                        );
+                                        let _ = kill(Pid::from_raw(raw_pid), Signal::SIGTERM);
+
+                                        match timeout(Duration::from_secs(2), child.wait()).await {
+                                            Ok(Ok(status)) => log::info!(
+                                                "Process {} exited gracefully: {}",
+                                                raw_pid,
+                                                status
+                                            ),
+                                            Ok(Err(e)) => log::error!(
+                                                "Error waiting after SIGTERM for {}: {}",
+                                                raw_pid,
+                                                e
+                                            ),
+                                            Err(_) => {
+                                                log::warn!(
+                                                    "SIGTERM timed out for PID {}; sending SIGKILL",
+                                                    raw_pid
+                                                );
+                                                let _ =
+                                                    kill(Pid::from_raw(raw_pid), Signal::SIGKILL);
+                                                let _ = child.wait().await;
+                                            }
+                                        }
+                                    }
+                                }
+
+                                #[cfg(windows)]
+                                {
+                                    use tokio::time::{timeout, Duration};
+                                    use windows_sys::Win32::Foundation::BOOL;
+                                    use windows_sys::Win32::System::Console::{
+                                        GenerateConsoleCtrlEvent, CTRL_C_EVENT,
+                                    };
+
+                                    if let Some(raw_pid) = child.id() {
+                                        log::info!(
+                                            "Sending Ctrl-C to PID {} during shutdown",
+                                            raw_pid
+                                        );
+                                        let ok: BOOL = unsafe {
+                                            GenerateConsoleCtrlEvent(CTRL_C_EVENT, raw_pid)
+                                        };
+                                        if ok == 0 {
+                                            log::error!("Failed to send Ctrl-C to PID {}", raw_pid);
+                                        }
+
+                                        match timeout(Duration::from_secs(2), child.wait()).await {
+                                            Ok(Ok(status)) => log::info!(
+                                                "Process {} exited after Ctrl-C: {}",
+                                                raw_pid,
+                                                status
+                                            ),
+                                            Ok(Err(e)) => log::error!(
+                                                "Error waiting after Ctrl-C for {}: {}",
+                                                raw_pid,
+                                                e
+                                            ),
+                                            Err(_) => {
+                                                log::warn!(
+                                                    "Timed out for PID {}; force-killing",
+                                                    raw_pid
+                                                );
+                                                let _ = child.kill().await;
+                                                let _ = child.wait().await;
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    });
                 }
                 let client = Client::new();
                 let url = "http://127.0.0.1:39291/processManager/destroy";