o402 0.1.5

//! Billing class for Bifrost / OpenAI-compatible HTTP paths.

use axum::http::Method;
use axum::http::header::{self, HeaderMap};

/// How a request is billed.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum Bill {
    /// Proxy with no x402 (`GET /v1/models`).
    Unpaid,
    /// Flat `pricing.default.price` (or model override).
    Exact,
    /// Token meter from JSON `usage` (chat / completions / embeddings / responses).
    Upto,
    /// Not exported (WebSocket realtime, Bifrost admin).
    Reject,
}

/// Classifies `method` + `path` + hop headers.
#[must_use]
pub(crate) fn classify(method: &Method, path: &str, headers: &HeaderMap) -> Bill {
    if is_websocket(headers) || is_realtime(path) {
        return Bill::Reject;
    }
    let path = openai_path(path);
    if *method == Method::GET && is_models(path) {
        return Bill::Unpaid;
    }
    if *method == Method::POST && is_upto_post(path) {
        return Bill::Upto;
    }
    if is_exported(path) {
        return Bill::Exact;
    }
    Bill::Reject
}

/// `/openai/v1/...` → `/v1/...`; other prefixes unchanged.
#[must_use]
pub(crate) fn openai_path(path: &str) -> &str {
    path.strip_prefix("/openai").unwrap_or(path)
}

fn is_websocket(headers: &HeaderMap) -> bool {
    headers
        .get(header::UPGRADE)
        .is_some_and(|value| value.as_bytes().eq_ignore_ascii_case(b"websocket"))
}

fn is_realtime(path: &str) -> bool {
    let path = openai_path(path);
    path == "/v1/realtime" || path.starts_with("/v1/realtime/")
}

fn is_models(path: &str) -> bool {
    path == "/v1/models" || path.starts_with("/v1/models/")
}

fn is_upto_post(path: &str) -> bool {
    matches!(
        path,
        "/v1/chat/completions" | "/v1/completions" | "/v1/embeddings" | "/v1/responses"
    )
}

fn is_exported(path: &str) -> bool {
    path == "/v1"
        || path.starts_with("/v1/")
        || path.starts_with("/anthropic/")
        || path.starts_with("/genai/")
        || path.starts_with("/bedrock/")
        || path.starts_with("/cohere/")
        || path.starts_with("/litellm/")
        || path.starts_with("/langchain/")
        || path.starts_with("/pydanticai/")
}

#[cfg(test)]
mod tests {
    use axum::http::header::HeaderValue;
    use axum::http::{HeaderMap, Method};

    use super::{Bill, classify};

    fn headers() -> HeaderMap {
        HeaderMap::new()
    }

    fn upgrade() -> HeaderMap {
        let mut headers = HeaderMap::new();
        headers.insert(
            axum::http::header::UPGRADE,
            HeaderValue::from_static("websocket"),
        );
        headers
    }

    #[test]
    fn upto_posts() {
        for path in [
            "/v1/chat/completions",
            "/openai/v1/chat/completions",
            "/v1/completions",
            "/v1/embeddings",
            "/v1/responses",
            "/openai/v1/responses",
        ] {
            assert_eq!(
                classify(&Method::POST, path, &headers()),
                Bill::Upto,
                "{path}"
            );
        }
    }

    #[test]
    fn exact_inference_and_media() {
        for path in [
            "/v1/images/generations",
            "/openai/v1/images/edits",
            "/v1/audio/speech",
            "/v1/audio/transcriptions",
            "/v1/files",
            "/v1/batches",
            "/v1/videos",
            "/v1/rerank",
            "/v1/responses/abc",
            "/anthropic/v1/messages",
            "/genai/v1beta/models/gemini-pro:generateContent",
        ] {
            assert_eq!(
                classify(&Method::POST, path, &headers()),
                Bill::Exact,
                "{path}"
            );
        }
        assert_eq!(
            classify(&Method::GET, "/v1/files", &headers()),
            Bill::Exact,
            "files list"
        );
        assert_eq!(
            classify(&Method::GET, "/v1/chat/completions", &headers()),
            Bill::Exact,
            "wrong method is exact not upto"
        );
    }

    #[test]
    fn unpaid_models() {
        assert_eq!(
            classify(&Method::GET, "/v1/models", &headers()),
            Bill::Unpaid,
            "models"
        );
        assert_eq!(
            classify(&Method::GET, "/openai/v1/models", &headers()),
            Bill::Unpaid,
            "openai models"
        );
        assert_eq!(
            classify(&Method::GET, "/v1/models/gpt-4o", &headers()),
            Bill::Unpaid,
            "retrieve"
        );
    }

    #[test]
    fn reject_realtime_and_admin() {
        assert_eq!(
            classify(&Method::GET, "/v1/realtime", &headers()),
            Bill::Reject,
            "realtime"
        );
        assert_eq!(
            classify(&Method::GET, "/openai/v1/realtime", &upgrade()),
            Bill::Reject,
            "upgrade"
        );
        assert_eq!(
            classify(&Method::POST, "/api/providers", &headers()),
            Bill::Reject,
            "admin"
        );
        assert_eq!(classify(&Method::GET, "/", &headers()), Bill::Reject, "ui");
    }
}