ai-app/server/src/main.rs

//! A phone interface to AI coding sessions -- the backend. See PLAN.md for
//! the whole picture; this is the entry point: config + session registry,
//! token bootstrap, and the one TLS listener.
//!
//! The listener binds the WireGuard interface's address only, and fails
//! closed -- if `wg0` is down the server refuses to start rather than
//! falling back to `0.0.0.0`, because this API *is* remote code execution
//! and the tunnel is what keeps its pre-auth surface (TLS handshake, HTTP
//! parsing, auth middleware) off the open internet. `--bind` overrides
//! explicitly for development; that is a deliberate, logged choice, never a
//! fallback.
//!
//! There is no plaintext listener at all, so the bearer token can't travel
//! unencrypted by misconfiguration -- even inside the tunnel.

mod auth;
mod config;
mod media;
mod models;
mod routes;
mod session;
mod setups;
mod ssh;
mod usage;

use std::net::{IpAddr, SocketAddr};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result};
use axum::middleware::Next;
use clap::Parser;
use tokio::signal::unix::{SignalKind, signal};

use wg_app_link::enroll;
use wg_app_link::netif::{self, WG_INTERFACE};
use wg_app_link::xdg::{config_home, data_home};

use config::TokenEntry;
use session::SessionManager;

const DEFAULT_PORT: u16 = 8443;

/// Serves AI coding sessions (Claude Code, llama.cpp) to the phone app.
#[derive(Parser)]
struct Args {
    /// TLS port for the whole API surface.
    #[arg(long, default_value_t = DEFAULT_PORT)]
    port: u16,

    /// Address to bind instead of the wg0 interface's -- a development
    /// override (e.g. 127.0.0.1 for curl, or a LAN address for a phone
    /// before the tunnel exists). Production runs without it and fails
    /// closed when wg0 is absent.
    #[arg(long)]
    bind: Option<IpAddr>,

    /// Where the token hashes, providers, hosts, and session list live.
    /// Defaults to `$XDG_CONFIG_HOME/ai-app/config.ron`.
    #[arg(long)]
    config: Option<PathBuf>,

    /// Directory for per-session data (transcripts, attachments, images).
    /// Defaults to `$XDG_DATA_HOME/ai-app/sessions`.
    #[arg(long)]
    data_dir: Option<PathBuf>,

    /// Directory for downloaded GGUF models. Defaults to
    /// `$XDG_DATA_HOME/ai-app/models`.
    #[arg(long)]
    models_dir: Option<PathBuf>,

    /// Directory holding the TLS certificates, generated here on first
    /// start. Defaults to `$XDG_CONFIG_HOME/ai-app/certs`.
    #[arg(long)]
    certs: Option<PathBuf>,

    /// Invalidate every enrolled token, generate a fresh one, and print
    /// its enrollment QR -- the whole lost-phone story.
    #[arg(long)]
    rotate_token: bool,

    /// Enroll one more device without touching the running server: mint a
    /// token, print its enrollment link (one line, stdout, nothing else)
    /// and exit. The server adopts the token the first time that device
    /// uses it. For a tool -- Dev Updater -- that opens the link on the
    /// phone, where a QR printed here cannot be scanned.
    #[arg(long)]
    enroll_link: bool,

    /// Hold every response back by this many milliseconds.
    ///
    /// A development aid, and a specific one: over the tunnel a phone's
    /// requests take tens to hundreds of milliseconds, and several faults
    /// live entirely in what the app does *while* one is outstanding --
    /// a page of history landing mid-fling, a screen drawn before its
    /// first answer arrives. On a loopback server every response is back
    /// within a millisecond or two, so those windows close before
    /// anything can be observed and the bug looks like it is not there.
    /// This reopens them on demand rather than by unplugging something.
    #[arg(long, default_value_t = 0, value_name = "MS")]
    delay: u64,

    /// Mark every session spawned here as throwaway: its process is
    /// stopped when this server exits, instead of being left running for
    /// the next start to adopt. On by default in a debug build.
    ///
    /// Sessions outlive the backend on purpose, which is right for the
    /// ones somebody is using and wrong for the ones a test made: a
    /// session spawned to check something leaves a `claude` behind that
    /// every later server adopts, and they accumulate silently -- twelve
    /// of them on this machine in a day, each holding a conversation open.
    /// So a development build cleans up after itself unless told not to
    /// (`--throwaway-sessions=false`), and a release build never does
    /// unless asked.
    ///
    /// The flag decides only what *new* sessions are marked as. What
    /// happens on the way out is decided by the mark, which is written
    /// into the session and outlives the server that made it -- so
    /// sessions spawned without it keep running, whichever server is up
    /// when one exits.
    #[arg(
        long,
        default_value_t = cfg!(debug_assertions),
        action = clap::ArgAction::Set,
        num_args = 0..=1,
        default_missing_value = "true",
        value_name = "BOOL",
    )]
    throwaway_sessions: bool,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Both rustls crypto providers are in the dependency graph (ureq
    // brings ring, axum-server brings aws-lc-rs), so rustls refuses to
    // pick one itself; choose before anything touches TLS.
    rustls::crypto::aws_lc_rs::default_provider()
        .install_default()
        .expect("no other TLS crypto provider is installed before main");

    // `info` unless RUST_LOG says otherwise. Written as a *fallback* rather than as the filter,
    // because `with_env_filter("info")` is a fixed directive that never reads the environment --
    // so the per-request diagnostics that AGENTS.md tells you to turn on with
    // `RUST_LOG=ai_server=debug` printed nothing, and the switch looked like the code it was
    // meant to instrument being wrong.
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .init();
    let args = Args::parse();

    let config_path = args
        .config
        .unwrap_or_else(|| config_home("ai-app").join("config.ron"));
    // Before the manager exists, on purpose: constructing it and seeding
    // setups touches sessions and subprocesses this invocation has no
    // business with while another instance is serving. Only the hash
    // reaches disk, in the spool `auth.rs` reads; the link itself goes to
    // stdout alone, because the caller opens whatever this prints.
    if args.enroll_link {
        let bind_ip = match args.bind {
            Some(ip) => ip,
            None => netif::wg_address("ai-server")?,
        };
        let token = enroll::generate_token();
        enroll::spool_pending(
            &config::pending_enrollments_dir(&config_path),
            "phone",
            &token,
        )?;
        println!(
            "{}",
            enroll::enrollment_uri("aiapp", bind_ip, args.port, &token)
        );
        return Ok(());
    }
    let data_dir = args
        .data_dir
        .unwrap_or_else(|| data_home("ai-app").join("sessions"));
    // Beside the session data rather than under it: models outlive every
    // session and are shared by all of them, so deleting a session must
    // never take a multi-gigabyte download with it.
    let models_dir = args
        .models_dir
        .unwrap_or_else(|| data_home("ai-app").join("models"));
    let models = Arc::new(models::ModelStore::new(models_dir.clone()));
    let manager = Arc::new(
        SessionManager::new(config_path.clone(), data_dir, models_dir.clone())
            .with_context(|| format!("failed to load {}", config_path.display()))?
            .marking_new_sessions_throwaway(args.throwaway_sessions),
    );
    if args.throwaway_sessions {
        tracing::warn!(
            "sessions spawned here are marked throwaway -- their processes are stopped when this \
             server exits rather than left running (--throwaway-sessions=false to keep them)"
        );
    }
    // After construction rather than inside it: seeding asks this machine
    // what it has, which is I/O, and a constructor that quietly runs a
    // subprocess is a surprise to every caller including the tests.
    manager.seed_setup().await?;

    tracing::info!("config: {}", config_path.display());
    tracing::info!("models: {}", models_dir.display());
    for setup in manager.setups() {
        match &setup.ssh {
            Some(ssh) => tracing::info!("  setup \"{}\" -> {}", setup.name, ssh.address),
            // No parenthetical naming the local machine: the default
            // setup is *called* "this machine", and the line read
            // "setup this machine (this machine)".
            None => tracing::info!("  setup \"{}\" runs here", setup.name),
        }
        for provider in &setup.providers {
            tracing::info!("    provider {} ({:?})", provider.name, provider.kind);
        }
    }
    for info in manager.sessions() {
        tracing::info!(
            "  session {} ({}, {:?})",
            info.id,
            info.provider,
            info.status
        );
    }

    // Before the interface check below, deliberately: the certificates are
    // also what the phone app embeds at build time, so they need to be
    // obtainable on a machine whose tunnel isn't up yet. The leaf is
    // reissued on every start, so once wg0 exists the next start covers it.
    let certs_dir = args
        .certs
        .unwrap_or_else(|| config_home("ai-app").join("certs"));
    let certificates = wg_app_link::certs::ensure("ai-app", &certs_dir, &netif::local_addresses())
        .with_context(|| format!("failed to prepare certificates in {}", certs_dir.display()))?;
    if certificates.ca_is_new {
        tracing::warn!(
            "a new CA was generated in {} -- any installed app pins the previous one and can no \
             longer reach this server. Rebuild it with app/build-apk.sh, which embeds this CA, \
             and reinstall through Dev Updater.",
            certs_dir.display(),
        );
    }

    let bind_ip = match args.bind {
        Some(ip) => {
            tracing::warn!(
                "binding {ip} by explicit --bind override -- production binds {WG_INTERFACE} only"
            );
            ip
        }
        None => netif::wg_address("ai-server")?,
    };

    // Token bootstrap: first run generates one; --rotate-token replaces
    // whatever exists. Either way the plaintext appears exactly once, in
    // the QR printed here.
    if args.rotate_token || manager.tokens().is_empty() {
        let rotating = args.rotate_token && !manager.tokens().is_empty();
        let token = enroll::generate_token();
        manager.set_tokens(vec![TokenEntry {
            name: "phone".to_string(),
            sha256: enroll::token_hash_hex(&token),
        }])?;
        if rotating {
            tracing::info!("rotated the enrolled token; the previous one is now invalid");
        }
        enroll::print_enrollment("aiapp", bind_ip, args.port, &token)?;
    }

    let tls_config = axum_server::tls_rustls::RustlsConfig::from_pem_file(
        &certificates.leaf_cert,
        &certificates.leaf_key,
    )
    .await
    .context("failed to load TLS cert/key")?;

    // No providers listed here any more: which machines can be asked, and
    // about what, comes from the setups at the moment the screen is opened
    // -- so a machine added from the phone reports its limits without a
    // restart, and the backend's own account stops standing in for every
    // machine's.
    let monitor = Arc::new(usage::UsageMonitor::new());

    // The bearer-token middleware wraps the entire router -- routes and
    // fallback alike -- here and only here, so a new route can't forget
    // auth. Zero unauthenticated endpoints.
    let app = routes::router(Arc::clone(&manager))
        .merge(routes::usage_router(monitor, Arc::clone(&manager)))
        .merge(routes::models_router(Arc::clone(&models)))
        .layer(axum::middleware::from_fn_with_state(
            Arc::clone(&manager),
            auth::require_token,
        ));

    // Outside the auth layer, so an unauthenticated request is refused at
    // the speed it always was: this is here to slow the app down, not to
    // widen the window on anything guessing at tokens.
    let app = match args.delay {
        0 => app,
        ms => {
            tracing::warn!("delaying every response by {ms}ms -- development override");
            app.layer(axum::middleware::from_fn(
                move |request, next: Next| async move {
                    tokio::time::sleep(Duration::from_millis(ms)).await;
                    next.run(request).await
                },
            ))
        }
    };

    let addr = SocketAddr::new(bind_ip, args.port);
    tracing::info!("serving https://{addr}");

    // Let go of the sessions on the way out rather than stopping them:
    // their processes are meant to outlive this one, so restarting the
    // backend does not end a turn somebody is waiting on. Each is recorded
    // in its session directory and adopted again on the way back up (see
    // `session::process`). The exception is the sessions marked throwaway,
    // which are stopped first -- see `--throwaway-sessions`. Both signals,
    // because systemd and OpenRC send TERM while a terminal sends INT.
    let serving = axum_server::bind_rustls(addr, tls_config)
        .serve(app.into_make_service_with_connect_info::<SocketAddr>());
    let mut terminate = signal(SignalKind::terminate()).context("listening for SIGTERM")?;
    tokio::select! {
        served = serving => served.context("TLS listener failed")?,
        _ = terminate.recv() => tracing::info!("SIGTERM -- letting go of sessions"),
        _ = tokio::signal::ctrl_c() => tracing::info!("interrupted -- letting go of sessions"),
    }
    // Stopped before the rest are let go of, and on every way out of the
    // select above: a throwaway session is one nobody meant to keep, and
    // the whole point is that nothing has to remember to clean it up.
    manager.stop_throwaway_sessions();
    manager.detach_all();

    Ok(())
}