o402 0.1.2

OpenAI-compatible gateway, paid with x402.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

//! `o402 serve --config PATH`

use std::net::SocketAddr;
use std::path::Path;
use std::time::Duration;

use axum::Router;
use tokio::net::TcpListener;

use crate::config::Config;
use crate::error::AppError;

/// Load `.env`, TOML, tracing, then serve.
///
/// # Errors
///
/// Returns [`AppError`] when config loading, tracing setup, bind, or serve fails.
pub(crate) async fn run(config_path: &Path) -> Result<(), AppError> {
    load_dotenv()?;
    let config = Config::from_path(config_path)?;
    crate::obs::init(&config.observability)?;
    listen(config).await
}

/// Bind `server.bind` (and optional `observability.metrics_bind`) and serve
/// until shutdown, then drain HTTP and spawned settles up to
/// `server.shutdown_timeout_secs`.
///
/// # Errors
///
/// Returns [`AppError::HttpClient`] or [`AppError::Config`] while building the
/// router, [`AppError::Bind`] if a listen socket fails, [`AppError::Metrics`]
/// if the recorder cannot be installed, or [`AppError::Server`] if an accept
/// loop fails.
pub(crate) async fn listen(config: Config) -> Result<(), AppError> {
    let addr = config.server.bind;
    let drain = Duration::from_secs(config.server.shutdown_timeout_secs);
    let payment_enabled = config.payment.enabled;
    let upstreams = config.upstreams.len();
    let metrics_bind = config.observability.metrics_bind;
    let (app, state) = crate::http::app_and_state(config)?;
    let listener = bind(addr).await?;
    let bound = listener.local_addr().unwrap_or(addr);
    let metrics = crate::obs::metrics::bind(metrics_bind).await?;
    let metrics_task = crate::obs::metrics::spawn(metrics);
    tracing::info!(
        %bound,
        payment_enabled,
        upstreams,
        "listening"
    );
    let result = serve_until_drain(listener, app, drain).await;
    state.inflight().wait_for_drain(drain).await;
    if let Some(task) = metrics_task {
        task.abort();
    }
    tracing::info!("shutdown complete");
    result
}

fn load_dotenv() -> Result<(), AppError> {
    match dotenvy::dotenv() {
        Ok(_) => Ok(()),
        Err(error) if error.not_found() => Ok(()),
        Err(error) => Err(AppError::Dotenv { source: error }),
    }
}

async fn bind(addr: SocketAddr) -> Result<TcpListener, AppError> {
    TcpListener::bind(addr)
        .await
        .map_err(|source| AppError::Bind { addr, source })
}

async fn serve_until_drain(
    listener: TcpListener,
    app: Router,
    drain: Duration,
) -> Result<(), AppError> {
    tokio::select! {
        result = axum::serve(listener, app).with_graceful_shutdown(crate::shutdown::signal()) => {
            result.map_err(|source| AppError::Server { source })
        }
        () = drain_deadline(drain) => {
            tracing::warn!(
                timeout_secs = drain.as_secs(),
                "shutdown timeout elapsed; dropping remaining connections"
            );
            Ok(())
        }
    }
}

/// First signal starts the clock; after `drain` the listen loop is aborted.
async fn drain_deadline(drain: Duration) {
    crate::shutdown::signal().await;
    tokio::time::sleep(drain).await;
}