sausage/cli/src/main.rs

//! Command parsing and the minimal line console.
use anyhow::{Result, bail};
use clap::{Parser, Subcommand};
use include_dir::{Dir, include_dir};
use sausage::{Host, world::World};
use std::path::{Path, PathBuf};

static AGENT: Dir<'_> = include_dir!("$CARGO_MANIFEST_DIR/../agent/src");
mod test_runner;
#[derive(Parser)]
#[command(version)]
struct Cli {
    #[arg(long, global = true)]
    headless: bool,
    #[arg(long, global = true)]
    cache_dir: Option<PathBuf>,
    #[command(subcommand)]
    command: Option<Cmd>,
}
#[derive(Subcommand)]
enum Cmd {
    Init {
        #[arg(default_value = "world")]
        dir: PathBuf,
    },
    Run {
        #[arg(default_value = "world")]
        dir: PathBuf,
    },
    /// Run a WASI libtest component in isolated instances.
    TestRunner {
        wasm: PathBuf,
        #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
        args: Vec<String>,
    },
}
fn headless(value: &str) -> std::result::Result<bool, String> {
    match value.to_ascii_lowercase().as_str() {
        "1" | "true" => Ok(true),
        "" | "0" | "false" => Ok(false),
        _ => Err("expected 1, true, 0, false, or an empty value".into()),
    }
}

// Readline can still wait on its detached thread when the host exits.
struct TerminalMode(Option<nix::sys::termios::Termios>);
impl Drop for TerminalMode {
    fn drop(&mut self) {
        if let Some(mode) = &self.0 {
            let _ = nix::sys::termios::tcsetattr(
                std::io::stdin(),
                nix::sys::termios::SetArg::TCSANOW,
                mode,
            );
        }
    }
}

/// Creates a world with editable configuration and builds its embedded agent.
async fn init(dir: &Path, cache: Option<&Path>) -> Result<()> {
    if dir.exists() {
        bail!("world already exists: {}", dir.display());
    }
    let parent = dir
        .parent()
        .filter(|p| !p.as_os_str().is_empty())
        .unwrap_or(Path::new("."));
    std::fs::create_dir_all(parent)?;
    let staging = tempfile::Builder::new()
        .prefix(".sausage-init-")
        .tempdir_in(parent)?;
    let world = World::create(&staging.path().join("world"))?;
    std::fs::write(
        world.root.join("config/agent.yaml"),
        include_str!("agent.yaml"),
    )?;
    std::fs::write(
        world.root.join("config/boot.yaml"),
        include_str!("boot.yaml"),
    )?;
    let cache = cache.map(Path::to_path_buf).unwrap_or_else(|| {
        let base = std::env::var_os("XDG_CACHE_HOME")
            .map(PathBuf::from)
            .unwrap_or_else(|| {
                PathBuf::from(std::env::var_os("HOME").unwrap_or_default()).join(".cache")
            });
        base.join("sausage")
    });
    let builder = sausage::cargo::Cargo::new(world.clone(), &cache)?;
    let source = world.plugin_dir("agent").join("source");
    std::fs::create_dir_all(source.join("src"))?;
    AGENT.extract(source.join("src"))?;
    std::fs::write(
        source.join("Cargo.toml"),
        include_str!("../../agent/Cargo.toml"),
    )?;
    std::fs::write(
        source.join("Cargo.lock"),
        include_bytes!("../../Cargo.lock"),
    )?;
    println!("The first agent build takes a few minutes.");
    let output = builder.build(&source).await?;
    std::fs::copy(output.wasm, world.plugin_dir("agent").join("plugin.wasm"))?;
    std::fs::rename(&world.root, dir)?;
    println!(
        "Edit {} before chatting.",
        dir.join("config/agent.yaml").display()
    );
    println!("Start the world with: sausage run {}", dir.display());
    Ok(())
}
#[tokio::main]
async fn main() -> Result<std::process::ExitCode> {
    let cli = Cli::parse();
    let _terminal = TerminalMode(nix::sys::termios::tcgetattr(std::io::stdin()).ok());
    let command = match cli.command {
        Some(command) => command,
        None => {
            let dir = PathBuf::from("world");
            if !dir.join("config/boot.yaml").is_file() {
                init(&dir, cli.cache_dir.as_deref()).await?;
            }
            Cmd::Run { dir }
        }
    };
    match command {
        Cmd::TestRunner { wasm, args } => {
            if !test_runner::run(&wasm, args).await? {
                return Ok(std::process::ExitCode::FAILURE);
            }
        }
        Cmd::Init { dir } => init(&dir, cli.cache_dir.as_deref()).await?,
        Cmd::Run { dir } => {
            let headless = cli.headless
                || std::env::var("SAUSAGE_HEADLESS")
                    .ok()
                    .map(|value| headless(&value).map_err(anyhow::Error::msg))
                    .transpose()?
                    .unwrap_or(false);
            let mut term =
                tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?;
            let mut interrupt =
                tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())?;
            let start = async {
                match cli.cache_dir {
                    Some(cache) => Host::start_with_cache(&dir, &cache).await,
                    None => Host::start(&dir).await,
                }
            };
            let host = tokio::select! {
                result=start=>result?,
                _=term.recv()=>return Ok(std::process::ExitCode::SUCCESS),
                _=interrupt.recv()=>return Ok(std::process::ExitCode::SUCCESS),
            };
            let mut notifications = host.subscribe();
            let notification_task = headless.then(|| {
                tokio::spawn(async move {
                    loop {
                        match notifications.recv().await {
                            Ok(notification) => println!("{notification}"),
                            Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
                            Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
                        }
                    }
                })
            });
            let console = sausage::console::run_input(host.clone());
            let mut console_result = Ok(());
            let exit = tokio::select! {
                _=host.exit_requested()=>true,
                _=term.recv()=>false,
                _=interrupt.recv(), if headless=>false,
                result = console, if !headless => {
                    console_result=result;
                    false
                },
            };
            let exit = exit || host.restart_requested();
            host.stop().await;
            if let Some(task) = notification_task {
                task.abort();
            }
            if exit {
                return Ok(std::process::ExitCode::from(75));
            }
            console_result?;
        }
    }
    Ok(std::process::ExitCode::SUCCESS)
}