steep/crates/steep-server/src/main.rs

use std::time::Duration;

use crate::{
    bot::{Bot, BotDeps},
    event_bus::EventBus,
    gitea::GiteaAPI,
    job_queue::JobQueue,
    open_router::OpenRouterClient,
    review_store::ReviewStore,
    sandbox::SandboxConfig,
    state::AppState,
};

use dotenvy::dotenv;
use tokio::signal::unix::{SignalKind, signal};
use tokio_util::sync::CancellationToken;
use tracing::{info, warn};
use tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt};

mod api;
mod bot;
mod bot_actions;
mod consts;
mod db;
mod domain;
mod env;
mod errors;
mod event_bus;
mod gitea;
mod job_queue;
mod metrics;
mod open_router;
mod review_store;
mod sandbox;
mod schemas;
mod state;
mod text;

/// Time left to the running reviews and the servers to stop cleanly.
///
/// Without it, a review stuck in an unresponsive sandbox would keep the process
/// alive until the supervisor kills it.
const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60);

fn main() -> anyhow::Result<()> {
    dotenv().ok();

    // TLS goes through rustls, which has no default crypto provider here: the
    // workspace compiles `ring` only, and it is installed once, for every client
    // (reqwest, sentry).
    rustls::crypto::ring::default_provider()
        .install_default()
        .map_err(|_| anyhow::anyhow!("a rustls crypto provider is already installed"))?;

    // Migration subcommand: `steep-server migration <generate|apply|...>`.
    if std::env::args().nth(1).as_deref() == Some("migration") {
        return tokio::runtime::Runtime::new()?.block_on(run_migrations());
    }

    tracing_subscriber::registry()
        .with(fmt::layer())
        .with(
            EnvFilter::try_from_default_env() // reads RUST_LOG from the environment
                .unwrap_or_else(|_| EnvFilter::new("info")),
        )
        .init();

    let _sentry_guard = if let Ok(sentry_dsn) = env::try_get_env("SENTRY_DSN") {
        info!("Initialize sentry");

        Some(sentry::init(
            sentry::ClientOptions::new()
                .dsn(&sentry_dsn)
                .maybe_release(sentry::release_name!())
                .send_default_pii(true),
        ))
    } else {
        warn!("SENTRY_DSN not set, sentry will not be initialized");
        None
    };

    tokio::runtime::Runtime::new()?.block_on(run())
}

async fn run_migrations() -> anyhow::Result<()> {
    let database_url = env::try_get_env("DATABASE_URL")?;
    let db = db::open(&database_url).await?;

    let migrations = match std::env::var("MIGRATIONS_DIR") {
        Ok(dir) => std::path::PathBuf::from(dir),
        Err(_) => std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("toasty"),
    };
    let config =
        toasty_cli::Config::new().migration(toasty_cli::MigrationConfig::new().path(migrations));

    toasty_cli::ToastyCli::with_config(db, config)
        .parse_and_run()
        .await?;

    Ok(())
}

async fn run() -> anyhow::Result<()> {
    let config = env::load_config()?;

    if let Some(metric_bind_addr) = &config.metrics_bind_addr {
        metrics::install(metric_bind_addr)?;
    }

    let gitea_api = GiteaAPI::new(&config.gitea_url, &config.gitea_token, config.gitea_timeout)?;
    let gitea_user = gitea_api.get_authorized_user().await?;

    info!(
        port = config.http_port,
        model = %config.open_router_model,
        gitea_url = %config.gitea_url,
        bot_name = %gitea_user.login,
        "Starting Steep"
    );

    let open_router_client = OpenRouterClient::new(
        &config.open_router_api_key,
        &config.open_router_model,
        config.open_router_timeout,
    )?;

    let shutdown = CancellationToken::new();

    let sandbox = SandboxConfig {
        runtime: devcontainer_rs::ContainerRuntime::connect()?,
        max_iterations: config.sandbox_max_iterations,
    };

    if !sandbox.runtime.available().await {
        warn!(
            endpoint = sandbox.runtime.endpoint(),
            "Container daemon is not reachable, every review will fail"
        );
    }

    let db = db::open(&config.database_url).await?;
    let reviews = ReviewStore::new(db.clone());

    let queue = JobQueue::new(db);
    queue.requeue_interrupted().await?;

    let event_bus = EventBus::new();

    let bot = Bot::new(
        BotDeps {
            gitea: gitea_api,
            open_router: open_router_client,
            sandbox,
            reviews: reviews.clone(),
            queue: queue.clone(),
            event_bus: event_bus.clone(),
        },
        config.bot_max_concurrent,
    );

    let app_state = AppState {
        queue,
        reviews,
        event_bus,
        bot_name: gitea_user.login.clone(),
        config,
        shutdown: shutdown.clone(),
    };

    let signal = async {
        let mut sigterm = signal(SignalKind::terminate())?;
        let mut sigint = signal(SignalKind::interrupt())?;
        tokio::select! {
            _ = sigterm.recv() => info!("Received SIGTERM"),
            _ = sigint.recv() => info!("Received SIGINT"),
        }

        info!("Shutting down...");
        shutdown.cancel();
        anyhow::Ok(())
    };

    let shutdown_deadline = async {
        shutdown.cancelled().await;
        tokio::time::sleep(SHUTDOWN_TIMEOUT).await;
    };

    tokio::select! {
        result = async {
            tokio::try_join!(
                bot.start(shutdown.clone()),
                api::start(app_state),
                signal
            )
        } => {
            result?;
        }
        () = shutdown_deadline => {
            warn!(
                timeout = ?SHUTDOWN_TIMEOUT,
                "Shutdown did not finish in time, exiting anyway"
            );
        }
    }

    info!("Shutdown complete");

    Ok(())
}