use axum::{
Router, middleware,
routing::{get, post},
};
use clap::Parser;
use octoroute::{
cli::{Cli, Command, generate_config_template},
config::Config,
error::AppError,
handlers::{self, AppState},
middleware::request_id_middleware,
telemetry,
};
use std::net::SocketAddr;
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let cli = Cli::parse();
if let Some(command) = cli.command {
match command {
Command::Config { output } => {
return handle_config_command(output).map_err(|e| e.into());
}
}
}
run_server(&cli.config).await
}
fn handle_config_command(output: Option<String>) -> Result<(), AppError> {
let template = generate_config_template();
match output {
Some(path) => {
if std::path::Path::new(&path).exists() {
return Err(AppError::ConfigFileExists { path });
}
std::fs::write(&path, template).map_err(|source| {
let remediation = match source.kind() {
std::io::ErrorKind::PermissionDenied => format!(
"\nPermission denied. Check that:\n\
1. Parent directory has write permissions\n\
2. Current user can write to: {}",
std::path::Path::new(&path)
.parent()
.map(|p| p.display().to_string())
.unwrap_or_else(|| ".".to_string())
),
std::io::ErrorKind::NotFound => format!(
"\nDirectory not found. Check that parent directory exists: {}",
std::path::Path::new(&path)
.parent()
.map(|p| p.display().to_string())
.unwrap_or_else(|| ".".to_string())
),
_ => String::new(),
};
AppError::ConfigFileWrite {
path: path.clone(),
source,
remediation,
}
})?;
eprintln!("Configuration template written to: {}", path);
eprintln!(
"Edit the file to configure your model endpoints, then run: octoroute --config {}",
path
);
}
None => {
print!("{}", template);
}
}
Ok(())
}
async fn run_server(config_path: &str) -> Result<(), Box<dyn std::error::Error>> {
let config = Config::from_file(config_path)?;
telemetry::init(&config.observability.log_level);
tracing::info!(
"Starting Octoroute server on {}:{}",
config.server.host,
config.server.port
);
let config = std::sync::Arc::new(config);
let state = AppState::new(config.clone())?;
let shutdown_state = state.clone();
let app = Router::new()
.route("/health", get(handlers::health::handler))
.route("/chat", post(handlers::chat::handler))
.route("/models", get(handlers::models::handler))
.route("/metrics", get(handlers::metrics::handler))
.route(
"/v1/chat/completions",
post(handlers::openai::completions::handler),
)
.route("/v1/models", get(handlers::openai::models::handler))
.with_state(state)
.layer(middleware::from_fn(request_id_middleware));
let ip_addr = config
.server
.host
.parse::<std::net::IpAddr>()
.map_err(|e| {
format!(
"Invalid IP address '{}' in config: {}. Expected format: 0.0.0.0 or 127.0.0.1",
config.server.host, e
)
})?;
let addr = SocketAddr::from((ip_addr, config.server.port));
tracing::info!("Listening on {}", addr);
tracing::info!("Health check available at http://{}/health", addr);
tracing::info!("Legacy chat endpoint at http://{}/chat", addr);
tracing::info!("Legacy models status at http://{}/models", addr);
tracing::info!("Metrics endpoint at http://{}/metrics", addr);
tracing::info!("OpenAI-compatible endpoints:");
tracing::info!(" POST http://{}/v1/chat/completions", addr);
tracing::info!(" GET http://{}/v1/models", addr);
let listener = tokio::net::TcpListener::bind(addr).await?;
axum::serve(listener, app)
.with_graceful_shutdown(shutdown_signal(shutdown_state))
.await?;
tracing::info!("Server shutdown complete");
Ok(())
}
async fn shutdown_signal(state: AppState) {
use tokio::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 signal handler")
.recv()
.await;
};
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
tokio::select! {
_ = ctrl_c => {
tracing::info!("Received SIGINT (Ctrl+C), starting graceful shutdown");
},
_ = terminate => {
tracing::info!("Received SIGTERM, starting graceful shutdown");
},
}
state.selector().health_checker().shutdown().await;
}