ForgeLM/tauri-app/src-tauri/src/runtime/mod.rs

pub mod config;
pub mod health;
pub mod process;
pub mod state;
pub mod types;

use std::time::Duration;

use tauri::{AppHandle, State};

use self::config::{load_runtime_config, save_runtime_config};
use self::health::wait_for_health;
use self::process::spawn_localllm;
use self::state::{
    begin_lifecycle_command, block_lifecycle_command, build_snapshot, finish_lifecycle_command,
    guard_start_request, guard_stop_request, refresh_process_state, requires_runtime_restart,
    runtime_error,
};
use self::types::{
    ConnectionMode, RuntimeConfig, RuntimeLifecycleCommand, RuntimeLifecycleOutcome,
    RuntimeSnapshot, RuntimeStatus, StartPolicy,
};

pub fn get_runtime_snapshot(state: State<'_, RuntimeState>) -> Result<RuntimeSnapshot, String> {
    let mut inner = state.0.lock().map_err(|e| e.to_string())?;
    refresh_process_state(&mut inner);
    Ok(build_snapshot(&inner))
}

pub fn save_runtime_config_command(
    app: AppHandle,
    state: State<'_, RuntimeState>,
    config: RuntimeConfig,
) -> Result<RuntimeSnapshot, String> {
    let mut inner = state.0.lock().map_err(|e| e.to_string())?;
    refresh_process_state(&mut inner);

    if inner.child.is_some() && requires_runtime_restart(&inner.config, &config) {
        let message =
            "Stop the local API before changing connection mode, model, address, or runtime sizing.";
        inner.last_error = Some(runtime_error("restart_required", message));
        block_lifecycle_command(&mut inner, RuntimeLifecycleCommand::Save, message);
        return Ok(build_snapshot(&inner));
    }

    begin_lifecycle_command(&mut inner, RuntimeLifecycleCommand::Save);
    if let Err(error) = save_runtime_config(&app, &config) {
        inner.last_error = Some(runtime_error("save_failed", &error));
        finish_lifecycle_command(
            &mut inner,
            RuntimeLifecycleCommand::Save,
            RuntimeLifecycleOutcome::Failed,
            format!("Failed to save runtime settings: {error}"),
        );
        return Err(error);
    }
    inner.config = config;
    inner.last_error = None;
    finish_lifecycle_command(
        &mut inner,
        RuntimeLifecycleCommand::Save,
        RuntimeLifecycleOutcome::Succeeded,
        "Runtime settings saved.",
    );
    Ok(build_snapshot(&inner))
}

fn start_runtime_inner(
    app: &AppHandle,
    state: &State<'_, RuntimeState>,
    override_config: Option<RuntimeConfig>,
    startup_status: RuntimeStatus,
) -> Result<RuntimeSnapshot, String> {
    let command = if override_config.is_some() {
        RuntimeLifecycleCommand::SaveAndStart
    } else {
        RuntimeLifecycleCommand::Start
    };

    let config = {
        let mut inner = state.0.lock().map_err(|e| e.to_string())?;
        refresh_process_state(&mut inner);

        if guard_start_request(&mut inner).is_err() {
            return Ok(build_snapshot(&inner));
        }

        begin_lifecycle_command(&mut inner, command);

        if let Some(config) = override_config {
            if let Err(error) = save_runtime_config(app, &config) {
                inner.last_error = Some(runtime_error("save_failed", &error));
                finish_lifecycle_command(
                    &mut inner,
                    command,
                    RuntimeLifecycleOutcome::Failed,
                    format!("Failed to save runtime settings: {error}"),
                );
                return Err(error);
            }
            inner.config = config;
        }

        if inner.config.connection_mode != ConnectionMode::Local {
            let message = "Switch to Local mode before starting the built-in API";
            inner.last_error = Some(runtime_error("remote_mode", message));
            inner.status = RuntimeStatus::Error;
            finish_lifecycle_command(
                &mut inner,
                command,
                RuntimeLifecycleOutcome::Failed,
                message,
            );
            return Ok(build_snapshot(&inner));
        }

        inner.status = startup_status;
        inner.last_error = None;
        inner.config.clone()
    };

    let child = match spawn_localllm(app, &config) {
        Ok(child) => child,
        Err(message) => {
            let mut inner = state.0.lock().map_err(|e| e.to_string())?;
            inner.status = RuntimeStatus::Error;
            inner.last_error = Some(runtime_error("spawn_failed", &message));
            finish_lifecycle_command(
                &mut inner,
                command,
                RuntimeLifecycleOutcome::Failed,
                message,
            );
            return Ok(build_snapshot(&inner));
        }
    };

    {
        let mut inner = state.0.lock().map_err(|e| e.to_string())?;
        inner.child = Some(child);
        inner.status = startup_status;
    }

    if wait_for_health(&config.addr, Duration::from_secs(15)) {
        let mut inner = state.0.lock().map_err(|e| e.to_string())?;
        inner.status = RuntimeStatus::Ready;
        inner.last_error = None;
        let message = if command == RuntimeLifecycleCommand::SaveAndStart {
            "Runtime settings saved and local API started."
        } else {
            "Local API started."
        };
        finish_lifecycle_command(
            &mut inner,
            command,
            RuntimeLifecycleOutcome::Succeeded,
            message,
        );
        Ok(build_snapshot(&inner))
    } else {
        let mut inner = state.0.lock().map_err(|e| e.to_string())?;
        if let Some(child) = inner.child.as_mut() {
            let _ = child.kill();
            let _ = child.wait();
        }
        inner.child = None;
        inner.status = RuntimeStatus::Error;
        let message = if startup_status == RuntimeStatus::Loading {
            format!(
                "Model loading timed out before the local API became ready on {}",
                config.addr
            )
        } else {
            format!(
                "Local API did not pass /health readiness on {} in time",
                config.addr
            )
        };
        inner.last_error = Some(runtime_error("health_timeout", &message));
        finish_lifecycle_command(
            &mut inner,
            command,
            RuntimeLifecycleOutcome::Failed,
            message,
        );
        Ok(build_snapshot(&inner))
    }
}

pub fn bootstrap_runtime(
    app: AppHandle,
    state: State<'_, RuntimeState>,
) -> Result<RuntimeSnapshot, String> {
    use crate::model_alias;
    use crate::model_library::paths::resolve_models_dir;

    let mut config = load_runtime_config(&app).unwrap_or_default();

    // Auto-resolve model if none is configured yet
    let model_path_empty = config
        .model_path
        .as_deref()
        .map(|p| p.trim().is_empty())
        .unwrap_or(true);

    if model_path_empty && config.connection_mode == ConnectionMode::Local {
        let models_dir = resolve_models_dir(config.models_dir.as_deref());
        let models = if models_dir.exists() {
            let dir_str = models_dir.to_string_lossy().to_string();
            crate::model_files::scan_models(dir_str).unwrap_or_default()
        } else {
            Vec::new()
        };

        if let Some(resolution) = model_alias::resolve_startup_model(
            config.preferred_startup_model_alias.as_deref(),
            &models,
        ) {
            eprintln!(
                "[bootstrap] auto-selected model: {} ({:?}) reason={:?}",
                resolution.alias.as_str(),
                resolution.model_path,
                resolution.reason
            );
            config.model_path = Some(resolution.model_path);
            // Auto-enable OnAppLaunch when we found a model automatically
            config.start_policy = StartPolicy::OnAppLaunch;
        }
    }

    let should_autostart = config.connection_mode == ConnectionMode::Local
        && config.start_policy == StartPolicy::OnAppLaunch
        && config
            .model_path
            .as_deref()
            .map(|value| !value.trim().is_empty())
            .unwrap_or(false);

    {
        let mut inner = state.0.lock().map_err(|e| e.to_string())?;
        inner.config = config.clone();
        refresh_process_state(&mut inner);
    }

    if should_autostart {
        start_runtime_inner(&app, &state, Some(config), RuntimeStatus::Loading)
    } else {
        get_runtime_snapshot(state)
    }
}

pub fn start_runtime(
    app: AppHandle,
    state: State<'_, RuntimeState>,
) -> Result<RuntimeSnapshot, String> {
    start_runtime_inner(&app, &state, None, RuntimeStatus::Starting)
}

pub fn start_runtime_with_status(
    app: AppHandle,
    state: State<'_, RuntimeState>,
    startup_status: RuntimeStatus,
) -> Result<RuntimeSnapshot, String> {
    start_runtime_inner(&app, &state, None, startup_status)
}

pub fn stop_runtime(state: State<'_, RuntimeState>) -> Result<RuntimeSnapshot, String> {
    let mut inner = state.0.lock().map_err(|e| e.to_string())?;
    refresh_process_state(&mut inner);

    if guard_stop_request(&mut inner).is_err() {
        return Ok(build_snapshot(&inner));
    }

    begin_lifecycle_command(&mut inner, RuntimeLifecycleCommand::Stop);
    inner.status = RuntimeStatus::Stopping;

    if let Some(child) = inner.child.as_mut() {
        if let Err(error) = child.kill() {
            inner.status = RuntimeStatus::Error;
            inner.last_error = Some(runtime_error(
                "stop_failed",
                format!("Failed to kill server: {error}"),
            ));
            finish_lifecycle_command(
                &mut inner,
                RuntimeLifecycleCommand::Stop,
                RuntimeLifecycleOutcome::Failed,
                format!("Failed to stop local API: {error}"),
            );
            return Ok(build_snapshot(&inner));
        }
        let _ = child.wait();
    }

    inner.child = None;
    inner.status = RuntimeStatus::Stopped;
    inner.last_error = None;
    finish_lifecycle_command(
        &mut inner,
        RuntimeLifecycleCommand::Stop,
        RuntimeLifecycleOutcome::Succeeded,
        "Local API stopped.",
    );
    Ok(build_snapshot(&inner))
}

pub use state::RuntimeState;