ForgeLM/tauri-app/src-tauri/src/lib.rs

use std::path::Path;
use std::process::{Child, Command};
use std::sync::Mutex;
use tauri::{Manager, State};

pub struct ServerState(pub Mutex<Option<Child>>);

/// Query GPU info via nvidia-smi.
#[tauri::command]
fn query_gpu() -> Result<GpuInfo, String> {
    let output = Command::new("nvidia-smi")
        .args([
            "--query-gpu=name,memory.total",
            "--format=csv,noheader,nounits",
        ])
        .output()
        .map_err(|_| "nvidia-smi not found — no NVIDIA GPU detected".to_string())?;

    if !output.status.success() {
        return Err("nvidia-smi failed — no NVIDIA GPU detected".into());
    }

    let stdout = String::from_utf8_lossy(&output.stdout);
    let line = stdout.lines().next().unwrap_or("").trim().to_string();
    if line.is_empty() {
        return Err("No GPU found".into());
    }

    let parts: Vec<&str> = line.splitn(2, ',').collect();
    if parts.len() < 2 {
        return Err(format!("Unexpected nvidia-smi output: {line}"));
    }

    let name = parts[0].trim().to_string();
    let vram_mib: u32 = parts[1].trim().parse().unwrap_or(0);
    let usable = vram_mib.saturating_sub(1024);
    let recommended_layers = (usable / 80).min(200) as i32;

    Ok(GpuInfo {
        name,
        vram_mib,
        recommended_layers,
    })
}

#[derive(serde::Serialize)]
pub struct GpuInfo {
    pub name: String,
    pub vram_mib: u32,
    pub recommended_layers: i32,
}

/// Start the inference server with full configuration.
#[tauri::command]
fn start_server(
    state: State<ServerState>,
    model_path: String,
    addr: String,
    gpu_layers: i32,
    ctx_size: u32,
) -> Result<String, String> {
    let mut guard = state.0.lock().map_err(|e| e.to_string())?;

    if guard.is_some() {
        return Err("Server is already running".into());
    }

    let binary = if cfg!(target_os = "windows") {
        "localllm.exe"
    } else {
        "localllm"
    };

    let child = Command::new(binary)
        .env("LOCALLLM_MODEL", &model_path)
        .env("LOCALLLM_ADDR", &addr)
        .env("LOCALLLM_GPU_LAYERS", gpu_layers.to_string())
        .env("LOCALLLM_CTX", ctx_size.to_string())
        .spawn()
        .map_err(|e| format!("Failed to start server: {e}"))?;

    let pid = child.id();
    *guard = Some(child);

    Ok(format!("Server started (pid={pid}) on {addr}"))
}

#[tauri::command]
fn stop_server(state: State<ServerState>) -> Result<String, String> {
    let mut guard = state.0.lock().map_err(|e| e.to_string())?;

    match guard.as_mut() {
        None => Err("Server is not running".into()),
        Some(child) => {
            child
                .kill()
                .map_err(|e| format!("Failed to kill server: {e}"))?;
            let _ = child.wait();
            *guard = None;
            Ok("Server stopped".into())
        }
    }
}

#[tauri::command]
fn server_status(state: State<ServerState>) -> String {
    let guard = state.0.lock().unwrap();
    if guard.is_some() {
        "running".into()
    } else {
        "stopped".into()
    }
}

/// Scan a directory for .gguf model files.
#[tauri::command]
fn scan_models(dir: String) -> Result<Vec<ModelFile>, String> {
    let path = Path::new(&dir);
    if !path.exists() {
        return Err(format!("Directory not found: {dir}"));
    }
    if !path.is_dir() {
        return Err(format!("Not a directory: {dir}"));
    }

    let entries = std::fs::read_dir(path).map_err(|e| e.to_string())?;
    let mut models = Vec::new();

    for entry in entries.flatten() {
        let p = entry.path();
        if p.is_file() {
            if let Some(ext) = p.extension() {
                if ext.eq_ignore_ascii_case("gguf") {
                    let name = p
                        .file_name()
                        .unwrap_or_default()
                        .to_string_lossy()
                        .to_string();
                    let full_path = p.to_string_lossy().to_string();
                    models.push(ModelFile {
                        name,
                        path: full_path,
                    });
                }
            }
        }
    }

    models.sort_by(|a, b| a.name.cmp(&b.name));
    Ok(models)
}

/// Copy a .gguf file into the models directory.
/// Returns the new ModelFile on success.
#[tauri::command]
fn import_model(src_path: String, dest_dir: String) -> Result<ModelFile, String> {
    let src = Path::new(&src_path);

    // Validate source
    if !src.exists() {
        return Err(format!("Source file not found: {src_path}"));
    }
    if !src.is_file() {
        return Err("Source is not a file".into());
    }
    let ext = src.extension().unwrap_or_default();
    if !ext.eq_ignore_ascii_case("gguf") {
        return Err("Only .gguf files are supported".into());
    }

    // Ensure destination directory exists
    let dest_dir_path = Path::new(&dest_dir);
    if !dest_dir_path.exists() {
        std::fs::create_dir_all(dest_dir_path)
            .map_err(|e| format!("Failed to create models directory: {e}"))?;
    }

    let file_name = src.file_name().unwrap_or_default();
    let dest_path = dest_dir_path.join(file_name);

    // If already in the same directory, just return it
    if src == dest_path {
        return Ok(ModelFile {
            name: file_name.to_string_lossy().to_string(),
            path: dest_path.to_string_lossy().to_string(),
        });
    }

    // Copy the file
    std::fs::copy(src, &dest_path).map_err(|e| format!("Failed to copy model: {e}"))?;

    Ok(ModelFile {
        name: file_name.to_string_lossy().to_string(),
        path: dest_path.to_string_lossy().to_string(),
    })
}

#[derive(serde::Serialize, Clone)]
pub struct ModelFile {
    pub name: String,
    pub path: String,
}

#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
    tauri::Builder::default()
        .plugin(tauri_plugin_shell::init())
        .manage(ServerState(Mutex::new(None)))
        .invoke_handler(tauri::generate_handler![
            query_gpu,
            scan_models,
            import_model,
            start_server,
            stop_server,
            server_status
        ])
        .setup(|app| {
            #[cfg(debug_assertions)]
            {
                let window = app.get_webview_window("main").unwrap();
                window.open_devtools();
            }
            Ok(())
        })
        .run(tauri::generate_context!())
        .expect("error while running tauri application");
}