o402 0.1.4

OpenAI-compatible gateway, paid with x402.
//! `GET /health` and `GET /ready`.

use axum::Json;
use axum::extract::State;
use axum::http::StatusCode;
use serde::Serialize;

use crate::state::AppState;

/// Liveness body.
#[derive(Serialize)]
pub(super) struct HealthBody {
    /// Always `ok` if the process is up.
    status: &'static str,
    /// Process name.
    service: &'static str,
    /// `CARGO_PKG_VERSION`.
    version: &'static str,
}

/// Readiness body.
#[derive(Serialize)]
pub(super) struct ReadyBody {
    /// `ok` or `unavailable`.
    status: &'static str,
}

/// Liveness. Always 200 while the HTTP server is accepting.
#[allow(clippy::unused_async, reason = "axum handler")]
pub(super) async fn health() -> Json<HealthBody> {
    Json(HealthBody {
        status: "ok",
        service: "o402",
        version: env!("CARGO_PKG_VERSION"),
    })
}

/// Readiness. 200 when payment is off, or when `/supported` succeeds.
pub(super) async fn ready(State(state): State<AppState>) -> (StatusCode, Json<ReadyBody>) {
    if !state.payment_enabled() {
        return (StatusCode::OK, Json(ReadyBody { status: "ok" }));
    }
    let Some(server) = state.resource_server() else {
        return (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(ReadyBody {
                status: "unavailable",
            }),
        );
    };
    match server.facilitator().supported().await {
        Ok(_) => (StatusCode::OK, Json(ReadyBody { status: "ok" })),
        Err(_) => (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(ReadyBody {
                status: "unavailable",
            }),
        ),
    }
}

#[cfg(test)]
mod tests {
    use axum::body::Body;
    use axum::http::{Request, StatusCode};
    use tower::ServiceExt;

    use crate::config::Config;
    use crate::http::app;

    const UNPAID: &str = r#"
[payment]
enabled = false

[[upstreams]]
name = "openai"
base_url = "https://api.openai.com"
api_key = "sk-test"
"#;

    fn unpaid_app() -> axum::Router {
        let cfg = Config::from_toml_str(UNPAID).expect("unpaid config");
        app(cfg).expect("router")
    }

    fn json_str<'a>(body: &'a serde_json::Value, key: &str) -> &'a str {
        body.get(key)
            .and_then(serde_json::Value::as_str)
            .unwrap_or("")
    }

    #[tokio::test]
    async fn health_ok() {
        let response = unpaid_app()
            .oneshot(
                Request::get("/health")
                    .body(Body::empty())
                    .expect("request"),
            )
            .await
            .expect("response");
        assert_eq!(response.status(), StatusCode::OK, "health status");
        let bytes = axum::body::to_bytes(response.into_body(), 4096)
            .await
            .expect("body");
        let body: serde_json::Value = serde_json::from_slice(&bytes).expect("json");
        assert_eq!(json_str(&body, "status"), "ok", "status");
        assert_eq!(json_str(&body, "service"), "o402", "service");
        assert_eq!(
            json_str(&body, "version"),
            env!("CARGO_PKG_VERSION"),
            "version"
        );
    }

    #[tokio::test]
    async fn ready_ok_when_payment_off() {
        let response = unpaid_app()
            .oneshot(Request::get("/ready").body(Body::empty()).expect("request"))
            .await
            .expect("response");
        assert_eq!(response.status(), StatusCode::OK, "ready status");
        let bytes = axum::body::to_bytes(response.into_body(), 4096)
            .await
            .expect("body");
        let body: serde_json::Value = serde_json::from_slice(&bytes).expect("json");
        assert_eq!(json_str(&body, "status"), "ok", "ready body");
    }

    const PAID: &str = r#"
[server]
base_url = "https://o402.example.com"

[payment]
enabled = true

[payment.facilitator]
url = "http://127.0.0.1:1"
timeout_secs = 1
supported_cache_ttl_secs = 0

[payment.pay_to]
"eip155:*" = "0x0000000000000000000000000000000000000001"

[[payment.accepts]]
scheme = "exact"
network = "eip155:8453"
asset = "usdc"

[pricing.default]
scheme = "exact"
price = "0.001"
request_floor = "0.00001"
input_per_million = "0.15"
output_per_million = "0.60"
max_input_tokens = 128000
default_max_output_tokens = 16384

[[upstreams]]
name = "openai"
base_url = "https://api.openai.com"
api_key = "sk-test"

[[models]]
id = "gpt-4o-mini"
upstream = "openai"
"#;

    #[tokio::test]
    async fn ready_unavailable_when_payment_on() {
        let cfg = Config::from_toml_str(PAID).expect("paid config");
        let response = app(cfg)
            .expect("router")
            .oneshot(Request::get("/ready").body(Body::empty()).expect("request"))
            .await
            .expect("response");
        assert_eq!(
            response.status(),
            StatusCode::SERVICE_UNAVAILABLE,
            "ready status"
        );
        let bytes = axum::body::to_bytes(response.into_body(), 4096)
            .await
            .expect("body");
        let body: serde_json::Value = serde_json::from_slice(&bytes).expect("json");
        assert_eq!(json_str(&body, "status"), "unavailable", "ready body");
    }
}