openlatch-client 0.6.12

//! Billing-mode detection (D-11).
//!
//! Inspect the caller's auth shape on the request — we already read the headers
//! to forward them, so this adds no work. **Default to `unknown`, never guess**
//! (D-20): an `unknown` suppresses dollar figures downstream (E-7) rather than
//! mispricing a call.
//!
//! A `Bearer` token is classified by its PREFIX or SHAPE alone. The token itself
//! is never logged, stored, emitted or put in an error — only [`BillingMode`]
//! leaves this module. The rules, in order:
//!
//! 1. Cursor's own wire format → `subscription`, whatever its token looks like.
//!    It is decided by the route-resolved [`WireFormat`], never by the token.
//! 2. `x-api-key` present → `api_key`.
//! 3. `Bearer sk-ant-oat…` → `subscription` (Claude OAuth).
//! 4. `Bearer` + a JWT (three non-empty base64url segments, starting `eyJ`) →
//!    `subscription` (ChatGPT / Codex OAuth), or `Bearer ya29.…` →
//!    `subscription` (a Google OAuth access token, e.g. Gemini CLI signed in
//!    with Google).
//! 5. `Bearer sk-…` (`sk-ant-api…`, `sk-proj-…`, any other `sk-`) → `api_key`.
//!    Checked AFTER rule 3, which is the more specific `sk-` prefix.
//! 6. Any other `Bearer`, or no/empty `Authorization` → `unknown`.
//!
//! ⚠️ **Unverified split (I-1 OQ3).** The `x-api-key` vs OAuth-`Bearer` mapping
//! has not been confirmed against a real Claude.ai Pro/Max subscription session.
//! When unsure the honest answer is `unknown`; track how often it fires (it
//! silently empties the money view — a logged risk).

use axum::http::header::{HeaderMap, AUTHORIZATION};

use super::wire_format::WireFormat;

/// Frozen enum `billing_mode = api_key | subscription | unknown`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BillingMode {
    /// `x-api-key`, or a `Bearer sk-…` key that is not a Claude OAuth token →
    /// pay-as-you-go API key billing.
    ApiKey,
    /// A Claude OAuth `Bearer sk-ant-oat…`, a JWT `Bearer` (ChatGPT / Codex
    /// OAuth), a Google OAuth `Bearer ya29.…`, or any Cursor wire-format
    /// request → a subscription session.
    Subscription,
    /// No recognised shape — never guessed.
    Unknown,
}

impl BillingMode {
    /// Frozen wire string.
    pub fn as_str(&self) -> &'static str {
        match self {
            BillingMode::ApiKey => "api_key",
            BillingMode::Subscription => "subscription",
            BillingMode::Unknown => "unknown",
        }
    }
}

/// Classify the caller's billing mode from the route-resolved wire format and
/// the request headers alone. See the module docs for the ordered rules.
pub fn detect_billing(wire: WireFormat, headers: &HeaderMap) -> BillingMode {
    // Cursor bills through its own plan whatever credential it carries, so its
    // label comes from the route, never from the token (user decision).
    if wire == WireFormat::CursorConnectRpc {
        return BillingMode::Subscription;
    }
    if headers.contains_key("x-api-key") {
        return BillingMode::ApiKey;
    }
    match bearer_token(headers) {
        Some(t) if t.starts_with("sk-ant-oat") => BillingMode::Subscription,
        Some(t) if is_jwt(t) || t.starts_with("ya29.") => BillingMode::Subscription,
        Some(t) if t.starts_with("sk-") => BillingMode::ApiKey,
        _ => BillingMode::Unknown,
    }
}

/// The token of an `Authorization: Bearer <token>` header, or `None` when the
/// header is absent, not `Bearer` (case-insensitive scheme), or empty. The
/// outer trim means a token, once split off, is never empty.
fn bearer_token(headers: &HeaderMap) -> Option<&str> {
    let v = headers.get(AUTHORIZATION)?.to_str().ok()?.trim();
    let (scheme, token) = v.split_once(' ')?;
    scheme
        .eq_ignore_ascii_case("bearer")
        .then_some(token.trim_start())
}

/// True for a JWT shape: three non-empty base64url segments separated by `.`,
/// the first starting `eyJ` (base64url of `{"`). Shape only — nothing is decoded.
fn is_jwt(token: &str) -> bool {
    let b64url = |s: &str| {
        !s.is_empty()
            && s.bytes()
                .all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_')
    };
    let mut parts = token.split('.');
    matches!(
        (parts.next(), parts.next(), parts.next(), parts.next()),
        (Some(h), Some(p), Some(s), None)
            if h.starts_with("eyJ") && b64url(h) && b64url(p) && b64url(s)
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::http::{HeaderName, HeaderValue};

    const JWT: &str = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ4In0.c2lnbmF0dXJl";

    fn hdr(pairs: &[(&str, &str)]) -> HeaderMap {
        let mut h = HeaderMap::new();
        for (k, v) in pairs {
            h.insert(
                HeaderName::from_bytes(k.as_bytes()).unwrap(),
                HeaderValue::from_str(v).unwrap(),
            );
        }
        h
    }

    fn anthropic(pairs: &[(&str, &str)]) -> BillingMode {
        detect_billing(WireFormat::AnthropicMessages, &hdr(pairs))
    }

    fn bearer(token: &str) -> BillingMode {
        anthropic(&[("authorization", &format!("Bearer {token}"))])
    }

    #[test]
    fn x_api_key_is_api_key() {
        let mode = anthropic(&[("x-api-key", "sk-ant-xyz")]);
        assert_eq!(mode, BillingMode::ApiKey);
        assert_eq!(mode.as_str(), "api_key");
    }

    #[test]
    fn claude_oauth_bearer_is_subscription() {
        assert_eq!(bearer("sk-ant-oat01-abc"), BillingMode::Subscription);
        assert_eq!(bearer("sk-ant-oat01-abc").as_str(), "subscription");
    }

    #[test]
    fn jwt_bearer_is_subscription() {
        assert_eq!(bearer(JWT), BillingMode::Subscription);
    }

    #[test]
    fn google_oauth_bearer_is_subscription() {
        assert_eq!(bearer("ya29.a0AfH6SMBabc"), BillingMode::Subscription);
    }

    #[test]
    fn anthropic_api_key_bearer_is_api_key() {
        assert_eq!(bearer("sk-ant-api03-abc"), BillingMode::ApiKey);
    }

    #[test]
    fn openai_project_key_bearer_is_api_key() {
        assert_eq!(bearer("sk-proj-abc"), BillingMode::ApiKey);
    }

    #[test]
    fn generic_sk_bearer_is_api_key() {
        assert_eq!(bearer("sk-abc123"), BillingMode::ApiKey);
    }

    #[test]
    fn opaque_bearer_is_unknown() {
        // Was `subscription` before: any bearer used to be read as OAuth.
        assert_eq!(bearer("oauth-token-value"), BillingMode::Unknown);
    }

    #[test]
    fn malformed_jwt_is_unknown() {
        assert_eq!(bearer("eyJabc"), BillingMode::Unknown);
        assert_eq!(bearer("eyJabc.def"), BillingMode::Unknown);
        assert_eq!(bearer("eyJabc..ghi"), BillingMode::Unknown);
        assert_eq!(bearer("eyJabc.def.ghi.jkl"), BillingMode::Unknown);
        assert_eq!(bearer("abc.def.ghi"), BillingMode::Unknown);
        assert_eq!(bearer("eyJa+c.def.ghi"), BillingMode::Unknown);
    }

    #[test]
    fn bearer_scheme_is_case_insensitive() {
        let h = [("authorization", "bearer sk-ant-oat01-abc")];
        assert_eq!(anthropic(&h), BillingMode::Subscription);
        let h = [("authorization", "BEARER sk-proj-abc")];
        assert_eq!(anthropic(&h), BillingMode::ApiKey);
    }

    #[test]
    fn no_recognised_shape_is_unknown() {
        assert_eq!(anthropic(&[]), BillingMode::Unknown);
        // A malformed/empty Authorization does not become subscription.
        assert_eq!(
            anthropic(&[("authorization", "Bearer")]),
            BillingMode::Unknown
        );
        assert_eq!(
            anthropic(&[("authorization", "Bearer ")]),
            BillingMode::Unknown
        );
        assert_eq!(
            anthropic(&[("authorization", "Basic sk-abc")]),
            BillingMode::Unknown
        );
    }

    #[test]
    fn x_api_key_wins_over_bearer() {
        // If both are present the pay-as-you-go key is authoritative.
        let h = [
            ("x-api-key", "sk-ant-xyz"),
            ("authorization", "Bearer sk-ant-oat01-t"),
        ];
        assert_eq!(anthropic(&h), BillingMode::ApiKey);
    }

    #[test]
    fn cursor_wire_format_is_subscription_whatever_the_token() {
        for h in [
            hdr(&[]),
            hdr(&[("authorization", "Bearer opaque-cursor-token")]),
            hdr(&[("authorization", "Bearer sk-proj-abc")]),
            hdr(&[("x-api-key", "sk-ant-xyz")]),
        ] {
            assert_eq!(
                detect_billing(WireFormat::CursorConnectRpc, &h),
                BillingMode::Subscription
            );
        }
    }

    #[test]
    fn non_cursor_opaque_bearer_is_unknown() {
        // The Cursor exemption is the route's, not the token's: the same opaque
        // bearer on any other format stays unknown.
        let h = hdr(&[("authorization", "Bearer opaque-cursor-token")]);
        for wire in WireFormat::ALL {
            if wire != WireFormat::CursorConnectRpc {
                assert_eq!(detect_billing(wire, &h), BillingMode::Unknown, "{wire:?}");
            }
        }
    }
}