boxmaker/crates/bxctl/src/main.rs

//! `bxctl`: the owner's chat client for `loopd`.

use std::io::{BufRead, BufReader, Read, Write};
use std::path::{Path, PathBuf};
use std::process::ExitCode;

use bxcore::{ErrorCode, SessionId, Timestamp, TurnDone};
use bxctl::admin::{self, AdminError};
use bxctl::chat::{Approvals, ChatError, OnPending, Printer, TurnIo, new_session_id, stream_turn};
use bxctl::cli::{self, ChatOptions, Command, USAGE};
use bxctl::escape::escape_model_text;

fn main() -> ExitCode {
    // `args_os`, because `args` panics on an argument that is not UTF-8. bxctl's arguments are
    // text (ids, messages, paths it prints back), so such an argument is a usage error.
    let args: Vec<String> = match std::env::args_os()
        .skip(1)
        .map(|a| a.into_string())
        .collect()
    {
        Ok(args) => args,
        Err(_) => {
            eprintln!("bxctl: an argument is not valid UTF-8\n{USAGE}");
            return ExitCode::from(2);
        }
    };
    // $BOXMAKER_HOME defaults to /var/lib/boxmaker; defaults for the sockets are read from it.
    let home = std::env::var_os("BOXMAKER_HOME")
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("/var/lib/boxmaker"));
    // Parse before connecting, so a usage error exits 2 even without a broker.
    let command = match cli::parse(&args, &home) {
        Ok(command) => command,
        Err(_) => {
            eprintln!("{USAGE}");
            return ExitCode::from(2);
        }
    };
    match command {
        Command::Chat(opts) => run(&opts),
        Command::Approvals { admin_socket } => exit(|| cmd_approvals(&admin_socket)),
        Command::Approve {
            admin_socket,
            approval,
        } => exit(|| cmd_approve(&admin_socket, approval)),
        Command::Refuse {
            admin_socket,
            approval,
            reason,
        } => exit(|| cmd_refuse(&admin_socket, approval, reason.as_deref())),
        Command::GrantsCheck { admin_socket } => exit(|| cmd_grants_check(&admin_socket)),
        Command::SanitizeCheck { text } => cmd_sanitize_check(text),
        Command::Memory { socket, ask } => {
            let mut out = std::io::stdout().lock();
            match bxctl::memory::run(&socket, &ask, bxctl::memory::MEMORY_TIMEOUT, &mut out) {
                Ok(()) => ExitCode::SUCCESS,
                Err(e) => {
                    eprintln!("bxctl: {e}");
                    ExitCode::from(1)
                }
            }
        }
        Command::AuditVerify { home } => {
            let mut out = std::io::stdout().lock();
            match bxctl::verify::run(&home, &mut out) {
                Ok(true) => ExitCode::SUCCESS,
                Ok(false) => ExitCode::from(1),
                Err(e) => {
                    eprintln!(
                        "bxctl: cannot read the audit log under {}: {e}",
                        home.display()
                    );
                    ExitCode::from(1)
                }
            }
        }
    }
}

// Run a command, mapping its result to an exit code: success for true, 1 for false or an error.
fn exit<F: FnOnce() -> Result<bool, AdminError>>(f: F) -> ExitCode {
    match f() {
        Ok(true) => ExitCode::SUCCESS,
        Ok(false) => ExitCode::from(1),
        Err(e) => {
            eprintln!("bxctl: {e}");
            ExitCode::from(1)
        }
    }
}

fn cmd_approvals(admin: &Path) -> Result<bool, AdminError> {
    let mut out = std::io::stdout().lock();
    admin::cmd_approvals(admin, Timestamp::now(), &mut out)
}

fn cmd_approve(admin: &Path, id: u64) -> Result<bool, AdminError> {
    let mut out = std::io::stdout().lock();
    admin::cmd_approve(admin, id, &mut out)
}

fn cmd_refuse(admin: &Path, id: u64, reason: Option<&str>) -> Result<bool, AdminError> {
    let mut out = std::io::stdout().lock();
    admin::cmd_refuse(admin, id, reason, &mut out)
}

fn cmd_grants_check(admin: &Path) -> Result<bool, AdminError> {
    let mut out = std::io::stdout().lock();
    admin::cmd_grants_check(admin, &mut out)
}

// `bxctl sanitize check`: 0 when no rule fires, 1 when one does. A text that cannot be read (not
// UTF-8, like an argument that is not) or an answer that cannot be written is 2, so 1 always
// means a rule fired.
fn cmd_sanitize_check(text: Option<String>) -> ExitCode {
    let text = match text {
        Some(text) => text,
        None => {
            let mut text = String::new();
            if let Err(e) = std::io::stdin().read_to_string(&mut text) {
                eprintln!("bxctl: cannot read standard input: {e}");
                return ExitCode::from(2);
            }
            text
        }
    };
    let mut out = std::io::stdout().lock();
    match bxctl::sanitize::run(&text, &mut out) {
        Ok(false) => ExitCode::SUCCESS,
        Ok(true) => ExitCode::from(1),
        Err(e) => {
            eprintln!("bxctl: cannot write the report: {e}");
            ExitCode::from(2)
        }
    }
}

fn on_pending(opts: &ChatOptions) -> OnPending {
    if opts.json {
        OnPending::EventOnly
    } else if opts.say.is_some() {
        OnPending::Show
    } else {
        OnPending::Ask
    }
}

fn run(opts: &ChatOptions) -> ExitCode {
    let stdin = std::io::stdin();
    let mut input = BufReader::new(stdin.lock());
    match &opts.say {
        Some(text) => run_say(opts, text, &mut input),
        None => run_interactive(opts, &mut input),
    }
}

fn run_say(opts: &ChatOptions, text: &str, input: &mut dyn BufRead) -> ExitCode {
    let session = match &opts.session {
        Some(s) => s.clone(),
        None => new_session_id(),
    };
    let mut printer = Printer::new(opts.show_thinking, opts.json);
    let code = run_turn_twice(
        opts,
        &session,
        text,
        opts.session.is_some(),
        &mut printer,
        input,
    );
    match code {
        Ok(()) => ExitCode::SUCCESS,
        Err(RunError::Io(e)) => {
            eprintln!("bxctl: {e}");
            ExitCode::from(1)
        }
        Err(RunError::Chat(e)) => {
            eprintln!("bxctl: {e}");
            ExitCode::from(1)
        }
    }
}

// Runs a --say turn, retrying once when the named session does not exist.
fn run_turn_twice(
    opts: &ChatOptions,
    session: &SessionId,
    text: &str,
    resume: bool,
    printer: &mut Printer,
    input: &mut dyn BufRead,
) -> Result<(), RunError> {
    let stderr = std::io::stderr();
    let mut handle = stderr.lock();
    let approvals = Approvals {
        admin_socket: &opts.admin_socket,
        on_pending: on_pending(opts),
    };
    let mut io = TurnIo {
        printer,
        input,
        out: &mut handle,
    };
    match stream_turn(
        &opts.socket,
        session,
        text,
        resume,
        opts.speaker.as_deref(),
        &approvals,
        &mut io,
    ) {
        Err(e) => Err(RunError::Io(e)),
        Ok(Ok(done)) => {
            write_answer(opts, &done).map_err(RunError::Io)?;
            Ok(())
        }
        Ok(Err(ChatError::Refused(w))) if w.code == ErrorCode::NoSuchSession => {
            // The named session does not exist: retry once, creating it (resume=false).
            let mut io = TurnIo {
                printer,
                input,
                out: &mut handle,
            };
            match stream_turn(
                &opts.socket,
                session,
                text,
                false,
                opts.speaker.as_deref(),
                &approvals,
                &mut io,
            ) {
                Err(e) => Err(RunError::Io(e)),
                Ok(Ok(done)) => {
                    write_answer(opts, &done).map_err(RunError::Io)?;
                    Ok(())
                }
                Ok(Err(e)) => Err(RunError::Chat(e)),
            }
        }
        Ok(Err(e)) => Err(RunError::Chat(e)),
    }
}

fn run_interactive(opts: &ChatOptions, reader: &mut dyn BufRead) -> ExitCode {
    let mut printer = Printer::new(opts.show_thinking, opts.json);
    let mut session = None;
    loop {
        eprint!("> ");
        let _ = std::io::stderr().flush();
        let mut line = String::new();
        match reader.read_line(&mut line) {
            Ok(0) => break,
            Ok(_) => {}
            Err(e) => {
                eprintln!("bxctl: {e}");
                return ExitCode::from(1);
            }
        }
        let text = line.trim();
        if text == "/quit" {
            break;
        }
        if text.is_empty() {
            continue;
        }
        let just_created = session.is_none();
        let this_session = session.get_or_insert_with(new_session_id);
        // The first turn creates the session (resume=false); later turns resume it.
        let resume = !just_created;
        if just_created {
            println!("session {}", this_session.as_str());
        }
        let stderr = std::io::stderr();
        let mut handle = stderr.lock();
        let approvals = Approvals {
            admin_socket: &opts.admin_socket,
            on_pending: on_pending(opts),
        };
        let mut io = TurnIo {
            printer: &mut printer,
            input: reader,
            out: &mut handle,
        };
        match stream_turn(
            &opts.socket,
            this_session,
            text,
            resume,
            opts.speaker.as_deref(),
            &approvals,
            &mut io,
        ) {
            Err(e) => {
                eprintln!("bxctl: {e}");
                return ExitCode::from(1);
            }
            Ok(Ok(done)) => {
                if let Err(e) = write_answer(opts, &done) {
                    eprintln!("bxctl: {e}");
                    return ExitCode::from(1);
                }
            }
            Ok(Err(e)) => {
                // A failed turn is reported and the loop goes on: the session still exists, so the
                // next line resumes it.
                eprintln!("bxctl: {e}");
            }
        }
    }
    ExitCode::SUCCESS
}

// The plain answer goes to stdout; in --json mode the TurnDone goes to stderr after the events.
fn write_answer(opts: &ChatOptions, done: &TurnDone) -> std::io::Result<()> {
    if opts.json {
        let line = serde_json::to_string(done).map_err(std::io::Error::other)?;
        let mut err = std::io::stderr().lock();
        err.write_all(line.as_bytes())?;
        err.write_all(b"\n")?;
    }
    let stdout = std::io::stdout();
    let mut out = stdout.lock();
    out.write_all(escape_model_text(&done.content).as_bytes())?;
    out.write_all(b"\n")?;
    out.flush()
}

enum RunError {
    Io(std::io::Error),
    Chat(ChatError),
}