use anyhow::Context; use axum::serve; use boosapi::{ config::AppConfig, db, services::{builtin_keys, runtime_auth}, AppState, }; use clap::{Parser, Subcommand}; use std::ffi::OsString; use std::net::SocketAddr; use tokio::signal; use tracing_subscriber::{prelude::*, EnvFilter}; #[derive(Parser, Debug)] #[command( name = "boosapi", version, author, about = "AI API Gateway — Rust rewrite of BoosAPI" )] struct Cli { #[command(subcommand)] command: Option, /// Config file path #[arg(short, long, default_value = "config.toml")] config: String, /// Host to bind to #[arg(long)] host: Option, /// Port to listen on #[arg(short, long)] port: Option, /// Log level (trace, debug, info, warn, error) #[arg(long)] log_level: Option, /// Only run database migrations and exit #[arg(long)] migrate_only: bool, /// Load and validate configuration, then exit #[arg(long)] check_config: bool, } #[derive(Subcommand, Debug)] enum Commands { /// Copy core data from SQLite into PostgreSQL and exit MigrateSqliteToPostgres { /// Source SQLite URL, e.g. sqlite:boosapi.db?mode=rwc #[arg(long)] sqlite_url: String, /// Target PostgreSQL URL, e.g. postgres://user:pass@host:5432/boosapi #[arg(long)] postgres_url: String, /// Truncate target business tables before importing #[arg(long)] force_truncate: bool, }, } #[tokio::main] async fn main() -> anyhow::Result<()> { let args = Cli::parse_from(normalize_legacy_args()); if let Some(Commands::MigrateSqliteToPostgres { sqlite_url, postgres_url, force_truncate, }) = &args.command { let sqlite = db::create_pool(sqlite_url, 1).await?; let pg = db::create_pg_pool(postgres_url, 2).await?; db::run_migrations(db::DbKind::Postgres, None, Some(&pg)).await?; let summary = db::sqlite_to_postgres::migrate(&sqlite, &pg, *force_truncate).await?; println!("{}", serde_json::to_string_pretty(&summary)?); return Ok(()); } // Load config let mut config = AppConfig::load_from_file(&args.config)?; // Apply CLI overrides if let Some(host) = args.host { config.server.host = host; } if let Some(port) = args.port { config.server.port = port; } if let Some(level) = args.log_level { config.server.log_level = level; } if args.check_config { println!("Configuration is valid."); println!("Server: {}:{}", config.server.host, config.server.port); println!("Database URL: {}", config.db.url); return Ok(()); } // Initialize tracing let filter = EnvFilter::try_from_default_env() .unwrap_or_else(|_| EnvFilter::new(&config.server.log_level)); tracing_subscriber::registry() .with(filter) .with(tracing_subscriber::fmt::layer()) .with(boosapi::services::runtime_logs::RuntimeLogLayer) .init(); tracing::info!("Loaded configuration from {}", args.config); let db_kind = db::DbKind::from_url(&config.db.url)?; // Create database pool let pool = db::create_pool(&config.db.url, config.db.max_connections).await?; tracing::info!("Database connected"); // Run migrations db::run_migrations(db_kind, Some(&pool), None).await?; db::schema_compat::ensure_sqlite_upstream_compat(&pool).await?; tracing::info!("Database migrations applied"); if args.migrate_only { tracing::info!("Migration only mode requested. Exiting."); return Ok(()); } // Bootstrap built-in system keys (idempotent). let comfyui_agent_key = builtin_keys::ensure_comfyui_agent_key(&pool) .await .context("failed to bootstrap built-in ComfyUI Agent key")?; tracing::info!( key_id = comfyui_agent_key.0, "Built-in ComfyUI Agent downstream key ready" ); let builtin = builtin_keys::BuiltinKeys { comfyui_agent: comfyui_agent_key, }; let runtime_auth = runtime_auth::RuntimeAuthSettings::load(&pool, config.auth.proxy_token.clone()) .await .context("failed to load runtime auth settings")?; // Build app state let state = AppState::new(config.clone(), pool, builtin, runtime_auth); if let Err(err) = state.comfyui.start().await { tracing::warn!(%err, "ComfyUI embedded process did not start"); } let comfyui_process = state.comfyui.clone(); // Build router let app = boosapi::routes::build_router(state); // Start server let addr = SocketAddr::new(config.server.host.parse()?, config.server.port); tracing::info!("Starting server on {addr}"); serve( tokio::net::TcpListener::bind(addr).await?, app.into_make_service_with_connect_info::(), ) .with_graceful_shutdown(shutdown_signal()) .await?; if let Err(err) = comfyui_process.stop().await { tracing::warn!(%err, "ComfyUI embedded process did not stop cleanly"); } Ok(()) } fn normalize_legacy_args() -> Vec { std::env::args_os() .map(|arg| { if arg == "-help" { OsString::from("--help") } else if arg == "-version" { OsString::from("--version") } else { arg } }) .collect() } async fn shutdown_signal() { let ctrl_c = async { signal::ctrl_c() .await .expect("failed to install Ctrl+C handler"); }; #[cfg(unix)] let terminate = async { signal::unix::signal(signal::unix::SignalKind::terminate()) .expect("failed to install SIGTERM handler") .recv() .await; }; #[cfg(not(unix))] let terminate = std::future::pending::<()>(); tokio::select! { _ = ctrl_c => {}, _ = terminate => {}, } tracing::info!("Shutdown signal received, starting graceful shutdown"); }