o402 0.1.5

//! `/health`, `/ready`, `/v1/pricing`, and billed catch-all proxy.

pub(crate) mod bill;
mod context;
mod health;
pub(crate) mod openai_error;
mod pricing;

use std::fmt;
use std::time::Duration;

use axum::http::header::HeaderName;
use axum::http::{HeaderValue, Request, Response, StatusCode};
use axum::routing::{any, get};
use axum::{Router, middleware};
use tower_http::cors::{AllowOrigin, Any, CorsLayer};
use tower_http::limit::RequestBodyLimitLayer;
use tower_http::timeout::TimeoutLayer;
use tower_http::trace::{
    DefaultOnFailure, DefaultOnRequest, DefaultOnResponse, MakeSpan, OnFailure, OnRequest,
    OnResponse, TraceLayer,
};
use tracing::Span;

pub(crate) use self::context::RequestContext;
use crate::config::{Config, ConfigError};
use crate::error::AppError;
use crate::state::AppState;

/// Build the Axum app from validated config.
///
/// # Errors
///
/// Returns [`AppError::Config`] when a CORS origin is not a valid header value,
/// or [`AppError::HttpClient`] when an upstream client cannot be built.
#[cfg(test)]
pub(crate) fn app(config: Config) -> Result<Router, AppError> {
    Ok(app_and_state(config)?.0)
}

/// Build the Axum app and the state handle (for inflight drain).
///
/// # Errors
///
/// Returns [`AppError::Config`] when a CORS origin is not a valid header value,
/// or [`AppError::HttpClient`] when an upstream client cannot be built.
pub(crate) fn app_and_state(config: Config) -> Result<(Router, AppState), AppError> {
    let body_limit = config.server.body_limit_bytes;
    let request_timeout = Duration::from_secs(config.server.request_timeout_secs);
    let cors = cors_layer(&config.server.cors_origins)?;
    let state = AppState::new(config)?;

    let router = Router::new()
        .route("/health", get(health::health))
        .route("/ready", get(health::ready))
        .route(
            "/v1/pricing",
            get(pricing::list).fallback(openai_error::unlisted),
        )
        .route("/v1/{*rest}", any(crate::proxy::handler))
        .route("/v1", any(crate::proxy::handler))
        .route("/openai/{*rest}", any(crate::proxy::handler))
        .route("/anthropic/{*rest}", any(crate::proxy::handler))
        .route("/genai/{*rest}", any(crate::proxy::handler))
        .route("/bedrock/{*rest}", any(crate::proxy::handler))
        .route("/cohere/{*rest}", any(crate::proxy::handler))
        .route("/litellm/{*rest}", any(crate::proxy::handler))
        .route("/langchain/{*rest}", any(crate::proxy::handler))
        .route("/pydanticai/{*rest}", any(crate::proxy::handler))
        .layer(middleware::from_fn_with_state(
            state.clone(),
            context::extract_context,
        ))
        .with_state(state.clone())
        .layer(RequestBodyLimitLayer::new(body_limit))
        .layer(TimeoutLayer::with_status_code(
            StatusCode::GATEWAY_TIMEOUT,
            request_timeout,
        ))
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(SkipProbes)
                .on_request(SkipProbeOnRequest)
                .on_response(SkipProbeOnResponse)
                .on_failure(SkipProbeOnFailure),
        );

    let router = match cors {
        Some(layer) => router.layer(layer),
        None => router,
    };
    Ok((router, state))
}

fn cors_layer(origins: &[String]) -> Result<Option<CorsLayer>, AppError> {
    if origins.is_empty() {
        return Ok(None);
    }
    // HeaderName::from_static requires lowercase. CorsLayer is outermost and
    // overwrites Access-Control-Expose-Headers from ensure_expose_headers.
    let expose = [
        HeaderName::from_static("payment-required"),
        HeaderName::from_static("payment-response"),
    ];
    if origins.iter().any(|origin| origin == "*") {
        tracing::warn!("server.cors.origins = [\"*\"] allows any Origin");
        return Ok(Some(
            CorsLayer::new()
                .allow_origin(Any)
                .allow_methods(Any)
                .allow_headers(Any)
                .expose_headers(expose),
        ));
    }
    let parsed = origins
        .iter()
        .map(|origin| HeaderValue::from_str(origin))
        .collect::<Result<Vec<_>, _>>()
        .map_err(|error| {
            AppError::Config(ConfigError::Validation(format!(
                "invalid CORS origin: {error}"
            )))
        })?;
    Ok(Some(
        CorsLayer::new()
            .allow_origin(AllowOrigin::list(parsed))
            .allow_methods(Any)
            .allow_headers(Any)
            .expose_headers(expose),
    ))
}

/// Skip `/health` and `/ready` so probes do not flood JSON logs.
#[derive(Clone, Copy, Debug)]
struct SkipProbes;

impl<B> MakeSpan<B> for SkipProbes {
    fn make_span(&mut self, request: &Request<B>) -> Span {
        let path = request.uri().path();
        if path == "/health" || path == "/ready" {
            return Span::none();
        }
        tracing::info_span!("http", method = %request.method(), path)
    }
}

/// `DefaultOnFailure` ignores the span and logs ERROR for 5xx, including `/ready` 503.
#[derive(Clone, Copy, Debug)]
struct SkipProbeOnFailure;

impl<FailureClass: fmt::Display> OnFailure<FailureClass> for SkipProbeOnFailure {
    fn on_failure(&mut self, class: FailureClass, latency: Duration, span: &Span) {
        if span.is_disabled() {
            return;
        }
        DefaultOnFailure::new().on_failure(class, latency, span);
    }
}

/// `DefaultOnResponse` logs DEBUG `"finished processing request"` for every probe.
#[derive(Clone, Copy, Debug)]
struct SkipProbeOnResponse;

impl<B> OnResponse<B> for SkipProbeOnResponse {
    fn on_response(self, response: &Response<B>, latency: Duration, span: &Span) {
        if span.is_disabled() {
            return;
        }
        DefaultOnResponse::new().on_response(response, latency, span);
    }
}

/// `DefaultOnRequest` logs DEBUG `"started processing request"` for every probe.
#[derive(Clone, Copy, Debug)]
struct SkipProbeOnRequest;

impl<B> OnRequest<B> for SkipProbeOnRequest {
    fn on_request(&mut self, request: &Request<B>, span: &Span) {
        if span.is_disabled() {
            return;
        }
        DefaultOnRequest::new().on_request(request, span);
    }
}