kcode-k1-daemon-lib 0.6.0

#![doc = include_str!("../Documentation.md")]

use axum::body::{Body, to_bytes};
use axum::extract::Request;
use axum::http::header::{CACHE_CONTROL, CONTENT_LENGTH, CONTENT_TYPE, HOST};
use axum::http::{HeaderValue, StatusCode};
use axum::middleware::{self, Next};
use axum::response::Response;
use axum::routing::get;
use axum::{Json, Router};
use kcode_codex_terra::CodexTerra;
use kcode_gemini_3_1_pro::Gemini31Pro;
use kcode_k1_access::K1Access;
use kcode_k1_access_full_audio::K1AccessFullAudio;
use kcode_k1_access_profiles::K1AccessProfiles;
use kcode_k1_accounting::Accounting;
use kcode_k1_accounts::K1Accounts;
use kcode_k1_audio_classification::AudioClassification;
use kcode_k1_daemon_files::DaemonFiles;
use kcode_k1_full_audio::K1FullAudio;
use kcode_k1_groups::{K1Groups, ModelId};
use kcode_k1_http::{Config, K1Http};
use kcode_k1_http_accounts::K1HttpAccounts;
use kcode_k1_http_people::K1HttpPeople;
use kcode_k1_http_replay::{ReplayConfig, ReplayWindow};
use kcode_k1_invites::K1Invites;
use kcode_k1_objects::K1Objects;
use kcode_k1_peering::K1Peering;
use kcode_k1_persons::K1Persons;
use kcode_k1_txn_ordering::K1TxnOrdering;
use kcode_k1_users::K1Users;
use kcode_k1_vault::{ExposeSecret, K1Vault, SecretString};
use kcode_speaker_v3_analysis::Analyzer;
use serde::Serialize;
use serde_json::Value;
use std::io::Write as _;
use std::path::{Path, PathBuf};
use std::process::ExitCode;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::net::TcpListener;
use tokio::signal::unix::{Signal, SignalKind, signal};

const LISTEN_ADDRESS: &str = "127.0.0.1:4450";
const PUBLIC_ORIGIN: &str = "http://localhost:4450";
const INVITE_LINK_URL: &str = "http://localhost:4321/lib/kcode-k1-ui/*/account.html";
const AUTHORITY: &str = "localhost:4450";
const STARTUP_BOUND: Duration = Duration::from_millis(100);
const PROVIDER_OPERATION_TIMEOUT: Duration = Duration::from_secs(30 * 60);
const GEMINI_API_KEY: &str = "gemini-api-key";
const GEMINI_MODEL_BYTES: [u8; 32] = *b"gemini-3.1-pro-preview..........";
const TERRA_MODEL_BYTES: [u8; 32] = *b"gpt-5.6-terra...................";
const API_OPERATION: &str = "serve API request";

#[derive(Clone, Serialize)]
struct PublicConfig {
    protocol: &'static str,
    server_id: String,
    public_origin: &'static str,
}

#[derive(Serialize)]
struct Ready {
    event: &'static str,
    public_origin: &'static str,
    unused_invites: usize,
}

struct Prepared {
    app: Router,
    listener: TcpListener,
    signals: Signals,
    unused_invites: usize,
    vault: Arc<K1Vault>,
}

struct Signals {
    interrupt: Signal,
    terminate: Signal,
}

pub fn run(k1_root: PathBuf) -> ExitCode {
    let runtime = match tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
    {
        Ok(runtime) => runtime,
        Err(_) => {
            eprintln!("kcode-k1-daemon: startup failed");
            return ExitCode::from(1);
        }
    };
    let passphrase = match rpassword::prompt_password("Unlock K1 vault: ") {
        Ok(passphrase) => match protect_passphrase(passphrase) {
            Ok(passphrase) => passphrase,
            Err(()) => {
                eprintln!("kcode-k1-daemon: startup failed");
                return ExitCode::from(1);
            }
        },
        Err(_) => {
            eprintln!("kcode-k1-daemon: startup failed");
            return ExitCode::from(1);
        }
    };
    runtime.block_on(run_async(k1_root, passphrase))
}

fn protect_passphrase(passphrase: String) -> Result<SecretString, ()> {
    (!passphrase.is_empty())
        .then(|| SecretString::from(passphrase))
        .ok_or(())
}

async fn run_async(k1_root: PathBuf, passphrase: SecretString) -> ExitCode {
    let started = Instant::now();
    let prepared = match startup(k1_root, passphrase).await {
        Ok(prepared) => prepared,
        Err(()) => {
            warn_if_slow(started.elapsed(), "error");
            eprintln!("kcode-k1-daemon: startup failed");
            return ExitCode::from(1);
        }
    };
    let elapsed = started.elapsed();
    if write_readiness(prepared.unused_invites).is_err() {
        warn_if_slow(elapsed, "error");
        eprintln!("kcode-k1-daemon: startup failed");
        return ExitCode::from(1);
    }
    warn_if_slow(elapsed, "ready");
    let Prepared {
        app,
        listener,
        signals,
        vault,
        ..
    } = prepared;
    let result = axum::serve(listener, app)
        .with_graceful_shutdown(signals.wait())
        .await;
    drop(vault);
    match result {
        Ok(()) => ExitCode::SUCCESS,
        Err(_) => {
            eprintln!("kcode-k1-daemon: listener failed");
            ExitCode::from(1)
        }
    }
}

async fn startup(k1_root: PathBuf, passphrase: SecretString) -> Result<Prepared, ()> {
    let state_root = state_root(&k1_root);
    let files = DaemonFiles::open(&state_root).map_err(|_| ())?;
    let ordering = Arc::new(K1TxnOrdering::open(&state_root.join("ordering")).map_err(|_| ())?);
    let peering = Arc::new(
        K1Peering::open(&state_root.join("peering"), Arc::clone(&ordering)).map_err(|_| ())?,
    );
    let vault = open_vault(
        &state_root,
        passphrase,
        Arc::clone(&ordering),
        Arc::clone(&peering),
    )?;
    let persons = Arc::new(
        K1Persons::open(
            &state_root.join("persons"),
            Arc::clone(&ordering),
            Arc::clone(&peering),
        )
        .map_err(|_| ())?,
    );
    let invites = Arc::new(
        K1Invites::open(
            &state_root.join("invites"),
            Arc::clone(&ordering),
            Arc::clone(&peering),
        )
        .map_err(|_| ())?,
    );
    let accounts = Arc::new(K1Accounts::open(Arc::clone(&invites)).map_err(|_| ())?);
    let users = Arc::new(K1Users::new(Arc::clone(&accounts), Arc::clone(&persons)));
    let groups = Arc::new(
        K1Groups::open(
            &state_root.join("groups"),
            Arc::clone(&ordering),
            Arc::clone(&peering),
        )
        .map_err(|_| ())?,
    );
    let profiles = Arc::new(
        K1AccessProfiles::open(
            &state_root.join("access-profiles"),
            Arc::clone(&ordering),
            Arc::clone(&peering),
        )
        .map_err(|_| ())?,
    );
    let gemini_key = vault.secret(GEMINI_API_KEY).map_err(|_| ())?.ok_or(())?;
    let accounting = Accounting::new();
    let gemini = Gemini31Pro::new(
        gemini_key.expose_secret().to_owned(),
        accounting.clone(),
        PROVIDER_OPERATION_TIMEOUT,
    )
    .map_err(|_| ())?;
    let terra = CodexTerra::new(
        accounting,
        "codex",
        std::env::current_dir().map_err(|_| ())?,
        PROVIDER_OPERATION_TIMEOUT,
    )
    .map_err(|_| ())?;
    let analyzer = Analyzer::new(gemini, terra);
    let objects =
        Arc::new(K1Objects::open(Arc::clone(&ordering), Arc::clone(&peering)).map_err(|_| ())?);
    let classification = Arc::new(
        AudioClassification::open(
            &state_root.join("audio-classification"),
            Arc::clone(&ordering),
            Arc::clone(&peering),
            Arc::clone(&objects),
            analyzer,
        )
        .map_err(|_| ())?,
    );
    let full_audio = Arc::new(
        K1FullAudio::open(
            resolve_ffmpeg()?,
            Arc::clone(&objects),
            Arc::clone(&classification),
        )
        .map_err(|_| ())?,
    );
    let access = Arc::new(
        K1Access::open(
            &state_root.join("access"),
            Arc::clone(&ordering),
            Arc::clone(&peering),
            Arc::clone(&groups),
        )
        .map_err(|_| ())?,
    );
    let audio = Arc::new(
        K1AccessFullAudio::open_for_models(
            access,
            Arc::clone(&profiles),
            full_audio,
            classification,
            Arc::clone(&groups),
            audio_models().to_vec(),
        )
        .map_err(|_| ())?,
    );
    let replay = ReplayWindow::open(ReplayConfig {
        epoch_file: files.replay_epoch_path().to_owned(),
        max_nonces_per_epoch: usize::MAX,
    })
    .await
    .map_err(|_| ())?;
    let unused_invites = kcode_k1_daemon_invite_stock::reconcile(
        &invites,
        files.invite_links_path(),
        INVITE_LINK_URL,
    )
    .map_err(|_| ())?;
    if unused_invites < 100 {
        return Err(());
    }
    let adapter = K1HttpAccounts::new(
        Arc::clone(&accounts),
        Arc::clone(&invites),
        Arc::clone(&users),
    );
    let people = K1HttpPeople::new(accounts, users, groups, profiles);
    let http = K1Http::new(
        Config {
            server_id: files.server_id().to_owned(),
            public_origin: PUBLIC_ORIGIN.to_owned(),
            max_body_bytes: usize::MAX,
        },
        replay,
        adapter.identity_provider(),
    )
    .map_err(|_| ())?;
    let authenticated = adapter
        .authenticated_routes()
        .merge(people.authenticated_routes())
        .merge(kcode_k1_http_audio::authenticated_routes(audio))
        .fallback(api_not_found);
    let api = http
        .router(
            adapter.registration_endpoint(),
            kcode_k1_terms::endpoint(),
            authenticated,
        )
        .layer(middleware::from_fn(contextualize_api_error));
    let config = PublicConfig {
        protocol: "K1-HTTP-1",
        server_id: files.server_id().to_owned(),
        public_origin: PUBLIC_ORIGIN,
    };
    let config_route = get(move || {
        let config = config.clone();
        async move { ([(CACHE_CONTROL, "no-store")], Json(config)) }
    });
    let app = Router::new()
        .route("/config.json", config_route)
        .merge(api)
        .layer(middleware::from_fn(require_authority));
    Ok(Prepared {
        app,
        listener: TcpListener::bind(LISTEN_ADDRESS).await.map_err(|_| ())?,
        signals: Signals::install()?,
        unused_invites,
        vault,
    })
}

fn audio_models() -> [ModelId; 2] {
    [
        ModelId::from_bytes(GEMINI_MODEL_BYTES),
        ModelId::from_bytes(TERRA_MODEL_BYTES),
    ]
}

fn resolve_ffmpeg() -> Result<PathBuf, ()> {
    let path = std::env::var_os("PATH").ok_or(())?;
    resolve_executable("ffmpeg", std::env::split_paths(&path))
}

fn resolve_executable(name: &str, paths: impl IntoIterator<Item = PathBuf>) -> Result<PathBuf, ()> {
    paths
        .into_iter()
        .find_map(|directory| {
            let candidate = directory.join(name);
            executable(&candidate)
                .then(|| std::fs::canonicalize(candidate).ok())
                .flatten()
                .filter(|path| path.is_absolute())
        })
        .ok_or(())
}

#[cfg(unix)]
fn executable(path: &Path) -> bool {
    use std::os::unix::fs::PermissionsExt as _;
    std::fs::metadata(path)
        .is_ok_and(|metadata| metadata.is_file() && metadata.permissions().mode() & 0o111 != 0)
}

#[cfg(not(unix))]
fn executable(path: &Path) -> bool {
    std::fs::metadata(path).is_ok_and(|metadata| metadata.is_file())
}

fn open_vault(
    state_root: &Path,
    passphrase: SecretString,
    ordering: Arc<K1TxnOrdering>,
    peering: Arc<K1Peering>,
) -> Result<Arc<K1Vault>, ()> {
    K1Vault::open(&state_root.join("vault"), passphrase, ordering, peering)
        .map(Arc::new)
        .map_err(|_| ())
}

fn state_root(k1_root: &Path) -> PathBuf {
    k1_root.join("state")
}

async fn api_not_found() -> Response {
    json_error(
        StatusCode::NOT_FOUND,
        "not_found",
        "authenticated API route not found",
    )
}

async fn contextualize_api_error(request: Request, next: Next) -> Response {
    let response = next.run(request).await;
    if !(response.status().is_client_error() || response.status().is_server_error()) {
        return response;
    }
    let (mut parts, body) = response.into_parts();
    let bytes = match to_bytes(body, usize::MAX).await {
        Ok(bytes) => bytes,
        Err(_) => return Response::from_parts(parts, Body::empty()),
    };
    let Some(contextualized) = contextualize_error_body(&bytes) else {
        return Response::from_parts(parts, Body::from(bytes));
    };
    parts.headers.remove(CONTENT_LENGTH);
    Response::from_parts(parts, Body::from(contextualized))
}

fn contextualize_error_body(bytes: &[u8]) -> Option<Vec<u8>> {
    let mut payload: Value = serde_json::from_slice(bytes).ok()?;
    let object = payload.as_object_mut()?;
    let code = object.get("error")?.as_str()?.to_owned();
    let source = object
        .get("message")
        .and_then(Value::as_str)
        .map(str::to_owned)
        .unwrap_or_else(|| format!("error code {code}"));
    object.insert(
        "message".to_owned(),
        Value::String(format!("{API_OPERATION}: {source}")),
    );
    Some(payload.to_string().into_bytes())
}

async fn require_authority(request: Request, next: Next) -> Response {
    let mut values = request.headers().get_all(HOST).iter();
    if values
        .next()
        .is_some_and(|value| value.as_bytes() == AUTHORITY.as_bytes())
        && values.next().is_none()
    {
        next.run(request).await
    } else {
        json_error(
            StatusCode::MISDIRECTED_REQUEST,
            "invalid_request_authority",
            "validate request authority: request authority is invalid",
        )
    }
}

fn json_error(status: StatusCode, code: &'static str, message: &'static str) -> Response {
    let mut response = Response::new(Body::from(
        serde_json::json!({"error": code, "message": message}).to_string(),
    ));
    *response.status_mut() = status;
    response
        .headers_mut()
        .insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
    response
        .headers_mut()
        .insert(CACHE_CONTROL, HeaderValue::from_static("no-store"));
    response
}

fn write_readiness(unused_invites: usize) -> Result<(), ()> {
    let stdout = std::io::stdout();
    let mut output = stdout.lock();
    serde_json::to_writer(
        &mut output,
        &Ready {
            event: "ready",
            public_origin: PUBLIC_ORIGIN,
            unused_invites,
        },
    )
    .map_err(|_| ())?;
    output.write_all(b"\n").map_err(|_| ())?;
    output.flush().map_err(|_| ())
}

fn warn_if_slow(elapsed: Duration, outcome: &'static str) {
    if elapsed > STARTUP_BOUND {
        eprintln!(
            "{{\"module\":\"kcode-k1-daemon\",\"operation\":\"startup\",\"elapsed_us\":{},\"outcome\":\"{outcome}\"}}",
            elapsed.as_micros()
        );
    }
}

impl Signals {
    fn install() -> Result<Self, ()> {
        Ok(Self {
            interrupt: signal(SignalKind::interrupt()).map_err(|_| ())?,
            terminate: signal(SignalKind::terminate()).map_err(|_| ())?,
        })
    }

    async fn wait(mut self) {
        tokio::select! {
            _ = self.interrupt.recv() => {}
            _ = self.terminate.recv() => {}
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn public_operation_accepts_only_the_state_root() {
        let _: fn(PathBuf) -> ExitCode = run;
    }

    #[test]
    fn accepted_passphrase_boundary_is_strict_and_protected() {
        assert!(protect_passphrase(String::new()).is_err());
        let text = "conspicuous-fake-passphrase-never-real";
        let protected = protect_passphrase(text.to_owned()).unwrap();
        assert!(!format!("{protected:?}").contains(text));
    }

    #[test]
    fn vault_composition_persists_at_the_fixed_path() {
        let root =
            std::env::temp_dir().join(format!("kcode-k1-daemon-vault-test-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&root);
        let state = state_root(&root);
        assert_eq!(state.join("vault"), root.join("state/vault"));
        let parts = || {
            let ordering = Arc::new(K1TxnOrdering::open(&state.join("ordering")).unwrap());
            let peering =
                Arc::new(K1Peering::open(&state.join("peering"), ordering.clone()).unwrap());
            (ordering, peering)
        };
        let password = || SecretString::from("fake-test-password-never-real");
        let (ordering, peering) = parts();
        let vault = open_vault(&state, password(), ordering.clone(), peering.clone()).unwrap();
        vault
            .set(
                "fake-provider-secret",
                SecretString::from("conspicuous-fake-value-never-real"),
            )
            .unwrap();
        drop((vault, peering, ordering));
        let (ordering, peering) = parts();
        let vault = open_vault(&state, password(), ordering.clone(), peering.clone()).unwrap();
        drop((vault, peering, ordering));
        let (ordering, peering) = parts();
        assert!(
            open_vault(
                &state,
                SecretString::from("wrong-fake-password-never-real"),
                ordering,
                peering
            )
            .is_err()
        );
        std::fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn audio_model_ids_are_fixed_distinct_and_in_order() {
        assert_eq!(GEMINI_MODEL_BYTES, *b"gemini-3.1-pro-preview..........");
        assert_eq!(TERRA_MODEL_BYTES, *b"gpt-5.6-terra...................");
        assert_eq!(GEMINI_MODEL_BYTES.len(), 32);
        assert_eq!(TERRA_MODEL_BYTES.len(), 32);
        let models = audio_models();
        assert_eq!(models[0].as_bytes(), &GEMINI_MODEL_BYTES);
        assert_eq!(models[1].as_bytes(), &TERRA_MODEL_BYTES);
        assert_ne!(models[0], models[1]);
    }

    #[test]
    fn only_the_fixed_gemini_vault_key_is_selected() {
        assert_eq!(GEMINI_API_KEY, "gemini-api-key");
    }

    #[cfg(unix)]
    #[test]
    fn executable_resolver_returns_an_absolute_path_for_a_fake_ffmpeg() {
        use std::os::unix::fs::PermissionsExt as _;
        let root = std::env::temp_dir().join(format!(
            "kcode-k1-daemon-ffmpeg-test-{}",
            std::process::id()
        ));
        let _ = std::fs::remove_dir_all(&root);
        std::fs::create_dir(&root).unwrap();
        let fake = root.join("ffmpeg");
        std::fs::write(&fake, b"#!/bin/sh\nexit 0\n").unwrap();
        std::fs::set_permissions(&fake, std::fs::Permissions::from_mode(0o700)).unwrap();
        assert_eq!(
            resolve_executable("ffmpeg", [root.clone()]).unwrap(),
            std::fs::canonicalize(&fake).unwrap()
        );
        std::fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn invite_link_and_backend_origins_remain_distinct() {
        assert_eq!(
            INVITE_LINK_URL,
            "http://localhost:4321/lib/kcode-k1-ui/*/account.html"
        );
        assert_eq!(PUBLIC_ORIGIN, "http://localhost:4450");
        assert_ne!(INVITE_LINK_URL, PUBLIC_ORIGIN);
    }

    #[test]
    fn existing_child_message_is_preserved_under_daemon_context() {
        let body = contextualize_error_body(
            br#"{"error":"group_failed","message":"load group: child failure","detail":7}"#,
        )
        .unwrap();
        let payload: Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(payload["error"], "group_failed");
        assert_eq!(payload["detail"], 7);
        assert_eq!(
            payload["message"],
            "serve API request: load group: child failure"
        );
    }

    #[test]
    fn missing_child_message_is_derived_from_stable_code() {
        let body = contextualize_error_body(br#"{"error":"invalid_signature"}"#).unwrap();
        let payload: Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(payload["error"], "invalid_signature");
        assert_eq!(
            payload["message"],
            "serve API request: error code invalid_signature"
        );
    }

    #[test]
    fn supplied_root_maps_only_to_state() {
        assert_eq!(
            state_root(Path::new("/trusted/k1")),
            PathBuf::from("/trusted/k1/state")
        );
    }
}