mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
3c1991887f
`classify_auth_signals` gated PAYG/OAUTH token-shape detection on
`authorization.startswith("Bearer ")`. RFC 7235 §2.1 defines the auth-scheme
token as case-insensitive, so a client sending `Authorization: bearer sk-...`
skipped the Bearer branch entirely and fell through to the `elif auth:`
OAUTH classification — misclassifying a PAYG API key as OAUTH.
Auth mode drives compression-policy routing and the cost/TOIN `auth_mode`
label, so the misclassification silently changes how a request is compressed
and accounted. Fold only the scheme for the comparison via `str.partition`;
the credential keeps its original case (tokens are case-sensitive).
93 lines
3.5 KiB
Python
93 lines
3.5 KiB
Python
"""Tests for pure auth and client classification policy."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from headroom.proxy.auth_policy import (
|
|
AuthMode,
|
|
AuthSignals,
|
|
classify_auth_signals,
|
|
classify_client_signals,
|
|
should_stamp_codex_client_signals,
|
|
)
|
|
|
|
|
|
def test_subscription_user_agent_wins_over_oauth_token() -> None:
|
|
signals = AuthSignals(
|
|
user_agent="claude-code/1.5.0 (linux; x86_64)",
|
|
authorization="Bearer sk-ant-oat01-abc123",
|
|
)
|
|
|
|
assert classify_auth_signals(signals) is AuthMode.SUBSCRIPTION
|
|
|
|
|
|
def test_oauth_bearer_token_shapes_are_oauth() -> None:
|
|
jwt = "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0In0.signature"
|
|
|
|
assert classify_auth_signals(AuthSignals(authorization="Bearer sk-ant-oat01-abc")) is (
|
|
AuthMode.OAUTH
|
|
)
|
|
assert classify_auth_signals(AuthSignals(authorization=f"Bearer {jwt}")) is AuthMode.OAUTH
|
|
|
|
|
|
def test_payg_key_shapes_are_payg() -> None:
|
|
assert classify_auth_signals(AuthSignals(authorization="Bearer sk-ant-api03-abc")) is (
|
|
AuthMode.PAYG
|
|
)
|
|
assert classify_auth_signals(AuthSignals(x_api_key="sk-ant-api03-abc")) is AuthMode.PAYG
|
|
assert classify_auth_signals(AuthSignals(x_goog_api_key="AIzaSyDUMMY")) is AuthMode.PAYG
|
|
|
|
|
|
def test_client_explicit_override_wins_over_user_agent() -> None:
|
|
signals = AuthSignals(user_agent="claude-code/1.2.3", x_client=" AIDER ")
|
|
|
|
assert classify_client_signals(signals) == "aider"
|
|
|
|
|
|
def test_grok_build_user_agent_is_subscription_client() -> None:
|
|
signals = AuthSignals(user_agent="grok/1.2.3")
|
|
|
|
assert classify_auth_signals(signals) is AuthMode.SUBSCRIPTION
|
|
assert classify_client_signals(signals) == "grok_build"
|
|
|
|
|
|
def test_codex_stamp_only_for_unidentified_responses_callers() -> None:
|
|
assert should_stamp_codex_client_signals("/v1/responses", AuthSignals()) is True
|
|
assert (
|
|
should_stamp_codex_client_signals(
|
|
"/v1/responses/foo",
|
|
AuthSignals(user_agent="codex-cli/0.5"),
|
|
)
|
|
is False
|
|
)
|
|
assert should_stamp_codex_client_signals("/v1/chat/completions", AuthSignals()) is False
|
|
|
|
|
|
def test_bearer_scheme_is_case_insensitive() -> None:
|
|
# RFC 7235 §2.1: the auth-scheme token is case-insensitive. A lowercase or
|
|
# mixed-case "bearer" carrying a PAYG key must classify as PAYG, not fall
|
|
# through to the non-Bearer OAUTH branch.
|
|
for scheme in ("bearer", "BEARER", "Bearer", "BeArEr"):
|
|
signals = AuthSignals(authorization=f"{scheme} sk-ant-api03-abc")
|
|
assert classify_auth_signals(signals) is AuthMode.PAYG, scheme
|
|
|
|
|
|
def test_bearer_case_insensitive_preserves_oauth_and_jwt_shapes() -> None:
|
|
jwt = "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0In0.signature"
|
|
assert classify_auth_signals(AuthSignals(authorization="bearer sk-ant-oat01-abc")) is (
|
|
AuthMode.OAUTH
|
|
)
|
|
assert classify_auth_signals(AuthSignals(authorization=f"bearer {jwt}")) is AuthMode.OAUTH
|
|
|
|
|
|
def test_credential_case_is_not_folded() -> None:
|
|
# Only the scheme is case-folded; the token keeps its case, so a PAYG
|
|
# `sk-...` prefix still matches exactly (it is lowercase by construction).
|
|
signals = AuthSignals(authorization="Bearer sk-PROJ-Abc123")
|
|
assert classify_auth_signals(signals) is AuthMode.PAYG
|
|
|
|
|
|
def test_non_bearer_scheme_still_oauth() -> None:
|
|
# AWS SigV4 (Bedrock) and any other non-Bearer scheme keep the OAUTH
|
|
# passthrough-prefer classification.
|
|
signals = AuthSignals(authorization="AWS4-HMAC-SHA256 Credential=AKIA/...")
|
|
assert classify_auth_signals(signals) is AuthMode.OAUTH
|