openlatch-client 0.6.19

//! 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:
//!
//! 0. **The route, before any token rule** (AGT-6 C-7). On `opencode.ai`, OpenCode Go's
//!    `/zen/go/v1/*` → `subscription` and Zen's `/zen/v1/*` (and its `/inference/v1/*` twin) →
//!    `api_key`. On a request attributed to OpenCode, `api.githubcopilot.com` → `subscription`
//!    (plan D-17). Decided from [`RouteFacts`] — the CONNECT host, the path, the attribution —
//!    never from a token.
//! 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::{self, ProviderHost, WireFormat};

/// OpenCode's own host, where Zen and OpenCode Go are served.
const OPENCODE_HOST: &str = "opencode.ai";

/// What the relay knows about a request's route before any token is read: the CONNECT host of
/// the tunnel it arrived in, its path, and whether it is attributed to OpenCode (AGT-6 C-7).
#[derive(Clone, Copy, Debug, Default)]
pub struct RouteFacts<'a> {
    /// The tunnel's CONNECT host; `None` outside a tunnel.
    pub host: Option<&'a str>,
    /// `uri.path()` — never the query.
    pub path: &'a str,
    /// The request is attributed to OpenCode: the CONNECT host's sole declarer, or OpenCode's
    /// claim on a host it shares (`proxy::opencode_claims_tunnel`). Gates the D-17 identities
    /// and Copilot's billing, so Cline's records on a shared host are unchanged.
    pub opencode: bool,
}

impl<'a> RouteFacts<'a> {
    /// A request outside any tunnel: no host, no attribution. Only the path speaks.
    pub fn path_only(path: &'a str) -> Self {
        RouteFacts {
            host: None,
            path,
            opencode: false,
        }
    }

    /// True on OpenCode's own host.
    fn on_opencode_host(&self) -> bool {
        self.host.is_some_and(is_opencode_host)
    }

    /// The new provider this request's CONNECT host belongs to, when it is attributed to
    /// OpenCode — never otherwise (plan D-17).
    fn opencode_provider(&self) -> Option<ProviderHost> {
        self.host
            .filter(|_| self.opencode)
            .and_then(ProviderHost::of)
    }
}

fn is_opencode_host(host: &str) -> bool {
    super::host_match::normalize(host) == OPENCODE_HOST
}

/// The provider identity the ROUTE decides, or `None` to leave it to the format (AGT-6 C-7,
/// plan D-17): `zen` for a Zen or OpenCode Go turn route on `opencode.ai`, and, on a request
/// attributed to OpenCode only, the identity of the provider its CONNECT host belongs to
/// (`github_copilot`, `groq`, `bedrock`, `azure`, `vertex_ai`). A route identity is
/// authoritative: `emit` writes it as-is and never runs the priced-identity rewrites over it.
pub fn route_provider(route: &RouteFacts) -> Option<&'static str> {
    if route.on_opencode_host() && wire_format::is_zen_route(route.path) {
        return Some("zen");
    }
    route.opencode_provider().map(ProviderHost::identity)
}

/// 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, the
/// route, and the request headers. See the module docs for the ordered rules.
pub fn detect_billing(wire: WireFormat, route: &RouteFacts, headers: &HeaderMap) -> BillingMode {
    // The route first, as for Cursor (AGT-6 C-7): Zen and OpenCode Go are told apart by their
    // path, whatever key the request carries, and Copilot bills through its plan.
    if route.on_opencode_host() {
        if wire_format::is_zen_go_route(route.path) {
            return BillingMode::Subscription;
        }
        if wire_format::is_zen_route(route.path) {
            return BillingMode::ApiKey;
        }
    }
    if route.opencode_provider() == Some(ProviderHost::GithubCopilot) {
        return BillingMode::Subscription;
    }
    // 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,
            &RouteFacts::path_only("/v1/messages"),
            &hdr(pairs),
        )
    }

    fn tunnel<'a>(host: &'a str, path: &'a str, opencode: bool) -> RouteFacts<'a> {
        RouteFacts {
            host: Some(host),
            path,
            opencode,
        }
    }

    /// AC-16, C-7: Zen bills `api_key` and OpenCode Go `subscription`, on all three formats,
    /// decided by the route before any token rule — whatever credential rides the request.
    #[test]
    fn billing_zen_and_go_are_decided_by_route() {
        let creds = [
            hdr(&[]),
            hdr(&[("x-api-key", "sk-zen-abc")]),
            hdr(&[("authorization", "Bearer sk-abc")]),
            hdr(&[("authorization", &format!("Bearer {JWT}"))]),
        ];
        for (path, wire) in [
            ("/zen/v1/messages", WireFormat::AnthropicMessages),
            ("/zen/v1/responses", WireFormat::OpenAiResponses),
            (
                "/zen/v1/chat/completions",
                WireFormat::OpenAiChatCompletions,
            ),
            ("/inference/v1/messages", WireFormat::AnthropicMessages),
        ] {
            for h in &creds {
                for opencode in [true, false] {
                    let r = tunnel("opencode.ai", path, opencode);
                    assert_eq!(detect_billing(wire, &r, h), BillingMode::ApiKey, "{path}");
                    assert_eq!(route_provider(&r), Some("zen"), "{path}");
                }
            }
        }
        for (path, wire) in [
            ("/zen/go/v1/messages", WireFormat::AnthropicMessages),
            ("/zen/go/v1/responses", WireFormat::OpenAiResponses),
            (
                "/zen/go/v1/chat/completions",
                WireFormat::OpenAiChatCompletions,
            ),
            ("/zen/go/v1/messages/", WireFormat::AnthropicMessages),
        ] {
            for h in &creds {
                let r = tunnel("opencode.ai", path, true);
                assert_eq!(
                    detect_billing(wire, &r, h),
                    BillingMode::Subscription,
                    "{path}"
                );
                assert_eq!(route_provider(&r), Some("zen"), "{path}");
            }
        }
    }

    /// A Zen path on any other host is that host's own: no route rule, the token decides.
    #[test]
    fn billing_zen_path_off_opencode_host_is_not_zen() {
        let h = hdr(&[("authorization", "Bearer sk-abc")]);
        for r in [
            tunnel("api.anthropic.com", "/zen/go/v1/messages", true),
            RouteFacts::path_only("/zen/go/v1/messages"),
        ] {
            assert_eq!(
                detect_billing(WireFormat::AnthropicMessages, &r, &h),
                BillingMode::ApiKey
            );
            assert_eq!(route_provider(&r), None);
        }
        // Any other `opencode.ai` path names no route identity.
        assert_eq!(route_provider(&tunnel("opencode.ai", "/share", true)), None);
    }

    /// D-17: Copilot bills `subscription` on an OpenCode request only; the same host from any
    /// other author keeps the token rules.
    #[test]
    fn billing_copilot_is_subscription_for_opencode_only() {
        let h = hdr(&[("authorization", "Bearer sk-abc")]);
        let wire = WireFormat::OpenAiChatCompletions;
        let oc = tunnel("api.githubcopilot.com", "/chat/completions", true);
        assert_eq!(detect_billing(wire, &oc, &h), BillingMode::Subscription);
        let other = tunnel("api.githubcopilot.com", "/chat/completions", false);
        assert_eq!(detect_billing(wire, &other, &h), BillingMode::ApiKey);
    }

    /// D-17: each new provider's identity on an OpenCode request, by CONNECT host; none on the
    /// same host from another author (a Cline Groq turn keeps today's identity).
    #[test]
    fn billing_route_provider_identities_are_opencode_only() {
        for (host, identity) in [
            ("api.githubcopilot.com", "github_copilot"),
            ("api.groq.com", "groq"),
            ("API.GROQ.COM.", "groq"),
            ("bedrock-runtime.us-east-1.amazonaws.com", "bedrock"),
            ("bedrock-runtime.cn-north-1.amazonaws.com.cn", "bedrock"),
            ("myres.openai.azure.com", "azure"),
            ("myres.cognitiveservices.azure.com", "azure"),
            ("myres.services.ai.azure.com", "azure"),
            ("aiplatform.googleapis.com", "vertex_ai"),
            ("us-east5-aiplatform.googleapis.com", "vertex_ai"),
        ] {
            assert_eq!(
                route_provider(&tunnel(host, "/x", true)),
                Some(identity),
                "{host}"
            );
            assert_eq!(route_provider(&tunnel(host, "/x", false)), None, "{host}");
        }
        for host in [
            "api.anthropic.com",
            "api.openai.com",
            "generativelanguage.googleapis.com",
            "openrouter.ai",
            "a.b.openai.azure.com",
            "bedrock-runtime..amazonaws.com",
        ] {
            assert_eq!(route_provider(&tunnel(host, "/x", true)), None, "{host}");
        }
        // Groq bills by the existing token rules, even for OpenCode.
        let h = hdr(&[("authorization", "Bearer gsk_abc")]);
        assert_eq!(
            detect_billing(
                WireFormat::OpenAiChatCompletions,
                &tunnel("api.groq.com", "/openai/v1/chat/completions", true),
                &h
            ),
            BillingMode::Unknown
        );
    }

    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,
                    &RouteFacts::path_only("/agent.v1.AgentService/Run"),
                    &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, &RouteFacts::path_only("/v1/x"), &h),
                    BillingMode::Unknown,
                    "{wire:?}"
                );
            }
        }
    }
}