sausage/cli/src/main.rs

use std::io::Write;
use std::path::PathBuf;

use anyhow::{Context, Result, bail};
use clap::{Parser, Subcommand};
use include_dir::{Dir, include_dir};
use serde_json::json;
use tokio::signal::unix::Signal;

use sausage::world::{ActorMeta, World};
use sausage::{Host, cargo_actor, console};

/// The source of the agent actor. `init` builds it like any other actor.
static AGENT_SOURCE: Dir<'_> = include_dir!("$CARGO_MANIFEST_DIR/../agent");

// Without a command, `sausage` runs the world in `./world` and creates it
// first (like `init`) if it does not exist yet.
#[derive(Parser)]
#[command(version, about)]
struct Cli {
    /// Do not read stdin, so end of input does not stop the world. It runs
    /// until SIGTERM (or SIGINT), or until an actor sends
    /// `@broker {"type":"exit"}`. Also set by SAUSAGE_HEADLESS=1.
    #[arg(long, global = true)]
    headless: bool,
    #[command(subcommand)]
    command: Option<Cmd>,
}

#[derive(Subcommand)]
enum Cmd {
    /// Create a new world and build the agent actor.
    ///
    /// Reads SAUSAGE_API, SAUSAGE_BASE_URL, SAUSAGE_MODEL and
    /// SAUSAGE_API_TOKEN, and asks for every value that is missing.
    Init {
        /// Folder of the new world.
        #[arg(default_value = "world")]
        dir: PathBuf,
    },
    /// Start the broker, all actors, and the console.
    Run {
        #[arg(default_value = "world")]
        dir: PathBuf,
    },
}

/// The value of `var`, or the answer to `question`.
fn setting(var: &str, question: &str, default: Option<&str>) -> Result<String> {
    if let Ok(v) = std::env::var(var)
        && !v.is_empty()
    {
        return Ok(v);
    }
    match default {
        Some(d) => print!("{question} [{d}]: "),
        None => print!("{question}: "),
    }
    std::io::stdout().flush()?;
    let mut line = String::new();
    std::io::stdin().read_line(&mut line)?;
    match (line.trim(), default) {
        ("", Some(d)) => Ok(d.to_string()),
        ("", None) => bail!("{question} is required (or set {var})"),
        (v, _) => Ok(v.to_string()),
    }
}

async fn init(dir: PathBuf) -> Result<()> {
    let api = setting("SAUSAGE_API", "API shape (openai-responses or anthropic-messages)", Some("openai-responses"))?;
    if api != "openai-responses" && api != "anthropic-messages" {
        bail!("unknown API shape {api:?}");
    }
    let base_url = setting("SAUSAGE_BASE_URL", "API base URL", None)?;
    let model = setting("SAUSAGE_MODEL", "Model", None)?;
    let token = setting("SAUSAGE_API_TOKEN", "API token", None)?;

    let world = World::create(&dir)?;
    let console = ActorMeta { subscriptions: ["agent".to_string()].into(), ..world.read_meta(console::ID) };
    world.write_meta(console::ID, &console)?;

    // The agent's configuration, in its /data.
    world.ensure_actor_dirs("agent")?;
    let config = json!({"api": api, "base_url": base_url.trim_end_matches('/'), "model": model, "token": token});
    let config_path = world.data_dir("agent").join("config.json");
    std::fs::write(&config_path, serde_json::to_string_pretty(&config)?)?;

    // The agent is an ordinary actor folder: `run` finds its `actor.wasm`,
    // tests it, and installs it as version 1.
    let source = world.source_dir("agent");
    std::fs::create_dir_all(&source)?;
    AGENT_SOURCE.extract(&source).context("extract agent source")?;
    println!("Building the agent actor (the first build takes a few minutes) …");
    let out = cargo_actor::Cargo::new(world.clone())?.build(&source, false).await.context("build agent")?;
    std::fs::copy(&out.wasm, world.wasm_path("agent"))?;
    println!("Built in {:.1} s. Start the world with: sausage run {}", out.seconds, dir.display());
    Ok(())
}

/// Install the SIGTERM and SIGINT handlers now (not when first polled).
fn termination_signals() -> Result<(Signal, Signal)> {
    use tokio::signal::unix::{SignalKind, signal};
    let term = signal(SignalKind::terminate()).context("install SIGTERM handler")?;
    let int = signal(SignalKind::interrupt()).context("install SIGINT handler")?;
    Ok((term, int))
}

/// Resolves on SIGTERM or SIGINT.
async fn termination((mut term, mut int): (Signal, Signal)) {
    tokio::select! {
        _ = term.recv() => {}
        _ = int.recv() => {}
    }
}

async fn run(dir: PathBuf, headless: bool) -> Result<()> {
    // Register the signal handlers before the host starts, so a signal that
    // arrives early is never the default (kill) action.
    let signals = if headless { Some(termination_signals()?) } else { None };
    let host = Host::start(&dir).await?;
    tokio::spawn(console::run_mailbox(host.mailbox(console::ID)));
    let broker = host.broker.clone();
    let exit = match signals {
        Some(signals) => {
            println!("Sausage headless: running until SIGTERM or an exit request.");
            tokio::select! {
                _ = termination(signals) => false,
                _ = broker.exit_requested() => true,
            }
        }
        None => tokio::select! {
            _ = console::run_input(host.broker.clone()) => false,
            _ = broker.exit_requested() => true,
        },
    };
    println!("Stopping actors …");
    host.stop().await;
    if exit {
        // The console still blocks on a stdin read that cannot be cancelled,
        // so the runtime would not shut down on its own. Exit code 75 tells
        // `nix run .#loop` to start the host again (0 = /quit = stop the loop).
        println!("Exiting, as an actor asked the broker to.");
        std::process::exit(75);
    }
    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    let headless = cli.headless || std::env::var("SAUSAGE_HEADLESS").is_ok_and(|v| !v.is_empty() && v != "0");
    match cli.command {
        Some(Cmd::Init { dir }) => init(dir).await,
        Some(Cmd::Run { dir }) => run(dir, headless).await,
        None => {
            let dir = PathBuf::from("world");
            if !dir.join("actor").is_dir() {
                println!("No world in ./world yet, creating it (same as `sausage init`).");
                init(dir.clone()).await?;
            }
            run(dir, headless).await
        }
    }
}