path: root/crates/asist-app/src/supervisor.rs

//! Starting and stopping the local servers.
//!
//! The assistant can start `llama-server` and `tts-server` itself, so getting
//! it running is a single command. The processes are launched in their own
//! group and stopped with SIGTERM before resorting to SIGKILL: killing
//! llama-server outright leaves the GPU busy until the driver recovers it.

use std::path::Path;
use std::process::{Child, Command, Stdio};
use std::time::{Duration, Instant};

use asist_core::config::{Config, LlamaProcess, TtsProcess};
use asist_core::error::{Error, Result};

pub struct Supervisor {
    children: Vec<Managed>,
    log_dir: std::path::PathBuf,
}

struct Managed {
    name: &'static str,
    child: Child,
}

impl Supervisor {
    pub fn new(log_dir: impl AsRef<Path>) -> Result<Self> {
        let log_dir = log_dir.as_ref().to_path_buf();
        std::fs::create_dir_all(&log_dir)?;
        Ok(Self {
            children: Vec::new(),
            log_dir,
        })
    }

    /// Starts the servers the configuration asks for that are not already up.
    ///
    /// Reusing one that is already listening is deliberate: during development
    /// the assistant is restarted many times, and reloading the models takes
    /// more than a minute.
    pub fn start(&mut self, config: &Config, llm_up: bool, tts_up: bool) -> Result<()> {
        if !config.supervisor.manage {
            return Ok(());
        }
        if llm_up {
            tracing::info!(target: "supervisor", "llama-server is already listening, reusing it");
        } else {
            let command = llama_command(&config.supervisor.llama, config)?;
            self.spawn("llama-server", command)?;
        }
        if tts_up {
            tracing::info!(target: "supervisor", "tts-server is already listening, reusing it");
        } else {
            let command = tts_command(&config.supervisor.tts, config)?;
            self.spawn("tts-server", command)?;
        }
        Ok(())
    }

    fn spawn(&mut self, name: &'static str, mut command: Command) -> Result<()> {
        let log_path = self.log_dir.join(format!("{name}.log"));
        let log = std::fs::File::create(&log_path)?;
        let errors = log.try_clone()?;

        // Output goes to the file: the models spit out hundreds of lines and
        // would bury the on-screen transcription. When something fails, the
        // error message points here.
        command
            .stdin(Stdio::null())
            .stdout(Stdio::from(log))
            .stderr(Stdio::from(errors));

        let child = command.spawn().map_err(|e| {
            Error::Config(format!(
                "could not start {name} ({:?}): {e}",
                command.get_program()
            ))
        })?;
        tracing::info!(
            target: "supervisor",
            process = name,
            pid = child.id(),
            registered = %log_path.display(),
            "lanzado"
        );
        self.children.push(Managed { name, child });
        Ok(())
    }

    /// Checks whether any of them died on its own, and returns its name.
    pub fn crashed(&mut self) -> Option<(&'static str, Option<i32>)> {
        for managed in &mut self.children {
            if let Ok(Some(status)) = managed.child.try_wait() {
                return Some((managed.name, status.code()));
            }
        }
        None
    }

    pub fn log_path(&self, name: &str) -> std::path::PathBuf {
        self.log_dir.join(format!("{name}.log"))
    }

    /// Stops everything that was launched. SIGTERM first so they release the GPU.
    pub fn shutdown(&mut self) {
        for managed in &mut self.children {
            terminate(&mut managed.child, managed.name);
        }
        self.children.clear();
    }
}

impl Drop for Supervisor {
    fn drop(&mut self) {
        self.shutdown();
    }
}

fn terminate(child: &mut Child, name: &str) {
    if matches!(child.try_wait(), Ok(Some(_))) {
        return;
    }
    // Direct SIGTERM: `Child::kill` sends SIGKILL, which gives the server no
    // chance to release GPU memory.
    #[cfg(unix)]
    unsafe {
        libc_kill(child.id() as i32, 15);
    }
    #[cfg(not(unix))]
    let _ = child.kill();

    let deadline = Instant::now() + Duration::from_secs(10);
    while Instant::now() < deadline {
        match child.try_wait() {
            Ok(Some(_)) => {
                tracing::info!(target: "supervisor", proc_handle = name, "detenido");
                return;
            }
            Ok(None) => std::thread::sleep(Duration::from_millis(100)),
            Err(_) => break,
        }
    }
    tracing::warn!(target: "supervisor", proc_handle = name, "did not respond to SIGTERM, forcing it");
    let _ = child.kill();
    let _ = child.wait();
}

#[cfg(unix)]
unsafe fn libc_kill(pid: i32, signal: i32) {
    // The only libc symbol needed is declared here, instead of pulling in
    // the whole crate for one call.
    unsafe extern "C" {
        fn kill(pid: i32, sig: i32) -> i32;
    }
    unsafe {
        kill(pid, signal);
    }
}

fn require(path: &Path, what: &str) -> Result<()> {
    if path.exists() {
        return Ok(());
    }
    Err(Error::Config(format!(
        "missing {what}: {}. Run scripts/bootstrap.sh to build the \
         engines and link the models",
        path.display()
    )))
}

fn llama_command(process: &LlamaProcess, config: &Config) -> Result<Command> {
    require(&process.binary, "the llama-server binary")?;
    require(&process.model, "the LLM model")?;

    let mut command = Command::new(&process.binary);
    command
        .arg("--model")
        .arg(&process.model)
        .arg("--host")
        .arg(&config.llm.host)
        .arg("--port")
        .arg(config.llm.port.to_string());

    if process.mmproj.exists() {
        command.arg("--mmproj").arg(&process.mmproj);
    }
    // Without this, the model template opens <think> and never closes it: the
    // assistant spends 7 to 9 s reasoning before the first word.
    if process.chat_template.exists() {
        command
            .arg("--jinja")
            .arg("--chat-template-file")
            .arg(&process.chat_template);
    } else {
        tracing::warn!(
            target: "supervisor",
            path = %process.chat_template.display(),
            "the no-reasoning template is missing: the model will take several \
             seconds to start talking"
        );
    }
    command.args(&process.extra_args);
    Ok(command)
}

fn tts_command(process: &TtsProcess, config: &Config) -> Result<Command> {
    require(&process.binary, "the tts-server binary")?;
    require(&process.model, "the TTS talker model")?;
    require(&process.codec, "the TTS codec")?;

    let mut command = Command::new(&process.binary);
    command
        .arg("--model")
        .arg(&process.model)
        .arg("--codec")
        .arg(&process.codec)
        .arg("--host")
        .arg(&config.tts.host)
        .arg("--port")
        .arg(config.tts.port.to_string())
        .arg("--lang")
        .arg(&config.tts.language)
        // The single most effective setting in the system: the stock value is
        // 24 s, which in practice means returning nothing until the sentence ends.
        .arg("--codec-chunk-dur")
        .arg(process.codec_chunk_dur.to_string());
    command.args(&process.extra_args);
    Ok(command)
}