agent-orchestrator/agor-pty/src/main.rs

mod auth;
mod daemon;
mod protocol;
mod session;

use std::path::PathBuf;

use tokio::signal::unix::{signal, SignalKind};
use tokio::sync::broadcast;

use auth::AuthToken;
use daemon::Daemon;

const VERSION: &str = "0.1.0";

// ---------------------------------------------------------------------------
// CLI argument parsing — no clap needed for 3 flags.
// ---------------------------------------------------------------------------

struct Cli {
    socket_dir: Option<PathBuf>,
    default_shell: Option<String>,
    verbose: bool,
}

impl Cli {
    fn parse() -> Result<Self, String> {
        let mut args = std::env::args().skip(1).peekable();
        let mut socket_dir = None;
        let mut default_shell = None;
        let mut verbose = false;

        while let Some(arg) = args.next() {
            match arg.as_str() {
                "--socket-dir" => {
                    socket_dir = Some(PathBuf::from(
                        args.next().ok_or("--socket-dir requires a value")?,
                    ));
                }
                "--shell" => {
                    default_shell = Some(
                        args.next().ok_or("--shell requires a value")?,
                    );
                }
                "--verbose" | "-v" => {
                    verbose = true;
                }
                "--help" | "-h" => {
                    print_usage();
                    std::process::exit(0);
                }
                other => {
                    return Err(format!("unknown argument: {other}"));
                }
            }
        }

        Ok(Self {
            socket_dir,
            default_shell,
            verbose,
        })
    }
}

fn print_usage() {
    eprintln!(
        "USAGE: agor-ptyd [OPTIONS]\n\
         \n\
         OPTIONS:\n\
           --socket-dir <PATH>   Socket directory\n\
                                 (default: /run/user/$UID/agor or ~/.local/share/agor/run)\n\
           --shell <PATH>        Default shell (default: $SHELL or /bin/bash)\n\
           --verbose             Enable debug logging\n\
           --help                Show this message"
    );
}

// ---------------------------------------------------------------------------
// Socket directory resolution
// ---------------------------------------------------------------------------

fn resolve_socket_dir(override_path: Option<PathBuf>) -> Result<PathBuf, String> {
    if let Some(p) = override_path {
        return Ok(p);
    }

    // Prefer XDG runtime dir.
    if let Ok(uid_str) = std::env::var("UID").or_else(|_| {
        // UID is not always exported; fall back to getuid().
        Ok::<_, std::env::VarError>(unsafe { libc_getuid() }.to_string())
    }) {
        let xdg = PathBuf::from(format!("/run/user/{uid_str}/agor"));
        if xdg.parent().map(|p| p.exists()).unwrap_or(false) {
            return Ok(xdg);
        }
    }

    // Fallback: ~/.local/share/agor/run
    let home = std::env::var("HOME").map_err(|_| "HOME not set".to_string())?;
    Ok(PathBuf::from(home).join(".local/share/agor/run"))
}

#[cfg(target_os = "linux")]
unsafe fn libc_getuid() -> u32 {
    // Safety: getuid() is always safe.
    extern "C" {
        fn getuid() -> u32;
    }
    getuid()
}

#[cfg(not(target_os = "linux"))]
unsafe fn libc_getuid() -> u32 {
    0
}

// ---------------------------------------------------------------------------
// Default shell resolution
// ---------------------------------------------------------------------------

fn resolve_shell(override_shell: Option<String>) -> String {
    override_shell
        .or_else(|| std::env::var("SHELL").ok())
        .filter(|s| !s.is_empty())
        .unwrap_or_else(|| "/bin/bash".into())
}

// ---------------------------------------------------------------------------
// Entry point
// ---------------------------------------------------------------------------

#[tokio::main]
async fn main() {
    let cli = match Cli::parse() {
        Ok(c) => c,
        Err(e) => {
            eprintln!("error: {e}");
            print_usage();
            std::process::exit(1);
        }
    };

    // Initialise logging.
    let log_level = if cli.verbose { "debug" } else { "info" };
    env_logger::Builder::from_env(
        env_logger::Env::default().default_filter_or(log_level),
    )
    .init();

    log::info!("agor-ptyd v{VERSION} starting");

    let socket_dir = match resolve_socket_dir(cli.socket_dir) {
        Ok(d) => d,
        Err(e) => {
            log::error!("cannot resolve socket directory: {e}");
            std::process::exit(1);
        }
    };

    // Ensure the directory exists.
    if let Err(e) = std::fs::create_dir_all(&socket_dir) {
        log::error!("cannot create socket directory {socket_dir:?}: {e}");
        std::process::exit(1);
    }

    let socket_path = socket_dir.join("ptyd.sock");
    let shell = resolve_shell(cli.default_shell);

    // Generate and persist auth token.
    let token = match AuthToken::generate_and_persist(&socket_dir) {
        Ok(t) => t,
        Err(e) => {
            log::error!("token generation failed: {e}");
            std::process::exit(1);
        }
    };

    // Shutdown broadcast channel — one sender, N receivers.
    let (shutdown_tx, shutdown_rx) = broadcast::channel::<()>(1);

    // Signal handlers for SIGTERM and SIGINT.
    let shutdown_tx_sigterm = shutdown_tx.clone();
    let shutdown_tx_sigint = shutdown_tx.clone();

    tokio::spawn(async move {
        let mut sigterm = signal(SignalKind::terminate()).expect("SIGTERM handler");
        sigterm.recv().await;
        log::info!("SIGTERM received");
        let _ = shutdown_tx_sigterm.send(());
    });

    tokio::spawn(async move {
        let mut sigint = signal(SignalKind::interrupt()).expect("SIGINT handler");
        sigint.recv().await;
        log::info!("SIGINT received");
        let _ = shutdown_tx_sigint.send(());
    });

    let daemon = Daemon::new(socket_path, token, shell);
    if let Err(e) = daemon.run(shutdown_rx).await {
        log::error!("daemon exited with error: {e}");
        std::process::exit(1);
    }

    log::info!("agor-ptyd shut down cleanly");
}