#![cfg_attr(
not(test),
deny(clippy::unwrap_used, clippy::panic, clippy::indexing_slicing)
)]
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};
use anyhow::Context;
use axum::middleware;
use kb::api::{self, AppState};
use kb::auth::{self, BearerToken};
use kb::embedding;
use kb::storage;
use tower_http::trace::{DefaultMakeSpan, DefaultOnResponse, TraceLayer};
use tracing::{Level, error, info};
use tracing_subscriber::fmt;
#[tokio::main]
#[allow(
clippy::disallowed_methods,
reason = "kb-server reads its KB_* runtime config once at the binary edge (REPO_INVARIANTS.md #5)"
)]
async fn main() -> anyhow::Result<()> {
fmt().json().with_target(false).init();
let token = require_env("KB_TOKEN");
let db_path = std::env::var("KB_DB_PATH")
.unwrap_or_else(|_| storage::default_db_path().to_string_lossy().into_owned());
let host = std::env::var("KB_HOST").unwrap_or_else(|_| "127.0.0.1".into());
let port: u16 = std::env::var("KB_PORT")
.unwrap_or_else(|_| "8080".into())
.parse()
.context("KB_PORT must be a valid port number")?;
let conn = storage::open_db(&db_path).context("failed to open database")?;
storage::init_db(&conn).context("failed to initialize database schema")?;
let (embedding_client, embedding_model) = match embedding::read_embedding_config_from_env() {
Some(cfg) => {
info!(model = %cfg.model, base_url = %cfg.base_url, "embedding hook enabled");
let model = cfg.model.clone();
let client: Arc<dyn kb::embedding::EmbeddingClient> =
Arc::new(embedding::http_embedding_client(cfg));
(Some(client), Some(model))
}
None => (None, None),
};
let state = Arc::new(AppState {
conn: Mutex::new(conn),
embedding_client,
embedding_model,
});
let trace_layer = TraceLayer::new_for_http()
.make_span_with(
DefaultMakeSpan::new()
.level(Level::INFO)
.include_headers(false),
)
.on_response(
DefaultOnResponse::new()
.level(Level::INFO)
.include_headers(false),
);
let bt = BearerToken::new(&token);
let app = api::build_router(state)
.layer(middleware::from_fn_with_state(bt, auth::bearer_auth))
.layer(trace_layer);
let addr: SocketAddr = format!("{host}:{port}")
.parse()
.context("invalid bind address")?;
info!("kb-server listening on {addr}");
let listener = tokio::net::TcpListener::bind(addr)
.await
.context("failed to bind listener")?;
axum::serve(listener, app)
.with_graceful_shutdown(shutdown_signal())
.await
.context("server error")?;
Ok(())
}
async fn shutdown_signal() {
match tokio::signal::ctrl_c().await {
Ok(()) => info!("shutting down"),
Err(e) => {
error!(error = %e, "failed to install Ctrl+C handler; graceful shutdown disabled");
std::future::pending::<()>().await;
}
}
}
#[allow(
clippy::disallowed_methods,
reason = "kb-server reads required runtime config once at the binary edge (REPO_INVARIANTS.md #5)"
)]
fn require_env(name: &str) -> String {
std::env::var(name).unwrap_or_else(|_| {
eprintln!("{name} environment variable is required");
std::process::exit(1);
})
}